One section per solver: its method and backend, what it took to run it in Mosaic, its checks on each setup, and how to call it from your own code. Solvers that share an input and output schema form a domain.
Effort levels
Every solver has an intervention record that lists everything done to run it beyond using the library as intended, with the reason and evidence for each item. Each intervention gets a level by its mechanism, and a solver’s effort on a setup is the highest level among the interventions it requires there.
| L0 |
Nothing beyond expected tweaks, meaning the public API with its defaults or documented options |
Setting a documented solver tolerance |
| L1 |
Undocumented configuration, found by reading the source |
An option that only the library’s source code mentions |
| L2 |
User code on the library’s public primitives |
A boundary condition or VJP rule that the library does not provide |
| L3 |
Changes to the solver itself |
A monkeypatch, a private API or a patched source file |
What counts as intended use depends on the solver’s abstraction level:
application: running a program as documented.framework: expressing the setup through the library’s API.kernel-toolkit: building on the library’s primitives.
The record also lists deviations, the places where a solver runs a substitute for part of the setup, with their effect on the results. An upstream fix or a better configuration can lower a solver’s effort, so contributions are welcome on every item.
Navier–Stokes (grid)
The solvers in this domain share the schema in mosaic/mosaic_shared/problems/navier_stokes_grid/schemas.py and are benchmarked on Periodic box, Channel flow past a cylinder, Forced Kolmogorov flow and Triply periodic box.
Exponax
Exponential time-differencing Runge–Kutta (ETDRK) spectral integrator for 2D and 3D Navier–Stokes, enforcing incompressibility to machine precision by construction. Gradients via JAX source-transformation AD
Abstraction level: framework.
Runs the 2-D flow in vorticity form L2 required
The 2-D path converts velocity to Exponax’s vorticity (q = -curl u) with a spectral curl, runs the vorticity stepper, and converts back with a spectral Biot-Savart inversion.
- Periodic box, Forced Kolmogorov flow (cannot express): Exponax offers only the vorticity form in 2-D, so the benchmark’s velocity state cannot be stepped directly.
Details
- Periodic box, Forced Kolmogorov flow: The benchmark’s state is velocity, and Exponax 0.2.0’s velocity steppers (NavierStokesVelocity, KolmogorovFlowVelocity) accept only three spatial dimensions. The mean velocity, which the vorticity form drops, is carried separately, under the mean force and drag when there are any. Evidence: https://github.com/Ceyron/exponax/blob/v0.2.0/exponax/stepper/_navier_stokes.py (num_spatial_dims checks).
- Record id:
vorticity-velocity-conversion
Sets the Kolmogorov forcing through a private attribute L3 required
With a body force, the 2-D path overwrites the injection term inside Exponax’s KolmogorovFlowVorticity stepper with the spectral curl of the force.
- Forced Kolmogorov flow (cannot express): The public forcing options place the force half a cell off the canonical grid, an error of about 20% of the laminar Kolmogorov solution at N = 64.
Details
- Forced Kolmogorov flow: The injection term is the private _integrator._nonlinear_fun.injection, set as complex coefficients. Exponax 0.2.0 builds the injection from injection_mode and injection_scale as a real coefficient on one mode, which samples sin(κy) at y = j·h. The canonical force is sampled at (j + ½)h, so the public options put it half a cell off, a phase error of κh/2 or about 20% of the laminar Kolmogorov solution at N = 64. Evidence: https://github.com/Ceyron/exponax/blob/v0.2.0/exponax/nonlin_fun/_vorticity_convection.py (the injection term); tests/test_ns_grid_forcing.py::test_laminar_kolmogorov_transient; #272.
- Record id:
canonical-force-injection
Setups and checks
Excluded experiments
forward/cylinder and optimization/drag_opt on Channel flow past a cylinder (not supported): a Fourier pseudo-spectral method represents only periodic fields, so it cannot express a channel with walls, an inflow and an outflow.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/exponax builds the image exponax_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
drag |
Array[(1,), Float32] |
✓ |
Linear drag coefficient |
order |
int |
|
ETDRK time integration order (0-4) |
kolmogorov_forcing |
bool |
|
Enable Kolmogorov sinusoidal body forcing for sustained turbulence |
injection_mode |
int |
|
Wavenumber at which energy is injected (Kolmogorov forcing) |
injection_scale |
Array[(1,), Float32] |
✓ |
Amplitude of the Kolmogorov body forcing |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
|
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
INS.jl
Julia staggered finite-volume solver with pressure projection and RK4 time integration, on the GPU through its CUDA backend (KernelAbstractions.jl). Periodic boxes use a spectral pressure solve; the channel a direct one, with the cylinder as a Brinkman penalization. Differentiates through the time loop via Zygote.jl reverse-mode AD.
Abstraction level: framework.
Hand-writes the RK4 rollout that Zygote differentiates L2 choice
The VJP differentiates a hand-written, non-mutating RK4 loop over create_right_hand_side with Zygote.pullback. Plain periodic runs of apply use the library’s in-place solve_unsteady instead.
- Periodic box, Forced Kolmogorov flow, Triply periodic box: Zygote cannot differentiate array mutation, so the VJP needs a non-mutating rollout. No reason is recorded for writing RK4 by hand instead of looping over INS.jl’s own timestep.
- Channel flow past a cylinder: The VJP needs the non-mutating rollout for the same reason as on the periodic boxes.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box: On periodic boxes apply runs solve_unsteady with RKMethods.RK44, except with a body force, a drag or snapshots, where it runs the hand-written loop too, so that with a force both integrate the same right-hand side. No comparison with the library’s force! hook (examples/Kolmogorov2D.jl) is recorded. INS.jl’s documented way to get a non-mutating rollout is a loop over its non-mutating timestep. Keeping the in-place stepper for apply was faster on the CPU, before the move to the GPU (#307): 0.56 s against 2.49 s for the non-mutating rollout at N = 256, 100 steps. The two paths agreed to 2.2e-4 (2-D) and 1.0e-4 (3-D) in float32. Evidence: https://github.com/agdestein/IncompressibleNavierStokes.jl/blob/v5.0.0/examples/Differentiability.jl; #103.
- Channel flow past a cylinder: In the channel apply runs the hand-written loop too, with the channel’s right-hand side, and records the force on the cylinder after every step. Because apply shares the loop, the drag gradient check compares the VJP with exactly the rollout it differentiates, and the per-step force comes from the same loop. solve_unsteady was not tried for the channel’s apply. Evidence: tests/test_ins_jl_cylinder.py::test_drag_gradient_matches_fd.
- Record id:
hand-written-rk4-rollout
Models the cylinder as a Brinkman penalization force L2 required
Represents the cylinder as a Brinkman penalization force −χu/η in the momentum right-hand side, with η = dt. The force on the cylinder is the penalization’s integral after every step.
- Channel flow past a cylinder (cannot express): INS.jl 5.0.0 has no obstacles or immersed boundaries, so the cylinder cannot be expressed without it.
Details
- Channel flow past a cylinder: χ = 1 on every face touching a cell of the shared staircase (solid_cells), and the force on the cylinder is the sum of χu/η Δx² over the faces. INS.jl’s documented hook for extra momentum terms is a body force (examples/Actuator2D.jl), and the penalization is one. The force is integrated explicitly: at η = dt / 2.5 a run at 20 cells per diameter went non-finite. Halving η from dt moved Cd by 0.1% at 10 cells per diameter. Evidence: #325, notes.
- Record id:
brinkman-cylinder
Adds the inflow after INS.jl’s boundary conditions L2 required
Gives the channel a no-slip DirichletBC at the inlet and adds the inflow profile to the inlet faces of the state after INS.jl applies the boundary conditions, in every right-hand side.
- Channel flow past a cylinder (no gradient): The inflow profile, which drag_opt optimizes, would get no gradient through INS.jl’s boundary conditions.
Details
- Channel flow past a cylinder: INS.jl’s rrule for apply_bc_u returns no tangent for the setup, which holds a DirichletBC’s value, and building a Setup inside Zygote.pullback fails with “Mutating arrays is not supported – called push!”. Added after the BCs, the drag gradient matched a central difference to 1e-5 (prototype, 40 steps at 10 cells per diameter) and 3e-5 over drag_opt’s horizon. Evidence: src/boundary_conditions.jl of IncompressibleNavierStokes.jl v5.0.0 (the rrule of apply_bc_u); tests/test_ins_jl_cylinder.py::test_drag_gradient_matches_fd.
- Record id:
differentiable-inlet
Solves the channel’s pressure equation by diagonalization L2 required
Solves the channel’s pressure Poisson equation by diagonalizing INS.jl’s Laplacian as a Kronecker sum of x and y operators, instead of with the default psolver_direct.
- Channel flow past a cylinder (crash, too slow): On the GPU the default direct solver fails to load, the transform solver rejects the outlet boundary, and CG was stopped after 600 s where this solver takes 2.7 s.
Details
- Channel flow past a cylinder: On the uniform grid the Laplacian is the Kronecker sum of an x and a y operator read off laplacian_mat(setup). psolver_direct is the solver default_psolver picks. The transforms run in Float64. On CuArrays psolver_direct needs INS.jl’s CUDSS extension, which imports CUDSS.CUDA. CUDSS 0.7 and later, the only versions compatible with CUDA.jl 6, no longer define it, so the extension fails to load, Pkg.precompile() exits 1 and psolver_direct raises “Unsupported array type”. psolver_transform asserts that every boundary is periodic or Dirichlet, which the PressureBC outlet is not. psolver_cg on the GPU was stopped after 600 s at 10 cells per diameter, where this solver’s run takes 2.7 s. Evidence: ext/CUDSSExt.jl and src/pressure.jl of IncompressibleNavierStokes.jl v5.0.0; CUDSS.jl’s registry compat (CUDSS 0.7 and later require CUDACore 6); #325, notes.
- Upstream: Not reported. INS.jl’s ext/CUDSSExt.jl would need to import CUDA (or CUDACore) directly instead of through CUDSS.
- Record id:
separable-pressure-solver
Runs the solve and its pullback in float32 L0 choice
Builds the Setup from Float32 grid vectors and keeps literals eltype-preserving, so the solve and its pullback run in float32.
- Periodic box, Forced Kolmogorov flow, Triply periodic box, Channel flow past a cylinder: With Float64 grid vectors the state and the Zygote tape were promoted to Float64 after the first projection, so INS.jl paid Float64 cost with float32 inputs and outputs.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box, Channel flow past a cylinder: Float32 matches the other float32 solvers. #103 offered keeping Float64 and disclosing it instead. No cost factor was measured separately, on the CPU or on the GPU. Evidence: #103, item 3; INS.jl’s manual on floating point precision (docs/src/manual/backends.md at v5.0.0).
- Record id:
float32-setup
Runs the solver on INS.jl’s CUDA backend L0 choice tuned
With a GPU present, builds the Setup with backend = CUDABackend() and moves the initial state to the device, so the solve and its Zygote pullback run on the GPU instead of the default CPU backend.
- Periodic box, Forced Kolmogorov flow (too slow): On the CPU, vjp_cost at N = 256 took 9.20 s against 0.96 to 2.59 s on a shared A100.
- Triply periodic box (too slow): On the CPU, vjp_cost at N = 64 took 54.6 s against 1.58 to 3.72 s.
- Channel flow past a cylinder (too slow): With one Julia thread on the CPU, a forward run to t = 8 at 20 cells per diameter took 25.5 s against 5.4 s on an A100.
Details
- Periodic box, Forced Kolmogorov flow: Small cases got slower on the GPU: spatial_cost at N = 64 went from 0.085 s to 0.16 to 0.64 s. Every status and forward error is unchanged. Evidence: #307, notes.
- Triply periodic box: 100 steps took 47.2 s on the CPU against 2.38 to 3.15 s. The 5120-step horizon_sweep_limits VJP needs 28.3 GB on the device and runs out of memory on the T4. Evidence: #307, notes.
- Channel flow past a cylinder: At 10 cells per diameter the same run took 3.4 s against 2.7 s. drag_opt’s VJP (10 cells per diameter, t = 4) took 8.2 s against 3.9 s. Warm calls. The A100 was not shared. Evidence: #325, notes.
- Record id:
cuda-backend
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/ins-jl builds the image ins_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
✓ |
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
jax-cfd
JAX-native incompressible Navier–Stokes solver on a staggered MAC grid with finite-difference advection and spectral FFT pressure projection. Gradients via JAX source-transformation AD through the time loop.
Abstraction level: framework.
Uses central advection with RK4 time stepping L0 choice tuned
Passes convect=advection.convect_linear and time_stepper=time_stepping.classic_rk4 to semi_implicit_navier_stokes instead of its defaults, van Leer advection with forward Euler.
- Periodic box, Forced Kolmogorov flow (inaccurate forward): With the default scheme the Taylor-Green errors were 5 to 6 times the peer median, and the convergence order was 1.25.
- Triply periodic box: Not compared with the default on this setup.
Details
- Periodic box, Forced Kolmogorov flow: With the default scheme the Taylor-Green errors were 1.3e-2 to 1.5e-2 (borderline) and the fd_check error 5.5e-4. With this scheme: 2.3e-3 to 2.4e-3 (1.0 times the median), order 2.00 and 9.3e-6. Central advection is unstable under forward Euler and RK2 in the PR’s probes, hence RK4. The forward pass is 2 to 3 times slower. Evidence: #225, before/after and scheme tables (merged through #232).
- Triply periodic box: jax-cfd joined ns-3d-grid with this scheme, and every cell was ok. Evidence: #283, notes.
- Upstream: https://github.com/google/jax-cfd/issues/259 asked which advection scheme the maintainers would pick; no reply, and the project says it is no longer maintained.
- Record id:
central-advection-rk4
Combines the body force and drag by hand L2 required
Adds the body force and the linear drag in one forcing function passed to semi_implicit_navier_stokes, in place of jax-cfd’s sum_forcings.
- Forced Kolmogorov flow (crash): Without it, combining the force and the drag through sum_forcings raises an AttributeError under the pinned JAX 0.10.0.
Details
- Forced Kolmogorov flow: sum_forcings adds the terms with equations.sum_fields, which calls jax.tree_map, and JAX 0.10.0 no longer has jax.tree_map. Reproduced on the CPU with jax-cfd 0.2.1 and JAX 0.10.0 by passing kolmogorov_forcing and linear_forcing to sum_forcings and calling the result on a periodic velocity: “module ‘jax’ has no attribute ‘tree_map’”. Evidence: https://github.com/google/jax-cfd/blob/0.2.1/jax_cfd/base/equations.py (sum_fields); #272.
- Upstream: No issue filed: the project says it is no longer maintained.
- Record id:
summed-forcing
Setups and checks
Excluded experiments
forward/cylinder and optimization/drag_opt on Channel flow past a cylinder (not supported): jax-cfd 0.2.1 cannot express the channel. Its advection raises NotImplementedError for any inflow (only no-penetration and periodic boundaries are supported), it pads a zero-gradient boundary only for cell-centered fields, so the outlet’s face-centered u_x cannot have one, its boundary values are scalars, so the inlet cannot carry a profile, and it has no obstacles.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/jax-cfd builds the image jax_cfd_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
|
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
density |
Array[(1,), Float32] |
✓ |
Density of the fluid |
inner_steps |
int |
|
Time steps of dt per output step, so a run advances steps * inner_steps * dt |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
|
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
OpenFOAM
OpenFOAM icoFoam PISO solver run as a forward-only reference baseline (C++ subprocess). No reverse-mode AD; used to anchor cross-solver comparisons.
Abstraction level: application.
Tightens the momentum solver tolerance to 1e-8 L0 required tuned
The periodic icoFoam case solves the momentum predictor to an absolute tolerance of 1e-8 instead of the cavity tutorial’s 1e-5. The 3-D periodic case uses the same file.
- Periodic box (inaccurate forward): At 1e-5 the 2-D Taylor-Green error stalls at 2.7e-4 between N = 64 and N = 128, which made forward/baseline borderline.
Details
- Periodic box: At 1e-8 it reaches 5.9e-5 at N = 128. Evidence: #222, table of U tolerances from 1e-5 to 1e-9 (merged through #232); 3-D results were unchanged.
- Record id:
momentum-tolerance
Solves the periodic pressure with GAMG L0 choice tuned
The periodic boxes solve the pressure with GAMG and the DICGaussSeidel smoother instead of the cavity tutorial’s PCG with DIC, at the tutorial’s tolerances (1e-6, relTol 0.05).
- Periodic box (too slow): With PCG/DIC, spatial_cost at N = 256 takes 37.4 s against 16.4 s.
- Triply periodic box: Cost moves by less than 8% either way.
Details
- Periodic box: temporal_cost at 1000 steps takes 44.7 s against 27.8 s. Every status is unchanged, and forward errors move in the fifth significant digit. Evidence: #306, before/after table.
- Triply periodic box: spatial_cost at N = 64 takes 28.7 s with PCG/DIC against 27.1 s. The 3-D box shares the 2-D box’s fvSolution. Evidence: #306, before/after table.
- Record id:
periodic-pressure-solver
Runs the periodic boxes in parallel with MPI L0 choice tuned
The periodic boxes are decomposed with scotch and run under mpirun, one rank per physical core and at least 10,000 cells per rank. Smaller cases run serially, as does the channel.
- Periodic box (too slow): Run serially, spatial_cost at N = 256 takes 9.90 s against 6.39 s decomposed.
- Triply periodic box (too slow): Run serially, spatial_cost at N = 64 takes 25.6 s against 12.9 s decomposed.
Details
- Periodic box: Statuses are unchanged. With the CPUs loaded by other solvers, decomposed runs were slower than serial ones (#306). The measurement assumes CI’s cost jobs run alone. Evidence: #314, before/after table; #306, notes.
- Triply periodic box: 100 steps at N = 48 take 18.2 s serially against 7.24 s decomposed. At 10 steps decomposing and reconstructing outweigh the solve (3.92 s serial against 4.33 s). On a loaded CPU, N = 64 took 150 to 280 s at 24 ranks against 29 s serial (#306). Evidence: #314, before/after table; #306, notes.
- Record id:
periodic-mpi-decomposition
Tightens the channel’s momentum tolerance to 1e-8 L0 choice
The cylinder channel solves the momentum predictor to 1e-8 instead of the tutorial’s 1e-5.
- Channel flow past a cylinder: None measured. At 1e-5, forward/cylinder gives the same Cd to five digits with no speed-up.
Details
- Channel flow past a cylinder: drag_opt was not run at 1e-5. Evidence: #269, notes.
- Record id:
channel-momentum-tolerance
Adds a non-orthogonal corrector on the cylinder mesh L0 required
One non-orthogonal corrector (nNonOrthogonalCorrectors 1) on the body-fitted O-grid, where the offsetCylinder tutorial uses 0.
- Channel flow past a cylinder (crash): Without it, the run at 20 cells per diameter fails (maximum Courant number 11.9).
Details
- Channel flow past a cylinder: The periodic case uses 0, as its tutorial does. Evidence: #269, before/after table.
- Record id:
channel-non-orthogonal-corrector
Grades the channel mesh toward the cylinder L0 choice tuned
The O-grid cells shrink radially to half the size of the box cells at the cylinder wall, and the wake block coarsens to 10 dx at the outlet.
- Channel flow past a cylinder (too slow): A uniform mesh takes 111 s for forward/cylinder against 82 s, though its Cd is closer to the reference.
Details
- Channel flow past a cylinder: The smallest graded cell is about 0.35 dx. Both times were measured with the tighter pressure tolerance and MPI decomposition the channel used before #269. The uniform mesh gives Cd 5.567318 and 5.578718 at 10 and 20 cells per diameter, against the graded mesh’s 5.563948 and 5.576608 and the reference 5.5795. Evidence: #269, before/after table; smallest cell from #263.
- Record id:
graded-channel-mesh
Setups and checks
Excluded experiments
cost/vjp_cost on Periodic box (not supported): OpenFOAM has no VJP to benchmark.gradient on Periodic box and Forced Kolmogorov flow (not supported): OpenFOAM is non-differentiable (C++, no AD path); DAFoam/OpenFOAM-AD exist but are not deployed in this tesseract.optimization on Channel flow past a cylinder (not supported): OpenFOAM is a non-differentiable, forward-only solver.forward/kolmogorov_laminar and forward/kolmogorov on Forced Kolmogorov flow (not implemented): this tesseract runs icoFoam, which has no momentum source term, so it cannot apply the body force or the drag. OpenFOAM’s foamRun incompressibleFluid solver supports both through an fvModel, but this tesseract does not use it.cost/vjp_cost on Triply periodic box (not supported): standard icoFoam has no VJP to benchmark.gradient on Triply periodic box (not supported): standard icoFoam is non-differentiable (C++, no AD path); DAFoam/OpenFOAM-AD exist but are not deployed in this tesseract.optimization on Triply periodic box (not supported): standard icoFoam is a non-differentiable, forward-only solver.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/openfoam builds the image openfoam_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
|
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
|
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
|
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
|
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
|
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
|
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
PhiFlow
Periodic boxes: 2nd-order central differences under phiflow’s incompressible RK4 with CG pressure projection. The cylinder channel: semi-Lagrangian advection with operator splitting, phiflow’s wind-tunnel boundaries and a staircase cylinder. Gradients via JAX autodiff with adjoint CG solves.
Abstraction level: framework.
Widens phiml’s sparse-matrix indices to int64 L3 required
Monkeypatches phiml’s private NumPyBackend.ravel_multi_index to widen flat indices to int64 when the index space exceeds int32.
- Periodic box, Forced Kolmogorov flow, Triply periodic box (crash): Building the pressure matrix fails at 2-D N >= 216 and 3-D N >= 36, because phiml casts the flat indices back to int32 and they wrap.
Details
Runs the channel’s pressure solve with CG L0 unknown
The channel’s pressure solve passes math.Solve(“CG”, …) instead of the “auto” method of make_incompressible’s default Solve().
- Channel flow past a cylinder: No reason is recorded.
Details
- Channel flow past a cylinder: The periodic boxes’ solve keeps incompressible_rk4’s default method (CG) and forward tolerance, with the changes recorded under periodic-cg-rank-deficiency and periodic-adjoint-relative-tolerance.
- Record id:
explicit-cg-solve
Runs CG on the singular periodic pressure system L0 required
The periodic boxes’ pressure solve passes rank_deficiency=0, so phiml runs CG on the singular pressure system, whose right-hand side make_incompressible balances to zero mean, instead of adding a Lagrange multiplier that makes the system indefinite.
- Forced Kolmogorov flow (non finite, incorrect gradient, too slow): Without it, states and the chain gradient went non-finite on CI, and forward/kolmogorov ran about six times slower (5,585 s against 936 s).
- Periodic box (too slow): Without it, the cost experiments run 1.3x to 7.5x slower, growing with N. No check changes.
- Triply periodic box (too slow): Without it, the cost experiments run 1.4x to 16x slower, growing with N. No check changes.
Details
- Forced Kolmogorov flow: CG is not guaranteed to converge on an indefinite system, and under jit phiml returns an unconverged solve without raising. At 256², 852 of 6400 solves to t = 4 stopped unconverged (128 of them as diverged, with residuals up to 10⁴ times the right-hand side), and two of three states went non-finite on CI. The chain gradient was non-finite on CI and off by 12% at the best ε on the A100. With rank_deficiency=0 every solve converged, in about 375 iterations against 570, both checks passed, and forward/kolmogorov took 936 s against 5,585 s. Evidence: forward/kolmogorov and gradient/chain_fd_check of #272; PR description.
- Periodic box: With it, all four PhiFlow cost anomalies went away. Evidence: #272, PR description.
- Triply periodic box: The largest factor is 15.7x, for the VJP at N = 64. Evidence: #272, PR description.
- Record id:
periodic-cg-rank-deficiency
Uses a relative tolerance in the adjoint solve L0 choice
The periodic boxes’ adjoint pressure solve (the Solve’s gradient_solve) stops at a relative residual of 1e-5 only, instead of inheriting the forward solve’s absolute tolerance of 1e-5.
- Forced Kolmogorov flow (incorrect gradient): The chain gradient misses the finite differences by 3.6e-3 instead of 1.1e-4, because the absolute tolerance stops the adjoint solves early for small cotangents. The check still passes.
- Periodic box, Triply periodic box: Its effect has not been measured separately. No check changed.
Details
- Forced Kolmogorov flow: The adjoint solve’s right-hand side scales with the cotangent. In the chain gradient check, whose cotangents are small, the absolute tolerance stopped the adjoint solves after 0 to 17 iterations, against about 100 at the relative tolerance. Both gradient errors are at the best ε. The peers miss by 7e-5 to 2e-3. Evidence: gradient/chain_fd_check of #272; PR description.
- Record id:
periodic-adjoint-relative-tolerance
Loosens the channel’s pressure CG tolerance L0 choice tuned
The channel’s pressure CG stops at a relative and absolute residual of 1e-2 instead of phiml’s float32 default of 1e-5.
- Channel flow past a cylinder (too slow): At a tolerance of 1e-5, forward/cylinder took 293 s against 43 s on a shared GPU.
Details
- Channel flow past a cylinder: On the same shared GPU, drag_opt took 10,782 s at 1e-5 against 1,271 s. Against 1e-7, Cd, Cl and the pressure drop move by at most 0.07%, and drag_opt reaches a drag reduction of 47.60% against 47.66%. Every status is unchanged. At 1e-1 the Cd error moves by 4%. Tuned with the channel-pressure-warm-start correction in place, so the tolerance applies to the correction’s residual. Evidence: #308, notes.
- Record id:
channel-cg-tolerance
Uses central differences with RK4 on periodic boxes L0 choice
Periodic boxes advect and diffuse with 2nd-order central differences under fluid.incompressible_rk4 instead of the default semi-Lagrangian advection with operator splitting.
- Periodic box, Forced Kolmogorov flow, Triply periodic box (inaccurate forward): The default semi-Lagrangian advection is first order, per PhiFlow’s docs, and was about 10x less accurate on the periodic box in an unpublished measurement.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box: PhiFlow’s docs point users who want accurate periodic flow to this pairing. The 10x figure comes from the cleanup notes, and that table is not published. Compared with the previous central-difference Euler step, errors change by a few percent and the cost is 2.4 to 4.4x higher, but the RK4 step stays stable where the Euler step went NaN (2-D, N = 32, nu = 0, 800 steps). The runs use pressure_order=2, with four pressure projections per step. pressure_order=2 follows the comparison notebook, and the function default (4) was not measured. The channel keeps the default, since incompressible_rk4 takes no obstacles. Evidence: #288; https://tum-pbs.github.io/PhiFlow/Fluid_Simulation.html (“Higher-Order Simulations”).
- Record id:
rk4-central-differences
Writes out the channel’s obstacle projection by hand L2 unknown
Writes the channel’s obstacle projection out instead of calling fluid.make_incompressible with an obstacle.
- Channel flow past a cylinder (cannot express): No reason is recorded.
Details
- Channel flow past a cylinder: The written-out projection builds face masks from the shared staircase solid_cells, zeroes solid faces after advection and after diffusion, and uses a masked_laplace CG solve over the fluid cells and a masked green-gauss pressure gradient. A plausible reason, not checked: make_incompressible rasterizes its own obstacle geometry, so it cannot take the staircase mask the other grid solvers share, and it does not return the momentum removed at solid faces, which body_force needs.
- Record id:
inlined-obstacle-projection
Solves for a pressure correction in the channel L2 unknown
Solves each step’s channel pressure as a correction to the previous step’s pressure, from a zero initial guess, instead of passing the previous pressure as the solve’s x0.
- Channel flow past a cylinder (too slow): Per the code comment, CG needs more iterations in the near-steady flow without it. This has not been measured for PhiFlow.
Details
- Channel flow past a cylinder: The channel’s CG tolerance was tuned with it in place. Passing the previous pressure as x0 fails under reverse mode with an UnexpectedTracerError inside lax.scan. Evidence: Comment in channel_fwd.step (#239).
- Upstream: https://github.com/tum-pbs/PhiML/issues/17 (x0 tracer error)
- Record id:
channel-pressure-warm-start
Computes the force on the cylinder by hand L2 required
Computes the per-step force on the cylinder as the momentum removed at solid faces in a step (advection and diffusion) plus the pressure on the faces between fluid and solid cells.
- Channel flow past a cylinder (cannot express): There is no body_force output for forward/cylinder and drag_opt to score, because PhiFlow has no force-on-obstacle routine.
Details
- Channel flow past a cylinder: The missing routine was checked for PhiFlow 3.2.0 and not re-checked for the pinned 3.4.0. Evidence: #239 makes body_force a required output.
- Record id:
body-force-output
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/phiflow builds the image phiflow_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
✓ |
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
PICT
GPU-accelerated differentiable PISO solver built on PyTorch with custom CUDA kernels (PISOtorch). VJP w.r.t. v0, inflow_profile, and viscosity flows through the PISO time loop via PISOtorch_diff (tracks VELOCITY, BOUNDARY_VELOCITY, and VISCOSITY gradients). dt is a Python scalar consumed at construction time; zero gradients are returned for it. The VJP scales the cotangents to a largest entry of 1 before PICT’s backward pass and scales the gradients back, because PICT’s linear solves stop at an absolute residual and would cut the adjoint solves short for small cotangents. PICT has no releases; this runs commit 72a6e86.
Abstraction level: framework.
Tightens the periodic solve tolerance to 1e-6 L0 required tuned
The 2-D periodic runs pass 1e-6 as the advection and pressure solve tolerance instead of PICT’s float32 default of 1e-5.
- Periodic box, Forced Kolmogorov flow (incorrect gradient): At 1e-5, fd_check cannot get below about 1e-3 at any step size, because of the residual error of the forward and adjoint solves.
Details
- Periodic box, Forced Kolmogorov flow: The reported error is 8.8e-4 at 1e-5 against 5.7e-5 at 1e-6, which becomes an anomaly against the peer median after #282. Measured without cotangent-scaling, which alone lowers the error at 1e-5 to 2.7e-4 (#294). The two together are not measured. Evidence: #287, before/after table; #294, A/B table.
- Record id:
periodic-solve-tolerance
Uses the same 1e-6 solve tolerance in 3-D L0 unknown
The 3-D periodic box passes the same 1e-6 advection and pressure solve tolerance, through the same code path as the 2-D box.
- Triply periodic box: Its effect on ns-3d-grid has not been measured.
Details
- Triply periodic box: #287 compared tolerances on ns-grid only.
- Record id:
periodic-3d-solve-tolerance
Scales the cotangents before PICT’s backward pass L2 choice
The VJP scales the cotangents to a largest entry of 1 before PICT’s backward pass and scales the gradients back, so that the absolute tolerance of PICT’s adjoint solves acts relative to the cotangent.
- Periodic box, Forced Kolmogorov flow (incorrect gradient): Without it, gradients through a chain of calls with small cotangents come out wrong, by up to 73% through 16 calls. Single calls with O(1) cotangents are barely affected.
- Triply periodic box: Statuses are unchanged, and fd_check moves within its spread.
Details
- Periodic box, Forced Kolmogorov flow: PICT’s adjoint solves stop at an absolute residual, so its VJP is not linear in the cotangent. The gradient suite checks single calls with O(1) cotangents. ns-grid fd_check at eps = 1 goes from 1.1e-3 to 2.7e-4 with it, and every status stays the same. The chain of 16 calls had cotangents of about 1e-4. Measured at the earlier solve tolerance of 1e-5, before periodic-solve-tolerance. Evidence: #294, A/B table; tests/test_pict_vjp.py.
- Triply periodic box: The recovery’s final IC error is 0.0022 to 0.0040 with it against 0.0013 to 0.0021 without, both far below the failure threshold. #294 has no verified explanation. Evidence: #294, A/B table.
- Upstream: None filed. A relative convergence criterion would need a change to PISOtorch, whose criteria are all absolute.
- Record id:
cotangent-scaling
Warm-starts the channel’s pressure solve L3 choice tuned
Overrides Simulation.linear_solve so that each channel pressure CG solve starts from the same corrector’s pressure one step earlier and stops at 1e-8.
- Channel flow past a cylinder (too slow): Without it, the channel forward runs 2.5x (10 cells per diameter) to 4.8x (20 cells per diameter) slower.
Details
- Channel flow past a cylinder: The override subclasses PISOtorch_simulation.Simulation. In differentiable mode the solve is a torch.autograd.Function on the private PISOtorch_diff._linear_solve_wrapper, with a warm-started adjoint. With it Cd stays within 0.05%. The code comment gives 2e-5 against a solve from zero at 1e-8, and at 1e-6 Cd at 20 cells per diameter still drifts by 1e-3 over the final 40% of the run. drag_opt runs the differentiable half, and its effect there has not been measured. Evidence: #239, summary and notes; comment above pressure_tol in _run_channel.
- Upstream: None filed.
- Record id:
warm-start-pressure
Loosens PICT’s startup flux check for the channel L1 unknown
Sets pressure_tol to twice the measured boundary flux imbalance, and at least 1e-6.
- Channel flow past a cylinder (crash): PICT refuses to start when the imbalance exceeds the tolerance.
Details
- Channel flow past a cylinder: Per the code comment, the float32 imbalance after balancing is about 4e-8 of the flux (1.4e-6 at 40 cells per diameter), which a fixed 1e-6 would fail. With the warm start active, pressure_tol only feeds the check at the start of Simulation.run(). PICT’s float32 default of 1e-5 may already be above the imbalance at the benchmark’s 10 and 20 cells per diameter, which has not been checked. Evidence: Comment above pressure_tol in _run_channel.
- Record id:
pressure-tol-flux-gate
Turns off PICT’s non-orthogonal correction in the channel L0 required
Runs the channel with non_orthogonal=False, which drops PICT’s skew correction although the O-grid cells near the box corners are skewed.
- Channel flow past a cylinder (crash): With the default non_orthogonal=True, CG and float32 BiCGStab break down in the first pressure solve.
Details
- Channel flow past a cylinder: Whether dropping the correction contributes to PICT’s Cd sitting about 3% high has not been tested. Evidence: #239, notes; comment at the _simulate call in _run_channel.
- Upstream: https://github.com/tum-pbs/PICT/issues/9
- Record id:
orthogonal-scheme-on-o-grid
Builds the channel in cell units L0 unknown
Builds the channel blocks in units of the canonical cell size dx, so that PICT’s absolute solver residual is measured on unit cells.
- Channel flow past a cylinder (inaccurate forward): In physical units the same absolute tolerance would mean a different relative accuracy. The effect has not been measured.
Details
- Channel flow past a cylinder: Vertices, viscosity and dt are divided by dx. The periodic box gets unit cells from PICT’s CreateBlockWithSize without asking. Evidence: Header comment of the channel section in tesseract_api.py.
- Record id:
channel-in-cell-units
Re-applies the inflow profile every step L0 unknown
Re-applies the inflow profile to the inlet faces in a PRE callback every step instead of once before the run.
- Channel flow past a cylinder (incorrect gradient): Per the code comment, the inflow_profile gradient would come back as zeros. This has not been reproduced.
Details
- Channel flow past a cylinder: Per the same comment, a profile set once is not in the autograd graph that PISOtorch_diff records. drag_opt uses this gradient, and with it PICT lowers Cd by 47% in 44 iterations (A100). Evidence: Comment in _pre_step in _run_channel.
- Record id:
inflow-in-autograd-graph
Integrates the wall force on the cylinder L2 required
Computes the force on the cylinder from wall pressure and shear, taken from second-order one-sided fits to the first two cells next to the wall and integrated over the arc faces every step.
- Channel flow past a cylinder (cannot express): PICT has no force or drag output, so body_force, which forward/cylinder scores, could not be returned.
Details
- Channel flow past a cylinder: The integration runs in a POST callback. #239 reports PICT’s control-volume Cd 7 to 11% off this surface Cd (about 1% for the other solvers), which has not been investigated. Evidence: #239; PISOtorch_simulation.py at PICT 72a6e86 has no force routine.
- Record id:
body-force-integration
Setup deviations
- Replaces the pressure outlet with advective outflow (Channel flow past a cylinder): PICT’s advective outflow replaces the prescribed p = 0 outlet: outlet faces start from the inlet profile, boundary fluxes are balanced, the outlet velocity is advected with the mean inflow speed every step, pressure is Neumann there, and the returned pressure is shifted to zero on the outlet face. Effect: Partly measured. In OpenFOAM, an advective velocity outlet with p = 0 leaves Cd, Cl and pressure drop unchanged to six digits, as expected for a steady flow; the missing pressure outlet could not be reproduced there. Its share of PICT’s 3% Cd excess (#239) is unmeasured. Evidence: PICT has no pressure outlet (comment in _run_channel); OpenFOAM A/B of 2026-10-29 in the intervention cleanup notes (policy 3), not yet published.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/pict builds the image pict_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
|
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
✓ |
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
Warp-NS
Chorin projection scheme with explicit Euler time stepping, implemented in NVIDIA Warp CUDA kernels with FFT pressure Poisson on periodic domains. The cylinder channel uses a staggered grid, a staircase cylinder and a conjugate-gradient pressure Poisson solve. Gradients via wp.Tape kernel-level VJP, with adjoint CG solves.
Abstraction level: kernel-toolkit.
Restructures the solver on a Warp developer’s advice L0 choice
The periodic solver was restructured following a Warp developer’s advice, with native tape adjoints, an on-device tile-FFT Poisson solve and fused kernels.
- Periodic box, Forced Kolmogorov flow, Triply periodic box (too slow): The earlier version was 15 to 30x slower in forward and VJP.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box: The native tape adjoints replaced hand-written NumPy adjoints. The fused kernels are the pack and unpack kernels, and the forward-only and taped paths got separate buffer policies. The slowdown was measured by direct apply and VJP calls on an A100. Evidence: #132; https://github.com/NVIDIA/warp/discussions/1636.
- Record id:
maintainer-advised-rewrite
Computes non-power-of-two FFTs with Bluestein’s algorithm L2 required
Grids whose size is not a power of two run their spectral Poisson solve through a Bluestein chirp-z transform built from power-of-two wp.tile_fft calls.
- Periodic box, Forced Kolmogorov flow, Triply periodic box (crash): wp.tile_fft (cuFFTDx) only compiles power-of-two lengths on the GPU, so the benchmark’s non-power-of-two sizes (N = 192 in 2-D, N = 48 in 3-D) failed at compile time.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box: Power-of-two grids keep the native path. The transform is exact to float32 round-off. Evidence: #132; section comment in tesseract_api.py. The scaling of this path was fixed in #232; tests/test_warp_ns_fft.py checks it against NumPy.
- Upstream: Not filed. Warp 1.15.0’s tile_fft documentation lists any size as supported on GPU builds with libmathdx, which is the default, so the failure in #132 contradicts it; it has not been reduced to a standalone reproducer.
- Record id:
bluestein-non-pow2-fft
Hand-writes the adjoint of the channel’s pressure solve L2 required
The backward pass of each channel pressure solve is a hand-written CG adjoint registered with tape.record_func, started from the previous adjoint solution.
- Channel flow past a cylinder (incorrect gradient): Without it, wp.Tape returns a zero gradient through the pressure solve, without an error.
Details
- Channel flow past a cylinder: The adjoint CG uses the same symmetric operator as the forward solve. The pressure solve is warp.optim.linear.cg, captured as a CUDA graph. Reproduced on the CPU with Warp 1.15.0 on a 32 × 32 symmetric positive definite system, without graph capture or a preconditioner: warp.optim.linear.cg solved to a relative error of 1e-7, and tape.backward gave a gradient of norm 0 against 0.13 for the exact one. The full solver path has not been run without the adjoint. Evidence: Docstring of channel_solve (#239); tests/test_warp_ns_cylinder.py checks the drag gradient against central finite differences.
- Record id:
channel-cg-adjoint
Preconditions the pressure CG with the cylinder-free operator L2 choice tuned
Preconditions the channel’s pressure CG with the exact inverse of the same operator without the cylinder, applied through dense 1-D eigenbases along x and y.
- Channel flow past a cylinder (too slow): Without it, CG needs about 150 iterations instead of 20 at 10 cells per diameter, and drag_opt takes twice as long. On finer grids the preconditioner stops paying off.
Details
- Channel flow past a cylinder: With the preconditioner, run time halves at 10 cells per diameter (drag_opt), breaks even at 20 (forward/cylinder) and doubles at 40. Measured at the earlier CG tolerance of 1e-5, not at the current 1e-3. Evidence: Docstring of _channel_pressure_solver (#239).
- Record id:
channel-cg-preconditioner
Runs the channel CG as a captured CUDA graph L0 choice
On CUDA the channel CG iterations run as a captured conditional graph, replayed per solution buffer, with non-convergence flagged on the device and raised after the run.
- Channel flow past a cylinder (too slow): Without it, each residual check would need a host sync. The effect has not been measured.
Details
- Channel flow past a cylinder: The saved host sync per residual check is stated in the docstring.
- Record id:
channel-cg-graph-capture
Loosens the channel’s pressure CG tolerance to 1e-3 L0 choice tuned
The channel’s pressure CG stops at a relative residual of 1e-3 (absolute tolerance 0), with a cap of 5000 iterations.
- Channel flow past a cylinder (too slow): At the earlier 1e-5, drag_opt took 407 s against 189 s on a shared GPU.
Details
- Channel flow past a cylinder: forward/cylinder took 21.6 s against 15.2 s at 1e-5, on the same shared GPU. Against 1e-7, Cd moves by 0.001%, Cl by 0.25% and the pressure drop by 0.015%, and drag_opt’s drag reduction by 0.01 points. Every status is unchanged. At 1e-2 Cl moves by 5%. Evidence: #308, notes.
- Record id:
channel-cg-tolerance
Setups and checks
Excluded experiments
gradient/chain_fd_check on Forced Kolmogorov flow (unstable): the solver advances with explicit Euler and central differences for advection, which is linearly unstable at the timestep of the chain (see forward/kolmogorov). Over the chain’s calls the unstable modes grow, so they dominate both the gradient and the finite differences, which then disagree at every ε a float32 run resolves. Where they grow fast enough, the chain from the ground-truth state diverges.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/warp-ns builds the image warp_ns_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
✓ |
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
XLB
JAX-accelerated lattice Boltzmann solver (D2Q9 / D3Q27) with BGK collision, running the streaming/collision loop on GPU. Gradients via source-transformation AD; recovers incompressible Navier–Stokes only to O(Ma²) via Chapman–Enskog expansion.
Abstraction level: framework.
Runs the solver and its gradients in float64 L0 unknown
Enables jax_enable_x64 and runs apply() and the VJP with XLB’s FP64FP64 precision policy, casting outputs to float32. XLB_VJP_FP32=1 switches everything to float32.
- Periodic box, Forced Kolmogorov flow, Channel flow past a cylinder, Triply periodic box (incorrect gradient): Per the code comment, float32 gradients fail gradient/fd_check. No float32 fd_check result is recorded.
Details
- Periodic box, Forced Kolmogorov flow, Channel flow past a cylinder, Triply periodic box: The comment attributes the failure to float32 cancellation near omega = 2 corrupting the gradients. Its reasoning predates the switch to BGK everywhere (#282). Evidence: Comment above _USE_F64_DIFF; https://github.com/pasteurlabs/mosaic/issues/125 (open) proposes measuring float32 against fd_check.
- Record id:
float64-everywhere
Reinitializes XLB for each velocity set L1 unknown
Calls xlb.init() again with each operator bundle’s velocity set and precision policy before constructing its operators, because XLB keeps both in global state.
- Periodic box, Forced Kolmogorov flow, Channel flow past a cylinder, Triply periodic box (crash): The recorded compile failure involved an operator that is no longer used. Whether the remaining operators need it has not been checked.
Details
- Periodic box, Forced Kolmogorov flow, Channel flow past a cylinder, Triply periodic box: The recorded failure was a D3Q27 bundle, built after xlb.init() with D2Q9, failing to compile. It involved the KBC operator, which #282 removed. Evidence: Docstring of _make_ops at 6517f75.
- Upstream: None filed.
- Record id:
xlb-init-per-velocity-set
Checkpoints the time loop in blocks of sqrt(steps) L0 required
Runs the time loop as nested jax.lax.scan blocks of about sqrt(steps) steps, with jax.checkpoint on each block and on each step inside it.
- Triply periodic box (out of memory): The 3-D VJP runs out of memory. A plain scan needs 110 GB at N = 64, against 0.3 GB for the forward.
- Periodic box, Forced Kolmogorov flow (too slow): Per #135, the 2-D VJP is slower without it. No number is recorded.
Details
- Triply periodic box: Reverse mode keeps the block-boundary populations and, while it works through a block, that block’s step inputs. The 110 GB for a plain scan is from #135. Checkpointing each step but not the blocks still keeps every step’s populations: on an A100 80 GB, the 3-D TGV VJP at N = 48 took 67 GB at 1000 steps and ran out of memory at 2000, against 5.6 GB at both with the blocks checkpointed. At 50 steps, N = 96 took 69 GB against 12 GB, and N = 128 ran out of memory against 38 GB. The block checkpoint makes the VJP 2-14% slower and changes the gradient by at most 5e-9 relative. Checkpointing the blocks without the steps inside them took about twice the memory and was no faster. Evidence: #135 (plain scan); #329 (A100, one container per point).
- Periodic box, Forced Kolmogorov flow: #135 says it fixes some of XLB’s VJP slowness in 2-D. The block checkpoint has not been timed in 2-D. Evidence: #135; https://github.com/pasteurlabs/mosaic/issues/125.
- Record id:
sqrt-checkpointing
Checkpoints the channel’s time loop in blocks L0 unknown
The channel runs its time loop through the same block- and step-checkpointed scan as the periodic boxes.
- Channel flow past a cylinder: No reason is recorded, and no run without it is recorded.
Details
- Channel flow past a cylinder: The channel reuses the periodic boxes’ helper since #239. drag_opt differentiates through the whole run.
- Record id:
channel-checkpointing
Starts periodic runs from populations consistent with v0 L2 required
Periodic runs start from populations consistent with v0, with density from a spectral pressure-Poisson solve plus the first-order Chapman-Enskog non-equilibrium part (Mei et al. 2006), instead of the equilibrium at density 1.
- Periodic box, Forced Kolmogorov flow (inaccurate forward): Start-up transients from the missing pressure and viscous stress raise the Taylor-Green error to 0.031 to 0.035, against 6e-4 to 8e-4 with consistent populations.
- Triply periodic box (inaccurate forward): Without it, a start-up transient made the N = 8 forward/baseline error 0.273, 80x the peers.
Details
Assembles the periodic time loop from XLB’s operators L2 unknown
The periodic time loop is assembled from XLB’s Stream, Macroscopic, QuadraticEquilibrium and BGK operators inside jax.lax.scan instead of using XLB’s stepper.
- Periodic box, Forced Kolmogorov flow, Triply periodic box (too slow): The recorded reason is compile time, which has not been measured.
Details
- Periodic box, Forced Kolmogorov flow, Triply periodic box: The initial tesseract description gave “O(1) compile time” as the reason. Nothing was measured, and whether XLB’s stepper works under scan and reverse mode was not recorded. Evidence: Initial commit (c47498a) description.
- Record id:
manual-lbm-loop
Applies exact-difference forcing by hand L2 required
Periodic runs with a body force or drag apply the exact difference method by hand in the manual LBM loop, with the drag evaluated semi-implicitly at the half-force velocity.
- Forced Kolmogorov flow (cannot express): XLB 0.3.1’s ExactDifference operator takes one constant force vector, so it cannot apply the Kolmogorov force sin(κy) or a drag that depends on the velocity.
Details
Writes the channel’s boundaries and loop by hand L2 choice
The channel is written by hand on XLB’s operators, with its own boundary conditions, momentum-exchange force and sub-stepped, checkpointed time loop.
- Channel flow past a cylinder: No reason is recorded.
Details
- Channel flow past a cylinder: The boundary conditions are boundary link masks for the staircase cylinder, halfway bounce-back on walls and cylinder, velocity bounce-back at the inlet and an anti-bounce-back outlet. XLB 0.3.1 ships halfway bounce-back, Zou-He, regularized and extrapolation-outflow boundary conditions and a momentum-transfer force operator. Whether they work under this tesseract’s scan and reverse mode was not checked. Evidence: https://github.com/Autodesk/XLB/tree/v0.3.1/xlb/operator/boundary_condition
- Record id:
channel-boundary-conditions
Lets the outlet density follow the outgoing characteristic L2 unknown
The outlet density follows the outgoing characteristic, 1 + (u_x - u_developed)/cs, so that p = 0 on the outlet face once the flow is steady.
- Channel flow past a cylinder (inaccurate forward): Per the docstring, a plain p = 0 face would reflect the start-up pressure waves, which then ring between inlet and outlet almost undamped at nu = 1e-3.
Details
- Channel flow past a cylinder: u_developed is the fully developed profile with the inlet’s flux. The reflection has not been measured. Evidence: Docstring of _channel_run (#239).
- Record id:
channel-characteristic-outlet
Setup deviations
- Sub-steps each step to bound the Mach number (Periodic box, Forced Kolmogorov flow, Channel flow past a cylinder, Triply periodic box): Each prescribed step runs as k = ceil(Ma / Ma_target) LBM steps of dt / k, never coarser than prescribed. Ma_target is 0.1, lowered in proportion to 1/N above N = 64 on periodic boxes. With a body force, Ma uses the larger of max|v0| and a bound on the speed the force can drive, so a flow it spins up from rest stays within the target. Effect: LBM ties the time step to the grid through the lattice Mach number, so XLB runs k times the prescribed number of steps (k = 9 at N = 256 on the 2-D box), and its cost includes them. Without sub-stepping the Taylor-Green case ran at Ma of about 0.88, with errors 84 to 114x the peer median. Evidence: https://github.com/pasteurlabs/mosaic/issues/38; #232.
Use it in your code
tesseract build mosaic/tesseracts/navier-stokes-grid/xlb builds the image xlb_navier_stokes_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
v0 |
GridVectorField |
✓ |
Initial velocity at the cell centers, shape (nx, ny, 1, 2) in 2-D or (nx, ny, nz, 3) in 3-D. Cell (i, j) is centered at ((i + 1/2) dx, (j + 1/2) dx). Default: 64×64 divergence-free random vortex field (seed=42). |
viscosity |
Array[(1,), Float32] |
✓ |
Kinematic viscosity ν. |
dt |
Array[(1,), Float32] |
✓ |
Timestep size. |
steps |
int |
|
Number of timesteps. Total simulated time = steps × dt. |
domain_extent |
float | list[float] |
|
Physical size of the domain. A scalar L is a square (cube) of side L, and a list gives the length per axis, [Lx, Ly] or [Lx, Ly, Lz]. Periodic boxes take a scalar. The grid spacing must be the same along every axis. |
boundary_conditions |
GridBC |
|
Per-face boundary conditions. Default: all periodic. |
obstacle |
GridObstacle | None |
|
Optional solid obstacle embedded in the domain (no-slip walls), in physical coordinates. |
inflow_profile |
Array[(None,), Float32] | None |
✓ |
Inlet velocity u_x(y) at the x_lo face, sampled at the cell-center heights y_j = (j + 1/2) dx, shape (ny,). u_y is zero on the inlet. Requires x_lo to be ‘dirichlet’ with no value. |
forcing |
GridVectorField | None |
|
Body force per unit mass (an acceleration) at the cell centers, constant in time, same shape as v0. 2-D periodic boxes only. |
linear_drag |
Array[(1,), Float32] | None |
|
Coefficient α of a linear drag -αu on the right-hand side of the momentum equation. 2-D periodic boxes only. |
snapshot_every |
int | None |
|
If set, also return the velocity after every snapshot_every steps as snapshots, so one run yields several times of a trajectory. steps must be a multiple of it. 2-D periodic boxes only. |
Outputs
result |
GridVectorField |
✓ |
Final velocity at the cell centers, same shape as v0. |
pressure |
GridScalarField | None |
|
Final kinematic pressure p/ρ at the cell centers, shape (nx, ny, nz). Values inside the obstacle are ignored. |
body_force |
Array[(None, None), Float32] | None |
✓ |
Force of the fluid on the obstacle per unit depth and density, shape (steps, ndim): row k is (F_x, F_y) at t = (k + 1) dt. Positive along +x and +y, so a flow in +x gives F_x > 0. |
snapshots |
Array[(None, None, None, None, None), Float32] | None |
|
Velocity at the cell centers after every snapshot_every steps, shape (steps / snapshot_every, nx, ny, 1, 2); the last equals result. Returned only when the inputs set snapshot_every. Not differentiable. |
Structural mechanics
The solvers in this domain share the schema in mosaic/mosaic_shared/problems/structural_mesh/schemas.py and are benchmarked on Linear-elastic cantilever.
deal.II
deal.II Q1 trilinear hexahedral FEM linear-elasticity solver via C++ subprocess (UMFPACK direct). Forward-only — reference implementation for cross-framework validation.
Abstraction level: framework.
Solves the system directly with UMFPACK L0 choice tuned
Solves the assembled system with deal.II’s SparseDirectUMFPACK instead of the CG with SSOR preconditioning of the step-8 tutorial the assembly follows.
- Linear-elastic cantilever (too slow): With the tutorial’s SSOR-preconditioned CG the solver runs 1.7 times slower at nx = 128, with the same compliance.
Details
- Linear-elastic cantilever: Evidence: #271, notes.
- Record id:
umfpack-direct-solve
Setups and checks
Excluded experiments
cost/vjp_cost and gradient (not supported): The deal.II structural tesseract is forward-only. It exposes no AD path (no VJP endpoint, no differentiable inputs), so cost/vjp and gradient experiments are categorically out of scope.optimization (not supported): The deal.II structural tesseract has no AD path, and topology optimization (compliance gradient-based ρ updates) requires a differentiable compliance output.
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/dealii-structural builds the image dealii_structural_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
|
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
E_max |
float |
|
Young’s modulus of the fully solid material [MPa]. |
nu |
float |
|
Poisson’s ratio. |
xmin |
float |
|
Void stiffness ratio (E_min = xmin * E_max). |
penal |
float |
|
SIMP penalisation exponent (E(ρ) = E_min + (E_max−E_min)·ρ^penal). |
Outputs
compliance |
Array[(), Float32] |
|
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
FEniCS
FEniCS/DOLFIN CG1 FEM linear-elasticity solver with SIMP penalisation. Gradient ∂C/∂ρ via dolfin-adjoint tape replay (discrete adjoint derived at the variational level via UFL).
Abstraction level: framework.
Solves the forward system with MUMPS L0 choice tuned
Passes solver_parameters={“linear_solver”: “mumps”} to the forward solve instead of using DOLFIN’s default LU backend (UMFPACK).
- Linear-elastic cantilever (too slow): With UMFPACK, apply at nx = 128 is about 1.15 times slower, 10.1 s against 8.8 s.
Details
- Linear-elastic cantilever: The times are from before and after the change that introduced it. Instrumentation put 99.6% of apply in the linear solve, so the gain is attributed to the solver switch. No failure without it is recorded on this setup. Evidence: #182, section 2 (measured together with tape replay).
- Record id:
mumps-forward-solver
Patches dolfin-adjoint to solve the adjoint with MUMPS L3 choice
Replaces dolfin-adjoint’s linear solve at import time with one that always uses MUMPS. Legacy dolfin-adjoint does not pass the forward solver_parameters to the adjoint solve, which then runs UMFPACK.
- Linear-elastic cantilever (too slow): Without the patch the VJP at nx = 128 takes 14.6 s against 11.1 s, with identical gradients.
Details
- Linear-elastic cantilever: At nx = 64 the VJP takes 4.7 s against 3.2 s. Small meshes show no difference. The patched function is fenics_adjoint.types.compat.linalg_solve. Evidence: A/B on an A100, 2026-10-26, in the intervention cleanup notes (item 3), not yet published.
- Upstream: Not filed: legacy dolfin-adjoint has had no release since 2019.1.0. Per #182, newer dolfin-adjoint passes linear_solver to the adjoint solve.
- Record id:
mumps-adjoint-monkeypatch
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/fenics-structural builds the image fenics_structural_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
E_max |
float |
|
Young’s modulus of the fully solid material. |
nu |
float |
|
Poisson’s ratio. |
xmin |
float |
|
Void stiffness ratio (E_min = xmin * E_max). |
penal |
float |
|
SIMP penalisation exponent (E(rho) = E_min + (E_max-E_min)*rho^penal). |
Outputs
compliance |
Array[(), Float32] |
✓ |
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
Firedrake
Firedrake CG1 FEM linear-elasticity solver with SIMP penalisation. Gradient ∂C/∂ρ via firedrake-adjoint (pyadjoint ReducedFunctional) — an independent tape-based discrete adjoint complement to FEniCS.
Abstraction level: framework.
No interventions are recorded.
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/firedrake-structural builds the image firedrake_structural_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
E_max |
float |
|
Young’s modulus of the fully solid material. |
nu |
float |
|
Poisson’s ratio. |
xmin |
float |
|
Void stiffness ratio (E_min = xmin * E_max). |
penal |
float |
|
SIMP penalisation exponent (E(rho) = E_min + (E_max-E_min)*rho^penal). |
Outputs
compliance |
Array[(), Float32] |
✓ |
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
JAX-FEM
JAX-FEM HEX8 linear-elasticity solver with SIMP stiffness E(ρ) = E_min + (E_max − E_min)·ρ^penal. Gradient ∂C/∂ρ via JAX automatic differentiation through the linear FE solve (hybrid implicit-function-theorem adjoint).
Abstraction level: framework.
Compiles JAX-FEM’s BiCGSTAB solve once per matrix size L2 choice tuned
Passes JAX-FEM’s own Jacobi-preconditioned BiCGSTAB, compiled once per matrix size, to ad_wrapper through the custom_solver option for the forward and adjoint solves.
- Linear-elastic cantilever (too slow): The default jax_solver recompiles on every call, and the umfpack_solver option used before is about 8 times slower at N = 128 (11.6 s against 1.49 s).
Details
- Linear-elastic cantilever: The custom solver uses the default jax_solver’s algorithm and tolerances. The default jax_solver runs bicgstab eagerly, so it compiles 2 XLA programs per forward call and 4 per VJP, a fixed overhead on every call. The umfpack_solver option gives identical outputs but is much slower on large meshes. At N = 128, spatial_cost took 11.6 s with umfpack_solver and 1.49 s with this solver, vjp_cost/by_N 23.5 s and 2.91 s. Evidence: #286, before/after table; comment above _bicgstab.
- Record id:
jitted-bicgstab-solver
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/jax-fem builds the image jax_fem_structural_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
Outputs
compliance |
Array[(), Float32] |
✓ |
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
TopOpt.jl
TopOpt.jl HEX8 FEM linear-elasticity solver with SIMP penalisation, Julia + Ferrite.jl backend. Gradient ∂C/∂ρ via analytical SIMP adjoint sensitivity computed alongside the forward solve.
Abstraction level: framework.
No interventions are recorded.
Setup deviations
- Applies the cantilever load as nodal forces (Linear-elastic cantilever): The harness gives TopOpt.jl the cantilever load as nodal forces (CLOAD) instead of a traction: the shares that integrating the traction over the bilinear face elements gives each node. Effect: None: on the uniform mesh the nodal forces equal the integrated traction. Evidence: TopOpt.jl 0.14.0’s INP reader takes surface loads only as a normal pressure (DLOAD P), and the cantilever load is tangential; _traction_solvers in mosaic/benchmarks/problems/structural_mesh/physics.py.
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/topopt-jl builds the image topopt_jl_structural_grid. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
E |
float |
|
Young’s modulus of the solid material. |
nu |
float |
|
Poisson’s ratio. |
xmin |
float |
|
Minimum density (void stiffness) to prevent singular stiffness matrix. |
Outputs
compliance |
Array[(), Float32] |
✓ |
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
torch-fem
torch-fem Solid HEX8 linear-elasticity solver with SIMP stiffness E(ρ) = E_min + (E_max − E_min)·ρ^p — the PyTorch backend in structural-mesh. Gradient ∂C/∂ρ via PyTorch autograd through the FE solve (implicit-function-theorem adjoint); GPU-native.
Abstraction level: framework.
Runs the solver on the GPU L0 choice
With a GPU present, runs torch-fem on it instead of torch’s default CPU device.
- Linear-elastic cantilever: No comparison with the CPU exists.
Details
- Linear-elastic cantilever: The tesseract has run on the GPU since its first version.
- Record id:
gpu-device
Sets torch’s default device and dtype globally L0 required
Sets torch’s default device and dtype (float64) process-wide, so the tensors torch-fem allocates internally land on the GPU with the mesh, material and load tensors.
- Linear-elastic cantilever (crash): torch-fem allocates auxiliary tensors without device=, so they land on the CPU and collide with the CUDA tensors.
Details
- Linear-elastic cantilever: No error log is recorded. Evidence: Code comment; torch-fem’s benchmarks/utils.py at v0.12.1 and its current installation docs set the default device the same way.
- Record id:
global-default-device
Solves with Jacobi-preconditioned CG on the GPU L0 choice
On the GPU, passes method=“cg” and preconditioner=“jacobi” to model.solve for the forward and adjoint solves. On the CPU torch-fem chooses.
- Linear-elastic cantilever (too slow): torch-fem 0.12.1 solves systems below 10,000 DOFs directly with SciPy on the CPU, even when the matrix is on the GPU.
Details
- Linear-elastic cantilever: Larger systems already get CG with Jacobi without AmgX, so the setting keeps the smaller meshes on the device. There is no isolated measurement. Evidence: #188; resolve_method and resolve_library in torch-fem v0.12.1 src/torchfem/sparse.py.
- Record id:
gpu-cg-jacobi
Loosens the CG stopping tolerance to 1e-6 L0 choice tuned
The GPU CG stops at a relative residual of 1e-6 (stol) instead of torch-fem’s 1e-10.
- Linear-elastic cantilever (too slow): At torch-fem’s default of 1e-10, spatial_cost at N = 128 took 0.435 s against 0.338 s, with every status unchanged.
Details
- Linear-elastic cantilever: vjp_cost took 0.717 s against 0.530 s. Measured on a shared GPU. Every status is unchanged from 1e-12 to 1e-6. Forward metrics are bit-identical and gradient checks move by less than 1e-9. Evidence: #301, notes.
- Record id:
cg-tolerance
Makes the CUDA assembly deterministic L0 choice
Enables torch.use_deterministic_algorithms(True, warn_only=True), so torch-fem’s CUDA scatter-adds in assembly sum in a fixed order.
- Linear-elastic cantilever: Run-to-run differences were shown on the heat setup only.
Details
- Linear-elastic cantilever: On this setup the setting is kept for consistency, with no measurable cost. Evidence: #200, notes.
- Record id:
deterministic-kernels
Use it in your code
tesseract build mosaic/tesseracts/structural-mesh/torch-fem-structural builds the image torch_fem_structural_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. Values in [0, 1]: 0 = void, 1 = solid. |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe displacement (values shape (n_groups, 3), None → zero); neumann.mask/values prescribe surface traction (values shape (n_groups, 3)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
E_max |
float |
|
Young’s modulus of fully solid material (ρ = 1). |
nu |
float |
|
Poisson’s ratio (density-independent). |
xmin |
float |
|
Void stiffness ratio E_min / E_max. |
p_exp |
float |
|
SIMP penalisation exponent p in E(ρ) = E_min + (E_max−E_min)·ρ^p. |
Outputs
compliance |
Array[(), Float32] |
✓ |
Structural compliance C = F^T U (scalar). Minimising compliance is equivalent to maximising stiffness. |
Heat transfer
The solvers in this domain share the schema in mosaic/mosaic_shared/problems/thermal_mesh/schemas.py and are benchmarked on Steady conduction in a slab.
deal.II
deal.II Q1 trilinear hexahedral FEM heat-conduction solver with SIMP conductivity (C++ subprocess, AMG-preconditioned CG). Forward-only — reference implementation for cross-framework validation.
Abstraction level: framework.
Preconditions the CG solve with algebraic multigrid L0 choice tuned
Preconditions the CG solve with Trilinos ML algebraic multigrid (PreconditionAMG, elliptic) instead of the SSOR of the step-5 tutorial the assembly follows.
- Steady conduction in a slab (too slow): With the tutorial’s SSOR preconditioner the largest mesh solves about 2.4 times slower (61.1 s against 25.6 s), with the same forward results.
Details
- Steady conduction in a slab: The times are spatial_cost at nx = 1024, where UMFPACK takes 35.6 s. Every forward metric is the same with SSOR. AMG is faster from nx = 256 on. At nx = 16 its setup makes it slower, 0.020 s against 0.014 s. Evidence: #305, notes table.
- Record id:
amg-preconditioner
Loosens the CG stopping tolerance to 1e-8 L0 choice tuned
The CG solve stops at a residual of 1e-8 times the right-hand-side norm, with a cap of 50000 iterations.
- Steady conduction in a slab (too slow): At the earlier tolerance of 1e-12 the largest mesh solves more slowly (60.1 s against 51.1 s), with identical forward metrics.
Details
- Steady conduction in a slab: The times are spatial_cost at nx = 1024, measured with SSOR. With AMG, #305 checked only that the forward metrics are unchanged at 1e-8. Evidence: #266, before/after table; #305, notes.
- Record id:
cg-tolerance
Setups and checks
Excluded experiments
cost/vjp_cost and gradient (not supported): The deal.II thermal tesseract is forward-only. It exposes no AD path (no VJP endpoint, no differentiable inputs), so cost/vjp and gradient experiments are categorically out of scope.optimization (not supported): The deal.II thermal tesseract exposes no AD path, and conductivity recovery (gradient-based ρ updates that minimize the misfit to a target temperature field) requires a differentiable identification-error output.
Use it in your code
tesseract build mosaic/tesseracts/thermal-mesh/dealii-heat builds the image dealii_heat_thermal_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
|
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. ρ=1 → fully conducting, ρ=0 → insulating. |
source |
Array[(None,), Float32] |
|
Per-element volumetric heat source (W/m³), shape (n_cells,). Active slice = source[:hex_mesh.n_faces]. Zero by default (topology-optimisation mode; no body heat source). |
target_temperature |
Array[(None,), Float32] |
|
Per-node target temperature for the inverse (source-identification) objective, shape (n_nodes,). identification_error = |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe temperature (values shape (n_groups, 1), None → zero); neumann.mask/values prescribe heat flux (values shape (n_groups, 1)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
k_max |
float |
|
Maximum thermal conductivity (fully solid material). |
p_exp |
float |
|
SIMP penalisation exponent p (k(ρ) = k_min + (k_max−k_min)·ρ^p). |
Outputs
thermal_compliance |
Array[(), Float32] |
|
Total thermal compliance C = ∮_Γ_N q_n · T dΓ (scalar). Equals the work done by the Neumann heat flux on the temperature field. |
identification_error |
Array[(), Float32] |
|
Source-identification objective: |
FEniCS
FEniCS P1 FEM heat-conduction solver with SIMP penalisation. Gradients (∂C/∂ρ, ∂C/∂source) via dolfin-adjoint tape replay.
Abstraction level: framework.
Solves the forward system with MUMPS L0 required
Passes solver_parameters={“linear_solver”: “mumps”} to the forward solve instead of using DOLFIN’s default LU backend (UMFPACK).
- Steady conduction in a slab (out of memory): With UMFPACK the forward solve runs out of memory at the largest cost-sweep size (nx = 1024).
Details
- Steady conduction in a slab: This showed as an error on the results page. No peak memory was recorded. Evidence: https://github.com/pasteurlabs/mosaic/issues/180 (comment by the FEniCS developer who proposed the change); merged in #182.
- Record id:
mumps-forward-solver
Patches dolfin-adjoint to solve the adjoint with MUMPS L3 required
Replaces dolfin-adjoint’s linear solve at import time with one that always uses MUMPS. Legacy dolfin-adjoint does not pass the forward solver_parameters to the adjoint solve, which then runs UMFPACK.
- Steady conduction in a slab (out of memory): With MUMPS in the forward solve only, the gradient computation still fails at nx = 1024, because the adjoint solve runs UMFPACK.
Details
- Steady conduction in a slab: The patched function is fenics_adjoint.types.compat.linalg_solve. Evidence: https://github.com/pasteurlabs/mosaic/issues/180 (comment by the FEniCS developer who proposed the patch); merged in #182.
- Upstream: Not filed: legacy dolfin-adjoint has had no release since 2019.1.0. Per #182, newer dolfin-adjoint passes linear_solver to the adjoint solve; per #180, DOLFINx-adjoint takes adjoint PETSc options.
- Record id:
mumps-adjoint-monkeypatch
Tapes the nodal misfit with a custom pyadjoint block L2 required
Puts the nodal identification error on the pyadjoint tape through a custom Block with hand-written adjoint and tangent-linear components, so pyadjoint differentiates it together with the solve.
- Steady conduction in a slab (cannot express, incorrect gradient): The nodal sum cannot be written as a UFL integral, because legacy FEniCS has no vertex quadrature on hexahedra. An earlier integral approximation gave gradients that did not match the reported value.
Details
- Steady conduction in a slab: The error is sum_v mask_v (T_v - T_target,v)^2. The approximation was an L2 integral rescaled by the node count. gradient/source_fd_check flagged FEniCS at 1753x the peer median, and conductivity_recovery_bfgs reached a final/initial misfit of 0.02779 against 0.00020 with the nodal adjoint. Evidence: #167; #184, notes; comment above NodalSquaredErrorBlock.
- Record id:
nodal-misfit-block
Use it in your code
tesseract build mosaic/tesseracts/thermal-mesh/fenics-heat builds the image fenics_heat_thermal_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. ρ=1 → fully conducting, ρ=0 → insulating. |
source |
Array[(None,), Float32] |
✓ |
Per-element volumetric heat source (W/m³), shape (n_cells,). Active slice = source[:hex_mesh.n_faces]. Zero by default (topology-optimisation mode; no body heat source). |
target_temperature |
Array[(None,), Float32] |
|
Per-node target temperature for the inverse (source-identification) objective, shape (n_nodes,). identification_error = |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe temperature (values shape (n_groups, 1), None → zero); neumann.mask/values prescribe heat flux (values shape (n_groups, 1)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
k_max |
float |
|
Maximum thermal conductivity (fully solid/conducting material). |
p_exp |
float |
|
SIMP penalisation exponent p (k(ρ) = k_min + (k_max−k_min)·ρ^p). |
Outputs
thermal_compliance |
Array[(), Float32] |
✓ |
Total thermal compliance C = ∮_Γ_N q_n · T dΓ (scalar). Equals the work done by the Neumann heat flux on the temperature field. |
identification_error |
Array[(), Float32] |
✓ |
Source-identification objective: |
Firedrake
Firedrake P1 FEM heat-conduction solver with SIMP penalisation. Gradients (∂C/∂ρ, ∂C/∂source) via firedrake-adjoint, complementing the FEniCS reference.
Abstraction level: framework.
Seeds the nodal misfit gradient by hand L2 unknown
Seeds the identification-error gradient with a hand-built cofunction at the nodes that have a target value, because the error is a nodal sum rather than an integral.
- Steady conduction in a slab (incorrect gradient): An earlier integral approximation gave gradients that did not match the reported value. It is not known whether the hand-built seed is needed.
Details
- Steady conduction in a slab: The cofunction is dI/dT = 2 (T - T_target) and seeds the derivative of the temperature functional. The approximation was an L2 integral rescaled by the node count. gradient/source_fd_check went from 4.58e-02 to 3.52e-05 with the nodal seed, and the rho gradient was 9% off central differences, with a conductivity_recovery_bfgs final/initial ratio of 0.00313 against 0.00020. Whether a vertex quadrature rule would let the nodal sum be written in UFL, which Firedrake’s custom quadrature may allow, was not checked. Evidence: #162; #167, summary; #184; current form from #229.
- Record id:
nodal-misfit-seed
Use it in your code
tesseract build mosaic/tesseracts/thermal-mesh/firedrake-heat builds the image firedrake_heat_thermal_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. ρ=1 → fully conducting, ρ=0 → insulating. |
source |
Array[(None,), Float32] |
✓ |
Per-element volumetric heat source (W/m³), shape (n_cells,). Active slice = source[:hex_mesh.n_faces]. Zero by default (topology-optimisation mode; no body heat source). |
target_temperature |
Array[(None,), Float32] |
|
Per-node target temperature for the inverse (source-identification) objective, shape (n_nodes,). identification_error = |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe temperature (values shape (n_groups, 1), None → zero); neumann.mask/values prescribe heat flux (values shape (n_groups, 1)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
k_max |
float |
|
Maximum thermal conductivity (fully solid/conducting material). |
p_exp |
float |
|
SIMP penalisation exponent p (k(ρ) = k_min + (k_max−k_min)·ρ^p). |
Outputs
thermal_compliance |
Array[(), Float32] |
✓ |
Total thermal compliance C = ∮_Γ_N q_n · T dΓ (scalar). Equals the work done by the Neumann heat flux on the temperature field. |
identification_error |
Array[(), Float32] |
✓ |
Source-identification objective: |
JAX-FEM
JAX-FEM HEX8 steady-state heat-conduction solver with SIMP conductivity and optional volumetric source. Gradients ∂C/∂ρ and ∂C/∂source via JAX autodiff through the linear FE solve (hybrid implicit-function-theorem adjoint)
Abstraction level: framework.
Compiles JAX-FEM’s BiCGSTAB solve once per matrix size L2 choice tuned
Passes JAX-FEM’s own Jacobi-preconditioned BiCGSTAB, compiled once per matrix size, to ad_wrapper through the custom_solver option for the forward and adjoint solves.
- Steady conduction in a slab (too slow): The default jax_solver recompiles on every call, and the umfpack_solver option used before is about 12 times slower at N = 1024 (60.8 s against 5.17 s).
Details
- Steady conduction in a slab: The custom solver uses the default jax_solver’s algorithm and tolerances. The default jax_solver runs bicgstab eagerly, so it compiles 2 XLA programs per forward call and 4 per VJP, a fixed overhead on every call. The umfpack_solver option gives identical outputs but is much slower on large meshes. At N = 1024, spatial_cost took 60.8 s with umfpack_solver and 5.17 s with this solver, vjp_cost/by_N 120 s and 10.8 s. Evidence: #286, before/after table; comment above _bicgstab.
- Record id:
jitted-bicgstab-solver
Use it in your code
tesseract build mosaic/tesseracts/thermal-mesh/jax-fem builds the image jax_fem_thermal_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. ρ=1 → fully conducting, ρ=0 → insulating. |
source |
Array[(None,), Float32] |
✓ |
Per-element volumetric heat source (W/m³), shape (n_cells,). Active slice = source[:hex_mesh.n_faces]. Zero by default (topology-optimisation mode; no body heat source). |
target_temperature |
Array[(None,), Float32] |
|
Per-node target temperature for the inverse (source-identification) objective, shape (n_nodes,). identification_error = |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe temperature (values shape (n_groups, 1), None → zero); neumann.mask/values prescribe heat flux (values shape (n_groups, 1)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
Outputs
thermal_compliance |
Array[(), Float32] |
✓ |
Total thermal compliance C = ∮_Γ_N q_n · T dΓ (scalar). Equals the work done by the Neumann heat flux on the temperature field. |
identification_error |
Array[(), Float32] |
✓ |
Source-identification objective: |
torch-fem
torch-fem SolidHeat HEX8 heat-conduction solver with IsotropicConductivity3D SIMP material — the first PyTorch backend in thermal-mesh. Gradients via PyTorch autograd through the FE solve; GPU-native.
Abstraction level: framework.
Runs the solver on the GPU L0 choice
With a GPU present, runs torch-fem on it instead of torch’s default CPU device.
- Steady conduction in a slab: No comparison with the CPU exists.
Details
- Steady conduction in a slab: The tesseract has run on the GPU since its first version.
- Record id:
gpu-device
Sets torch’s default device and dtype globally L0 required
Sets torch’s default device and dtype (float64) process-wide, so the tensors torch-fem allocates internally land on the GPU with the mesh, material and load tensors.
- Steady conduction in a slab (crash): torch-fem allocates auxiliary tensors without device=, so they land on the CPU and collide with the CUDA tensors.
Details
- Steady conduction in a slab: No error log is recorded. Evidence: Code comment; torch-fem’s benchmarks/utils.py at v0.12.1 and its current installation docs set the default device the same way.
- Record id:
global-default-device
Solves with Jacobi-preconditioned CG on the GPU L0 choice
On the GPU, passes method=“cg” and preconditioner=“jacobi” to model.solve for the forward and adjoint solves. On the CPU torch-fem chooses.
- Steady conduction in a slab (too slow): torch-fem 0.12.1 solves systems below 10,000 DOFs directly with SciPy on the CPU, even when the matrix is on the GPU.
Details
- Steady conduction in a slab: Larger systems already get CG with Jacobi without AmgX, so the setting keeps the smaller meshes on the device. There is no isolated measurement. Evidence: #187; resolve_method and resolve_library in torch-fem v0.12.1 src/torchfem/sparse.py.
- Record id:
gpu-cg-jacobi
Loosens the CG stopping tolerance to 1e-6 L0 choice tuned
The GPU CG stops at a relative residual of 1e-6 (stol) instead of torch-fem’s 1e-10.
- Steady conduction in a slab: No measurable gain here: cost moved by -16% to +11% against 1e-10, within the run-to-run spread on a shared GPU.
Details
- Steady conduction in a slab: Every status is unchanged from 1e-12 to 1e-6. Forward metrics are bit-identical. Evidence: #301, notes.
- Record id:
cg-tolerance
Makes the CUDA assembly deterministic L0 choice
Enables torch.use_deterministic_algorithms(True, warn_only=True), so torch-fem’s CUDA scatter-adds in assembly sum in a fixed order.
- Steady conduction in a slab (non reproducible): Without it, identical GPU runs differed: the conductivity_recovery_bfgs final loss varied between 241 and 5.7e3.
Details
- Steady conduction in a slab: The spread depended on whether the run left a plateau before the iteration cap. With the setting, reruns are identical (#200, and again at every CG tolerance in #301). It cost about 5 to 15% in #200, on an RTX 2000 Ada. Evidence: #200, notes; #301, notes.
- Record id:
deterministic-kernels
Use it in your code
tesseract build mosaic/tesseracts/thermal-mesh/torch-fem-thermal builds the image torch_fem_thermal_mesh. To load it and call it from JAX, follow Use a solver in your code.
Fields marked ∂ are differentiable.
Inputs
rho |
Array[(None,), Float32] |
✓ |
Per-cell density field, shape (n_max_cells,). Active slice = rho[:hex_mesh.n_faces]. ρ=1 → fully conducting, ρ=0 → insulating. |
source |
Array[(None,), Float32] |
✓ |
Per-element volumetric heat source (W/m³), shape (n_cells,). Active slice = source[:hex_mesh.n_faces]. Zero by default (topology-optimisation mode; no body heat source). |
target_temperature |
Array[(None,), Float32] |
|
Per-node target temperature for the inverse (source-identification) objective, shape (n_nodes,). identification_error = |
boundary_conditions |
MeshBC |
|
Mesh boundary conditions. dirichlet.mask/values prescribe temperature (values shape (n_groups, 1), None → zero); neumann.mask/values prescribe heat flux (values shape (n_groups, 1)). |
hex_mesh |
HexMesh |
|
Hexahedral mesh (8-node hex elements). n_points / n_faces give the active slice sizes. |
k_max |
float |
|
Maximum thermal conductivity (fully solid material). |
p_exp |
float |
|
SIMP penalisation exponent p (k(ρ) = k_min + (k_max−k_min)·ρ^p). |
Outputs
thermal_compliance |
Array[(), Float32] |
✓ |
Total thermal compliance C = ∮_Γ_N q_n · T dΓ (scalar). Equals the work done by the Neumann heat flux on the temperature field. |
identification_error |
Array[(), Float32] |
✓ |
Source-identification objective: |