Merge pull request numpy#103 from rossbar/reuse-fractal-masks

mattip · web-flow · commit 1fecd1946bab · 2021-08-23T13:37:18.000+03:00
ENH: compute mask only once per iteration in fractals notebook
diff --git a/content/tutorial-plotting-fractals.md b/content/tutorial-plotting-fractals.md
@@ -198,10 +198,11 @@ def divergence_rate(mesh, num_iter=10, radius=2):
     z = mesh.copy()
     diverge_len = np.zeros(mesh.shape)  # Keep tally of the number of iterations
 
+    # Iterate on element if and only if |element| < radius (Otherwise assume divergence)
     for i in range(num_iter):
-        # Iterate on element if and only if |element| < radius (Otherwise assume divergence)
-        diverge_len[np.abs(z) < radius] += 1
-        z[np.abs(z) < radius] = f(z[np.abs(z) < radius])
+        conv_mask = np.abs(z) < radius
+        diverge_len[conv_mask] += 1
+        z[conv_mask] = f(z[conv_mask])
 
     return diverge_len
 ```
@@ -252,8 +253,9 @@ def julia(mesh, c=-1, num_iter=10, radius=2):
     diverge_len = np.zeros(z.shape)
 
     for i in range(num_iter):
-        z[np.abs(z) < radius] = np.square(z[np.abs(z) < radius]) + c
-        diverge_len[np.abs(z) < radius] += 1
+        conv_mask = np.abs(z) < radius
+        z[conv_mask] = np.square(z[conv_mask]) + c
+        diverge_len[conv_mask] += 1
 
     return diverge_len
 ```
@@ -325,8 +327,9 @@ def mandelbrot(mesh, num_iter=10, radius=2):
     diverge_len = np.zeros(z.shape)
 
     for i in range(num_iter):
-        z[np.abs(z) < radius] = np.square(z[np.abs(z) < radius]) + c[np.abs(z) < radius]
-        diverge_len[np.abs(z) < radius] += 1
+        conv_mask = np.abs(z) < radius
+        z[conv_mask] = np.square(z[conv_mask]) + c[conv_mask]
+        diverge_len[conv_mask] += 1
 
     return diverge_len
 ```
@@ -349,8 +352,9 @@ def general_julia(mesh, c=-1, f=np.square, num_iter=100, radius=2):
     diverge_len = np.zeros(z.shape)
 
     for i in range(num_iter):
-        z[np.abs(z) < radius] = f(z[np.abs(z) < radius]) + c
-        diverge_len[np.abs(z) < radius] += 1
+        conv_mask = np.abs(z) < radius
+        z[conv_mask] = f(z[conv_mask]) + c
+        diverge_len[conv_mask] += 1
 
     return diverge_len
 ```
@@ -391,10 +395,11 @@ def newton_fractal(mesh, f, df, num_iter=10, r=2):
     diverge_len = np.zeros(z.shape)
 
     for i in range(num_iter):
-        pz = f(z[np.abs(z) < r])
-        dp = df(z[np.abs(z) < r])
-        z[np.abs(z) < r] = z[np.abs(z) < r] - pz/dp
-        diverge_len[np.abs(z) < r] += 1
+        conv_mask = np.abs(z) < r
+        pz = f(z[conv_mask])
+        dp = df(z[conv_mask])
+        z[conv_mask] = z[conv_mask] - pz/dp
+        diverge_len[conv_mask] += 1
 
     return diverge_len
 ```
