Solver¶
In this framework's pipeline (Problem → Domain → Model → Physics → Solver
→ Backend → Researcher), "Solver" means the training policy: which
optimizer, how many stages, in what order, and with what scheduling —
distinct from pinneapple_simulation.numerical_solvers, which are classical
FEM/FDM/CFD solvers used to generate reference data (see
Package Layers for that distinction).
Where it lives¶
The training-policy layer is pinneapple_neural.trainer (re-exported at
pinneapple_neural, and via the legacy alias pinneapple_train):
Trainer+TrainConfig— the baseline policy: an Adam optimizer (torch.optim.Adam, configured byTrainConfig.lr/weight_decay), with optional gradient clipping, AMP, early stopping, checkpointing, and aphysics_aware_validationflag that runs validation undertorch.enable_grad()so PDE-residual validation losses (which need second derivatives) don't break inside ano_gradblock.TwoPhaseTrainer+TwoPhaseConfig— a two-stage policy: Phase 1 fits the model to reference/measured data (supervised MSE), Phase 2 switches to the physics (PDE residual + BC) loss;combined=Trueinstead ramps a physics weight in during Phase 1 rather than switching abruptly.TimeMarchingTrainer— splits[0, T]into sequential windows and trains a fresh copy of the network per window, using the previous window's prediction at its end time as the next window's initial condition (Wight & Zhao 2020) — for stiff or long-horizon temporal PDEs.CausalPINNTrainer+CausalWeightScheduler— weights collocation points by a decaying function of the cumulative residual at earlier times, so the network is pushed to get early times right before late ones (Wang et al. 2022, "Respecting causality...").DDPPINNTrainer— the same training loop wrapped fortorch.nn.parallel.DistributedDataParallel.MultiRestartTrainer— runs multiple random restarts and keeps the best.- Loss-weight balancers, usable independently of which
Traineryou pick:WeightScheduler/SelfAdaptiveWeights,GradNormBalancer,LossRatioBalancer,NTKWeightBalancer,ReLoBRaLo,SoftAdapt,AugmentedLagrangian,InverseDirichlet,PCGrad,JointAdaptiveWeights,AutoBalancer.
Using it¶
from pinneapple_neural import train_model # Trainer shortcut
result = train_model(model, losses, epochs=5000)
train_model is a thin wrapper over Trainer(model, loss_fn).fit(...); for
anything beyond the default single-stage Adam policy (two-phase, causal,
time-marching, distributed), construct the specific trainer class directly.
What Solver does not do¶
It has no opinion on which physics or architecture it's optimizing — it receives a model and a loss function and drives weight updates. It also doesn't decide where those updates run (CPU/GPU, PyTorch/JAX); that's Backend.
One solve(problem) for every kind of solver¶
pp.solve(problem, method) takes a PhysicalProblem, a ProblemSpec or a preset name and returns a Solution with
predict(X), wall_time_s and metadata() (the same keys for every backend: problem fingerprint, method, kind,
fields, coordinates, and the backend's own details). pp.compare scores several methods against a reference at the
same points and reports each method's kind.
| Method | Kind | What it is |
|---|---|---|
pinn |
neural | physics-informed network trained with solve_pde |
fem |
classical | P1 finite elements for steady poisson / laplace on a box with Dirichlet conditions (1D-3D) |
reference |
classical | the solver named in the problem's reference_solver |
external |
external | any code returning grid fields, runner=callable(problem) |
analytic, exact |
analytic | closed form known to the library, or a function you pass |
Register more with @register_method("name", kind="classical"); list_methods(kind=...) and list_kinds() show what is
available. A method that cannot handle a problem raises MethodNotAvailable with the reason, never a partial answer.
The source term of poisson is passed as ctx={"source_fn": fn} and read by both pinn and fem, so the two solve the
same problem. Example: examples/solver_api/01_classical_vs_pinn.py (FEM relative L2 4.2e-3 in 0.02 s, PINN 1.3e-3 in
30 s on the manufactured Poisson problem).
Not yet behind this contract: FNO, SPH, LBM and FVM backends, and time-dependent classical solves of arbitrary presets.
pp.compile: solve the same problem with less work¶
compiled = pp.compile("burgers_1d", optimize="physics")
sol = compiled.solve("pinn", epochs=2000)
print(compiled.benchmark(epochs=100)) # baseline vs compiled CPU time and how far the weights differ
optimize="physics" memoizes first derivatives within one loss evaluation. The PDE residual is built from several
derivative calls on the same field (the time derivative and the convective term of Burgers both differentiate u), and
each call builds an autograd graph; the compiled problem builds it once. The losses are the same, and the trained
weights agree to floating-point summation order. Measured with benchmark (80 epochs, 2048 collocation points,
modified_mlp 64x4, one CPU thread, best of 3 alternating runs):
| Problem | First-derivative graphs built | Baseline | Compiled | Speedup | Max weight difference |
|---|---|---|---|---|---|
| burgers_1d | 1 of 2 | 3.33 s | 2.84 s | 1.18x | 3e-8 |
| crystal_phonon | 1 of 3 | 3.26 s | 2.52 s | 1.29x | 6e-8 |
| heston_pde_2d | 4 of 6 | 7.49 s | 6.31 s | 1.19x | 5e-8 |
| sod_shock_tube_astro | 5 of 6 | 4.21 s | 3.57 s | 1.18x | 0 |
| threaded_coupling_tc50_rotating | 10 of 11 | 9.58 s | 9.11 s | 1.05x | 2e-8 |
The gain is modest and depends on how many derivatives the PDE repeats. Where the time goes: about half of a PINN
step is the backward pass through the second-order graph, which caching does not touch. Two other ideas were measured
on CPU and rejected: forward-mode (nested jvp) derivatives for the Burgers residual were slower (59 ms against 35 ms
per step), and a batched evaluation of the diagonal second derivatives changed the time by less than 10% (41.4 to
41.8 ms in 2D, 57.3 to 51.6 ms in 3D). Not done: kernel-level operator fusion, batching of the boundary and initial
forward passes, memory planning, and compiling a bare model (pp.compile(model) raises TypeError, since the
optimisations act on the PDE residual a model alone does not define). solve_pde(..., cache_derivatives=True) and
compile_problem(spec, cache_derivatives=True) expose the same switch.