feat(solver): update torch-fem to 0.12.1, solve with CG+Jacobi on GPU

[meyer-nils]

Summary

This updates the torch-fem thermal-mesh backend to torch-fem 0.12.1. On GPU, the linear solve now runs as Jacobi-preconditioned CG in torch, so the forward and adjoint solves stay on the device.

Type of change

• New solver backend
• Solver tuning / improvement
• New benchmark domain
• Harness / infrastructure change
• Documentation
• Bug fix

For new solvers

• tesseract_config.yaml has a metadata.mosaic: block with at least name and backend
• tesseract build mosaic/tesseracts/<domain>/<solver> succeeds
• mosaic run -p <domain> --suites forward -s <solver> --debug completes
• mosaic status -p <domain> -f output pasted below
• Exclusions and explained anomalies documented (if any)

For solver tuning

• Before/after mosaic status --format json snapshots compared
• No regressions to other solvers

Benchmark label

benchmark:solver

Notes

• torch-fem 0.12.1: boundary conditions are now set through SolidHeat's temperatures / heat_flux. SolidHeat no longer inherits the mechanics setters displacements / forces, so assigning those left the applied load at zero. The local tensors are renamed to match.
• sparse_solve monkeypatch removed: torch-fem now moves all solver inputs to the solve device itself.
• Solver choice: GPU uses CG+Jacobi in torch. CPU-only hosts keep torch-fem's automatic solver choice. The numerics metadata is updated to CG+Jacobi.
• LD_LIBRARY_PATH: now points at the Python version the image actually uses.
• Testing: pytest passes. mosaic run -p thermal-mesh --suites forward,gradient --debug -s torch-fem passes on GPU, and its results match the previous torch-fem version. The VJP agrees with a finite-difference check on CPU.

[PasteurBot]

📊 View the full benchmark results

The rendered docs preview has every plot for this run (forward accuracy, gradients, cost, optimization) merged with existing baseline results on main. The summary below reports pass/fail status.

---

Coverage: measured torch-fem on thermal-mesh this run. Other cells are shown from the baseline.

Status diff vs base

Legend · ✅ ok · 🟠 anom · ❌ fail · · missing · 🚫 excluded (permanent — out of score denominator) · ⚪ excluded (work-to-do) · * stale — result predates current benchmark run

0 regression(s) · 0 improvement(s) · 3 metric change(s) · 0 other transition(s) · 0 resource-frontier shift(s) · 0 new row(s) · 0 removed row(s) · score 0.93 → 0.93

📊 Metric changes

• 🟢 thermal-mesh · optimization/conductivity_recovery_bfgs · torch-fem · final/initial loss ratio 5.4e-05 → 3.1e-05 (-43%)
• 🟢 thermal-mesh · optimization/conductivity_recovery_bfgs · torch-fem · final loss 1.02e+05 → 5.83e+04 (-43%)

⏱ Timing (indicative — shared-runner contention)

Wall-clock varies with CI-runner load, so treat these as indicative rather than controlled measurements.

• 🟢 thermal-mesh · cost/spatial_cost · torch-fem · median time 0.297s → 0.187s (-37%)

Full Mosaic status

Mosaic status

Legend · ✅ ok · 🟠 anom · ❌ fail · · missing · 🚫 excluded (permanent — out of score denominator) · ⚪ excluded (work-to-do) · * stale — result predates current benchmark run

Each solver is run against every experiment in the suite. ok = produced valid results; fail = crashed or returned invalid data; anom = ran successfully but tripped an automated quality check (e.g. poor gradient accuracy, outlier wall-clock time, or diverged optimisation). Thresholds are defined per-problem in the problem config.

problem	ok	anom	fail	missing	excl (work)	excl (perm)	stale	score
ns-3d-grid	90	6	7	0	0	9	0	🟢 0.90
ns-grid	106	16	2	0	0	16	0	🟢 0.92
structural-mesh	44	1	0	0	0	5	0	🟢 0.99
thermal-mesh	58	4	1	0	0	7	0	🟢 0.95
overall	298	27	10	0	0	37	0	🟢 0.93

Failures & anomalies

