fix(solvers): correct reverse rules of the projected damage solve

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,

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)