Better initialization for wave and shallow water benchmarks (#19)

Author: tkarna

examples/shallow_water.py

@@ -54,25 +54,26 @@ 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)
+
+        def transpose(a):
+            return np.permute_dims(a, [1, 0])
 
         all_axes = [0, 1]
         init(False)
 
     elif backend == "numpy":
         import numpy as np
-        from numpy import fromfunction
 
         if comm is not None:
             assert (
                 comm.Get_size() == 1
             ), "Numpy backend only supports serial execution."
 
         create_full = np.full
+        transpose = np.transpose
 
         fini = sync = lambda x=None: None
         all_axes = None
@@ -110,34 +111,32 @@ 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,
-    )
-    x_u_2d = fromfunction(lambda i, j: xmin + i * dx, (nx + 1, ny), dtype=dtype)
-    y_u_2d = fromfunction(
-        lambda i, j: ymin + j * dy + dy / 2,
-        (nx + 1, ny),
-        dtype=dtype,
-    )
-    x_v_2d = fromfunction(
-        lambda i, j: xmin + i * dx + dx / 2,
-        (nx, ny + 1),
-        dtype=dtype,
-    )
-    y_v_2d = fromfunction(lambda i, j: ymin + j * dy, (nx, ny + 1), 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 = transpose(np.reshape(ind, (ny, nx)))
+        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)
     F_shape = (nx + 1, ny + 1)
+    sync()
+    x_t_2d = xmin + ind_arr(T_shape, True) * dx + dx / 2
+    y_t_2d = ymin + ind_arr(T_shape) * dy + dy / 2
+
+    x_u_2d = xmin + ind_arr(U_shape, True) * dx
+    y_u_2d = ymin + ind_arr(U_shape) * dy + dy / 2
+
+    x_v_2d = xmin + ind_arr(V_shape, True) * dx + dx / 2
+    y_v_2d = ymin + ind_arr(V_shape) * dy
+    sync()
 
     dofs_T = int(numpy.prod(numpy.asarray(T_shape)))
     dofs_U = int(numpy.prod(numpy.asarray(U_shape)))
@@ -205,14 +204,6 @@ def bathymetry(x_t_2d, y_t_2d, lx, ly):
         bath = 1.0
         return bath * create_full(T_shape, 1.0, dtype)
 
-    # inital elevation
-    u0, v0, e0 = exact_solution(
-        0, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d
-    )
-    e[:, :] = e0
-    u[:, :] = u0
-    v[:, :] = v0
-
     # set bathymetry
     h[:, :] = bathymetry(x_t_2d, y_t_2d, lx, ly)
     # steady state potential energy
@@ -329,6 +320,18 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
         v[:, 1:-1] = v[:, 1:-1] / 3.0 + 2.0 / 3.0 * (v2[:, 1:-1] + dt * dvdt)
         e[:, :] = e[:, :] / 3.0 + 2.0 / 3.0 * (e2[:, :] + dt * dedt)
 
+    # warm up jit cache
+    step(u, v, e, u1, v1, e1, u2, v2, e2)
+    sync()
+
+    # initial solution
+    u0, v0, e0 = exact_solution(
+        0, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d
+    )
+    e[:, :] = e0
+    u[:, :] = u0
+    v[:, :] = v0
+
     t = 0
     i_export = 0
     next_t_export = 0

examples/wave_equation.py

@@ -54,25 +54,26 @@ 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)
+
+        def transpose(a):
+            return np.permute_dims(a, [1, 0])
 
         all_axes = [0, 1]
         init(False)
 
     elif backend == "numpy":
         import numpy as np
-        from numpy import fromfunction
 
         if comm is not None:
             assert (
                 comm.Get_size() == 1
             ), "Numpy backend only supports serial execution."
 
         create_full = np.full
+        transpose = np.transpose
 
         fini = sync = lambda x=None: None
         all_axes = None
@@ -110,17 +111,23 @@ 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 = transpose(np.reshape(ind, (ny, nx)))
+        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)))
@@ -162,9 +169,6 @@ def exact_elev(t, x_t_2d, y_t_2d, lx, ly):
         sol_t = numpy.cos(2 * omega * t)
         return amp * sol_x * sol_y * sol_t
 
-    # inital elevation
-    e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly)
-
     # compute time step
     alpha = 0.5
     c = (g * h) ** 0.5
@@ -215,6 +219,16 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
         v[:, 1:-1] = v[:, 1:-1] / 3.0 + 2.0 / 3.0 * (v2[:, 1:-1] + dt * dvdt)
         e[:, :] = e[:, :] / 3.0 + 2.0 / 3.0 * (e2[:, :] + dt * dedt)
 
+    # warm up jit cache
+    step(u, v, e, u1, v1, e1, u2, v2, e2)
+    sync()
+
+    # initial solution
+    e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly)
+    u[:, :] = create_full(U_shape, 0.0, dtype)
+    v[:, :] = create_full(V_shape, 0.0, dtype)
+    sync()
+
     t = 0
     i_export = 0
     next_t_export = 0