Convert fill_complex_plane() to Cython.

distarray/examples/julia_set/benchmark_julia.py

@@ -88,16 +88,8 @@ def create_complex_plane(context, resolution, dist, re_ax, im_ax):
     import numpy as np
 
     def fill_complex_plane(arr, re_ax, im_ax, resolution):
-        """Fill in points on the complex coordinate plane."""
-        # Drawing the coordinate plane directly like this is currently much
-        # faster than trying to do it by indexing a distarray.
-        # This may not be the most DistArray-thonic way to do this.
-        re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
-        im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
-        for i in arr.distribution[0].global_iter:
-            for j in arr.distribution[1].global_iter:
-                arr.global_index[i, j] = complex(re_ax[0] + re_step * i,
-                                                 im_ax[0] + im_step * j)
+        from distarray.examples.julia_set.kernel import fill_complex_plane
+        fill_complex_plane(arr, re_ax, im_ax, resolution)
 
     # Create an empty distributed array.
     distribution = Distribution(context, (resolution[0], resolution[1]),

distarray/examples/julia_set/kernel.pyx

@@ -1,11 +1,43 @@
+# cython: boundscheck=False
+# cython: wraparound=False
+
 cimport cython
 import numpy as np
 
+def fill_complex_plane(arr, re_ax, im_ax, resolution):
+    """Fill in points on the complex coordinate plane."""
+    # Drawing the coordinate plane directly like this is currently much
+    # faster than trying to do it by indexing a distarray.
+    # This may not be the most DistArray-thonic way to do this.
+
+    cdef:
+        float _re_ax0 = re_ax[0]
+        float _im_ax0 = im_ax[0]
+        float re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
+        float im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
+        unsigned int row_start, col_start, row_step, col_step, nrow, ncol, i, j, glb_i, glb_j
+        float complex[:,::1] _arr = arr.ndarray
+
+    s0, s1 = arr.distribution.global_slice
+    row_start, row_step = s0.start, (s0.step or 1)
+    d1 = arr.distribution[1]
+    col_start, col_step = s1.start, (s1.step or 1)
+
+    nrow = _arr.shape[0]
+    ncol = _arr.shape[1]
+
+    for i in range(nrow):
+        glb_i = row_start + i * row_step
+        for j in range(ncol):
+            glb_j = col_start + j * col_step
+            _arr[i, j].real = _re_ax0 + re_step * glb_i
+            _arr[i, j].imag = _im_ax0 + im_step * glb_j
+
+
 cdef inline float cabs2(float complex z):
     return z.real * z.real + z.imag * z.imag
 
-@cython.boundscheck(False)
-@cython.wraparound(False)
+
 def cython_julia_calc(float complex[:,::1] z, float complex c,
                       float z_max, unsigned int n_max):
 
@@ -25,4 +57,4 @@ def cython_julia_calc(float complex[:,::1] z, float complex c,
                 count += 1
             counts[i,j] = count
 
-    return counts
+    return np.asarray(counts)