Quick start

Quick start#

The C1 case of the test matrix: shallow water (SME(level=0)) on $[0, 10]^2$ with walls, at rest with a Gaussian hump in the depth, second order, until $t = 2.5$. The scripts are in docs/book/examples/; run them from there after source env.sh (it puts this checkout’s libraries on PYTHONPATH).

The case#

"""The C1 bump of the test matrix: shallow water (SME level 0) on [0, 10]^2,
walls, at rest with a Gaussian hump in the depth.

Model -> SystemModel -> NumericalSystemModel; every backend takes the NSM."""
import numpy as np

import zoomy_core.model.initial_conditions as IC
from zoomy_core.mesh import LSQMesh
from zoomy_core.model.boundary_conditions import BoundaryConditions, FromModel
from zoomy_core.model.models import SME
from zoomy_core.model.models.closures import ManningFriction, Newtonian, StressFree
from zoomy_core.numerics import (
    Limiter, NumericalSystemModel, ReconstructedVariables, ReconstructionSpec)
from zoomy_core.systemmodel.system_model import SystemModel


def bump(order=2, nx=40):
    """Returns ``(mesh, nsm)``."""
    # 1. Model: SME level 0 = shallow water; dimension 3 = (x, y) + vertical.
    walls = [FromModel(tag=t, definition="wall") for t in ("left", "right", "bottom", "top")]
    model = SME(level=0, dimension=3,
                closures=[Newtonian(), ManningFriction(), StressFree()],
                boundary_conditions=BoundaryConditions(walls))

    # 2. SystemModel: the operator slots (flux, nonconservative_matrix, source, ...).
    sm = SystemModel.from_model(model)

    def ic(x):                                   # state (b, h, q_x, q_y)
        h = 0.5 + 0.05 * np.exp(-((x[0] - 5.0) ** 2 + (x[1] - 5.0) ** 2) / 2.0)
        return np.array([0.0, h, 0.0, 0.0])
    sm.initial_conditions = IC.UserFunction(function=ic)

    # 3. NumericalSystemModel: reconstruction, Riemann solver (default: Rusanov).
    nsm = NumericalSystemModel.from_system_model(
        sm, reconstruction=ReconstructionSpec(order=order, limiter=Limiter.minmod,
                                              variables=ReconstructedVariables.conservative))
    mesh = LSQMesh.create_2d(domain=(0, 10, 0, 10), nx=nx, ny=nx)
    return mesh, nsm

JAX#

"""Quick start: the C1 bump on JAX (CPU or GPU)."""
import jax

jax.config.update("jax_enable_x64", True)

import numpy as np  # noqa: E402

from bump import bump  # noqa: E402
from zoomy_jax.symsolver.solver import SymbolicHyperbolicSolverJax  # noqa: E402

mesh, nsm = bump(order=2, nx=40)
solver = SymbolicHyperbolicSolverJax(time_end=2.5, cfl=0.9, mood=False)
Q, Qaux = solver.setup(mesh, nsm)            # prints the template, compiles
time, Q, Qaux = solver.run(Q, Qaux)

h = np.asarray(Q)[1]
print(f"t = {float(time):.3f}, cells = {h.size}, h in [{h.min():.5f}, {h.max():.5f}]")
np.save("quickstart_jax.npy", np.asarray(Q))

t = 2.500, cells = 1600, h in [0.49671, 0.51884]

About 20 s on 8 CPU cores, most of it printing and compiling the template.

foam and AMReX: the same NSM#

The C++ backends take the NSM, the mesh and the initial state; they print the same template as C++, build it in their container (cached by a hash of the sources) and run it.

"""Quick start: the same NSM on foam (OpenFOAM 13 in apptainer)."""
from pathlib import Path

import numpy as np

from bump import bump
from zoomy_foam.symsolver import SymbolicHyperbolicSolverFoam

mesh, nsm = bump(order=2, nx=40)
Q0 = nsm.sm.initial_conditions.apply(
    np.asarray(mesh.cell_centers)[:, :mesh.n_inner_cells], np.zeros((4, mesh.n_inner_cells)))

solver = SymbolicHyperbolicSolverFoam(time_end=2.5, cfl=0.9, mood=False)   # nprocs=4: MPI
time, Q = solver.run(mesh, nsm, Q0, Path("work_foam").resolve())          # prints, builds (cached), runs

h = Q[1]
print(f"t = {time:.3f}, cells = {h.size}, h in [{h.min():.5f}, {h.max():.5f}]")
ref = np.load("quickstart_jax.npy")
print(f"max |h_foam - h_jax| = {np.abs(h - ref[1]).max():.2e}")

t = 2.500, cells = 1600, h in [0.49669, 0.51856]
max |h_foam - h_jax| = 3.51e-04

"""Quick start: the same NSM on AMReX (apptainer; ranks=N for MPI, gpu=True for CUDA)."""
from pathlib import Path

import numpy as np

from bump import bump
from zoomy_amrex.symsolver import SymbolicHyperbolicSolverAmrex

mesh, nsm = bump(order=2, nx=40)
Q0 = nsm.sm.initial_conditions.apply(
    np.asarray(mesh.cell_centers)[:, :mesh.n_inner_cells], np.zeros((4, mesh.n_inner_cells)))

solver = SymbolicHyperbolicSolverAmrex(time_end=2.5, cfl=0.9, mood=False)
time, Q = solver.run(mesh, nsm, Q0, Path("work_amrex").resolve())         # prints, builds (cached), runs

h = Q[1]
print(f"t = {time:.3f}, cells = {h.size}, h in [{h.min():.5f}, {h.max():.5f}]")
ref = np.load("quickstart_jax.npy")
print(f"max |h_amrex - h_jax| = {np.abs(h - ref[1]).max():.2e}")

t = 2.500, cells = 1600, h in [0.49671, 0.51884]
max |h_amrex - h_jax| = 1.82e-14

AMReX runs the same arithmetic as JAX and agrees to round-off. foam computes second-order gradients with OpenFOAM’s own stencils, so it agrees with JAX to the discretization error (at first order: to round-off; matrix case C1). The first AMReX run compiles AMReX itself (about a minute here); later runs reuse the build.

Outputs above: library commits core 09d1584, jax 5fd4252, foam e17026f, amrex ce3f72a; containers OpenFOAM 13 and AMReX as in Installation.

Next#

Contents