fix(interpolation): P2 phase 0 prerequisites (bugs 3, 4, 5, 6, 11, F1)

packages/loop_common/src/loop_common/supports/_3d_p2_tetra.py

@@ -25,7 +25,14 @@ def __init__(
                     "neighbours arrays, or origin/step_vector/nsteps_cells to build a "
                     "quadratic tetrahedral mesh over a bounding box"
                 )
+            origin = np.asarray(origin, dtype=float)
+            step_vector = np.asarray(step_vector, dtype=float)
+            nsteps_cells = np.asarray(nsteps_cells)
             nodes, elements, neighbours = self._build_from_bbox(origin, step_vector, nsteps_cells)
+            # kept so that inside() and the setup diagnostics work on a bbox-built mesh
+            self.origin = origin
+            self.step_vector = step_vector
+            self.nsteps_cells = nsteps_cells
         UnStructuredTetMesh.__init__(self, nodes, elements, neighbours, aabb_nsteps)
         self.type = SupportType.P2UnstructuredTetMesh
         if self.elements.shape[1] != 10:

packages/loop_common/src/loop_common/supports/_3d_unstructured_tetra.py

@@ -186,8 +186,11 @@ def _init_face_table(self):
         )
         shared_faces = faces[np.array(np.sum(faces, axis=1) == 3).flatten(), :]
         row, col = shared_faces.nonzero()
-        row = row[row.argsort()]
-        col = col[row.argsort()]
+        # one stable sort for both arrays: the default sort is not stable, so ties
+        # (three nodes per face) came out in a numpy/CPU dependent order
+        order = row.argsort(kind="stable")
+        row = row[order]
+        col = col[order]
         shared_face_index = np.zeros((shared_faces.shape[0], 3), dtype=int)
         shared_face_index[:] = -1
         shared_face_index[row.reshape(-1, 3)[:, 0], :] = col.reshape(-1, 3)

packages/loop_interpolation/src/loop_interpolation/_discrete_fold_interpolator.py

@@ -53,14 +53,17 @@ def update_fold(self, fold):
         self.fold = fold
 
     def setup_interpolator(self, **kwargs):