• 🟠 ns-3d-grid · forward/baseline · INS.jl — error 0.147 at sweep=8 is 43.1× peer median (0.00341); threshold k=3.0
• 🟠 ns-3d-grid · forward/baseline · PhiFlow — error 0.147 at sweep=8 is 43.1× peer median (0.00341); threshold k=3.0
• 🟠 ns-3d-grid · forward/baseline · XLB — error 0.273 at sweep=8 is 80.2× peer median (0.00341); threshold k=3.0
• ❌ ns-3d-grid · cost/spatial_cost · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• ❌ ns-3d-grid · cost/temporal_cost · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• ❌ ns-3d-grid · cost/vjp_cost/by_N · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• 🟠 ns-3d-grid · cost/vjp_cost/by_N · XLB — median time 173.9s is 65× peer median (2.67s); threshold k=20.0
• ❌ ns-3d-grid · cost/vjp_cost/by_steps · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• 🟠 ns-3d-grid · cost/vjp_cost/by_steps · XLB — median time 569.8s is 210× peer median (2.72s); threshold k=20.0
• 🟠 ns-3d-grid · gradient/fd_check · INS.jl — best-ε median FD rel_err 2.36e-03 > max_rel_err=1e-03
• ❌ ns-3d-grid · gradient/horizon_sweep_limits · PhiFlow — nan
• ❌ ns-3d-grid · gradient/horizon_sweep_limits · PICT — error
• ❌ ns-3d-grid · gradient/horizon_sweep_limits · Warp-NS — nan
• 🟠 ns-grid · forward/agreement/multimode · XLB — error 0.0725 at sweep=0.005 is 16.9× peer median (0.00429); threshold k=3.0
• 🟠 ns-grid · forward/agreement/tgv · jax-cfd — error 0.0145 at sweep=0.001 is 6.0× peer median (0.00241); threshold k=3.0
• 🟠 ns-grid · forward/agreement/tgv · PhiFlow — phiflow's double CenteredGrid↔StaggeredGrid resampling gives 4.18% amplitude damping (ratio=0.9582); cosine=0.9999924 (pattern correct); arithmetic-average output conversion fix worsened error 9×; upstream library change required
• 🟠 ns-grid · forward/agreement/tgv · XLB — error 0.274 at sweep=0.001 is 113.7× peer median (0.00241); threshold k=3.0
• 🟠 ns-grid · forward/baseline · INS.jl — staggered MAC grid double-interpolation: collocated TGV IC -> staggered faces -> collocated output gives sin^2(pi/N) round-trip error at all N; 35-40x above collocated peers
• 🟠 ns-grid · forward/baseline · jax-cfd — staggered MAC grid double-interpolation: collocated TGV IC -> staggered faces -> collocated output gives sin^2(pi/N) round-trip error at all N; 35-40x above collocated peers
• 🟠 ns-grid · forward/baseline · XLB — error 0.00704 at sweep=128 is 11.7× peer median (0.000602); threshold k=3.0
• 🟠 ns-grid · forward/cylinder · OpenFOAM — error 0.694 at sweep=0.05 > max_error=0.5
• 🟠 ns-grid · forward/cylinder · PICT — error 0.672 at sweep=0.05 > max_error=0.5
• 🟠 ns-grid · forward/tgv_nu_sweep · jax-cfd — error 0.0145 at sweep=0.0001 is 6.0× peer median (0.00241); threshold k=3.0
• 🟠 ns-grid · forward/tgv_nu_sweep · XLB — error 0.275 at sweep=0.0001 is 114.2× peer median (0.00241); threshold k=3.0
• 🟠 ns-grid · cost/spatial_cost · OpenFOAM — median time 10.7s is 36× peer median (0.30s); threshold k=20.0
• ❌ ns-grid · cost/spatial_cost · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• 🟠 ns-grid · cost/temporal_cost · PhiFlow — median time 11.7s is 39× peer median (0.30s); threshold k=20.0
• ❌ ns-grid · cost/vjp_cost/by_N · PhiFlow — error: RuntimeError: Error 500 from Tesseract: Traceback (most recent call last):
• 🟠 ns-grid · optimization/drag_opt · PhiFlow — final/initial = 0.85 (> 0.5)
• 🟠 ns-grid · optimization/drag_opt · PICT — final/initial = 0.74 (> 0.5)
• 🟠 ns-grid · optimization/drag_opt · XLB — final/initial = 0.69 (> 0.5)
• 🟠 structural-mesh · cost/temporal_cost · deal.II — median time 1.5s is 30× peer median (0.05s); threshold k=20.0
• 🟠 thermal-mesh · forward/baseline · deal.II — error 1.87e-06 at sweep=12 is 10.5× peer median (1.78e-07); threshold k=3.0
• 🟠 thermal-mesh · forward/baseline · torch-fem — error 3.56e-07 at sweep=6 is 4.0× peer median (8.91e-08); threshold k=3.0
• 🟠 thermal-mesh · forward/source_baseline · deal.II — error 9.99e-07 at sweep=6 is 13.0× peer median (7.68e-08); threshold k=3.0
• 🟠 thermal-mesh · forward/source_baseline · torch-fem — error 3.84e-07 at sweep=6 is 5.0× peer median (7.68e-08); threshold k=3.0
• ❌ thermal-mesh · cost/vjp_cost/by_N · JAX-FEM — error: ConnectionError: ('Connection aborted.', RemoteDisconnected('Remote end closed connection without response'))

ns-3d-grid — 16 experiment(s)

