Fix examples with device arg (#90)

Author: tkarna

.github/workflows/ci.yml

@@ -147,17 +147,20 @@ jobs:
           mpirun -n 2 -genv SHARPY_FALLBACK=numpy python -u ./stencil-2d.py 5 2048 star 2
           mpirun -n 3 -genv SHARPY_FALLBACK=numpy python -u ./stencil-2d.py 5 2048 star 2
           mpirun -n 4 -genv SHARPY_FALLBACK=numpy python -u ./stencil-2d.py 5 2048 star 2
+          python -u ./wave_equation.py -b numpy
           python -u ./wave_equation.py -d f32
           python -u ./wave_equation.py -d f64
           SHARPY_FORCE_DIST=1 python -u ./wave_equation.py
           mpirun -n 2 python -u ./wave_equation.py
           mpirun -n 3 python -u ./wave_equation.py
           mpirun -n 4 python -u ./wave_equation.py
+          python -u ./shallow_water.py -b numpy
           python -u ./shallow_water.py -d f32
           python -u ./shallow_water.py -d f64
           mpirun -n 2 python -u ./shallow_water.py
           mpirun -n 3 python -u ./shallow_water.py
           mpirun -n 4 python -u ./shallow_water.py
+          python -u ./black_scholes.py -b numpy
           python -u ./black_scholes.py -d f32
           python -u ./black_scholes.py -d f64
           SHARPY_FORCE_DIST=1 python -u ./black_scholes.py

examples/black_scholes.py

@@ -23,7 +23,9 @@
 
 """
 import argparse
+import os
 import time as time_mod
+from functools import partial
 
 import numpy
 
@@ -45,7 +47,28 @@ def info(s):
         print(s)
 
 
-def initialize(np, nopt, seed, dtype):
+def naive_erf(x):
+    """
+    Error function (erf) implementation
+
+    Adapted from formula 7.1.26 in
+    Abramowitz and Stegun, "Handbook of Mathematical Functions", 1965.
+    """
+    y = numpy.abs(x)
+
+    a1 = 0.254829592
+    a2 = -0.284496736
+    a3 = 1.421413741
+    a4 = -1.453152027
+    a5 = 1.061405429
+    p = 0.3275911
+
+    t = 1.0 / (1.0 + p * y)
+    f = (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t
+    return numpy.sign(x) * (1.0 - f * numpy.exp(-y * y))
+
+
+def initialize(create_full, nopt, seed, dtype):
     """
     Initialize arrays.
 
@@ -85,15 +108,15 @@ def initialize(np, nopt, seed, dtype):
     # t = random.uniform(TL, TH, nopt)
 
     # constant values
-    price = np.full((nopt,), 49.81959369901096, dtype)
-    strike = np.full((nopt,), 40.13264789835957, dtype)
-    t = np.full((nopt,), 1.8994311692123782, dtype)
+    price = create_full((nopt,), 49.81959369901096, dtype)
+    strike = create_full((nopt,), 40.13264789835957, dtype)
+    t = create_full((nopt,), 1.8994311692123782, dtype)
     # parameters
     rate = RISK_FREE
     volatility = VOLATILITY
     # output arrays
-    call = np.zeros((nopt,), dtype)  # 16.976097804669887
-    put = np.zeros((nopt,), dtype)  # 0.34645174725098116
+    call = create_full((nopt,), 0.0, dtype)  # 16.976097804669887
+    put = create_full((nopt,), 0.0, dtype)  # 0.34645174725098116
 
     return (price, strike, t, rate, volatility, call, put)
 
@@ -153,17 +176,23 @@ def run(nopt, backend, iterations, datatype):
         import sharpy as np
         from sharpy import fini, init, sync
 
-        init(False)
+        device = os.getenv("SHARPY_USE_GPU", "")
+        create_full = partial(np.full, device=device)
         erf = np.erf
+
+        init(False)
     elif backend == "numpy":
         import numpy as np
-        from scipy.special import erf
 
-        fini = sync = lambda x=None: None
+        erf = naive_erf
+
         if comm is not None:
             assert (
                 comm.Get_size() == 1
             ), "Numpy backend only supports serial execution."
+
+        create_full = np.full
+        fini = sync = lambda x=None: None
     else:
         raise ValueError(f'Unknown backend: "{backend}"')
 
@@ -179,7 +208,7 @@ def run(nopt, backend, iterations, datatype):
     info(f"Datatype: {datatype}")
 
     # initialize
-    args = initialize(np, nopt, seed, dtype)
+    args = initialize(create_full, nopt, seed, dtype)
     sync()
 
     def eval():

examples/shallow_water.py

@@ -27,45 +27,57 @@
 import math
 import os
 import time as time_mod
+from functools import partial
 
 import numpy
 
-device = os.getenv("SHARPY_USE_GPU", "")
+try:
+    import mpi4py
+
+    mpi4py.rc.finalize = False
+    from mpi4py import MPI
+
+    comm_rank = MPI.COMM_WORLD.Get_rank()
+    comm = MPI.COMM_WORLD
+except ImportError:
+    comm_rank = 0
+    comm = None
+
+
+def info(s):
+    if comm_rank == 0:
+        print(s)
 
 
 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
+        from sharpy.numpy import fromfunction as _fromfunction
+
+        device = os.getenv("SHARPY_USE_GPU", "")
+        create_full = partial(np.full, device=device)
+        fromfunction = partial(_fromfunction, device=device)
 
         all_axes = [0, 1]
         init(False)
 
-        try:
-            import mpi4py
-
-            mpi4py.rc.finalize = False
-            from mpi4py import MPI
-
-            comm_rank = MPI.COMM_WORLD.Get_rank()
-        except ImportError:
-            comm_rank = 0
-
     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
+
         fini = sync = lambda x=None: None
         all_axes = None
-        comm_rank = 0
     else:
         raise ValueError(f'Unknown backend: "{backend}"')
 
-    def info(s):
-        if comm_rank == 0:
-            print(s)
-
     info(f"Using backend: {backend}")
 
     dtype = {
@@ -102,32 +114,24 @@ def info(s):
         lambda i, j: xmin + i * dx + dx / 2,
         (nx, ny),
         dtype=dtype,
-        device=device,
     )
     y_t_2d = fromfunction(
         lambda i, j: ymin + j * dy + dy / 2,
         (nx, ny),
         dtype=dtype,
-        device=device,
-    )
-    x_u_2d = fromfunction(
-        lambda i, j: xmin + i * dx, (nx + 1, ny), dtype=dtype, device=device
     )
+    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,
-        device=device,
     )
     x_v_2d = fromfunction(
         lambda i, j: xmin + i * dx + dx / 2,
         (nx, ny + 1),
         dtype=dtype,
-        device=device,
-    )
-    y_v_2d = fromfunction(
-        lambda i, j: ymin + j * dy, (nx, ny + 1), dtype=dtype, device=device
     )
