Use a solver in your code

Every Mosaic solver is a Tesseract that you can call from your own JAX or PyTorch code, without the benchmark harness. This page shows how to install what you need, load a solver, differentiate through it, and set up mesh solvers and GPUs.

Setup

The canonical schemas and input helpers live in mosaic_shared, which is not on PyPI yet, so install it from a clone of the repository. For automatic differentiation, add tesseract-jax or tesseract-torch, both of which need Python 3.12 or later:

git clone https://github.com/pasteurlabs/mosaic && cd mosaic
pip install -e mosaic/mosaic_shared   # schemas and helpers (pydantic and tesseract-core only)
pip install tesseract-jax             # or: pip install tesseract-torch

There are two ways to load a solver.

From a built image. This works for every solver, whatever language it is written in, and pins all of its dependencies. tesseract build builds the image once. The image is named after the name field in the solver’s tesseract_config.yaml and tagged latest:

tesseract build mosaic/tesseracts/navier-stokes-grid/exponax

Tesseract.from_image then starts the container, and stops it when the with block ends:

from tesseract_core import Tesseract

with Tesseract.from_image("exponax_navier_stokes_grid:latest") as t:
    ...

Without a container. Tesseract.from_tesseract_api loads the solver’s tesseract_api.py into your Python process, with no build step:

t = Tesseract.from_tesseract_api("mosaic/tesseracts/navier-stokes-grid/exponax/tesseract_api.py")

This requires the solver’s dependencies, listed in its tesseract_requirements.txt, to be installed in your environment, which is practical for solvers built on pip-installable Python libraries. Solvers that need a separate toolchain, system packages or a conda environment, such as OpenFOAM, deal.II, FEniCS, Firedrake and JAX-FEM, are easier to use from the image.

Inputs and outputs

Each solver exposes the Tesseract endpoints apply, which runs the forward simulation, and, if it is differentiable, vector_jacobian_product, which computes the reverse-mode gradient. Both take a dictionary of inputs and return a dictionary of outputs.

The inputs and outputs follow the canonical schema of the solver’s domain, defined in mosaic/mosaic_shared/problems/<domain>/schemas.py. For the navier-stokes-grid domain, for example, the main inputs are the initial velocity v0 with shape (nx, ny, 1, 2) in 2D or (nx, ny, nz, 3) in 3D, the kinematic viscosity, the timestep dt and the number of steps, and the main output is the final velocity result. All of its inputs have defaults, so t.apply({}) runs a default simulation. The mesh domains have required inputs (see Mesh solvers).

A solver may add inputs of its own on top of the canonical schema, and it declares which inputs and outputs it can differentiate. The solver reference lists both for every solver. On a loaded solver, t.available_endpoints lists the endpoints it implements, and t.openapi_schema returns its full schema.

OpenFOAM and the two deal.II solvers are forward-only. They implement apply but not vector_jacobian_product, so you can run them but cannot differentiate through them.

Differentiate through a solver

With tesseract-jax, apply_tesseract calls the solver as a JAX function that works with jax.grad and jax.jit, and taking a gradient calls the solver’s vector_jacobian_product endpoint:

import jax
import jax.numpy as jnp
from tesseract_core import Tesseract
from tesseract_jax import apply_tesseract
from mosaic_shared.problems.navier_stokes_grid.schemas import make_vortex_ic

ic = make_vortex_ic(N=64, seed=42)
inputs = {"v0": ic, "viscosity": jnp.array([0.01]), "steps": 50}

with Tesseract.from_image("exponax_navier_stokes_grid:latest") as t:
    outputs = apply_tesseract(t, inputs)

    def loss(v0):
        out = apply_tesseract(t, {**inputs, "v0": v0})
        return jnp.mean(out["result"] ** 2)

    grad_v0 = jax.grad(loss)(inputs["v0"])

With tesseract-torch, the solver joins PyTorch’s autograd graph, and .backward() calls the same endpoint:

import torch
from tesseract_core import Tesseract
from tesseract_torch import apply_tesseract
from mosaic_shared.problems.navier_stokes_grid.schemas import make_vortex_ic

v0 = torch.tensor(make_vortex_ic(N=64, seed=42), requires_grad=True)
inputs = {"v0": v0, "viscosity": torch.tensor([0.01]), "steps": 50}

with Tesseract.from_image("exponax_navier_stokes_grid:latest") as t:
    loss = apply_tesseract(t, inputs)["result"].pow(2).mean()
    loss.backward()

grad_v0 = v0.grad

tesseract-torch supports PyTorch’s standard autograd API but not the torch.func transforms. See its documentation for details.

Mesh solvers

Mesh schemas carry hexahedral connectivity and per-face boundary conditions, and these inputs have no defaults. Each mesh domain’s schema module provides make_default_inputs(...), which builds a minimal valid input dictionary you can then modify. For the structural cantilever with a FEniCS image built from mosaic/tesseracts/structural-mesh/fenics-structural:

import jax
import jax.numpy as jnp
import numpy as np
from tesseract_core import Tesseract
from tesseract_jax import apply_tesseract
from mosaic_shared.problems.structural_mesh import make_default_inputs

inputs = make_default_inputs(nx=16, ny=8, nz=1)
rho = np.random.uniform(0.2, 0.8, size=16 * 8).astype(np.float32)  # one density per cell

with Tesseract.from_image("fenics_structural_mesh:latest") as t:
    def compliance(rho):
        return apply_tesseract(t, {**inputs, "rho": rho})["compliance"]

    print("Compliance:", compliance(rho))
    grad_rho = jax.grad(compliance)(jnp.asarray(rho))

The thermal domain follows the same pattern with mosaic_shared.problems.thermal_mesh.make_default_inputs(...).

GPU and memory

Tesseract.from_image gives the container no GPU unless you ask for one. For solvers that use a GPU (shown in the solver reference), pass the device, for example Tesseract.from_image(image, gpus=["0"]). Without it, some of them fall back to the CPU, which can be much slower.

Solvers that differentiate with reverse-mode AD usually keep every state of the rollout in memory, so memory use grows with the number of steps and can exhaust RAM or VRAM on long rollouts. The solver reference lists each solver’s AD strategy, and the results show where each one ran out of memory or time in the benchmark. from_image also accepts memory to cap the container’s memory.