AMReX: from symbolic model to generated C++

AMReX: from symbolic model to generated C++#

AMReX is Zoomy’s block-structured adaptive-mesh backend. The model you write is the same one the NumPy tutorial uses — only the final step differs: instead of lambdifying into NumPy, the operators are printed as C++ and compiled against AMReX.

Before you start — you need the AMReX container#

The codegen cells below run anywhere zoomy_core is installed. Compiling and running needs the AMReX toolchain, which ships as a container:

# Apptainer (recommended)
apptainer pull zoomy_amrex.sif oras://ghcr.io/zoomylab/zoomy_amrex_sif:latest
apptainer run --bind $PWD:/workspace zoomy_amrex.sif jupyter
# -> open the printed http://127.0.0.1:8888/?token=... and run this notebook there

# Docker
docker run --rm -p 8888:8888 -v "$PWD:/workspace" \
    ghcr.io/zoomylab/zoomy_amrex:latest jupyter

See Installation for all containers and the three run modes. Nothing on this page needs a local AMReX build.

1. The model#

Identical to the NumPy tutorial — that is the point. Shallow water as SME(level=0), wall boundaries, dam-break initial state.

import numpy as np

from zoomy_core.model.models import SME, Newtonian, NavierSlip, StressFree
import zoomy_core.model.boundary_conditions as BC
import zoomy_core.model.initial_conditions as IC

model = SME(
    level=0,
    closures=[Newtonian(), NavierSlip(), StressFree()],
    boundary_conditions=BC.BoundaryConditions([
        BC.Wall(tag="left"), BC.Wall(tag="right"),
    ]),
    initial_conditions=IC.RP(
        high=lambda n: np.array([0.0, 2.0] + [0.0] * (n - 2)),
        low=lambda n: np.array([0.0, 1.0] + [0.0] * (n - 2)),
        jump_position_x=5.0,
    ),
)

2. Freeze and inspect#

describe() is the cheapest sanity check there is. Read the slots before you spend a compile cycle: an empty hydrostatic_pressure row on a free-surface model means the derivation went wrong, and you can see that here in seconds rather than after a build.

from zoomy_core.systemmodel import SystemModel

sm = SystemModel.from_model(model)

# Display, do NOT print: describe() has a _repr_markdown_ that renders the
# operator matrices as real LaTeX. print() would dump the raw markup.
sm.describe()

SystemModel — 3 equations, 1 spatial dimension

Parameters: \(g = 9.81\), \(\rho = 1.0\), \(\nu = 0.0\), \(\lambda_{s} = 0.0\), \(e_{x} = 0.0\)

System form:

\[M(Q)\,\partial_t Q + \nabla\cdot\!\big(F(Q) + P(Q)\big) + \sum_{d} B_{d}(Q)\,\partial_{d} Q - \nabla\cdot\!\big(A(Q)\,\nabla Q\big) - S(Q) = 0\]

State \(Q\):

\[\begin{split} \left[\begin{matrix}b\\h\\q_{0}\end{matrix}\right] \end{split}\]

Auxiliary \(Q_{aux}\):

\[\begin{split} \left[\begin{matrix}dq0dx\\dhdx\\dbdx\end{matrix}\right] \end{split}\]

Mass matrix \(M\):

\[\begin{split} \left[\begin{matrix}1 & 0 & 0\\0 & 1 & 0\\0 & 0 & 1\end{matrix}\right] \end{split}\]

Flux \(F\):

\[\begin{split} \left[\begin{matrix}0\\q_{0}\\\frac{q_{0}^{2}}{h}\end{matrix}\right] \end{split}\]

Hydrostatic pressure \(P\):

\[\begin{split} \left[\begin{matrix}0\\0\\\frac{g h^{2}}{2}\end{matrix}\right] \end{split}\]

Diffusive flux \(A\) \(= 0\)

NCP \(B\):

\[\begin{split} \left[\begin{matrix}0 & 0 & 0\\0 & 0 & 0\\g h & 0 & 0\end{matrix}\right] \end{split}\]

Source \(S\) (implicit):

\[\begin{split} \left[\begin{matrix}0\\0\\e_{x} g h - \frac{\lambda_{s} q_{0}}{h \rho}\end{matrix}\right] \end{split}\]

Operations: from_model

3. Attach the numerics#

Same NumericalSystemModel as every other backend. The scheme lives here, at the symbolic level — not in the AMReX code. That is deliberate: every backend gets the same well-balancing, positivity and tolerances, emitted from one place, so they cannot drift apart.

from zoomy_core.numerics import NumericalSystemModel, ReconstructionSpec

nsm = NumericalSystemModel.from_system_model(
    sm, reconstruction=ReconstructionSpec(order=1))

riemann = nsm.riemann
print("Riemann solver :", getattr(riemann, "__name__", None) or type(riemann).__name__)
print("reconstruction :", nsm.reconstruction)

Riemann solver : NonconservativeRusanov
reconstruction : ReconstructionSpec(order=1, limiter='venkatakrishnan', free_surface_aware=False, positivity='')

5. Running it#

Writing the headers to a case directory and building is a compile step, so it belongs in the container rather than in a rendered notebook page. Inside the AMReX container:

from pathlib import Path
out = Path("case"); out.mkdir(exist_ok=True)
(out / "Model.H").write_text(model_src)
(out / "Numerics.H").write_text(numerics_src)

cd case && make && ./run

The maintained end-to-end examples are the AMReX test suite itself — library/zoomy_amrex/tests/ — which builds and runs these same generated headers:

cd library/zoomy_amrex && python3 -m pytest -m small -q

That suite covers a wet dam break (Stoker), a dry one (Ritter), 2-D, a lake-at-rest well-balancing check, NumPy parity, and a two-rank parallel run — each a better starting point for a real case than a hand-assembled directory.

What carries over#

Nothing about the model changed to target AMReX. Swap the last step and the same nsm drives NumPy, JAX, or OpenFOAM — see Backends.

Contents