Skip to content

Commit

Permalink
Better initialization for wave and shallow water benchmarks (#19)
Browse files Browse the repository at this point in the history
  • Loading branch information
@tkarna
tkarna authored Nov 11, 2024
1 parent 4520d97 commit 9506c28
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 47 deletions.
71 changes: 37 additions & 34 deletions examples/shallow_water.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
40 changes: 27 additions & 13 deletions examples/wave_equation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 9506c28

Please sign in to comment.