SOFAx

SOFAx v2 is a differentiable 3D FEM dynamics framework designed for matrix-free implicit time integration, iterative linear solvers, and end-to-end automatic differentiation.

The design goal is to preserve the numerical structure of constrained mechanics—Newton–Krylov iterations, coupled primal–dual systems, and projection operators—while remaining fully compatible with JAX/XLA compilation and multi-GPU execution.

Key ideas

Matrix-free formulation
No global sparse matrix assembly. All operators are applied as linear maps (JVP/VJP or custom LinearOp abstractions). This enables scalability to large systems and seamless integration with automatic differentiation.
Implicit time integration
Each time step is formulated as a nonlinear root-finding problem, providing unconditional stability for stiff systems and exact constraint enforcement.
Iterative solvers
Krylov methods (CG/GMRES) are used throughout, with planned support for block preconditioning and Schur complement strategies. The matrix-free structure naturally accommodates learned preconditioners and solver assistants.
Differentiable by construction
Residuals and linear actions are defined so that JAX can differentiate through the solver, either by unrolling iterations or via implicit differentiation. This enables inverse problems, parameter estimation, and end-to-end learning.

Why this structure enables AI integration

The operator-based, matrix-free design creates natural integration points for learning-based components:

Residual evaluation can incorporate learned constitutive laws without changing the solver structure
Linear operator actions (JVP/VJP) enable learned preconditioners to be inserted as pure functions
Iterative solver structure allows AI models to predict warm starts, correction directions, or approximate inverse actions
End-to-end differentiability supports training AI components through the entire simulation pipeline

The key insight is that AI components act as operators within the solver loop, preserving the physical structure while accelerating or enriching local operations. For detailed discussion of integration opportunities, see AI for Simulation.

Stack

The implementation is fully JAX-based. Typical building blocks include:

equinox: PyTree modules and parameter handling
optimistix: Newton-style nonlinear solvers
lineax: Krylov linear solvers (CG, GMRES)
optax: optimizers for inverse problems and learning
diffrax: ODE tooling when continuous-time formulations are required

This stack provides the foundation for both physics simulation and learning-based extensions, with shared abstractions (PyTrees, operators, transformations) enabling seamless integration.

Modeling assumptions

Current scope and assumptions:

No contacts — Contact mechanics is not yet implemented
Fixed topology — No remeshing or dynamic mesh adaptation
Static boundary conditions — Boundary conditions do not change during simulation
Fixed holonomic constraints — Constraints are implemented via static indices/slices
Shallow scene graph — A solver node typically has only one generation of children

These assumptions simplify the implementation while maintaining the core structure needed for differentiable constrained mechanics. Future extensions may relax some of these constraints.

Target application domains:

large-deformation hyperelasticity
electro-mechanical coupling
cardiac muscle modeling
volume constraints

Documentation structure

This section covers:

Mathematical Foundations:
- Problem Formulation — Dynamic and discrete formulations from Newton's second law
- Constraints & Physical Interactions — Geometric decomposition and exact constraint enforcement
- Preconditioning — Block preconditioners and Schur complement methods
Numerical Methods:
- Time Integration — Implicit schemes and residual formulation
- Residual Operator — Tree-based evaluation and matrix-free operations
- Newton–Krylov Solver — Matrix-free nonlinear solver structure
Implementation:
- Code Architecture — Scene graph and PyTree structure
- CPU Setup — Implementation details

For AI integration strategies and learning-based extensions, see the AI for Simulation section.