From a8dd39cae43af8b1f35c4cb93d135c1f3dd5ac09 Mon Sep 17 00:00:00 2001
From: Nicholas Sharp <nmwsharp@gmail.com>
Date: Sun, 29 Dec 2024 16:10:39 -0500
Subject: [PATCH] attempt to shade based on inf/nan values

---
 include/polyscope/scalar_quantity.h   |  8 +++++++
 include/polyscope/scalar_quantity.ipp | 32 +++++++++++++++++++++++++++
 src/render/opengl/shaders/rules.cpp   | 17 +++++++++++---
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/include/polyscope/scalar_quantity.h b/include/polyscope/scalar_quantity.h
index e999d039..ad47ee1a 100644
--- a/include/polyscope/scalar_quantity.h
+++ b/include/polyscope/scalar_quantity.h
@@ -46,6 +46,12 @@ class ScalarQuantity {
   QuantityT* setColorMap(std::string val);
   std::string getColorMap();
 
+  // Substitute colors for inf and nan
+  QuantityT* setInfColor(glm::vec3 val);
+  glm::vec3 getInfColor();
+  QuantityT* setNaNColor(glm::vec3 val);
+  glm::vec3 getNaNColor();
+
   // Data limits mapped in to colormap
   QuantityT* setMapRange(std::pair<double, double> val);
   std::pair<double, double> getMapRange();
@@ -74,6 +80,8 @@ class ScalarQuantity {
 
   // Parameters
   PersistentValue<std::string> cMap;
+  PersistentValue<glm::vec3> infColor;
+  PersistentValue<glm::vec3> nanColor;
   PersistentValue<bool> isolinesEnabled;
   PersistentValue<ScaledValue<float>> isolineWidth;
   PersistentValue<float> isolineDarkness;
diff --git a/include/polyscope/scalar_quantity.ipp b/include/polyscope/scalar_quantity.ipp
index 0c5fba9d..a52041a4 100644
--- a/include/polyscope/scalar_quantity.ipp
+++ b/include/polyscope/scalar_quantity.ipp
@@ -11,6 +11,8 @@ ScalarQuantity<QuantityT>::ScalarQuantity(QuantityT& quantity_, const std::vecto
       vizRangeMin(quantity.uniquePrefix() + "vizRangeMin", -777.), // set later,
       vizRangeMax(quantity.uniquePrefix() + "vizRangeMax", -777.), // including clearing cache
       cMap(quantity.uniquePrefix() + "cmap", defaultColorMap(dataType)),
+      infColor(quantity.uniquePrefix() + "infColor", glm::vec3(1., 0., 1.)),
+      nanColor(quantity.uniquePrefix() + "nanColor", glm::vec3(1., 0., 0.)),
       isolinesEnabled(quantity.uniquePrefix() + "isolinesEnabled", false),
       isolineWidth(quantity.uniquePrefix() + "isolineWidth",
                    absoluteValue((dataRange.second - dataRange.first) * 0.02)),
@@ -166,6 +168,12 @@ template <typename QuantityT>
 void ScalarQuantity<QuantityT>::buildScalarOptionsUI() {
   if (ImGui::MenuItem("Reset colormap range")) resetMapRange();
   if (ImGui::MenuItem("Enable isolines", NULL, isolinesEnabled.get())) setIsolinesEnabled(!isolinesEnabled.get());
+  if (ImGui::ColorEdit3("Color for +-Inf values", &infColor.get()[0], ImGuiColorEditFlags_NoInputs)) {
+    setInfColor(getInfColor());
+  }
+  if (ImGui::ColorEdit3("Color for NaN values", &nanColor.get()[0], ImGuiColorEditFlags_NoInputs)) {
+    setNaNColor(getNaNColor());
+  }
 }
 
 template <typename QuantityT>
@@ -182,6 +190,8 @@ template <typename QuantityT>
 void ScalarQuantity<QuantityT>::setScalarUniforms(render::ShaderProgram& p) {
   p.setUniform("u_rangeLow", vizRangeMin.get());
   p.setUniform("u_rangeHigh", vizRangeMax.get());
+  p.setUniform("u_infColor", infColor.get());
+  p.setUniform("u_nanColor", nanColor.get());
 
   if (isolinesEnabled.get()) {
     p.setUniform("u_modLen", getIsolineWidth());
@@ -236,6 +246,28 @@ std::string ScalarQuantity<QuantityT>::getColorMap() {
   return cMap.get();
 }
 
+template <typename QuantityT>
+QuantityT* ScalarQuantity<QuantityT>::setInfColor(glm::vec3 val) {
+  infColor = val;
+  requestRedraw();
+  return &quantity;
+}
+template <typename QuantityT>
+glm::vec3 ScalarQuantity<QuantityT>::getInfColor() {
+  return infColor.get();
+}
+
+template <typename QuantityT>
+QuantityT* ScalarQuantity<QuantityT>::setNaNColor(glm::vec3 val) {
+  nanColor = val;
+  requestRedraw();
+  return &quantity;
+}
+template <typename QuantityT>
+glm::vec3 ScalarQuantity<QuantityT>::getNaNColor() {
+  return nanColor.get();
+}
+
 template <typename QuantityT>
 QuantityT* ScalarQuantity<QuantityT>::setMapRange(std::pair<double, double> val) {
   vizRangeMin = val.first;
diff --git a/src/render/opengl/shaders/rules.cpp b/src/render/opengl/shaders/rules.cpp
index da510e4c..31e9622a 100644
--- a/src/render/opengl/shaders/rules.cpp
+++ b/src/render/opengl/shaders/rules.cpp
@@ -122,16 +122,27 @@ const ShaderReplacementRule SHADE_COLORMAP_VALUE(
           uniform float u_rangeHigh;
           uniform float u_rangeLow;
           uniform sampler1D t_colormap;
+          uniform vec3 u_infColor;
+          uniform vec3 u_nanColor;
         )"},
       {"GENERATE_SHADE_COLOR", R"(
-          float rangeTVal = (shadeValue - u_rangeLow) / (u_rangeHigh - u_rangeLow);
-          rangeTVal = clamp(rangeTVal, 0.f, 1.f);
-          vec3 albedoColor = texture(t_colormap, rangeTVal).rgb;
+          vec3 albedoColor = vec3(0., 0., 0.);
+          if (isnan(shadeValue)) {
+            albedoColor = u_nanColor;
+          } else if(isinf(shadeValue)) {
+            albedoColor = u_infColor;
+          } else {
+            float rangeTVal = (shadeValue - u_rangeLow) / (u_rangeHigh - u_rangeLow);
+            rangeTVal = clamp(rangeTVal, 0.f, 1.f);
+            albedoColor = texture(t_colormap, rangeTVal).rgb;
+          }
       )"}
     },
     /* uniforms */ {
         {"u_rangeLow", RenderDataType::Float},
         {"u_rangeHigh", RenderDataType::Float},
+        {"u_infColor", RenderDataType::Vector3Float},
+        {"u_nanColor", RenderDataType::Vector3Float},
     },
     /* attributes */ {},
     /* textures */ {