Code architecture

The architecture encodes the problem formulation as a scene graph represented by an immutable Equinox PyTree. This structure separates the dynamic unknowns \((u, v, a, \lambda)\) from static operators \((M, B, f_{\mathrm{int}}, \Phi)\), enabling matrix-free residual evaluation and automatic differentiation.

Mathematical structure

The discrete problem from Problem Formulation is:

\[ \begin{cases} M(u)\,a + B(u)\,v + f_{\mathrm{int}}(u) + G(u)^\top \lambda = f_{\mathrm{ext}} \\ \Phi(u) = 0 \end{cases} \]

The architecture represents this as:

Dynamic state: \((u, v, a) \in \mathbb{R}^n\) and \(\lambda \in \mathbb{R}^{n_c}\) stored as PyTree leaves
Static operators: \(M(u)\), \(B(u)\), \(f_{\mathrm{int}}(u)\), \(\Phi(u)\) encoded as pure functions
Derived operators: \(G(u) = \partial \Phi / \partial u\) computed via automatic differentiation
Residual evaluation: \(F(u, v, a, \lambda) \mapsto (r_u, r_c)\) computed via tree traversal

This separation enables XLA to compile static operators once while tracing only dynamic state through time steps.

Core design principles

Immutability: Static operators are frozen in the PyTree; only dynamic leaves \((u, v, a, \lambda)\) change
Matrix-free: All operators act as functions \(x \mapsto M(u) \cdot x\), never assembled
Differentiable: Pure functions enable differentiation w.r.t. state \(\partial F / \partial (u, a, \lambda)\) and parameters \(\partial F / \partial \theta\) (material properties, learnable weights) via automatic differentiation
Hierarchical: Scene graph enables composition via linear mappings \(P\)

Class Diagram

