still graph breaks

e2e/mandelbrot/mandelbrot.py

@@ -30,22 +30,56 @@ def mandelbrot(xmin, xmax, ymin, ymax, xn, yn, maxiter, horizon=2.0):
     return Z, N
 
 
-if __name__ == '__main__':
-    # Benchmark
-    xmin, xmax, xn = -2.25, +0.75, int(3000/3)
-    ymin, ymax, yn = -1.25, +1.25, int(2500/3)
-    maxiter = 200
+def abs2(a):
+    r"""abs(a) replacement"""
+    return a[..., 0]**2 + a[..., 1]**2
+
+
+def sq2(a):
+    """a**2 replacement."""
+    z = np.empty_like(a)
+    z[..., 0] = a[..., 0]**2 - a[..., 1]**2
+    z[..., 1] = 2 * a[..., 0] * a[..., 1]
+    return z
+
 
+@torch.compile
+def step(n, CC, Z, N):
+    I = abs2(Z) < horizon**2
+    N = np.where(I, n, N)                       # N[I] = n
+    Z = np.where(I[..., None], sq2(Z) + CC, Z)  # Z[I] = Z[I]**2 + C[I]
+    return N, Z
+
+
+def mandelbrot_c(xmin, xmax, ymin, ymax, xn, yn, maxiter, horizon=2.0):
+    # Adapted from https://www.ibm.com/developerworks/community/blogs/jfp/...
+    #              .../entry/How_To_Compute_Mandelbrodt_Set_Quickly?lang=en
+    X = np.linspace(xmin, xmax, xn, dtype='float32')
+    Y = np.linspace(ymin, ymax, yn, dtype='float32')
+    CC = np.stack(np.broadcast_arrays(X[None, :], Y[:, None]), axis=-1)
+
+    N = np.zeros(CC.shape[:-1], dtype='int')
+    Z = np.zeros(CC.shape, 'float32')
+    for n in range(maxiter):
+        N, Z = step(n, CC, N, Z)
+    N = np.where(N == maxiter -1, 0, N)             # N[N == maxiter-1] = 0
+    return Z, N
+
+
+if __name__ == '__main__':
     # Visualization
     xmin, xmax, xn = -2.25, +0.75, int(3000/2)
     ymin, ymax, yn = -1.25, +1.25, int(2500/2)
     maxiter = 200
     horizon = 2.0 ** 40
     log_horizon = np.log(np.log(horizon))/np.log(2)
-    Z, N = mandelbrot(xmin, xmax, ymin, ymax, xn, yn, maxiter, horizon)
+    Z, N = mandelbrot_c(xmin, xmax, ymin, ymax, xn, yn, maxiter, horizon)
 
     # Normalized recount as explained in:
     # http://linas.org/art-gallery/escape/smooth.html
+
+    Z = Z[..., 0] + 1j*Z[..., 1]
+
     M = np.nan_to_num(N + 1 - np.log(np.log(abs(Z)))/np.log(2) + log_horizon)
 
     dpi = 72