"""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()