experiment	Exponax	INS.jl	OpenFOAM	PhiFlow	PICT	Warp-NS	XLB
forward/agreement	✅	✅	✅	✅	✅	✅	✅
forward/baseline	✅	🟠	✅	🟠	✅	✅	🟠
forward/physical_laws/vs_N	✅	✅	✅	✅	✅	✅	✅
forward/physical_laws/vs_nu	✅	✅	✅	✅	✅	✅	✅
forward/physical_laws/vs_steps	✅	✅	✅	✅	✅	✅	✅
cost/spatial_cost	✅	✅	✅	❌	✅	✅	✅
cost/temporal_cost	✅	✅	✅	❌	✅	✅	✅
cost/vjp_cost/by_N	✅	✅	🚫	❌	✅	✅	🟠
cost/vjp_cost/by_steps	✅	✅	🚫	❌	✅	✅	🟠
gradient/fd_check	✅	🟠	🚫	✅	✅	✅	✅
gradient/horizon_sweep_limits	✅	✅	🚫	❌	❌	❌	✅
gradient/jacobian_svd	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_nu01	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_steps20	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_steps40	✅	✅	🚫	✅	✅	✅	✅
optimization/recovery_constant_ic_bfgs_proj	✅	✅	🚫	✅	✅	✅	✅

ns-grid — 20 experiment(s)

experiment	INS.jl	jax-cfd	OpenFOAM	PhiFlow	PICT	Warp-NS	XLB
forward/agreement/multimode	✅	✅	✅	✅	✅	✅	🟠
forward/agreement/tgv	✅	🟠	✅	🟠	✅	✅	🟠
forward/baseline	🟠	🟠	✅	✅	✅	✅	🟠
forward/cylinder	🚫	🚫	🟠	✅	🟠	🚫	✅
forward/physical_laws/vs_N	✅	✅	✅	✅	✅	✅	✅
forward/physical_laws/vs_nu	✅	✅	✅	✅	✅	✅	✅
forward/physical_laws/vs_steps	✅	✅	✅	✅	✅	✅	✅
forward/tgv_nu_sweep	✅	🟠	✅	✅	✅	✅	🟠
cost/spatial_cost	✅	✅	🟠	❌	✅	✅	✅
cost/temporal_cost	✅	✅	✅	🟠	✅	✅	✅
cost/vjp_cost/by_N	✅	✅	🚫	❌	✅	✅	✅
cost/vjp_cost/by_steps	✅	✅	🚫	✅	✅	✅	✅
gradient/fd_check	✅	✅	🚫	✅	✅	✅	✅
gradient/horizon_sweep	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_nu01	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_steps20	✅	✅	🚫	✅	✅	✅	✅
gradient/jacobian_svd_steps40	✅	✅	🚫	✅	✅	✅	✅
gradient/param_sweep	✅	✅	🚫	✅	✅	✅	✅
optimization/drag_opt	🚫	🚫	🚫	🟠	🟠	🚫	🟠

structural-mesh — 10 experiment(s)

experiment	deal.II	FEniCS	Firedrake	JAX-FEM	TopOpt.jl
forward/agreement	✅	✅	✅	✅	✅
forward/baseline	✅	✅	✅	✅	✅
forward/physical_laws	✅	✅	✅	✅	✅
cost/spatial_cost	✅	✅	✅	✅	✅
cost/temporal_cost	🟠	✅	✅	✅	✅
cost/vjp_cost/by_N	🚫	✅	✅	✅	✅
cost/vjp_cost/by_steps	🚫	✅	✅	✅	✅
gradient/fd_check	🚫	✅	✅	✅	✅
gradient/param_sweep	🚫	✅	✅	✅	✅
optimization/topopt	🚫	✅	✅	✅	✅

thermal-mesh — 14 experiment(s)

experiment	deal.II	FEniCS	Firedrake	JAX-FEM	torch-fem
forward/agreement	✅	✅	✅	✅	✅
forward/baseline	🟠	✅	✅	✅	🟠
forward/physical_laws	✅	✅	✅	✅	✅
forward/source_baseline	🟠	✅	✅	✅	🟠
forward/source_linearity	✅	✅	✅	✅	✅
cost/spatial_cost	✅	✅	✅	✅	✅
cost/temporal_cost	✅	✅	✅	✅	✅
cost/vjp_cost/by_N	🚫	✅	✅	❌	✅
cost/vjp_cost/by_steps	🚫	✅	✅	✅	✅
gradient/fd_check	🚫	✅	✅	✅	✅
gradient/param_sweep	🚫	✅	✅	✅	✅
gradient/source_fd_check	🚫	✅	✅	✅	✅
gradient/source_width_sweep	🚫	✅	✅	✅	✅
optimization/conductivity_recovery_bfgs	🚫	✅	✅	✅	✅

[meyer-nils]

The two 🟠 on torch-fem (forward/baseline, forward/source_baseline) are likely float32 rounding issues: every thermal solver casts thermal_compliance to float32 before returning it. An absolute floor in the outlier check would probably clear these (deal.II's included).

[dionhaefner]

Will do (in a separate PR). Thanks @meyer-nils !