Updates to poly_fractal shaders

manimlib/shaders/meta_poly_fractal/frag.glsl

@@ -0,0 +1,130 @@
+#version 330
+
+uniform vec3 light_source_position;
+uniform float gloss;
+uniform float shadow;
+uniform float focal_distance;
+
+uniform vec4 color0;
+uniform vec4 color1;
+uniform vec4 color2;
+
+uniform vec2 root0;
+uniform vec2 root1;
+
+// uniform vec2 coef0;
+// uniform vec2 coef1;
+// uniform vec2 coef2;
+// uniform vec2 coef3;
+
+uniform vec2 z0;
+
+uniform float n_roots;
+uniform float n_steps;
+uniform float julia_highlight;
+uniform float color_mult;
+
+uniform vec2 frame_shape;
+
+in vec3 xyz_coords;
+
+out vec4 frag_color;
+
+#INSERT finalize_color.glsl
+
+const int MAX_DEGREE = 3;
+
+vec2 complex_mult(vec2 z, vec2 w){
+    return vec2(z.x * w.x - z.y * w.y, z.x * w.y + z.y * w.x);
+}
+
+vec2 complex_div(vec2 z, vec2 w){
+    return complex_mult(z, vec2(w.x, -w.y)) / (w.x * w.x + w.y * w.y);
+}
+
+vec2 complex_pow(vec2 z, int n){
+    vec2 result = vec2(1.0, 0.0);
+    for(int i = 0; i < n; i++){
+        result = complex_mult(result, z);
+    }
+    return result;
+}
+
+vec2 poly(vec2 z, vec2[MAX_DEGREE + 1] coefs){
+    vec2 result = vec2(0.0);
+    for(int n = 0; n < int(n_roots) + 1; n++){
+        result += complex_mult(coefs[n], complex_pow(z, n));
+    }
+    return result;
+}
+
+vec2 dpoly(vec2 z, vec2[MAX_DEGREE + 1] coefs){
+    vec2 result = vec2(0.0);
+    for(int n = 1; n < int(n_roots) + 1; n++){
+        result += n * complex_mult(coefs[n], complex_pow(z, n - 1));
+    }
+    return result;
+}
+
+vec2 seek_root(vec2 z, vec2[MAX_DEGREE + 1] coefs, int max_steps, out float n_iters){
+    float last_len;
+    float curr_len;
+    float threshold = 1e-3;
+
+    for(int i = 0; i < max_steps; i++){
+        last_len = curr_len;
+        n_iters = float(i);
+        vec2 step = complex_div(poly(z, coefs), dpoly(z, coefs));
+        curr_len = length(step);
+        if(curr_len < threshold){
+            break;
+        }
+        z = z - step;
+    }
+    n_iters -= clamp((threshold - curr_len) / (last_len - curr_len), 0.0, 1.0);
+
+    return z;
+}
+
+void main() {
+    vec2 root2 = xyz_coords.xy;
+
+    vec2 coef0 = -complex_mult(complex_mult(root0, root1), root2);
+    vec2 coef1 = complex_mult(root0, root1) + complex_mult(root0, root2) + complex_mult(root1, root2);
+    vec2 coef2 = -(root0 + root1 + root2);
+    vec2 coef3 = vec2(1.0, 0.0);
+
+    vec2[MAX_DEGREE + 1] coefs = vec2[MAX_DEGREE + 1](coef0, coef1, coef2, coef3);
+    vec2[MAX_DEGREE] roots = vec2[MAX_DEGREE](root0, root1, root2);
+    vec4[MAX_DEGREE] colors = vec4[MAX_DEGREE](color0, color1, color2);
+
+    // vec2 z = z0;
+    vec2 z = -coef2 / 3.0;
+    float n_iters;
+    vec2 found_root = seek_root(z, coefs, int(n_steps), n_iters);
+
+    vec4 color = vec4(0.0);
+    float min_dist = 1e10;
+    float dist;
+    for(int i = 0; i < int(n_roots); i++){
+        dist = distance(roots[i], found_root);
+        if(dist < min_dist){
+            min_dist = dist;
+            color = colors[i];
+        }
+    }
+    // color *= (1.0 + (color_mult - 1) * (n_iters - 5));
+
+    if (min_dist > 1e-2){
+        color = vec4(0.0, 0.0, 0.0, 1.0);
+    }
+
+    frag_color = finalize_color(
+        color,
+        xyz_coords,
+        vec3(0.0, 0.0, 1.0),
+        light_source_position,
+        gloss,
+        shadow
+    );
+ }

manimlib/shaders/meta_poly_fractal/vert.glsl

@@ -0,0 +1,17 @@
+#version 330
+
+#INSERT camera_uniform_declarations.glsl
+
+in vec3 point;
+out vec3 xyz_coords;
+
+uniform float scale_factor;
+uniform vec3 offset;
+
+#INSERT position_point_into_frame.glsl
+#INSERT get_gl_Position.glsl
+
+void main(){
+    xyz_coords = (point - offset) / scale_factor;
+    gl_Position = get_gl_Position(position_point_into_frame(point));
+}

manimlib/shaders/poly_fractal/frag.glsl

@@ -1,5 +1,3 @@
-// Copied from surface/frag.glsl
-
 #version 330
 
 uniform vec3 light_source_position;
@@ -28,6 +26,8 @@ uniform vec2 root4;
 
 uniform float n_roots;
 uniform float n_steps;
