ENH: compute mask only once per iteration.

Mukulikaa · Aug 15, 2021 · 06b270d · 06b270d
1 parent d50f534
commit 06b270d
Showing 1 changed file with 18 additions and 13 deletions.
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
 ```