-        return setup_with_fold_constraints(
+        report = setup_with_fold_constraints(
             fold=self.fold,
             kwargs=kwargs,
             interpolator_name=self.__class__.__name__,
             base_setup=super().setup_interpolator,
             add_fold_constraints=self.add_fold_constraints,
             finalize_report=self.finalize_setup_diagnostics_report,
         )
+        # includes fold_weights, which the base setup does not see
+        self._setup_kwargs = dict(kwargs)
+        return report
 
     def add_fold_constraints(
         self,

packages/loop_interpolation/src/loop_interpolation/_discrete_interpolator.py

@@ -26,6 +26,23 @@
 logger = getLogger(__name__)
 
 
+def _drop_nan_rows(A, B, idc, w, name):
+    """Remove constraint rows with a nan in A, B or idc, and log how many were dropped.
+
+    A, B, idc are 2d, 1d, 2d with one row per constraint. ``w`` is filtered
+    too when it is a per-row array.
+    """
+    bad = np.isnan(A).any(axis=1) | np.isnan(B) | np.isnan(idc).any(axis=1)
+    n_bad = int(bad.sum())
+    if n_bad == 0:
+        return A, B, idc, w
+    logger.warning(f"Dropping {n_bad} of {bad.shape[0]} constraint rows containing nan: {name}")
+    keep = ~bad
+    if isinstance(w, np.ndarray) and w.shape[0] == bad.shape[0]:
+        w = w[keep]
+    return A[keep], B[keep], idc[keep], w
+
+
 class DiscreteInterpolator(GeologicalInterpolator):
     """ """
 
@@ -41,6 +58,7 @@ def __init__(self, support, data=None, c=None, up_to_date=False):
         """
         GeologicalInterpolator.__init__(self, data=data, up_to_date=up_to_date)
         self.B = []
+        self._setup_kwargs = {}
         self.support = support
         self.dimensions = support.dimension
         self.c = (
@@ -352,11 +370,12 @@ def add_constraints_to_least_squares(self, A, B, idc, w=1.0, name="undefined"):
             B = B.reshape(A.shape[0])
             # w = w.reshape((A.shape[0]))
 
-        # Check for nan before any row normalisation, which would otherwise
+        # Drop nan rows before any row normalisation, which would otherwise
         # zero out nan rows in A and mask this check further down.
-        if np.any(np.isnan(idc)) or np.any(np.isnan(A)) or np.any(np.isnan(B)):
-            logger.warning(f"Constraints contain nan not adding constraints: {name}")
+        A, B, idc, w = _drop_nan_rows(A, B, idc, w, name)
+        if A.shape[0] == 0:
             return
+        n_rows = A.shape[0]
 
         # normalise by rows of A
         # Should this be done? It should make the solution more stable
@@ -1334,7 +1353,7 @@ def update(self) -> bool:
             logger.debug("Cannot rerun interpolator")
             return False
         if not self.up_to_date:
-            self.setup_interpolator()
+            self.setup_interpolator(**self._setup_kwargs)
             self.up_to_date = self.solve_system(self.solver)
             return self.up_to_date
         return bool(self.up_to_date)

packages/loop_interpolation/src/loop_interpolation/_fd_fold_interpolator.py

@@ -76,14 +76,17 @@ def setup_interpolator(self, **kwargs):
         All other kwargs are forwarded to the parent
         :meth:`FiniteDifferenceInterpolator.setup_interpolator`.
         """
-        return setup_with_fold_constraints(
+        report = setup_with_fold_constraints(
             fold=self.fold,
             kwargs=kwargs,
             interpolator_name=self.__class__.__name__,
             base_setup=super().setup_interpolator,
             add_fold_constraints=self.add_fold_constraints,
             finalize_report=self.finalize_setup_diagnostics_report,
         )
+        # includes fold_weights, which the base setup does not see
+        self._setup_kwargs = dict(kwargs)
+        return report
 
     # ------------------------------------------------------------------
     # Fold constraints

packages/loop_interpolation/src/loop_interpolation/_finite_difference_interpolator.py

@@ -156,6 +156,8 @@ def setup_interpolator(self, **kwargs):
         -------
 
         """
+        # kept so that update() can rebuild the system with the same options
+        self._setup_kwargs = dict(kwargs)
         self.reset()
         regularisation_config = self.resolve_regularisation_config(
             regularisation=kwargs.get("regularisation", None),

packages/loop_interpolation/src/loop_interpolation/_p1interpolator.py

@@ -193,6 +193,8 @@ def setup_interpolator(self, **kwargs):
         -------
 
         """
+        # kept so that update() can rebuild the system with the same options
+        self._setup_kwargs = dict(kwargs)
         # can't reset here, clears fold constraints
         self.reset()
         regularisation_config = self.resolve_regularisation_config(

packages/loop_interpolation/src/loop_interpolation/_p2interpolator.py

@@ -60,6 +60,8 @@ def setup_interpolator(self, **kwargs):
         -------
 
         """
+        # kept so that update() can rebuild the system with the same options
+        self._setup_kwargs = dict(kwargs)
         self.reset()
         regularisation_config = self.resolve_regularisation_config(
             regularisation=kwargs.get("regularisation", None),
@@ -109,17 +111,16 @@ def add_gradient_constraints(self, w: float = 1.0):
         if points.shape[0] > 0:
             grad, elements = self.support.evaluate_shape_derivatives(points[:, : self.dimensions])
             inside = elements > -1
-            area = self.support.element_size[elements[inside]]
-            wt = np.ones(area.shape[0])
-            wt *= w * area
+            # element size cancels in the row normalisation, w is applied after it
+            wt = np.full(int(np.sum(inside)), float(w))
             A = np.einsum(
                 "ikj,ij->ik",
                 grad[inside, :],
                 points[inside, self.dimensions : self.dimensions * 2],
             )
             B = np.zeros(A.shape[0])
             elements = self.support.elements[elements[inside]]
-            self.add_constraints_to_least_squares(A * wt[:, None], B, elements, name="gradient")
+            self.add_constraints_to_least_squares(A, B, elements, w=wt, name="gradient")
 
     def add_gradient_orthogonal_constraints(
         self, points: np.ndarray, vector: np.ndarray, w=1.0, B=0
@@ -141,30 +142,27 @@ def add_gradient_orthogonal_constraints(
         if points.shape[0] > 0:
             grad, elements = self.support.evaluate_shape_derivatives(points[:, : self.dimensions])
             inside = elements > -1
-            area = self.support.element_size[elements[inside]]
-            wt = np.ones(area.shape[0])
-            wt *= w * area
+            wt = np.full(int(np.sum(inside)), float(w))
             A = np.einsum("ijk,ij->ik", grad[inside, :], vector[inside, :])
             B = np.zeros(A.shape[0])
             elements = self.support.elements[elements[inside]]
             self.add_constraints_to_least_squares(
-                A * wt[:, None], B, elements, name="gradient orthogonal"
+                A, B, elements, w=wt, name="gradient orthogonal"
             )
 
     def add_norm_constraints(self, w: float = 1.0):
         points = self.get_norm_constraints()
         if points.shape[0] > 0:
             grad, elements = self.support.evaluate_shape_derivatives(points[:, : self.dimensions])
             inside = elements > -1
-            area = self.support.element_size[elements[inside]]
-            wt = np.ones(area.shape[0])
-            wt *= w * area
+            wt = np.full(int(np.sum(inside)), float(w))
             elements = np.tile(self.support.elements[elements[inside]], (self.dimensions, 1, 1))
             elements = elements.swapaxes(0, 1)
             self.add_constraints_to_least_squares(
-                grad[inside, :, :] * wt[:, None, None],
-                points[inside, self.dimensions : self.dimensions * 2] * wt[:, None],
+                grad[inside, :, :],
+                points[inside, self.dimensions : self.dimensions * 2],
                 elements,
+                w=wt,
                 name="norm",
             )
 
@@ -207,7 +205,7 @@ def minimise_grad_steepness(
         elements = np.arange(0, len(self.support.elements))
         mask = np.ones(self.support.neighbours.shape[0], dtype=bool)
         if not maskall:
-            mask[:] = np.all(self.support.neighbours > 3, axis=1)
+            mask[:] = np.all(self.support.neighbours >= 0, axis=1)
 
         d2 = self.support.evaluate_shape_d2(elements[mask])
         # d2 shape is [ele_idx, deriv, node]

packages/loop_interpolation/src/loop_interpolation/_solver_strategy.py

@@ -11,6 +11,9 @@
 import numpy as np
 from scipy import sparse
 
+# scipy's default rtol (1e-5) stops P2 solves up to 0.6 from the converged field
+DEFAULT_CG_RTOL = 1e-8
+
 
 def resolve_solver_choice(
     solver: Callable[[sparse.csr_matrix, np.ndarray], np.ndarray] | str | None, logger
@@ -58,6 +61,8 @@ def solve_with_cg(
     logger.info("Solving using cg")
     if ("atol" not in solver_kwargs or "rtol" not in solver_kwargs) and tol is not None:
         solver_kwargs["atol"] = tol
+    elif tol is None and "rtol" not in solver_kwargs and "atol" not in solver_kwargs:
+        solver_kwargs["rtol"] = DEFAULT_CG_RTOL
 
     logger.info(f"Solver kwargs: {solver_kwargs}")
     solve_step_started = perf_counter()
@@ -68,7 +73,7 @@ def solve_with_cg(
             "CG reached iteration limit (%s) and did not converge; using last iteration",
             res[1],
         )
-    return np.asarray(res[0]), True
+    return np.asarray(res[0]), res[1] == 0
 
 
 def solve_with_cg_normal_equations(
@@ -92,6 +97,8 @@ def solve_with_cg_normal_equations(
     logger.info("Solving using cg (matrix-free fused regularisation normal equations)")
     if ("atol" not in solver_kwargs or "rtol" not in solver_kwargs) and tol is not None:
         solver_kwargs["atol"] = tol
+    elif tol is None and "rtol" not in solver_kwargs and "atol" not in solver_kwargs:
+        solver_kwargs["rtol"] = DEFAULT_CG_RTOL
 
     logger.info(f"Solver kwargs: {solver_kwargs}")
     solve_step_started = perf_counter()
@@ -102,7 +109,7 @@ def solve_with_cg_normal_equations(
             "CG reached iteration limit (%s) and did not converge; using last iteration",
             res[1],
         )
-    return np.asarray(res[0]), True
+    return np.asarray(res[0]), res[1] == 0
 
 
 def solve_with_lsmr(

packages/loop_interpolation/tests/golden/README.md

@@ -65,20 +65,22 @@ and by how much, against the golden file. A failing constraint test names the fa
 
 ## Findings from building these tests (all confirmed by running code)
 
-Known bugs from the review that the suite now demonstrates (`test_known_*`, strict xfail):
+Bugs from the review that the suite demonstrated. The tests are now normal tests (`test_default_solve_is_converged`,
+`test_p2_normal_weight_has_effect`, `test_update_keeps_fold_weights`, `test_dfi_is_repeatable_*`).
 
-| Test | Review item | Observed |
+| Test | Review item | Status |
 |---|---|---|
-| `test_known_default_solve_is_converged` | bug 3 | Default CG on P2 `planar_dip` is wrong by up to 0.64 in a unit box. It still returns `True`. P1 is off by about 2e-3. |
-| `test_known_p2_normal_weight_has_effect` | bug 4 | P2: changing `npw` from 1 to 100 moves the field by 1e-7. P1 moves by 7e-4, FD by 9e-3. |
-| `test_known_update_keeps_fold_weights` | bug 11 | After `update()` the fold orientation weight goes from 50 back to 10. The field moves by 0.28. DFI and FDFold. |
-| `test_known_dfi_is_repeatable` | section 3, item 6 | Two unseeded DFI builds differ by 6.6e-3. The harness reseeds `loop_common.math.rng`. |
+| `test_default_solve_is_converged` | bug 3 | Fixed (Phase 0.4). CG returns `False` at the iteration limit. Default `rtol` is 1e-8: P2 `planar_dip` error 0.65 -> 2.5e-4, P1 1.8e-2 -> 5e-7. |
+| `test_p2_normal_weight_has_effect` | bug 4 | Fixed (Phase 0.3). `npw` 1 -> 100 moves the P2 field by 7.4e-4 (was 1e-7). |
+| `test_update_keeps_fold_weights` | bug 11 | Fixed (Phase 0.5). `update()` replays the kwargs of the last `setup_interpolator` call. |
+| `test_dfi_is_repeatable_without_reseed` | section 3, item 6 | Fixed. DFI no longer shuffles rows. |
+| (no golden test; see `test_p0_nan_constraints_skipped.py`) | bug 6 | Fixed (Phase 0.6). NaN rows are dropped, not the whole family. |
+| `TestP2GradSteepnessMask` (`test_p2_interpolator.py`) | bug 5 | Fixed (Phase 0.2). Mask is `neighbours >= 0`. P2 fields moved by at most 2.3e-3. |
 
 New findings, not in the review:
 
-- **F1.** `P2UnstructuredTetMesh(origin=..., step_vector=..., nsteps_cells=...)` never stores those attributes.
-  `P2Interpolator.setup_interpolator` then raises `AttributeError: ... has no attribute 'origin'` from
-  `inside()` in the diagnostics report. `engines.py` sets them by hand. Remove that workaround when fixed.
+- **F1.** (Fixed, Phase 0.1.) `P2UnstructuredTetMesh(origin=..., step_vector=..., nsteps_cells=...)` did not store those
+  attributes, so `inside()` and the diagnostics report raised `AttributeError`.
 - **F2.** `TetMesh` (P1 and DFI) returns NaN for points on the lower x, y or z faces (strict `>` in `inside`).
   631 of 3375 probe points. FD and P2 return values. The golden files keep the NaN. Metrics use the inner box.
 - **F3.** Fold amplitude is damped. The median ratio of solved to true gradient magnitude is 0.66 (`fold_section`),

packages/loop_interpolation/tests/golden/engines.py

@@ -72,11 +72,6 @@ def make_interpolator(engine: str, case):
         return DiscreteFoldInterpolator(mesh, fold=make_fold_event(case))
     if engine == "P2":
         mesh = P2UnstructuredTetMesh(origin=origin, step_vector=np.full(3, 1.0 / P2_CELLS), nsteps_cells=np.full(3, P2_CELLS))
-        # KNOWN BUG (golden finding F1): a bbox-built P2 mesh never stores these, so
-        # `inside()` and the setup diagnostics raise AttributeError. Set them here.
-        for attr, val in (("origin", origin), ("step_vector", np.full(3, 1.0 / P2_CELLS)), ("nsteps_cells", np.full(3, P2_CELLS))):
-            if not hasattr(mesh, attr):
-                setattr(mesh, attr, val)
         return P2Interpolator(mesh)
     raise ValueError(engine)