zoomy_jax.symsolver.solver

`zoomy_jax.symsolver.solver`#

Driver of the JAX backend: prints the solver template (hyperbolic, Chorin) of a NumericalSystemModel and runs it against zoomy_jax.symsolver.runtime. ad_mode=AdMode.reverse checkpoints the time loop for jax.grad, couple= makes the march a preCICE participant, comm= (mpi4py) runs it distributed, one rank per device.

References

Bradbury et al. 2018, https://github.com/jax-ml/jax

Module Contents#

Classes#

`SymbolicHyperbolicSolverJax`	Finite-volume solver on JAX: prints `HyperbolicTemplate` for a NumericalSystemModel, builds the runtime and marches the jitted program. Use `setup(mesh, model)`, then `run(Q, Qaux)`.
`SymbolicChorinSolverJax`	The Chorin projection solver on JAX (`ChorinTemplate`): `setup(mesh, pred, press, corr)` takes the predictor NSM on the full state and the pressure and corrector systems. The split must use zoomy’s time-step symbol `dt` (positive).

API#

class zoomy_jax.symsolver.solver.SymbolicHyperbolicSolverJax(**kwargs)#

Bases: zoomy_core.symsolver.settings.HyperbolicSettings

Finite-volume solver on JAX: prints HyperbolicTemplate for a NumericalSystemModel, builds the runtime and marches the jitted program. Use setup(mesh, model), then run(Q, Qaux).

Initialization

fold_geometry#: ‘Boolean(…)’

on_write#: ‘Callable(…)’

output#: ‘String(…)’

ad_mode#: ‘OptionParam(…)’

ad_checkpoints#: ‘Integer(…)’

couple#: ‘Dict(…)’

comm#: ‘Parameter(…)’

subcycling#: ‘Boolean(…)’

setup(mesh, model)#

Lift model to an NSM if needed, then build the solver on mesh.

Args

mesh: a mesh, or the path of a mesh file
model: SystemModel or NumericalSystemModel

Returns

initial (Q, Qaux) (distributed: the local piece)

initial_aux(Q, Qaux, parameters=None, X=None)#: The aux of the initial state: the pointwise formula plus the derivative aux with the BC ghosts and the boundary closure.

property parameters#: The runtime’s parameter vector.

run(Q, Qaux, parameters=None, time0=None, X=None)#: March from time0 (default time_start; a restart: the saved state) to time_end; returns (time, Q, Qaux). parameters and the vertex coordinates X (dim, n_vertices) may be traced for AD.

class zoomy_jax.symsolver.solver.SymbolicChorinSolverJax(**kwargs)#

Bases: zoomy_jax.symsolver.solver.SymbolicHyperbolicSolverJax, zoomy_core.symsolver.settings.ChorinSettings

The Chorin projection solver on JAX (ChorinTemplate): setup(mesh, pred, press, corr) takes the predictor NSM on the full state and the pressure and corrector systems. The split must use zoomy’s time-step symbol dt (positive).