diff --git a/src/underworld3/meshing/smoothing/mmpde.py b/src/underworld3/meshing/smoothing/mmpde.py index 00b22324..bb5ace24 100644 --- a/src/underworld3/meshing/smoothing/mmpde.py +++ b/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)) diff --git a/tests/test_0850_mesh_smoothing.py b/tests/test_0850_mesh_smoothing.py index 7dccc99f..037227b3 100644 --- a/tests/test_0850_mesh_smoothing.py +++ b/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])