File size: 4,037 Bytes
411b9de | 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 | """Unified-constitutive plasticity module (design.md 'UNIFIED CONSTITUTIVE FRAMEWORK').
Box plasticity on principal stretches of the elastic part (the Genesis-family form the data engine
uses): trial F_e = F F_p^{-1}; SVD; clamp singular values to [y_lo, y_hi]; excess folds into F_p.
Gradient safety: torch.linalg.svd's backward divides by (sigma_i^2 - sigma_j^2) and explodes on the
(common!) repeated-stretch states — uniform squash has sigma_1=sigma_2. SafeSVD3 keeps torch's
forward but hand-codes the standard SVD adjoint with the denominator clamped away from zero
(|den| >= EPS * scale, sign preserved). Verified against finite differences in verify_plas.py,
including repeated-sigma cases.
Unclamped points NEVER route through the SVD graph (exact identity path), so with the yield window
wide open the whole module is a bit-exact no-op vs the elastic-only code — the regression test."""
import torch
EPS_DEN = 1e-4
class SafeSVD3(torch.autograd.Function):
"""batched 3x3 SVD with a repeated-singular-value-safe adjoint."""
@staticmethod
def forward(ctx, A):
U, S, Vh = torch.linalg.svd(A)
# sign-fix to proper rotations is NOT needed for the clamp map (U diag(S) Vh reconstructs A
# regardless); keep torch's convention.
ctx.save_for_backward(U, S, Vh)
return U, S, Vh
@staticmethod
def backward(ctx, gU, gS, gVh):
"""first-order SVD perturbation, derived from
P = U^T dA V, dsigma_i = P_ii,
Omega_U[ij] = (s_j P_ij + s_i P_ji)/(s_j^2 - s_i^2), dU = U Omega_U,
Omega_V[ij] = (s_i P_ij + s_j P_ji)/(s_j^2 - s_i^2), dVh = -Omega_V V^T.
Collecting <gU,dU> + <gS,dS> + <gVh,dVh> = sum_ij W_ij P_ij gives gA = U W V^T with
W_ij = ( s_j (a_ij - a_ji) + s_i (b_ij - b_ji) ) / (s_j^2 - s_i^2) (i != j)
W_ii = gS_i, a = U^T gU, b = -gVh V.
(FD-validated in verify_plas.py on sign-invariant observables, incl. repeated sigmas.)
Denominator clamped away from zero (repeated-sigma safety), sign preserved."""
U, S, Vh = ctx.saved_tensors
V = Vh.transpose(-1, -2)
S2 = S * S
den = S2.unsqueeze(-2) - S2.unsqueeze(-1) # den[i,j] = s_j^2 - s_i^2
scale = S2.max(dim=-1, keepdim=True).values.unsqueeze(-1) + 1e-30
sgn = torch.where(den >= 0, torch.ones_like(den), -torch.ones_like(den))
den = torch.where(den.abs() < EPS_DEN * scale, EPS_DEN * scale * sgn, den)
a = U.transpose(-1, -2) @ gU if gU is not None else torch.zeros_like(U)
b = -(gVh @ V) if gVh is not None else torch.zeros_like(U)
Si = S.unsqueeze(-1) # s_i down rows
Sj = S.unsqueeze(-2) # s_j across cols
W = (Sj * (a - a.transpose(-1, -2)) + Si * (b - b.transpose(-1, -2))) / den
W = W - torch.diag_embed(torch.diagonal(W, dim1=-2, dim2=-1))
if gS is not None:
W = W + torch.diag_embed(gS)
return U @ W @ Vh
def return_map(Fe_trial, y_lo, y_hi):
"""box return map on principal stretches. Returns (Fe, need_mask (N,) bool).
Points whose stretches sit inside [y_lo, y_hi] take the EXACT identity path (no SVD in their
graph, bit-identical passthrough). y_lo/y_hi are 0-dim tensors (backbone params) — gradients
flow into them from the clamped points only."""
with torch.no_grad(): # trigger test detached
St = torch.linalg.svdvals(Fe_trial)
need = ((St < y_lo.detach()) | (St > y_hi.detach())).any(-1)
if not bool(need.any()):
return Fe_trial, need
idx = need.nonzero(as_tuple=True)[0]
U, S, Vh = SafeSVD3.apply(Fe_trial[idx])
Sc = torch.clamp(S, min=y_lo, max=y_hi) # grads reach y_lo/y_hi here
Fe = Fe_trial.clone()
Fe[idx] = U @ torch.diag_embed(Sc) @ Vh
return Fe, need
|