File size: 7,775 Bytes
24c2ab9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """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()
|