zoomy_jax.symsolver.runtime

`zoomy_jax.symsolver.runtime`#

JAX runtime of the symbolic solvers: what each name of a printed solver means on JAX. The printed template (solver template -> PyPrinter) is plain Python calling this name table: nsm is the vmapped JaxRuntime of the kernels, every other entry a zoomy_core.symsolver block.

References

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

Module Contents#

Classes#

`AdMode`	How the JAX march is differentiated.
`Geometry`	The `mesh` argument of a printed solver, computed from the vertex coordinates inside the traced function so that `jax.grad` reaches the vertices. Topology (indices, stencils) is static and lives in the runtime’s closures.
`PreciceAdapter`	Host side of the `couple_*` effects: the preCICE API calls only. `comm` (mpi4py) makes a distributed participant, one preCICE process per rank, so it must span the world (one `mpiexec` per participant).

Functions#

`geometry_of`	All solver geometry from vertex coordinates `X` `(dim, n_vertices)`.
`runtime`	Build the name table (dict of callables) a printed solver runs against: the JAX implementation of every `zoomy_core.symsolver` block on this mesh. Called by the JAX solver driver’s `_setup`.
`coupled_face_indices`	Boundary faces carrying boundary tag `tag` (by name).
`hoisted_jit`	`jax.jit(fn)` with the constants `fn` captures (the mesh index arrays) passed as arguments, so compile time and program size stay independent of the mesh.
`build`	Execute printed source against the name table and return the function `name`; raises if the source uses a block the table lacks.

API#

class zoomy_jax.symsolver.runtime.AdMode#

Bases: zoomy_core.misc.options.Option

How the JAX march is differentiated.

Initialization

Initialize self. See help(type(self)) for accurate signature.

forward#: ‘forward’

reverse#: ‘reverse’

class zoomy_jax.symsolver.runtime.Geometry#

Bases: typing.NamedTuple

The mesh argument of a printed solver, computed from the vertex coordinates inside the traced function so that jax.grad reaches the vertices. Topology (indices, stencils) is static and lives in the runtime’s closures.

cell_volumes: jax.numpy.ndarray#: None

cell_centers: jax.numpy.ndarray#: None

cell_inradius: jax.numpy.ndarray#: None

face_normals: jax.numpy.ndarray#: None

face_volumes: jax.numpy.ndarray#: None

face_centers: jax.numpy.ndarray#: None

lsq_gradQ: jax.numpy.ndarray#: None

lsq_closure: jax.numpy.ndarray#: None

zoomy_jax.symsolver.runtime.geometry_of(topo, X)#: All solver geometry from vertex coordinates X (dim, n_vertices).

zoomy_jax.symsolver.runtime.runtime(nsm, rt, mesh_np, mesh_j, topo, *, ad_mode=AdMode.forward, checkpoints=64, on_write=None, on_save=None, on_load=None, couple=None, implicit=None, systems=None, partition=None)#

Build the name table (dict of callables) a printed solver runs against: the JAX implementation of every zoomy_core.symsolver block on this mesh. Called by the JAX solver driver’s _setup.

Args

topo: static topology (geometry_xp.build_topology) with the boundary-closure classes
ad_mode: AdMode.forward, or AdMode.reverse for a time loop with checkpoints checkpoints
on_write: on_write(time, Q, Qaux) snapshot callback (forward mode)
on_save, on_load: the restart pair (host), on_save(time, Q, Qaux) (forward mode) and on_load(time, Q, Qaux) -> (time, Q, Qaux); None: no restart file / not restarting
couple: preCICE settings dict(adapter=, tag=, side=)
implicit: Newton and linear-solver settings of root and linsolve
systems: further {name: (nsm, JaxRuntime)} a template uses besides nsm
partition: this rank’s piece of a distributed run; its halo and global reductions replace the serial ones

Returns

dict name -> callable

class zoomy_jax.symsolver.runtime.PreciceAdapter(name, config, mesh_name, write, read, comm=None)#

Host side of the couple_* effects: the preCICE API calls only. comm (mpi4py) makes a distributed participant, one preCICE process per rank, so it must span the world (one mpiexec per participant).