classDiagram class Legend { + : dynamic pytree leaf #40;JAX-traced#41; - : static over steps #40;GPU constant#41; DOFs x = #123;u,v,a#125; } %% ======================= %% Scene Graph Hierarchy %% ======================= class SceneNode { -NodeDOFs dofs -MeshTopology mesh -Tuple~SceneNode~ children -Tuple~PhysicalField~ fields } class SolverNode { -SolverState solver -ProjectiveBC bc } class MappedNode { -Mapping mapping_to_parent } SceneNode <|-- SolverNode SceneNode <|-- MappedNode %% ======================= %% Dynamic State %% ======================= class NodeDOFs { +jnp.ndarray u +jnp.ndarray v +jnp.ndarray a } class SolverState { -float dt +float time +int step_id -int max_iters -float tol } %% ======================= %% Topology & Geometry %% ======================= class MeshTopology { -jnp.ndarray points -jnp.ndarray connectivity -Dict~str,jnp.ndarray~ point_data -Dict~str,jnp.ndarray~ cell_data -GraphData graph -FiniteElementData fe } class GraphData { -jnp.ndarray adjacency } class FiniteElementData { -jnp.ndarray jacobian_det -jnp.ndarray shape_fn_derivatives } %% ======================= %% Mappings & Constraints %% ======================= class Mapping { -jnp.ndarray parent_node_indices -jnp.ndarray weights apply_P(x) jnp.ndarray apply_PT(x) jnp.ndarray } class ProjectiveBC { -jnp.ndarray mask_free_dofs -jnp.ndarray fixed_values apply_Pc(x) jnp.ndarray } %% ======================= %% Physical Fields %% ======================= class PhysicalField { +jnp.ndarray parameters add_M(a) LinearOp add_B(v) LinearOp add_K(u) LinearOp add_f() jnp.ndarray evaluate(node, x) FieldContribution } class ConstraintField { +jnp.ndarray lambda add_c(u) LinearOp } class FieldContribution { +jnp.ndarray Ma +jnp.ndarray Bv +jnp.ndarray Ku +jnp.ndarray f } PhysicalField <|-- ConstraintField PhysicalField ..> FieldContribution : evaluates ConstraintField ..> FieldContribution : evaluates %% ======================= %% Relationships %% ======================= SceneNode *-- NodeDOFs SceneNode *-- MeshTopology SceneNode *-- PhysicalField SolverNode *-- SolverState SolverNode *-- ProjectiveBC MappedNode *-- Mapping MeshTopology *-- GraphData MeshTopology *-- FiniteElementData

Scene graph abstraction

The diagram above shows how the problem structure maps to code:

SceneNode: Base node that may own mesh topology, physical fields, and constraints
SolverNode: Owns the primary unknowns \((u, v, a) \in \mathbb{R}^n\) and \(\lambda \in \mathbb{R}^{n_c}\) that solve the coupled system
MappedNode: Expresses child DOFs as \(u_{\mathrm{child}} = P \cdot u_{\mathrm{parent}}\) where \(P\) is a static interpolation operator

Each node contributes to the residual:

DOFs: Dynamic unknowns \((u, v, a)\) or mapped via \(P\)
Fields: Encode operators \(M(u)\), \(B(u)\), \(f_{\mathrm{int}}(u)\) that contribute to \(r_u\)
Constraints: Encode \(\Phi(u)\) and \(G(u)\) that contribute to \(r_c\) and \(G(u)^\top \lambda\)

The graph is typically shallow: one solver node with one level of mapped children, reflecting the structure of the coupled primal–dual system.

Dynamic state

Degrees of freedom

Each node carrying mechanical state owns the dynamic unknowns from the problem formulation:

\(u \in \mathbb{R}^{n_u}\): displacement DOFs
\(v \in \mathbb{R}^{n_u}\): velocity DOFs (\(v = \dot{u}\))
\(a \in \mathbb{R}^{n_u}\): acceleration DOFs (\(a = \ddot{u}\))

These are mutable PyTree leaves updated by the time integrator. At each Newton step, we solve for \((a_{n+1}, \lambda_{n+1})\) given \((u_n, v_n)\).

Solver state

Solver nodes additionally carry:

\(t, \Delta t\): current time and step size (used in time integration)
\(\lambda \in \mathbb{R}^{n_c}\): Lagrange multipliers (constraint force intensities)

The multipliers \(\lambda\) appear in the dynamics residual as \(G(u)^\top \lambda\), representing constraint reaction forces. They are solved simultaneously with accelerations in the coupled system \(\begin{bmatrix} A & G^\top \\ G & 0 \end{bmatrix}\).

Static data

Mesh and topology

Each node may own mesh data that defines the discretization:

Point coordinates \(X \in \mathbb{R}^{n_{\mathrm{points}} \times 3}\): reference configuration
Connectivity: element topology for finite element evaluation
Graph adjacency: mesh structure for GNN-based learned components

This data is frozen in the PyTree and transferred once to GPU. It defines the domain \(\Omega\) over which the continuous problem is discretized.

Finite element precomputations

Precomputed quantities used during residual evaluation:

Jacobian determinants \(J(\xi)\) at quadrature points \(\xi\): for volume integration
Shape function derivatives \(\nabla N(\xi)\): for strain computation \(\varepsilon = \nabla N \cdot u\)
Element geometric factors: for efficient \(f_{\mathrm{int}}(u)\) evaluation

These enable efficient evaluation of \(f_{\mathrm{int}}(u) = -\partial W(u)/\partial u\) where \(W(u) = \int_\Omega \psi(\mathbf{F}(u)) \, dV\) and \(\mathbf{F} = \mathbf{I} + \nabla u\).

Mappings

Mapped nodes express child DOFs as linear interpolations of parent DOFs:

\[ u_{\mathrm{child}} = P \cdot u_{\mathrm{parent}} \]

where \(P \in \mathbb{R}^{n_{\mathrm{child}} \times n_{\mathrm{parent}}}\) is a sparse interpolation matrix (static, frozen in PyTree).

Physical interpretation: The mapping \(P\) represents kinematic constraints or reduced-order approximations. The transpose \(P^\top\) projects child forces onto the parent, preserving virtual work:

\[ \delta W = \delta u_{\mathrm{child}}^\top \cdot f_{\mathrm{child}} = \delta u_{\mathrm{parent}}^\top \cdot P^\top f_{\mathrm{child}} \]

This ensures energy and power flow are correctly transferred through the hierarchy. Mappings are applied during the residual evaluation traversal.

Boundary conditions

Dirichlet conditions \(u|_{\partial \Omega} = u_0\) are enforced via a projective operator \(P_c\):

\[ P_c \cdot u = \begin{cases} u_i & \text{if DOF } i \text{ is free} \\ u_{0,i} & \text{if DOF } i \text{ is constrained} \end{cases} \]

Mathematical role: The projector \(P_c\) restricts the trial space to the admissible manifold \(\{u : u|_{\partial \Omega} = u_0\}\). It is applied to:

Trial states: \(P_c \cdot u\) ensures \(u\) satisfies boundary conditions
Residuals: \(P_c \cdot r\) projects \(r\) onto free DOFs only
Linear solver directions: \(P_c \cdot \delta x\) ensures search directions respect boundaries

This is matrix-free: no modification of \(A\) or \(G\), only operator composition. The multipliers \(\lambda\) remain unconstrained, consistent with the Lagrange multiplier formulation where boundary reactions are not explicitly tracked.

Physical fields

Physical behavior is encoded through fields that contribute to the residual \(r_u\):

Inertia field: Evaluates \(M(u) \cdot a\) where \(M(u)\) is the mass operator
Damping field: Evaluates \(B(u) \cdot v\) where \(B(u)\) is the damping operator
Stiffness field: Evaluates \(f_{\mathrm{int}}(u) = -\partial W(u)/\partial u\) where \(W(u)\) is strain energy
External forces: Evaluates \(f_{\mathrm{ext}}(t)\) (time-dependent loading)
Constraint field: Evaluates \(\Phi(u)\) and \(G(u) = \partial \Phi / \partial u\) (computed via automatic differentiation)

Operator structure: Each field exposes:

Residual contribution: \(r_{\mathrm{field}}(u, v, a) \in \mathbb{R}^n\)
Linearized action: \(\partial r_{\mathrm{field}} / \partial (u, v, a)\) applied via JVP

Constraint Jacobian via AD: The constraint Jacobian \(G(u) = \partial \Phi / \partial u\) is computed via automatic differentiation of \(\Phi(u)\), applied as JVP operations \(G(u) \cdot v\) in the Krylov solver. This is consistent with the matrix-free approach: \(G\) is never assembled, only its action on vectors is computed.

Field parameters: Physical fields depend on simulation parameters \(\theta\) (e.g., material properties, learnable weights). The architecture supports differentiation w.r.t. these parameters for inverse problems:

Parameter derivatives: \(\partial F / \partial \theta\) enables gradient-based parameter estimation
State derivatives: \(\partial F / \partial (u, a, \lambda)\) enables sensitivity analysis and adjoint methods

This differentiability enables parameter estimation, material identification, and optimization-based inverse problems.

The tangent operator \(A = \partial r_u / \partial (u, a)\) required by Newton is assembled via automatic differentiation of field contributions, never explicitly stored.

Connection to problem formulation

The architecture directly encodes the mathematical structure from Problem Formulation:

Dynamic unknowns: \((u, v, a, \lambda)\) stored as PyTree leaves, updated by solver
Static operators: \(M(u)\), \(B(u)\), \(f_{\mathrm{int}}(u)\), \(G(u)\), \(\Phi(u)\) encoded as pure functions
Residual evaluation: \(F(u, v, a, \lambda) = (r_u, r_c)\) computed via tree traversal
Coupled structure: The block system \(\begin{bmatrix} A & G^\top \\ G & 0 \end{bmatrix}\) emerges from residual linearization

This separation enables:

Efficient XLA compilation: Static operators compiled once, dynamic state traced
Matrix-free operations: All operators act as functions, never assembled
Automatic differentiation: \(\partial F / \partial (u, a, \lambda)\) computed via JAX
Scalable execution: GPU-friendly structure with minimal memory allocation

The scene graph structure ensures that residual evaluation respects the hierarchical composition while preserving the mathematical structure of constrained mechanics.