Use Mosaic solvers elsewhere

You do not need the full benchmark harness to use a Mosaic solver. This page introduces the Tesseract interface every solver exposes, then shows the two ways to call one from your own code.

How Tesseracts work

Each solver is wrapped as a Tesseract: a Docker container exposing two entry points (see endpoint contract):

The container isolates all dependencies (language, libraries, GPU runtime). A Python client calls the container over a local socket via tesseract-core; tesseract-jax wraps those calls as native JAX functions with full jit and grad support. Containerization adds a fixed per-call overhead (a few ms of IPC) that is independent of solver complexity.

Forward-only solvers (OpenFOAM, deal.II) implement apply only. Mosaic’s finite-difference-through-Tesseract mode provides gradients automatically at O(N) cost.

Each domain defines a canonical InputSchema / OutputSchema (Pydantic, all fields with defaults) in mosaic/mosaic_shared/problems/<domain>/schemas.py. Field shapes use the Array[(shape,), dtype] annotation and the Differentiable[...] wrapper from tesseract-core. Solvers subclass the canonical schema and may add solver-specific parameters:

from mosaic_shared.problems.navier_stokes_grid import InputSchema as _Base

class InputSchema(_Base):
    inner_steps: int = Field(default=1, description="Sub-steps per outer step")

Prerequisites

# Schemas and types for all Mosaic domains (only deps: pydantic + tesseract-core)
pip install -e mosaic/mosaic_shared

# Option A --- also install the solver's native dependencies
pip install tesseract-core exponax jax

# Option B --- container path (also enables jax.grad through containers)
pip install tesseract-core tesseract-jax jax

mosaic_shared alone is enough to import schema types (InputSchema, OutputSchema, GridBC, MeshBC, HexMesh, make_differentiable, …) without Docker or any specific solver installed.

Option A: Local usage (no Docker)

Point Tesseract.from_tesseract_api at the solver’s tesseract_api.py and call it directly — no container, no build step. Requires the solver’s native dependencies, so this path is limited to Python-only solvers (no Julia, no C++).

import numpy as np
from tesseract_core import 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": np.array([0.01], dtype=np.float32), "steps": 50}

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

Option B: Via Tesseract container

The canonical path: works for every solver regardless of language, pins all dependencies, integrates with jax.grad through tesseract-jax.

tesseract build compiles the solver into a Docker image (one-time):

tesseract build mosaic/tesseracts/navier-stokes-grid/exponax
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"])

Tesseract.from_image starts and tears down the container. For GPU solvers (marked GPU in the Solver Reference), pass gpus=["0"] to from_image; CPU-only solvers ignore the argument. A CUDA-capable GPU is recommended for the GPU-enabled solvers to run efficiently. Autodiff without gradient checkpointing materialises the entire trajectory in memory, so RAM/VRAM can saturate quickly on longer rollouts — more memory is better.

Mesh solvers

Mesh schemas carry hexahedral / tetrahedral connectivity and per-face boundary conditions, so they take more setup than grid solvers. Each mesh schema module ships a make_default_inputs(...) helper that builds a minimal valid input dict you can then mutate:

from tesseract_core import Tesseract
from mosaic_shared.problems.structural_mesh import make_default_inputs

inputs = make_default_inputs(nx=16, ny=8, nz=1)
inputs["rho"] = np.random.uniform(0.2, 0.8, size=16 * 8).astype(np.float32)

t = Tesseract.from_tesseract_api(
    "mosaic/tesseracts/structural-mesh/jax-fem/tesseract_api.py"
)
print("Compliance:", t.apply(inputs)["compliance"])

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