diff --git a/src/phast/solvers/damage_solver.py b/src/phast/solvers/damage_solver.py index 40e5ea723..44327a13a 100644 --- a/src/phast/solvers/damage_solver.py +++ b/src/phast/solvers/damage_solver.py @@ -217,7 +217,7 @@ def _max_or_zero(values: torch.Tensor) -> float: def make_gamma_corrected_Gc_field(solver, Gc_scalar): - """Build a per-element Gc field from a scalar, baking in gamma correction. + """Apply the elementwise gamma correction to nominal ``Gc`` values. The phase-field solver (when built with ``material.gamma_correction=True``) uses an effective ``Gc_eff_e = Gc * gamma_factor_e`` per element, where @@ -235,13 +235,14 @@ def make_gamma_corrected_Gc_field(solver, Gc_scalar): Must have been constructed with ``gamma_correction=True`` so that the element sizes and ``c_w`` coefficient are available. Gc_scalar : torch.Tensor - 0-d tensor (optionally ``requires_grad=True``) for the nominal Gc. + A scalar or per-element tensor (optionally ``requires_grad=True``) + containing nominal physical toughness values. Returns ------- Gc_field : torch.Tensor, shape ``(n_elems,)`` ``Gc_scalar * gamma_factor_e`` on the solver's compute device/dtype. - Differentiable wrt ``Gc_scalar`` via autograd. + Differentiable with respect to the supplied nominal values. Example ------- @@ -265,6 +266,21 @@ def make_gamma_corrected_Gc_field(solver, Gc_scalar): return Gc_scalar * gamma_factor_e +def _history_pullback_damage(solver, d_unc, d_sol): + """Damage at which dA/dH is evaluated in the history pullback. + + Solve-then-clamp routes ('post_clamp', 'direct') differentiate the linear + solve A(H) d_unc = b(H), so the matrix derivative is taken at the + unconstrained solution d_unc. The projected-CG route solves a + bound-constrained problem; its exact sensitivity needs a reduced + active-set system that is not implemented, and it keeps the projected + damage here (an approximation). + """ + if getattr(solver, '_bounds_method', 'post_clamp') in ('post_clamp', 'direct'): + return d_unc + return d_sol + + def _adjoint_quadratic_forms_perelem(solver, lam, d_sol): """Per-element version of _adjoint_quadratic_forms for spatial Gc(x). @@ -785,7 +801,8 @@ def backward(ctx, grad_output): lam_node = lam.to(d_sol.device, d_sol.dtype) elements = solver.fem.mesh.elements areas = solver.fem.mesh.areas - d_e = d_sol[elements] + d_e = _history_pullback_damage(solver, d_unc, d_sol).to( + d_sol.device, d_sol.dtype)[elements] lam_e = lam_node[elements] d_sum = d_e.sum(dim=1, keepdim=True) per_node = 2.0 * areas.unsqueeze(1) * ( @@ -998,7 +1015,8 @@ def backward(ctx, grad_output): lam_node = lam.to(d_sol.device, d_sol.dtype) elements = solver.fem.mesh.elements areas = solver.fem.mesh.areas - d_e = d_sol[elements] + d_e = _history_pullback_damage(solver, d_unc, d_sol).to( + d_sol.device, d_sol.dtype)[elements] lam_e = lam_node[elements] d_sum = d_e.sum(dim=1, keepdim=True) per_node = 2.0 * areas.unsqueeze(1) * ( @@ -1076,12 +1094,20 @@ def backward(ctx, grad_output): solver, pf_mask, pf_vals) # Adjoint RHS: grad_output (upstream gradient dL/dd_new) - g, _grad_d_prev_unused = _masked_adjoint_inputs( + g, grad_d_prev = _masked_adjoint_inputs( solver, grad_output, d_sol, d_prev, fixed, vals) if need_transfer: g = g.to( dtype=solver._cg_dtype, device=solver._cg_device) + init = None + if fixed is not None: + init = torch.zeros_like(b) + init[fixed] = vals[fixed] + d_unc = _pcg_solve_adjoint_system( + solver, b, reaction_coeff, M_inv, fixed, vals, + initial=init, rel_tol=1e-24) + # Solve A lambda = g using simple Jacobi-preconditioned CG # (reuse the same _Ax operator with same reaction_coeff) lam = _pcg_solve_adjoint_system( @@ -1111,7 +1137,8 @@ def backward(ctx, grad_output): lam_np = lam.to(d_sol.device, d_sol.dtype) elements = solver.fem.mesh.elements areas = solver.fem.mesh.areas - d_e = d_sol[elements] # (E, 3) + d_e = _history_pullback_damage(solver, d_unc, d_sol).to( + d_sol.device, d_sol.dtype)[elements] # (E, 3) lam_e = lam_np[elements] # (E, 3) d_sum = d_e.sum(dim=1, keepdim=True) # (E, 1) @@ -1123,8 +1150,7 @@ def backward(ctx, grad_output): 1.0/3.0 - (d_e + d_sum) / 12.0) # (E, 3) grad_H = (lam_e * per_node).sum(dim=1) # (E,) - # No gradient for solver itself or d_prev - grads = (None, grad_H, None, None, None) + grads = (None, grad_H, grad_d_prev, None, None) return grads[:ctx.n_inputs] @@ -2566,7 +2592,13 @@ def solve(self, H_input: torch.Tensor, "(per-element), not both.") if l0 is None: l0 = torch.tensor(self._l0, dtype=torch.float64) - if not Gc_field_needs_grad and not l0_needs_grad: + if not ( + getattr(self, 'differentiable', False) + or Gc_field_needs_grad + or l0_needs_grad + or self._tensor_needs_grad(H_input) + or self._tensor_needs_grad(d_prev) + ): return self._solve_forward_with_material_overrides( H_input, d_prev, l0=l0, Gc_field=Gc_field, pf_dirichlet_mask=pf_dirichlet_mask, @@ -2581,7 +2613,13 @@ def solve(self, H_input: torch.Tensor, Gc = torch.tensor(self._Gc, dtype=torch.float64) if l0 is None: l0 = torch.tensor(self._l0, dtype=torch.float64) - if not Gc_needs_grad and not l0_needs_grad: + if not ( + getattr(self, 'differentiable', False) + or Gc_needs_grad + or l0_needs_grad + or self._tensor_needs_grad(H_input) + or self._tensor_needs_grad(d_prev) + ): return self._solve_forward_with_material_overrides( H_input, d_prev, Gc=Gc, l0=l0, pf_dirichlet_mask=pf_dirichlet_mask, diff --git a/tests/test_damage_adjoint_regressions.py b/tests/test_damage_adjoint_regressions.py new file mode 100644 index 000000000..775d797c2 --- /dev/null +++ b/tests/test_damage_adjoint_regressions.py @@ -0,0 +1,170 @@ +"""Regression tests for the reverse rules of the projected damage solve. + +The damage update solves a linear system for the unconstrained damage d_unc +and then projects it, d = clamp(max(d_unc, d_prev), 0, 1). Three earlier +defects broke the reverse pass of that composition: + +1. the history (H) pullback evaluated the matrix derivative at the projected + damage instead of d_unc; +2. the history-only route returned no gradient for d_prev, so the + irreversibility branch (d = d_prev) cut the chain to earlier steps; +3. with fixed material values the solve was dispatched to a forward-only path, + detaching the damage from H and d_prev. + +Tests 1 and 2 drive the production ``_AdjointDamageSolve`` on a two-triangle +operator assembled independently from linear basis functions and compare +against a dense solve-and-project reference differentiated by autograd and +against central differences. Test 3 uses the full solver on a small structured mesh. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from phast.solvers.damage_solver import _AdjointDamageSolve + + +class TwoTriangleDamage: + """Minimal stand-in exposing the attributes the reverse rule reads.""" + + def __init__(self): + dtype = torch.float64 + xy = torch.tensor([[0., 0.], [1., 0.], [1., 1.], [0., 1.]], dtype=dtype) + elements = torch.tensor([[0, 1, 2], [0, 2, 3]]) + areas = torch.tensor([.5, .5], dtype=dtype) + self.fem = SimpleNamespace(mesh=SimpleNamespace(elements=elements, areas=areas)) + self.masses = [] + self.K = torch.zeros(4, 4, dtype=dtype) + for nodes, area in zip(elements, areas): + interpolation = torch.cat([torch.ones(3, 1, dtype=dtype), xy[nodes]], dim=1) + gradients = torch.linalg.inv(interpolation)[1:, :].T + local_mass = area / 12 * (torch.eye(3, dtype=dtype) + torch.ones(3, 3, dtype=dtype)) + scatter = torch.eye(4, dtype=dtype)[:, nodes] + self.masses.append(scatter @ local_mass @ scatter.T) + self.K = self.K + scatter @ (area * gradients @ gradients.T) @ scatter.T + self._cg_device = torch.device("cpu") + self._cg_dtype = dtype + self.max_iter = 50 + # Node 0 carries a phase-field Dirichlet value, as on a pre-cracked edge. + self.mask = torch.tensor([True, False, False, False]) + self.values = torch.tensor([.8, 0., 0., 0.], dtype=dtype) + + def system(self, H): + A = .7 * self.K + sum((2 * h + .6) * M for h, M in zip(H, self.masses)) + b = sum(2 * h * M.sum(dim=1) for h, M in zip(H, self.masses)) + return A, b + + def unconstrained(self, H): + A, b = self.system(H) + free = ~self.mask + xf = torch.linalg.solve(A[free][:, free], + b[free] - A[free][:, self.mask] @ self.values[self.mask]) + return self.values.clone().masked_scatter(free, xf) + + def projected(self, H, d_prev): + x = self.unconstrained(H) + return torch.where(self.mask, self.values, torch.maximum(x, d_prev).clamp(0, 1)) + + # Interface used by _AdjointDamageSolve. + def _solve_dispatch(self, H, d_prev, *ignored): + return self.projected(H, d_prev) + + def _prepare_cg(self, H, d_prev): + A, b = self.system(H) + return (d_prev.clone(), d_prev.clone(), H, b, 1 / A.diag(), d_prev.device, d_prev.dtype, False) + + def _prepare_pf_dirichlet(self, mask, values, *ignored): + return mask, values + + def _Ax(self, x, H): + return self.system(H)[0] @ x + + +def _irreversibility_active_state(model): + """H, d_prev with node 1 held at d_prev and nodes 2, 3 strictly free.""" + H = torch.tensor([.3, .8], dtype=torch.float64, requires_grad=True) + x = model.unconstrained(H).detach() + d_prev = (.5 * x).clone() + d_prev[1] = .5 * (1 + x[1]) + assert d_prev[1] > x[1] + 1e-3 and bool(torch.all(x[2:] > d_prev[2:] + 1e-3)) + return H, d_prev.requires_grad_(True) + + +@pytest.mark.parametrize("weights", [[.5, -.7, 1.2, -.4], [-.3, .8, -.6, 1.1], [1., 1., 1., 1.]]) +def test_history_pullback_uses_unconstrained_damage(weights): + model = TwoTriangleDamage() + H, d_prev = _irreversibility_active_state(model) + w = torch.tensor(weights, dtype=torch.float64) + reference = torch.autograd.grad(w @ model.projected(H, d_prev), H)[0] + out = _AdjointDamageSolve.apply(model, H, d_prev, model.mask, model.values) + torch.testing.assert_close(out, model.projected(H, d_prev), rtol=0, atol=1e-12) + grad_H = torch.autograd.grad(w @ out, H)[0] + torch.testing.assert_close(grad_H, reference, rtol=1e-9, atol=1e-11) + h = 1e-5 # small enough to keep every node on its projection branch + for i in range(2): + e = torch.zeros_like(H) + e[i] = 1 + fd = w @ (model.projected(H + h * e, d_prev) - model.projected(H - h * e, d_prev)) / (2 * h) + torch.testing.assert_close(grad_H[i], fd, rtol=1e-8, atol=1e-10) + + +def test_previous_damage_receives_gradient_where_irreversibility_is_active(): + model = TwoTriangleDamage() + H, d_prev = _irreversibility_active_state(model) + w = torch.tensor([.5, -.7, 1.2, -.4], dtype=torch.float64) + reference = torch.autograd.grad(w @ model.projected(H, d_prev), d_prev)[0] + out = _AdjointDamageSolve.apply(model, H, d_prev, model.mask, model.values) + (grad_prev,) = torch.autograd.grad(w @ out, d_prev, allow_unused=True) + assert grad_prev is not None, "irreversibility branch must pass the gradient to d_prev" + assert abs(float(reference[1])) > .1 + torch.testing.assert_close(grad_prev, reference, rtol=1e-12, atol=1e-12) + + +def _unit_square_mesh(n=8): + """Structured T3 mesh of the unit square (no mesh generator needed).""" + from phast.mesh import FEMMesh + + xs = torch.linspace(0.0, 1.0, n + 1, dtype=torch.float64) + yy, xx = torch.meshgrid(xs, xs, indexing="ij") + nodes = torch.stack([xx.reshape(-1), yy.reshape(-1)], dim=1) + tris = [] + for j in range(n): + for i in range(n): + a = j * (n + 1) + i + b, c, d = a + 1, a + n + 2, a + n + 1 + tris += [[a, b, c], [a, c, d]] + return FEMMesh.from_tensors(nodes, torch.tensor(tris, dtype=torch.long), device="cpu") + + +def test_history_gradient_is_not_detached_with_fixed_material(): + from phast.device import DeviceContext + from phast.fem_operators import FEMOperators + from phast.material import Material + from phast.damage_solver import PhaseFieldDamageSolver + + mesh = _unit_square_mesh() + material = Material(E=210000.0, nu=0.3, Gc=2.7, l0=0.1, rho=7.8e-9, + energy_split="spectral", pf_model="AT2", eta_residual=1e-7) + ctx = DeviceContext(device="cpu", profile=False, compile_solvers=False) + solver = PhaseFieldDamageSolver(FEMOperators(mesh, material, ctx=ctx), tol=1e-12, max_iter=2000, ctx=ctx) + + torch.manual_seed(0) + n_elem = mesh.elements.shape[0] + H = (5.0 + 3.0 * torch.rand(n_elem, dtype=torch.float64)).requires_grad_(True) + d_prev = torch.zeros(mesh.n_nodes, dtype=torch.float64) + Gc = torch.tensor(2.7, dtype=torch.float64) # fixed material value, no gradient requested + + loss = lambda d: (d ** 2).sum() # noqa: E731 + out = solver.solve(H, d_prev, Gc=Gc) + assert out.requires_grad, "damage must stay on the graph when H requires a gradient" + grad_H = torch.autograd.grad(loss(out), H)[0] + assert float(grad_H.abs().max()) > 0.0 + + direction = torch.randn(n_elem, dtype=torch.float64) + h = 1e-6 + with torch.no_grad(): + plus = loss(solver.solve(H.detach() + h * direction, d_prev, Gc=Gc)) + minus = loss(solver.solve(H.detach() - h * direction, d_prev, Gc=Gc)) + fd = (plus - minus) / (2 * h) + torch.testing.assert_close(grad_H @ direction, fd, rtol=1e-5, atol=1e-10)