wave-equation: remove fromfunction initialization

Author: tkarna

examples/wave_equation.py

@@ -54,18 +54,15 @@ def run(n, backend, datatype, benchmark_mode):
     if backend == "sharpy":
         import sharpy as np
         from sharpy import fini, init, sync
-        from sharpy.numpy import fromfunction as _fromfunction
 
         device = os.getenv("SHARPY_DEVICE", "")
         create_full = partial(np.full, device=device)
-        fromfunction = partial(_fromfunction, device=device)
 
         all_axes = [0, 1]
         init(False)
 
     elif backend == "numpy":
         import numpy as np
-        from numpy import fromfunction
 
         if comm is not None:
             assert (
@@ -110,17 +107,24 @@ def run(n, backend, datatype, benchmark_mode):
     t_export = 0.02
     t_end = 1.0
 
-    # coordinate arrays
-    x_t_2d = fromfunction(
-        lambda i, j: xmin + i * dx + dx / 2, (nx, ny), dtype=dtype
-    )
-    y_t_2d = fromfunction(
-        lambda i, j: ymin + j * dy + dy / 2, (nx, ny), dtype=dtype
-    )
+    def ind_arr(shape, columns=False):
+        """Construct an (nx, ny) array where each row/col is an arange"""
+        nx, ny = shape
+        if columns:
+            ind = np.arange(0, nx * ny, 1, dtype=np.int32) % nx
+            ind = np.reshape(ind, (ny, nx))
+            ind = np.permute_dims(ind, [1, 0])
+        else:
+            ind = np.arange(0, nx * ny, 1, dtype=np.int32) % ny
+            ind = np.reshape(ind, (nx, ny))
+        return ind.astype(dtype)
 
+    # coordinate arrays
     T_shape = (nx, ny)
     U_shape = (nx + 1, ny)
     V_shape = (nx, ny + 1)
+    x_t_2d = xmin + ind_arr(T_shape, True) * dx + dx / 2
+    y_t_2d = ymin + ind_arr(T_shape) * dy + dy / 2
 
     dofs_T = int(numpy.prod(numpy.asarray(T_shape)))
     dofs_U = int(numpy.prod(numpy.asarray(U_shape)))