Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Better initialization for wave and shallow water benchmarks #19

Merged
merged 4 commits into from
Nov 11, 2024
Merged

Better initialization for wave and shallow water benchmarks #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0d34e27
wave-equation: remove fromfunction initialization
tkarna Nov 7, 2024
7fb1742
shallow-water: remove fromfunction initialization
tkarna Nov 7, 2024
048f1c6
wave-equation: exclude jit compilation from timing
tkarna Nov 7, 2024
aba4ce8
shallow-water: exclude jit compilation from timing
tkarna Nov 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 37 additions & 34 deletions examples/shallow_water.py
View file
Open in desktop
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
Open in desktop
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
Loading