+uniform float julia_highlight;
+uniform float color_mult;
 
 uniform vec2 frame_shape;
 
@@ -44,8 +44,7 @@ vec2 complex_mult(vec2 z, vec2 w){
 }
 
 vec2 complex_div(vec2 z, vec2 w){
-    float w_norm_squared = w.x * w.x + w.y * w.y;
-    return complex_mult(z, vec2(w.x, -w.y)) / w_norm_squared;
+    return complex_mult(z, vec2(w.x, -w.y)) / (w.x * w.x + w.y * w.y);
 }
 
 vec2 complex_pow(vec2 z, int n){
@@ -58,7 +57,7 @@ vec2 complex_pow(vec2 z, int n){
 
 vec2 poly(vec2 z, vec2[MAX_DEGREE + 1] coefs){
     vec2 result = vec2(0.0);
-    for(int n = 0; n < MAX_DEGREE + 1; n++){
+    for(int n = 0; n < int(n_roots) + 1; n++){
         result += complex_mult(coefs[n], complex_pow(z, n));
     }
     return result;
@@ -66,20 +65,28 @@ vec2 poly(vec2 z, vec2[MAX_DEGREE + 1] coefs){
 
 vec2 dpoly(vec2 z, vec2[MAX_DEGREE + 1] coefs){
     vec2 result = vec2(0.0);
-    for(int n = 1; n < MAX_DEGREE + 1; n++){
+    for(int n = 1; n < int(n_roots) + 1; n++){
         result += n * complex_mult(coefs[n], complex_pow(z, n - 1));
     }
     return result;
 }
 
-vec2 seek_root(vec2 z, vec2[MAX_DEGREE + 1] coefs){
-    for(int i = 0; i < int(n_steps); i++){
+vec2 seek_root(vec2 z, vec2[MAX_DEGREE + 1] coefs, int max_steps, out float n_iters){
+    float last_len;
+    float curr_len;
+    float threshold = 1e-3;
+
+    for(int i = 0; i < max_steps; i++){
+        last_len = curr_len;
+        n_iters = float(i);
         vec2 step = complex_div(poly(z, coefs), dpoly(z, coefs));
-        if(length(step) < 1e-2){
+        curr_len = length(step);
+        if(curr_len < threshold){
             break;
         }
         z = z - step;
     }
+    n_iters -= clamp((threshold - curr_len) / (last_len - curr_len), 0.0, 1.0);
 
     return z;
 }
@@ -91,17 +98,52 @@ void main() {
     vec4[MAX_DEGREE] colors = vec4[MAX_DEGREE](color0, color1, color2, color3, color4);
 
     vec2 z = xyz_coords.xy;
-    vec2 found_root = seek_root(z, coefs);
+    float n_iters;
+    vec2 found_root = seek_root(z, coefs, int(n_steps), n_iters);
 
     vec4 color = vec4(0.0);
-    float curr_min = 1e10;
-    for(int i = 0; i < MAX_DEGREE; i++){
-        float dist = distance(roots[i], found_root);
-        if(dist < curr_min){
-            curr_min = dist;
+    float min_dist = 1e10;
+    float dist;
+    for(int i = 0; i < int(n_roots); i++){
+        dist = distance(roots[i], found_root);
+        if(dist < min_dist){
+            min_dist = dist;
             color = colors[i];
         }
     }
+    color *= (1.0 + (color_mult - 1) * (n_iters - 5));
+
+    // if(julia_highlight > 0.0){
+    //     float factor = min_dist / distance(z, found_root);
+    //     factor *= pow(2.0, n_iters);
+    //     float t = smoothstep(0, 5, factor);
+    //     t *= 2.0;
+    //     // color = vec4(0.0) * (1 - t) * (1 - t) + 2 * color * (1 - t) * t + vec4(1.0) * t * t;
+    //     color *= t;
+    // }
+
+    if(julia_highlight > 0.0){
+        float radius = julia_highlight;  // TODO
+        vec2[4] samples = vec2[4](
+            z + vec2(radius, 0.0),
+            z + vec2(-radius, 0.0),
+            z + vec2(0.0, radius),
+            z + vec2(0.0, -radius)
+        );
+        for(int i = 0; i < 4; i++){
+            for(int j = 0; j < n_steps; j++){
+                vec2 z = samples[i];
+                z = z - complex_div(poly(z, coefs), dpoly(z, coefs));
+                samples[i] = z;
+            }
+        }
+        float dist_sum = 0.0;
+        for(int i = 0; i < 4; i++){
+            dist_sum += distance(samples[i], samples[(i + 1) % 4]);
+        }
+        float factor = dist_sum / (4 * sqrt(2) * radius);
+        color *= 1.0 * smoothstep(0.0, 50, factor);
+    }
 
     frag_color = finalize_color(
         color,
@@ -111,4 +153,4 @@ void main() {
         gloss,
         shadow
     );
-}
+ }

manimlib/shaders/poly_fractal/vert.glsl

@@ -1,5 +1,3 @@
-// Copied from surface/vert.glsl
-
 #version 330
 
 #INSERT camera_uniform_declarations.glsl
@@ -12,7 +10,6 @@ uniform vec3 offset;
 
 #INSERT position_point_into_frame.glsl
 #INSERT get_gl_Position.glsl
-#INSERT get_rotated_surface_unit_normal_vector.glsl
 
 void main(){
     xyz_coords = (point - offset) / scale_factor;