fix(mmpde): finite SPD fallback for fully-degenerate metric evaluations (#352)

src/underworld3/meshing/smoothing/mmpde.py

@@ -14,6 +14,75 @@
                     _global_sum, _global_min, _global_max, _global_mean)
 
 
+def _spd_sanitise(M):
+    """Project a batch of metric tensors ``(N, d, d)`` onto SPD.
+
+    The MMPDE functional uses fractional powers (``sqrt(detM)``,
+    ``detM**((1-p)/2)``, ``S**q``) that are defined only for a
+    symmetric-positive-definite metric. FE extrapolation outside the current
+    mesh can hand back a non-SPD tensor; this projection floors its
+    eigenvalues at a small RELATIVE level so a genuine SPD metric passes
+    through unchanged while garbage becomes a benign "coarsen here" tensor
+    (tiny positive eigenvalues) instead of a NaN in the energy.
+
+    The floor is relative to the batch's largest finite eigenvalue with an
+    absolute minimum of 1e-8, so the pass-through guarantee holds for the
+    O(1)-normalised metrics the mover uses; a metric batch whose valid
+    eigenvalues sit far below unit scale would be floored too.
+    """
+    # Symmetrise: the metric is symmetric by construction, so for a valid
+    # tensor this is an exact no-op (M_ij == M_ji bit-for-bit).
+    Ms = 0.5 * (M + np.swapaxes(M, -1, -2))
+    if Ms.shape[0] == 0:
+        return Ms                                   # rank owns no cells
+    try:
+        w, Vc = np.linalg.eigh(Ms)
+    except np.linalg.LinAlgError:
+        # Degenerate-input eigh behaviour is LAPACK-path dependent: the 2x2
+        # kernel returns quiet NaNs, the general kernel raises — and a
+        # batched eigh raises for the WHOLE batch if any one tensor fails.
+        # Retry per tensor so one degenerate tensor cannot take down its
+        # neighbours; a tensor that still fails is marked non-finite and
+        # rebuilt as the isotropic fallback below (#352).
+        w = np.empty(Ms.shape[:-1])
+        Vc = np.empty_like(Ms)
+        for i, Mi in enumerate(Ms):
+            try:
+                w[i], Vc[i] = np.linalg.eigh(Mi)
+            except np.linalg.LinAlgError:
+                w[i] = np.nan
+                Vc[i] = np.eye(Ms.shape[-1])
+    wmax = float(np.nanmax(w[np.isfinite(w)], initial=-np.inf))
+    if not np.isfinite(wmax):
+        # Every eigenvalue on this rank is NaN/inf: the metric carries no
+        # information here, so anchor the floor at the unit scale — the
+        # projection below then rebuilds every tensor as the same benign
+        # isotropic fallback instead of propagating NaN (#352).
+        wmax = 1.0
+    floor = max(wmax, 1.0) * 1.0e-8
+    # Per-tensor SPD test: a cell is "bad" only if one of its OWN
+    # eigenvalues is non-finite or below the floor. Project just those
+    # cells; every already-SPD tensor is returned untouched (bit-identical
+    # to the symmetrised input), so one bad point cannot perturb the rest.
+    bad = ~np.isfinite(w).all(axis=1) | (w.min(axis=1) < floor)
+    if not bad.any():
+        return Ms
+    out = Ms.copy()
+    wf = np.clip(np.nan_to_num(w[bad], nan=floor, posinf=wmax, neginf=floor),
+                 floor, None)
+    # The eigenvector basis itself can be NaN for a fully-degenerate input
+    # tensor; fall back to the identity basis there so the rebuilt tensor
+    # is the isotropic floor metric, not NaN (#352).
+    Vb = Vc[bad]
+    nan_basis = ~np.isfinite(Vb).all(axis=(1, 2))
+    if nan_basis.any():
+        Vb = Vb.copy()
+        Vb[nan_basis] = np.eye(Ms.shape[-1])
+    out[bad] = np.einsum('nij,nj,nkj->nik', Vb, wf, Vb)
+    return out
+
+
+
 def _mmpde_mover(mesh, metric, pinned_labels, verbose,
                    n_outer=150, p=1.5, theta=1.0 / 3.0, tau=1.0,
                    step_frac=0.2, area_floor_frac=0.01,
@@ -219,51 +288,16 @@ def _eval_M(pts):
         _eval_M = _eval_M_analytic
 
     # --- SPD sanitiser on the evaluated metric -------------------------
-    # The MMPDE functional G uses fractional powers that are defined ONLY
-    # for an SPD metric: sqrt(detM), detM**((1-p)/2), and S**q with
-    # S = tr(J M⁻¹ Jᵀ). The metric is a guide field FE-evaluated at the
-    # FIXED reference cloud; once the interior has deformed, a reference
-    # point can fall OUTSIDE the current mesh and the P1 metric field is
-    # then evaluated by FE EXTRAPOLATION (out-of-cell basis functions go
-    # negative), yielding a non-SPD tensor — e.g. a scalar density ρ·I
-    # with ρ<0. Its determinant ρ² stays positive (so a detM>0 test
-    # passes) but M is negative-definite, so S<0 and S**q = NaN → the
-    # energy is non-finite and the mover bails with zero displacement
-    # (no adaptation). Project every evaluated tensor onto SPD with a
-    # small RELATIVE eigenvalue floor: a genuine SPD metric is returned
-    # unchanged (no-op), while extrapolation garbage becomes a benign
-    # "coarsen here" (tiny positive eigenvalues) instead of a NaN.
+    # The metric is a guide field FE-evaluated at the FIXED reference
+    # cloud; once the interior has deformed, a reference point can fall
+    # OUTSIDE the current mesh and the P1 metric field is then evaluated
+    # by FE EXTRAPOLATION (out-of-cell basis functions go negative),
+    # yielding a non-SPD tensor — e.g. a scalar density ρ·I with ρ<0
+    # whose determinant ρ² still passes a detM>0 test. Project every
+    # evaluated tensor onto SPD (module-level _spd_sanitise) so the
+    # fractional powers in the energy stay finite.
     _eval_M_raw = _eval_M
 
-    def _spd_sanitise(M):
-        # Symmetrise: the metric is symmetric by construction, so for a valid
-        # tensor this is an exact no-op (M_ij == M_ji bit-for-bit).
-        Ms = 0.5 * (M + np.swapaxes(M, -1, -2))
-        if Ms.shape[0] == 0:
-            return Ms                                   # rank owns no cells
-        w, Vc = np.linalg.eigh(Ms)
-        # TODO(BUG): if EVERY eigenvalue on this rank is NaN/inf (e.g. a
-        # fully-degenerate metric evaluation), np.nanmax(w) is non-finite,
-        # max(wmax, 1.0) propagates the NaN, the floor is non-finite and
-        # the nan_to_num/clip projection below emits NaN tensors instead
-        # of the intended benign SPD fallback. Pre-existing (preserved
-        # verbatim by the Wave D split); guard wmax with a finite default
-        # (e.g. wmax = 1.0 when not finite) when next touching numerics.
-        wmax = float(np.nanmax(w))
-        floor = max(wmax, 1.0) * 1.0e-8
-        # Per-tensor SPD test: a cell is "bad" only if one of its OWN
-        # eigenvalues is non-finite or below the floor. Project just those
-        # cells; every already-SPD tensor is returned untouched (bit-identical
-        # to the symmetrised input), so one bad point cannot perturb the rest.
-        bad = ~np.isfinite(w).all(axis=1) | (w.min(axis=1) < floor)
-        if not bad.any():
-            return Ms
-        out = Ms.copy()
-        wf = np.clip(np.nan_to_num(w[bad], nan=floor, posinf=wmax, neginf=floor),
-                     floor, None)
-        out[bad] = np.einsum('nij,nj,nkj->nik', Vc[bad], wf, Vc[bad])
-        return out
-
     def _eval_M(pts):
         return _spd_sanitise(_eval_M_raw(pts))
 

tests/test_0850_mesh_smoothing.py

@@ -236,3 +236,71 @@ def test_mmpde_2d_smoke_still_works(self):
         after = np.asarray(mesh.X.coords)
         assert np.all(np.isfinite(after))
         assert not np.allclose(before, after)   # the mover actually moved
+
+
+class TestSPDSanitise:
+    """Regression (#352): the SPD projection must return a finite SPD
+    fallback even when a rank's ENTIRE metric evaluation is degenerate —
+    previously an all-NaN eigenvalue set made the relative floor itself
+    NaN and the "sanitised" output was NaN."""
+
+    def _assert_finite_spd(self, out):
+        assert np.all(np.isfinite(out))
+        w = np.linalg.eigvalsh(out)
+        assert np.all(w > 0)
+
+    def test_all_nan_batch_returns_finite_spd(self):
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        out = _spd_sanitise(np.full((5, 2, 2), np.nan))
+        self._assert_finite_spd(out)
+
+    def test_all_inf_batch_returns_finite_spd(self):
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        out = _spd_sanitise(np.full((3, 2, 2), np.inf))
+        self._assert_finite_spd(out)
+
+    def test_valid_spd_input_is_returned_unchanged(self):
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        rng = np.random.default_rng(7)
+        A = rng.standard_normal((10, 2, 2))
+        M = np.einsum('nij,nkj->nik', A, A) + 0.5 * np.eye(2)  # SPD by construction
+        out = _spd_sanitise(M)
+        # Symmetrisation of an exactly-symmetric tensor is a no-op, so a
+        # genuine SPD metric must pass through bit-identical.
+        assert np.array_equal(out, M)
+
+    def test_one_degenerate_tensor_does_not_perturb_the_rest(self):
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        rng = np.random.default_rng(11)
+        A = rng.standard_normal((6, 2, 2))
+        M = np.einsum('nij,nkj->nik', A, A) + 0.5 * np.eye(2)
+        M[2] = np.nan
+        out = _spd_sanitise(M)
+        self._assert_finite_spd(out)
+        good = np.ones(6, dtype=bool)
+        good[2] = False
+        assert np.array_equal(out[good], M[good])
+
+    def test_3x3_degenerate_batch_returns_finite_spd(self):
+        # Degenerate-input eigh is LAPACK-path dependent: the 2x2 kernel
+        # returns quiet NaNs, the general (3x3) kernel raises LinAlgError.
+        # The fallback must hold for both — this is the path the 3D MMPDE
+        # capstone will stand on.
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        out = _spd_sanitise(np.full((4, 3, 3), np.nan))
+        self._assert_finite_spd(out)
+
+    def test_3x3_mixed_batch_survives_batched_eigh_failure(self):
+        # A single non-converging tensor makes the BATCHED eigh raise for
+        # the whole batch; the per-tensor retry must preserve the valid
+        # neighbours bit-identical and rebuild only the degenerate one.
+        from underworld3.meshing.smoothing.mmpde import _spd_sanitise
+        rng = np.random.default_rng(3)
+        A = rng.standard_normal((5, 3, 3))
+        M = np.einsum('nij,nkj->nik', A, A) + 0.5 * np.eye(3)
+        M[1] = np.inf
+        out = _spd_sanitise(M)
+        self._assert_finite_spd(out)
+        good = np.ones(5, dtype=bool)
+        good[1] = False
+        assert np.array_equal(out[good], M[good])