Quazim0t0/neural-physics-engine

🌊 Neural Physics Engine — learned constraint projectors

A projective-dynamics (PD) engine in which the per-element local constraint projections are learned neural networks, while rotations and the global solve stay exactly analytic. One tiny network, shared across every element and across constraint types via material tokens, sits inside an exact reduced global solve.

Status: work in progress. This is a research artifact from an 8-week build; the local-projector results are strong, the global-reduction results are proven in-distribution, and the fully-learned-solver path is deliberately left for later. It was not finished — see Roadmap & what's next below.

Try it here

What's in here

File	What it is	Params
unified_projector.pt	A tied constraint projector serving 5 solid materials and the fluid via material tokens. Encoder→latent→decoder; rotation stays analytic. Ships state_dict, materials, k, and fluid_scales (the water-token calibration).	~10k
warm_start_net.pt	A rotation-equivariant net that predicts the PD solver's converged correction as a residual on classical extrapolation — cuts iterations at fixed tolerance.	~1.5k
engine3d/	The engine: solver.py (PD solid solver, exact reduced global step), fluid.py (PBF-3D with a lambda_fn hook), neural.py (NeuralTiedProjector, WarmStartNet), rotation.py, strain.py, materials.py, mesh.py.
images/	Validation figures (below).
*.html	Standalone browser demos (WebGPU).
neural_physics_engine_roadmap.md	The full design/scaling roadmap this was built against.

---

The core idea

Every system here follows one skeleton:

per-element LOCAL PROJECTION   →   GLOBAL RECONCILIATION
(tied / shared across elements)     (exact reduced solve)

The bet: replace each hand-derived local projection with a learned one, keep the skeleton, and keep the analytic symmetry handling. Rotations are not learned — polar decomposition (closed-form 2D, Müller-iterative 3D) is exact and cheap, and removing it from the learning problem is what lets a tiny latent suffice. Two guards are baked into the architecture (not patched on):

• Rejection form f(e) = e − r(e): the network learns what to remove, so admissible strains pass at gain ≈ 1 (mirrors PCA's within-subspace gain of exactly 1; avoids artificial damping inside the solver loop).
• Zero-anchoring r(0) = 0: rest strain maps exactly to rest.

A new material is a new token row, not a new network — the tied-embedding thesis. The fluid is just a 6th token: PBF's density constraint is another tied local projector, so the same weights that serve the solids also compute the fluid's λ multiplier.

---

Measured results

• One tied projector matched five per-material PCAs. Local-step accuracy, train and OOD: steel 0.99901 · rubber 0.99693 · foam 0.99901 · composite 0.99913 · gel 0.99244 — the water token added later caused no interference with the solids.
• Fluid unification held in-loop. Running the full dam break on the neural λ held the same density as the exact analytic solver: 0.0121 (neural) vs 0.0124 (analytic).
• Learned warm start cut solver work by ~32% fewer PD iterations at fixed tolerance, with zero correctness risk (it only moves the loop's starting iterate; the fixed point is unchanged — classical warm starts are exact special cases of its form).
• Earlier 2D findings that motivated the scale-up: one 8-float strain matrix transferred to unseen load cases at 99%+, co-rotation reduced strain to an exactly 3-D space, and constraint folding was algebraically exact (free 3.25×).

Figures

Neural tied projector vs per-material PCA — one token-conditioned network matches five separate bases (train + OOD).

Fluid token, in the loop — dam break driven by the neural λ vs the analytic rule; density maintenance tracks.

PBF-3D validation — density error over a dam break.

Learned warm start — PD iterations saved at fixed tolerance.

3D tied-basis solid — co-rotated strain in a reduced basis.

Cantilever beam — tip settle + energy signature (correctness/invariant check).

---

How it was trained

Small models, cheap runs. The 2D corpus and all solid experiments ran on CPU in seconds; the 3D neural projector trained on a single GPU.

1. Tied strain bases (3D) — replicate the 2D PCA results in 3D: co-rotated 6-D symmetric strain compressed to k≈3, shared across all tets (experiments/w3_tied_basis.py).
2. Neural tied projector with material tokens — one network over 5 materials, per-material input scaling, rejection + zero-anchor guards; trained to beat the per-material PCAs (w45_neural_projector.py, corpus w45_corpus.npz).
3. Learned warm start — trained to predict the converged correction as a residual on linear extrapolation, from per-vertex history invariants (w6_warm_start.py).
4. Fluid token — a 6th "water" token learns the PBF density-constraint rule λ = −C/(Σ‖∇C‖² + ε); blanketed with synthetic samples across the plausible (C, g²) domain to survive in-loop distribution shift, then validated in-loop against the analytic solver (w7_fluid_token.py).

Losses prioritized trajectory matching (positions and velocities — velocity error caught an over-damping bug), constraint residuals, long-horizon stability, and an energy-gain penalty.

---

Demos

• Live (server-side, this repo's Space):
• Standalone (browser, WebGPU): dam_break_gpu.html, tied_subspace_water.html — open in Chrome/Edge 113+. These parse the .pt files in-browser and run the learned λ in a WebGPU compute shader (no server, no ONNX).

---

Roadmap & what's next (unfinished)

The local-projector thesis has strong evidence; the global reduction is proven in-distribution with a known literature fix (CROM-style continuous fields) for its OOD weakness; the fully-learned-solver path is the speculative tail, sequenced last. Remaining steps from the roadmap: CROM-style global decoder, a unified constraint zoo (strain + density + volume + contact as tokens through one network), a full GPU/Warp training port, and nested-latent LOD. See neural_physics_engine_roadmap.md.

Usage

import torch
from engine3d.neural import NeuralTiedProjector
from engine3d.fluid import PBF3D, dam_break_block
import numpy as np

ck = torch.load("unified_projector.pt", map_location="cpu", weights_only=False)
net = NeuralTiedProjector(n_materials=6, k=ck["k"], hidden=64)
net.load_state_dict(ck["state_dict"]); net.eval()
# see the Space's app.py for the water-token λ rule and the dam-break rollout.

Citation

@misc{byrne2026neuralphysics,
  title  = {Neural Physics Engine: learned constraint projectors in an exact
            projective-dynamics solve},
  author = {Byrne, Dean},
  year   = {2026},
  note   = {Work in progress. https://huggingface.co/Quazim0t0}
}

Adjacent work: CROM, differentiable/neural Projective Dynamics, subspace neural physics (Holden et al.), GNS/MeshGraphNets, NCLaw. The under-explored middle ground here — weight-tied per-element projectors in co-rotated frames, shared across constraint types via tokens, inside an exact reduced global solve — is the lane.