+    y_v_2d = fromfunction(lambda i, j: ymin + j * dy, (nx, ny + 1), dtype=dtype)
 
     T_shape = (nx, ny)
     U_shape = (nx + 1, ny)
@@ -144,32 +148,32 @@ def info(s):
     info(f"Total     DOFs: {dofs_T + dofs_U + dofs_V}")
 
     # prognostic variables: elevation, (u, v) velocity
-    e = np.full(T_shape, 0.0, dtype, device=device)
-    u = np.full(U_shape, 0.0, dtype, device=device)
-    v = np.full(V_shape, 0.0, dtype, device=device)
+    e = create_full(T_shape, 0.0, dtype)
+    u = create_full(U_shape, 0.0, dtype)
+    v = create_full(V_shape, 0.0, dtype)
 
     # potential vorticity
-    q = np.full(F_shape, 0.0, dtype, device=device)
+    q = create_full(F_shape, 0.0, dtype)
 
     # bathymetry
-    h = np.full(T_shape, 0.0, dtype, device=device)
+    h = create_full(T_shape, 0.0, dtype)
 
-    hu = np.full(U_shape, 0.0, dtype, device=device)
-    hv = np.full(V_shape, 0.0, dtype, device=device)
+    hu = create_full(U_shape, 0.0, dtype)
+    hv = create_full(V_shape, 0.0, dtype)
 
-    dudy = np.full(F_shape, 0.0, dtype, device=device)
-    dvdx = np.full(F_shape, 0.0, dtype, device=device)
+    dudy = create_full(F_shape, 0.0, dtype)
+    dvdx = create_full(F_shape, 0.0, dtype)
 
     # vector invariant form
-    H_at_f = np.full(F_shape, 0.0, dtype, device=device)
+    H_at_f = create_full(F_shape, 0.0, dtype)
 
     # auxiliary variables for RK time integration
-    e1 = np.full(T_shape, 0.0, dtype, device=device)
-    u1 = np.full(U_shape, 0.0, dtype, device=device)
-    v1 = np.full(V_shape, 0.0, dtype, device=device)
-    e2 = np.full(T_shape, 0.0, dtype, device=device)
-    u2 = np.full(U_shape, 0.0, dtype, device=device)
-    v2 = np.full(V_shape, 0.0, dtype, device=device)
+    e1 = create_full(T_shape, 0.0, dtype)
+    u1 = create_full(U_shape, 0.0, dtype)
+    v1 = create_full(V_shape, 0.0, dtype)
+    e2 = create_full(T_shape, 0.0, dtype)
+    u2 = create_full(U_shape, 0.0, dtype)
+    v2 = create_full(V_shape, 0.0, dtype)
 
     def exact_solution(t, x_t_2d, y_t_2d, x_u_2d, y_u_2d, x_v_2d, y_v_2d):
         """
@@ -198,7 +202,7 @@ def bathymetry(x_t_2d, y_t_2d, lx, ly):
         Water depth at rest
         """
         bath = 1.0
-        return bath * np.full(T_shape, 1.0, dtype, device=device)
+        return bath * create_full(T_shape, 1.0, dtype)
 
     # inital elevation
     u0, v0, e0 = exact_solution(