From 92b0953081b1649803cb101c7d28451ed21e9eb6 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Sun, 10 Jan 2021 08:35:06 -0800
Subject: [PATCH] Move around where colormap code lives

---
 manimlib/mobject/mobject.py                  |  9 ++----
 manimlib/shader_wrapper.py                   | 29 +++++---------------
 manimlib/shaders/inserts/finalize_color.glsl | 11 +++++++-
 3 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/manimlib/mobject/mobject.py b/manimlib/mobject/mobject.py
index 6863d021..b3bac0ac 100644
--- a/manimlib/mobject/mobject.py
+++ b/manimlib/mobject/mobject.py
@@ -1277,13 +1277,10 @@ class Mobject(object):
         """
         for char in "xyz":
             glsl_snippet = glsl_snippet.replace(char, "point." + char)
-        self.replace_shader_code(
-            "///// INSERT COLOR_MAP FUNCTION HERE /////",
-            get_colormap_code(colormap)
-        )
         self.set_color_by_code(
-            "color.rgb = colormap({}, {}, {});".format(
-                glsl_snippet, float(min_value), float(max_value)
+            "color.rgb = float_to_color({}, {}, {}, {});".format(
+                glsl_snippet, float(min_value), float(max_value),
+                get_colormap_code(colormap)
             )
         )
         return self
diff --git a/manimlib/shader_wrapper.py b/manimlib/shader_wrapper.py
index cef266f4..1195efd7 100644
--- a/manimlib/shader_wrapper.py
+++ b/manimlib/shader_wrapper.py
@@ -154,28 +154,13 @@ def get_shader_code_from_file(filename):
 
 
 def get_colormap_code(colormap="viridis"):
-    code = """
-        const vec3[9] COLOR_MAP_DATA = vec3[9](
-            // INSERT DATA //
-        );
-        vec3 colormap(float value, float min_val, float max_val){
-            float alpha = smoothstep(min_val, max_val, value);
-            int disc_alpha = min(int(alpha * 8), 7);
-            return mix(
-                COLOR_MAP_DATA[disc_alpha],
-                COLOR_MAP_DATA[disc_alpha + 1],
-                8.0 * alpha - disc_alpha
-            );
-        }
-    """
-    colors = get_cmap(colormap).colors
-    sparse_colors = [
-        colors[int(n)]
-        for n in np.linspace(0, len(colors) - 1, 9)
+    rgbs = get_cmap(colormap).colors  # Make more general?
+    sparse_rgbs = [
+        rgbs[int(n)]
+        for n in np.linspace(0, len(rgbs) - 1, 9)
     ]
-    insertion = "".join(
+    data = ",".join(
         "vec3({}, {}, {}),".format(*color)
-        for color in sparse_colors
+        for color in sparse_rgbs
     )
-    insertion = insertion[:-1]  # Remove final comma
-    return code.replace("// INSERT DATA //", insertion)
+    return f"vec3[9]({data})"
diff --git a/manimlib/shaders/inserts/finalize_color.glsl b/manimlib/shaders/inserts/finalize_color.glsl
index d72eee59..16406991 100644
--- a/manimlib/shaders/inserts/finalize_color.glsl
+++ b/manimlib/shaders/inserts/finalize_color.glsl
@@ -1,4 +1,13 @@
-///// INSERT COLOR_MAP FUNCTION HERE /////
+vec3 float_to_color(float value, float min_val, float max_val, vec3[9] colormap_data){
+    float alpha = smoothstep(min_val, max_val, value);
+    int disc_alpha = min(int(alpha * 8), 7);
+    return mix(
+        colormap_data[disc_alpha],
+        colormap_data[disc_alpha + 1],
+        8.0 * alpha - disc_alpha
+    );
+}
+
 
 vec4 add_light(vec4 color,
                vec3 point,