Multi-physics Extensions
The SOFAx v2 architecture, built on a matrix-free Newton–Krylov solver and a primal–dual formulation, is naturally extensible to multi-physics problems. This document details how to integrate domains like electrophysiology and fluid dynamics into the existing solver structure without architectural changes.
1. Unified State Vector (PyTree)
The core principle is to extend the state vector \(y\) to include the physical fields of all coupled domains. Since SOFAx treats states as PyTrees, this extension is seamless and does not require a monolithic array structure.
The unified state \(y\) becomes:
Breaking this down into primal and dual components:
- Mechanics:
- Primal: \(a\) (acceleration)
- Dual: \(\lambda_{\text{mech}}\) (constraint multipliers)
- Electrophysiology:
- Primal: \(V\) (transmembrane potential)
- Dual: \(\emptyset\) (typically purely primal, unless flux constraints are enforced)
- Fluids (Incompressible):
- Primal: \(u_f\) (velocity)
- Dual: \(p\) (pressure, acting as a Lagrange multiplier for incompressibility)
The total system state for the solver is:
2. Residual Formulation
The Newton solver seeks the root of a global residual \(R(y) = 0\). This global residual is the concatenation of residuals from each physical domain, including coupling terms.
Electrophysiology (Reaction-Diffusion)
The monodomain equation typically reads: \(\beta C_m \dot{V} + \beta I_{\text{ion}}(V) - \nabla \cdot (\sigma \nabla V) = I_{\text{stim}}\).
Discretized residual contribution:
Coupling: - The diffusion tensor \(K_{\text{diff}}(u)\) depends on the mechanical deformation (mech \(\to\) elec). - This residual is purely primal.
Fluid Dynamics (Incompressible Navier-Stokes)
Modeled as a saddle-point problem naturally fitting the primal–dual structure.
-
Momentum (Primal):
\[ R_{\text{mom}}(u_f, p) = \rho (\dot{u}_f + u_f \cdot \nabla u_f) - \nabla \cdot \mu \nabla u_f + \nabla p - f_{\text{ext}} \] -
Incompressibility (Dual): $$ R_{\text{div}}(u_f) = \nabla \cdot u_f = 0 $$
Coupling: - FSI (Fluid-Structure Interaction): Interface conditions appear as boundary terms in both \(R_{\text{mech}}\) and \(R_{\text{fluid}}\). - ALE formulation can be used if the fluid mesh moves with the solid.
Coupling Terms
Coupling manifests as dependencies between residuals:
- Electromechanical:
- \(f_{\text{active}}(V)\) is added to the mechanical residual \(R_{\text{mech}}\).
- Deformation gradients from \(u\) modify conductivity in \(R_{\text{elec}}\).
- Fluid-Structure:
- Traction continuity: Fluid pressure forces added to solid boundary.
- Velocity continuity: Solid velocity acts as boundary condition for fluid.
3. Solver Compatibility
The Newton–Krylov solver (see Newton–Krylov Solver) remains unchanged. It simply operates on a larger PyTree \(y\) and evaluates a more complex functional \(R(y)\).
Block Structure of the Jacobian
The Jacobian \(J = \partial R / \partial y\) inherits a block structure reflecting the multi-physics coupling:
- Diagonal Blocks (\(A\)): Standard tangent operators for each physics.
- Off-diagonal Blocks (\(K\)): Coupling terms (e.g., \(K_{\text{em}} = \partial f_{\text{active}} / \partial V\)).
- Constraint Blocks (\(G\)): Geometric constraints and incompressibility.
Matrix-Free Advantage
Assembling this global Jacobian would be prohibitively expensive and complex. SOFAx's matrix-free approach avoids this entirely:
- Residual Evaluation: We only implement the function \(R(y)\).
- Jacobian Action (JVP): JAX automatic differentiation computes \(J(y) \cdot \delta y\) on the fly.
- Krylov Solve: GMRES solves the linearized system using only these vector products.
This allows adding new physics by simply defining their residual contributions, without deriving or assembling coupling matrices.
4. Implementation Strategy
Preconditioning
While the solver is matrix-free, convergence requires good preconditioning. The block structure suggests operator splitting or block-diagonal preconditions:
Each block can be preconditioned independently (e.g., Algebraic Multigrid for diffusion, Schur complement for Stokes), ignoring weak couplings during the preconditioner step.
Differentiation & Inverse Problems
Since the formulation remains a root-finding problem \(R(y, \theta) = 0\), the entire coupled simulation is differentiable with respect to parameters \(\theta\) (e.g., conductivities, tissue stiffness).
- Parameter Identification: Adjust material properties to match observed data (e.g., ECG or motion).
- Control: Optimize stimulation protocols \(I_{\text{stim}}\) to achieve desired activation patterns.
Summary
| Physics | Primal Variable | Dual Variable | Residual Type |
|---|---|---|---|
| Mechanics | Acceleration \(a\) | Multipliers \(\lambda\) | \(Ma - f + G^T\lambda\) |
| Electrophysiology | Potential \(V\) | - | Reaction-Diffusion |
| Fluid (Stokes) | Velocity \(u_f\) | Pressure \(p\) | Momentum + Divergence |
By mapping these domains to the unified \((y, R(y))\) abstraction, SOFAx can solve strongly coupled multi-physics problems using the same robust, differentiable numerical core.