engine3d/solver.py

"""3D tetrahedral Projective Dynamics solver.

Skeleton (same as every 2D system in this project):
    per-tet LOCAL PROJECTION (tied)  ->  GLOBAL RECONCILIATION (prefactored)

Local step: co-rotated constraint per tet — project F onto SO(3) via
Mueller extraction (warm-started quaternion cache, SVD fallback guard).
Global step: one prefactored sparse Cholesky-like solve (scipy splu),
identical scalar system applied per coordinate.

Lessons baked in from the 2D work:
- residual-based iteration stopping, never a fixed budget
- inversion guard inside rotation extraction, not a patch
- pinned vertices removed from the system (exact Dirichlet), not penalized
"""
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import splu

from .rotation import extract_rotations, mat_to_quat
from .strain import corotated_strain, vec6_to_sym


class PDSolver3D:
    def __init__(self, verts, tets, density=1000.0, stiffness=1e5,
                 dt=1e-3, gravity=(0.0, -9.81, 0.0), damping=0.999,
                 pinned=None, strain_limit=None, strain_basis=None,
                 incompressible=False):
        # strain_limit: admissible co-rotated strain. None -> pure corotated
        #   elasticity (target strain 0). Scalar -> radial clamp on ||e||_F.
        #   Length-6 vector -> anisotropic per-component clamp (materials with
        #   direction-dependent compliance, e.g. fiber composites or foams).
        # strain_basis: optional projector with a .project(e) method
        #   (TiedStrainBasis or a neural MaterialProjector); the projection
        #   target is reconstructed through it (the tied-subspace local step).
        # incompressible: remove the volumetric part (trace) of the target
        #   strain — near-incompressible materials like rubber.
        if strain_limit is not None and np.ndim(strain_limit) > 0:
            strain_limit = np.asarray(strain_limit, dtype=np.float64)
        self.strain_limit = strain_limit
        self.incompressible = incompressible
        self.strain_basis = strain_basis
        self.x = verts.astype(np.float64).copy()
        self.x0 = self.x.copy()
        self.v = np.zeros_like(self.x)
        self.tets = tets
        self.dt = dt
        self.gravity = np.asarray(gravity, dtype=np.float64)
        self.damping = damping
        V, T = len(verts), len(tets)

        # rest-shape matrices
        d = verts[tets]
        Dm = np.stack([d[:, 1] - d[:, 0], d[:, 2] - d[:, 0], d[:, 3] - d[:, 0]], axis=-1)
        self.vol = np.abs(np.linalg.det(Dm)) / 6.0
        Dminv = np.linalg.inv(Dm)
        # G (T,4,3): F = X^T G with X the (4,3) tet vertex positions
        self.G = np.empty((T, 4, 3))
        self.G[:, 1:, :] = Dminv
        self.G[:, 0, :] = -Dminv.sum(axis=1)

        self.w = stiffness * self.vol  # per-tet constraint weight

        # lumped mass
        self.mass = np.zeros(V)
        np.add.at(self.mass, tets.ravel(), np.repeat(density * self.vol / 4.0, 4))

        # rotation warm-start cache
        self.q = np.tile(np.array([1.0, 0, 0, 0]), (T, 1))
        self.inverted_count = 0

        # global matrix  A = M/h^2 + sum_t w_t G_t G_t^T  (V x V, per-coordinate)
        GGt = np.einsum("tic,tjc->tij", self.G, self.G)  # (T,4,4)
        rows = np.repeat(tets, 4, axis=1).ravel()
        cols = np.tile(tets, (1, 4)).ravel()
        vals = (self.w[:, None, None] * GGt).ravel()
        L = sp.coo_matrix((vals, (rows, cols)), shape=(V, V)).tocsc()
        A = L + sp.diags(self.mass / dt**2)

        self.pinned = np.zeros(V, dtype=bool)
        if pinned is not None:
            self.pinned[np.asarray(pinned)] = True
        self.free = np.flatnonzero(~self.pinned)
        f = self.free
        self.A_ff = A[np.ix_(f, f)].tocsc()
        self.A_fp = A[np.ix_(f, np.flatnonzero(self.pinned))].tocsc()
        self.lu = splu(self.A_ff)

    def _project_local(self, F, R):
        """Local constraint projection: F -> nearest admissible target.

        Pure corotated (no limit): target = R.
        Strain-limited: target = R (I + e*), e* the radially clamped
        co-rotated strain, optionally reconstructed through the tied basis.
        """
        if self.strain_limit is None and not self.incompressible:
            return R
        e = corotated_strain(F, R)
        if self.strain_limit is None:
            e_star = e.copy()
        elif np.ndim(self.strain_limit) > 0:
            lim = self.strain_limit
            e_star = np.clip(e, -lim, lim)
        else:
            norm = np.linalg.norm(e, axis=1)
            scale = np.minimum(1.0, self.strain_limit / np.maximum(norm, 1e-12))
            e_star = e * scale[:, None]
        if self.incompressible:
            e_star[:, :3] -= e_star[:, :3].mean(axis=1, keepdims=True)
        if self.strain_basis is not None:
            e_star = self.strain_basis.project(e_star)
        return np.einsum("tij,tjk->tik", R, np.eye(3) + vec6_to_sym(e_star))

    def _deformation_gradients(self):
        return np.einsum("tvd,tvc->tdc", self.x[self.tets], self.G)

    def inertia_prediction(self, ext_accel=None):
        """The PD momentum target s (what x converges toward absent elasticity).
        Exposed so warm-start policies can form corrections relative to it."""
        h = self.dt
        accel = self.gravity if ext_accel is None else self.gravity + ext_accel
        s = self.x + h * self.v + h * h * accel
        s[self.pinned] = self.x[self.pinned]
        return s

    def step(self, max_iters=20, tol=1e-6, ext_accel=None, x_init=None):
        """One PD step. Returns iterations used.

        x_init: optional initial iterate (warm start). Only the starting
        point of the local/global loop changes — the fixed point does not,
        so a warm start can save iterations but never change the physics.
        """
        h = self.dt
        s = self.inertia_prediction(ext_accel)

        x_prev_step = self.x.copy()
        x = s.copy() if x_init is None else x_init.copy()
        x[self.pinned] = self.x[self.pinned]
        rhs_inertia = (self.mass[:, None] / h**2) * s

        it = 0
        for it in range(1, max_iters + 1):
            # local: project each F onto SO(3) (tied across all tets)
            F = np.einsum("tvd,tvc->tdc", x[self.tets], self.G)
            R, self.q, bad = extract_rotations(F, self.q)
            self.inverted_count += int(bad.sum())
            target = self._project_local(F, R)

            # global: rhs and prefactored solve, per coordinate
            rhs = rhs_inertia.copy()
            contrib = self.w[:, None, None] * np.einsum("tvc,tdc->tvd", self.G, target)
            np.add.at(rhs, self.tets.ravel(),
                      contrib.reshape(-1, 3))
            b = rhs[self.free] - self.A_fp @ x[self.pinned]
            x_new = x.copy()
            x_new[self.free] = self.lu.solve(b)

            delta = np.abs(x_new - x).max()
            x = x_new
            if delta < tol:
                break

        self.v = self.damping * (x - x_prev_step) / h
        self.v[self.pinned] = 0.0
        self.x = x
        return it

    # --- diagnostics (the invariant suite reads these) ---
    def momentum(self):
        return (self.mass[:, None] * self.v).sum(axis=0)

    def kinetic_energy(self):
        return 0.5 * (self.mass * (self.v**2).sum(axis=1)).sum()

    def strain6(self):
        """Current per-tet co-rotated strain (T,6) — corpus sampling hook."""
        F = self._deformation_gradients()
        R, _, _ = extract_rotations(F, self.q.copy())
        return corotated_strain(F, R)

    def elastic_energy(self):
        F = self._deformation_gradients()
        R, _, _ = extract_rotations(F, self.q.copy())
        return 0.5 * (self.w * ((F - R)**2).sum(axis=(1, 2))).sum()