Download this example.
Lock-exchange testcase
This example illustrates the classical lock-exchange problem in a closed basin. Two fluids of different density are initially separated in a domain; when the barrier is removed, the dense fluid flows along the bottom and the light fluid flows above in the opposite direction, forming gravity currents and a propagating density front. It is commonly used to test an ocean model’s ability to represent stratified flow, sharp interfaces, and mixing.
import os
import time
import numpy as np
from slim4 import mpi, slim, slim2d, slim3d, pre3d
from slim4.tools import visualize_prepro_data
We optionally write preprocessing checks to disk. These files can be useful when validating the bathymetry, vertical extrusion, and initial tracer fields.
check_preprocessing = False
First, we load the 2D mesh and define the bathymetry and free surface used to build the 3D domain.
print("Loading hydro setup ...")
mesh2d = slim2d.Mesh("lock_exchange_structured.msh")
num_2d_elements = mesh2d.n_triangles
prepro_dir = "prepro"
os.makedirs(prepro_dir, exist_ok=True)
print(f"Prepro dir: {prepro_dir}")
bathymetry_2d = slim.DataStatic.create_constant(mesh2d, [20])
free_surface_2d = slim.DataStatic.create_constant(mesh2d, [0])
slim.DataFile.write(
f"./{prepro_dir}/bathymetry", [("Depth", bathymetry_2d.eval_mesh())], "DG"
)
Next, we define the vertical discretization. This example uses sigma layers We extrude the 2D mesh vertically to obtain the 3D computational mesh.
n_layers = 20
mesh, z0, bathymetry_2d = pre3d.extrude(mesh2d, bathymetry_2d, pre3d.VMapProportional(np.ones(n_layers)))
Now we create the 3D hydrostatic problem and add the tracers used in the testcase. Temperature controls the density anomaly, while salinity is kept as a passive background field in this configuration.
g = 9.80616
rho_0 = 1000
num_3d_elements = mesh.z_layers[-1]
problem = slim3d.Problem(mesh, g, bathymetry_2d, rho_0)
temperature_equation = problem.add_tracer(1, "temperature")
salinity_equation = problem.add_tracer(1, "salinity")
problem.set_wall_boundary("Wall", slip=False)
We initialize the prognostic fields and moving vertical coordinates used by the 3D model.
z = slim.ScalarArrayDeviceMPI.create_from_host(z0)
uv = problem.problemUV.new_vector()
temperature = temperature_equation.new_vector()
salinity = salinity_equation.new_vector()
The initial temperature field defines the lock exchange: colder, denser water is placed on the left side of the domain, while warmer, lighter water fills the rest of the basin.
temp_init = np.ones((1, 6, num_3d_elements)) * 30
x_threshold = 32e3
left_columns = np.all(mesh2d.x[0, :, :] <= x_threshold, axis=0)
left_triangles = np.arange(num_2d_elements)[left_columns]
for triangle in left_triangles:
column = np.arange(mesh.z_layers[triangle], mesh.z_layers[triangle + 1])
temp_init[0, :, column] = 5
temperature.set(temp_init)
slim.DataFile.write(f"./{prepro_dir}/temperature", [("Temperature", temp_init)], "DG")
Salinity is initialized uniformly. The density is then linked to temperature through a linear equation of state.
salinity_initial = np.ones((1, 6, num_3d_elements)) * 35
salinity.set(salinity_initial)
problem.set_state_equation_linear([(temperature, -0.2)])
Optionally, we write preprocessing figures and 3D fields to help verify that the bathymetry and initial condition were generated correctly.
if check_preprocessing:
visualize_prepro_data(
f"./{prepro_dir}/bathymetry",
"Depth",
f"./{prepro_dir}/check/bathymetry",
mesh2d,
)
mesh.create_xdmf(z0, f"./{prepro_dir}/check/temperature")
mesh.append_to_xdmf(
f"./{prepro_dir}/check/temperature", 0, 0, z0, [("temperature", temp_init)]
)
Next, we define the viscous and diffusive parameters. This testcase uses no explicit horizontal tracer diffusion, a fixed vertical viscosity, and a small horizontal viscosity.
vertical_viscosity = np.full((1, 1, mesh.n_elements), 1e-4)
problem.problemUV.set_vertical_viscosity(
slim.ScalarArrayDevice.create_from_host(vertical_viscosity)
)
horizontal_viscosity = slim.DataStatic.create_constant(mesh, value=1e-2)
problem.problemUV.set_horizontal_viscosity(horizontal_viscosity, "direct")
We prepare the output files and choose the time-stepping parameters.
nu_h_value = horizontal_viscosity.eval_mesh()[0, 0, 0]
print(
f" -> Running lock-exchange testcase with nu_h = {nu_h_value:.3g} m^2/s and {n_layers} vertical layers"
)
output_dir = "out/"
path2d = output_dir + "slim2D"
path3d = output_dir + "slim3D"
mesh2d.create_xdmf(path2d)
mesh.create_xdmf(z0, path3d)
problem.problemUV.update_density(z, problem.get_gpu_tmpvars().rho)
t_end = 17 * 3600
export_dt = 900
rho = problem.get_gpu_tmpvars().rho
problem.problemUV.update_density(z, rho)
dt3d, dt2d = problem.compute_cfl(2, rho.get())
nsub = int(np.ceil(export_dt / dt3d))
dt = export_dt / nsub
M = dt / dt2d
nstep = int(np.ceil(t_end / dt))
export_stride = int(np.ceil(export_dt / dt))
print(f"dt3d = {dt:.3f}, dt2d = {dt2d:.2f}, M = {M:.2f}, nstep = {nstep}")
Finally, we run the simulation and export the solution at regular intervals. The exported fields include density diagnostics in addition to temperature and salinity.
t = 0
iout = 0
export_kwargs = dict(turbulence=False, Huv=False, rho=True)
problem.do_export(
path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
)
start_time = time.time()
loop_start = start_time
iteration = 0
while t + dt <= t_end:
t = slim3d.iterate_impl(problem, t, dt, M, 1, z, uv, [temperature, salinity])
# t = slim3d.iterate_midpoint(problem, t, dt, M, 1, z, uv, [temperature, salinity])
export_index, remainder = divmod(iteration + 1, export_stride)
if remainder == 0:
iout += 1
problem.do_export(
path2d, path3d, iout, t, z, uv, [temperature, salinity], **export_kwargs
)
uv_max = mpi.max_scalar(np.linalg.norm(uv.get(), axis=0).max())
w_max = mpi.max_scalar(np.abs(problem.get_gpu_tmpvars().w_gpu.get()).max())
print(
f"Output {iout} at t = {t / 3600:.2f}h, "
f"uv_max = {uv_max:.2e}, w_max = {w_max:.2e}, "
f"{slim.si_format((time.time() - loop_start) / (iteration + 1), 3)}s/iter"
)
iteration += 1
print(f"Hydro finished in {(time.time() - start_time) / 60:.2f} minutes")