feat(interpolation): P2 phase 1 - P2FoldInterpolator

packages/loop_common/src/loop_common/supports/_2d_p2_unstructured.py

@@ -211,7 +211,31 @@ def evaluate_shape_d2(self, indexes):
                 ii += 1
         return d2
 
-    def evaluate_shape_derivatives(self, locations, elements=None):
+    def evaluate_shape_derivatives(self, locations: np.ndarray, elements=None):
+        """Evaluate P2 shape derivatives in global coordinates.
+
+        Returns
+        -------
+        dN, elements, inside
+            dN has shape (n, 2, n_nodes_per_element). Points outside the mesh
+            have element index -1, ``inside`` False and NaN rows in dN. They do
+            not affect other points in the batch.
+        """
+        locations = np.asarray(locations)
+        if elements is None:
+            _verts, _c, elements, inside = self.get_element_for_location(locations)
+            elements = np.asarray(elements)
+            valid = (elements >= 0) & np.asarray(inside, dtype=bool)
+        else:
+            elements = np.asarray(elements)
+            valid = elements >= 0
+        out_elements = np.where(valid, elements, -1)
+        dN = np.full((locations.shape[0], 2, self.elements.shape[1]), np.nan)
+        if np.any(valid):
+            dN[valid], _ = self._evaluate_shape_derivatives_valid(locations[valid], elements[valid])
+        return dN, out_elements, valid
+
+    def _evaluate_shape_derivatives_valid(self, locations, elements=None):
         """
         Compute dN/ds (1st row), dN/dt(2nd row)
         """
@@ -221,7 +245,7 @@ def evaluate_shape_derivatives(self, locations, elements=None):
         else:
             tri = elements
             M = np.ones((elements.shape[0], 3, 3))
-            M[:, :, 1:] = self.vertices[self.elements[elements], :][:, :3, :]
+            M[:, :, 1:] = self.nodes[self.elements[elements], :][:, :3, :]
             points_ = np.ones((locations.shape[0], 3))
             points_[:, 1:] = locations
             minv = np.linalg.inv(M)
@@ -334,8 +358,7 @@ def evaluate_gradient(self, pos: np.ndarray, property_array: np.ndarray) -> np.n
             raise ValueError("property array must have same length as nodes")
         values = np.zeros(pos.shape)
         values[:] = np.nan
-        element_gradients, tri = self.evaluate_shape_derivatives(pos[:, :2])
-        inside = tri >= 0
+        element_gradients, tri, inside = self.evaluate_shape_derivatives(pos[:, :2])
         values[inside, :] = (
             element_gradients[inside, :, :] * property_array[self.elements[tri[inside], None, :]]
         ).sum(2)

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

@@ -191,7 +191,31 @@ def evaluate_shape_d2(self, indexes: np.ndarray) -> np.ndarray:
                 ii += 1
         return d2
 
-    def evaluate_shape_derivatives(
+    def evaluate_shape_derivatives(self, locations: np.ndarray, elements=None):
+        """Evaluate P2 shape derivatives in global coordinates.
+
+        Returns
+        -------
+        dN, elements, inside
+            dN has shape (n, 3, n_nodes_per_element). Points outside the mesh
+            have element index -1, ``inside`` False and NaN rows in dN. They do
+            not affect other points in the batch.
+        """
+        locations = np.asarray(locations)
+        if elements is None:
+            _verts, _c, elements, inside = self.get_element_for_location(locations)
+            elements = np.asarray(elements)
+            valid = (elements >= 0) & np.asarray(inside, dtype=bool)
+        else:
+            elements = np.asarray(elements)
+            valid = elements >= 0
+        out_elements = np.where(valid, elements, -1)
+        dN = np.full((locations.shape[0], 3, self.elements.shape[1]), np.nan)
+        if np.any(valid):
+            dN[valid], _ = self._evaluate_shape_derivatives_valid(locations[valid], elements[valid])
+        return dN, out_elements, valid
+
+    def _evaluate_shape_derivatives_valid(
         self, locations: np.ndarray, elements: np.ndarray = None
     ) -> np.ndarray:
         """
@@ -372,8 +396,7 @@ def evaluate_gradient(self, pos: np.ndarray, property_array: np.ndarray) -> np.n
             raise ValueError("property array must have same length as nodes")
         values = np.zeros(pos.shape)
         values[:] = np.nan
-        element_gradients, tetra = self.evaluate_shape_derivatives(pos[:, :3])
-        inside = tetra >= 0
+        element_gradients, tetra, inside = self.evaluate_shape_derivatives(pos[:, :3])
         values[inside, :] = (
             element_gradients[:, :, :] * property_array[self.elements[tetra[inside], None, :]]
         ).sum(2)

packages/loop_common/src/loop_common/supports/_p2_structured_tetra.py

@@ -260,6 +260,30 @@ def evaluate_shape(self, locations: np.ndarray):
         return N, elements, inside
 
     def evaluate_shape_derivatives(self, locations: np.ndarray, elements=None):
+        """Evaluate P2 shape derivatives in global coordinates.
+
+        Returns
+        -------
+        dN, elements, inside
+            dN has shape (n, 3, n_nodes_per_element). Points outside the mesh
+            have element index -1, ``inside`` False and NaN rows in dN. They do
+            not affect other points in the batch.
+        """
+        locations = np.asarray(locations)
+        if elements is None:
+            _verts, _c, elements, inside = self.get_element_for_location(locations)
+            elements = np.asarray(elements)
+            valid = (elements >= 0) & np.asarray(inside, dtype=bool)
+        else:
+            elements = np.asarray(elements)
+            valid = elements >= 0
+        out_elements = np.where(valid, elements, -1)
+        dN = np.full((locations.shape[0], 3, self.elements.shape[1]), np.nan)
+        if np.any(valid):
+            dN[valid], _ = self._evaluate_shape_derivatives_valid(locations[valid], elements[valid])
+        return dN, out_elements, valid
+
+    def _evaluate_shape_derivatives_valid(self, locations: np.ndarray, elements=None):
         """Evaluate quadratic tetrahedral shape derivatives at locations."""
         locations = np.array(locations)
         if elements is None:
@@ -480,8 +504,7 @@ def evaluate_gradient(self, pos: np.ndarray, property_array: np.ndarray) -> np.n
         values = np.zeros(pos.shape)
         values[:] = np.nan
 
-        element_gradients, tetras = self.evaluate_shape_derivatives(pos)
-        _, _, inside = self.evaluate_shape(pos)
+        element_gradients, tetras, inside = self.evaluate_shape_derivatives(pos)
         values[inside, :] = np.einsum(
             "ijk,ik->ij",
             element_gradients[inside, :, :],

packages/loop_common/tests/test_p2_shape_derivatives_outside.py

@@ -0,0 +1,39 @@
+"""evaluate_shape_derivatives must handle points outside the mesh."""
+
+import numpy as np
+import pytest
+from loop_common.supports._3d_p2_tetra import P2UnstructuredTetMesh
+from loop_common.supports._p2_structured_tetra import P2TetMesh
+
+
+def _unstructured():
+    return P2UnstructuredTetMesh(
+        origin=np.zeros(3), step_vector=np.ones(3) / 4, nsteps_cells=np.array([4, 4, 4])
+    )
+
+
+def _structured():
+    return P2TetMesh(origin=np.zeros(3), nsteps=np.array([3, 3, 3]), step_vector=np.ones(3))
+
+
+@pytest.mark.parametrize(
+    "make,inside_pt,outside_pt",
+    [
+        (_unstructured, [0.41, 0.52, 0.33], [5.0, 5.0, 5.0]),
+        (_structured, [1.1, 1.2, 1.3], [50.0, 50.0, 50.0]),
+    ],
+)
+def test_outside_point_gives_nan_and_does_not_poison_batch(make, inside_pt, outside_pt):
+    mesh = make()
+    ref = mesh.evaluate_shape_derivatives(np.array([inside_pt]))
+    pts = np.array([inside_pt, outside_pt, inside_pt])
+    res = mesh.evaluate_shape_derivatives(pts)
+    assert len(res) == 3
+    dN, elements, inside = res
+    assert dN.shape == (3, 3, 10)
+    assert elements[1] == -1
+    assert not inside[1]
+    assert np.all(np.isnan(dN[1]))
+    assert inside[0] and inside[2]
+    assert np.all(np.isfinite(dN[0])) and np.all(np.isfinite(dN[2]))
+    np.testing.assert_allclose(dN[0], ref[0][0])

packages/loop_common/tests/test_p2_structured_tetra.py

@@ -183,7 +183,7 @@ def test_p2tetmesh_shape_function_derivatives_exist(self):
         mesh = P2TetMesh(nsteps=np.array([3, 3, 3]))
 
         centroid = np.array([[1.0, 1.0, 1.0]])
-        dN, _elem_ids = mesh.evaluate_shape_derivatives(centroid)
+        dN, _elem_ids, _inside = mesh.evaluate_shape_derivatives(centroid)
 
         # dN should have shape (n_points, 3, 10)
         assert dN.shape == (1, 3, 10)

packages/loop_interpolation/src/loop_interpolation/_discrete_interpolator.py

@@ -364,7 +364,8 @@ def add_constraints_to_least_squares(self, A, B, idc, w=1.0, name="undefined"):
         if len(A.shape) > 2:
             n_rows = A.shape[0] * A.shape[1]
             if isinstance(w, np.ndarray):
-                w = np.tile(w, (A.shape[1]))
+                # reshape below puts row (i, j) at i * A.shape[1] + j, so repeat, not tile
+                w = np.repeat(w, A.shape[1])
             A = A.reshape((A.shape[0] * A.shape[1], A.shape[2]))
             idc = idc.reshape((idc.shape[0] * idc.shape[1], idc.shape[2]))
             B = B.reshape(A.shape[0])

packages/loop_interpolation/src/loop_interpolation/_p2_fold_interpolator.py

@@ -0,0 +1,129 @@
+"""
+Piecewise quadratic interpolator using fold constraints.
+
+Summary of findings_1_1.md (what exists / what was added)
+---------------------------------------------------------
+Existing:
+- ``P2Interpolator`` and the P2 supports. ``evaluate_shape_derivatives``
+  gives dN of shape (n, 3, 10), node order [v0..v3, e23, e03, e01, e12, e13, e02].
+- ``setup_with_fold_constraints`` (_fold_setup.py) and
+  ``resolve_fold_norm_target`` (_fold_norm_alignment.py), both reused as is.
+- ``add_constraints_to_least_squares`` drops NaN rows, normalises the rows of A,
+  then applies the weight ``w``. A per-row ``w`` is therefore applied after
+  normalisation.
+- ``get_quadrature_points`` on the support is a triangle (face) rule for edge
+  jumps. It does not fit a volume rule.
+Added (Phase 1):
+- ``_quadrature.py``: ``tet_quadrature`` / ``tet_quadrature_points``
+  (centroid, gauss4, weights sum to 1, tet index kept).
+- P2 ``evaluate_shape_derivatives`` returns (dN, elements, inside) with NaN
+  rows for outside points.
+- This class. ``copy()`` is overridden because the base returns P2Interpolator.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+from loop_common.logging import get_logger as getLogger
+
+from ._fold_event import FoldEvent
+from ._fold_norm_alignment import resolve_fold_norm_target
+from ._fold_setup import setup_with_fold_constraints
+from ._p2interpolator import P2Interpolator
+from ._quadrature import tet_quadrature_points
+
+logger = getLogger(__name__)
+
+_SAMPLE_RULES = ("centroid", "gauss4")
+
+
+class P2FoldInterpolator(P2Interpolator):
+    """P2 interpolator with fold constraints (Laurent et al., 2016)."""
+
+    def __init__(self, support, fold: FoldEvent | None = None):
+        P2Interpolator.__init__(self, support)
+        self.fold = fold
+
+    def copy(self):
+        return P2FoldInterpolator(self.support, fold=self.fold)
+
+    def setup_interpolator(self, **kwargs):
+        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,
+        )
+        # set after the call: includes fold_weights, which the base setup does not see
+        self._setup_kwargs = dict(kwargs)
+        return report
+
+    def add_fold_constraints(
+        self,
+        fold_orientation=10.0,
+        fold_axis_w=10.0,
+        fold_normalisation=1.0,
+        fold_norm=-1.0,
+        dgz_alignment="warn",
+        sample_rule="gauss4",
+        mask_fn: Callable | None = None,
+    ):
+        """Add the three fold constraint families at quadrature points.
+
+        Parameters
+        ----------
+        fold_orientation, fold_axis_w, fold_normalisation : float or None
+            weights of the families. None skips a family.
+        fold_norm : float
+            norm target, see ``resolve_fold_norm_target``
+        dgz_alignment : {"none", "warn", "correct"}
+        sample_rule : {"gauss4", "centroid"}
+            quadrature rule used to place sample points in every tetrahedron
+        mask_fn : callable, optional
+            called with the sample points; points where it returns True get zero weight
+        """
+        if sample_rule not in _SAMPLE_RULES:
+            raise ValueError(f"Unknown sample_rule '{sample_rule}', use one of {_SAMPLE_RULES}")
+        support = self.support
+        points, tet_idx, q = tet_quadrature_points(support.nodes, support.elements, sample_rule)
+        dN, _out_elements, inside = support.evaluate_shape_derivatives(points, elements=tet_idx)
+        deformed_orientation, fold_axis, dgz = self.fold.get_deformed_orientation(points)
+        # energy per tet must not depend on the rule: weight the residual by w*sqrt(q)
+        sq = np.sqrt(q)
+        if mask_fn is not None:
+            sq = sq * (~np.asarray(mask_fn(points), dtype=bool))
+        # keep finite rows only, so A, B, idc and w stay aligned
+        valid = np.asarray(inside, dtype=bool) & np.all(np.isfinite(dN), axis=(1, 2))
+        idc_all = support.elements[tet_idx]
+
+        def _add(vectors, target, w, name):
+            v = valid & np.all(np.isfinite(vectors), axis=1)
+            # no volume factor: row normalisation would cancel it
+            A = np.einsum("ij,ijk->ik", vectors[v], dN[v])
+            B = np.full(A.shape[0], float(target))
+            self.add_constraints_to_least_squares(A, B, idc_all[v], w=float(w) * sq[v], name=name)
+
+        if fold_orientation is not None:
+            logger.info(f"Adding fold orientation constraint to w = {fold_orientation}")
+            _add(deformed_orientation, 0.0, fold_orientation, "fold orientation")
+        if fold_axis_w is not None:
+            logger.info(f"Adding fold axis constraint to w = {fold_axis_w}")
+            _add(fold_axis, 0.0, fold_axis_w, "fold axis")
+        if fold_normalisation is not None:
+            logger.info(f"Adding fold normalisation constraint to w = {fold_normalisation}")
+            target_norm = resolve_fold_norm_target(
+                fold=self.fold,
+                normal_constraints=self.get_norm_constraints(),
+                fold_norm=fold_norm,
+                dgz_alignment=dgz_alignment,
+                logger=logger,
+                default_fold_norm=-1.0,
+            )
+            # Decision: unlike DiscreteFoldInterpolator, B is NOT multiplied by
+            # fold_normalisation (nor by volume). That scaling changes the target
+            # norm with the weight, which is a quirk. Here the weight only sets
+            # the importance and the target stays fold_norm.
+            _add(dgz, target_norm, fold_normalisation, "fold normalisation")

packages/loop_interpolation/src/loop_interpolation/_p2interpolator.py

@@ -109,8 +109,9 @@ def copy(self):
     def add_gradient_constraints(self, w: float = 1.0):
         points = self.get_gradient_constraints()
         if points.shape[0] > 0:
-            grad, elements = self.support.evaluate_shape_derivatives(points[:, : self.dimensions])
-            inside = elements > -1
+            grad, elements, inside = self.support.evaluate_shape_derivatives(
+                points[:, : self.dimensions]
+            )
             # element size cancels in the row normalisation, w is applied after it
             wt = np.full(int(np.sum(inside)), float(w))
             A = np.einsum(
@@ -140,8 +141,9 @@ def add_gradient_orthogonal_constraints(
 
         """
         if points.shape[0] > 0:
-            grad, elements = self.support.evaluate_shape_derivatives(points[:, : self.dimensions])
-            inside = elements > -1
+            grad, elements, inside = self.support.evaluate_shape_derivatives(
+                points[:, : self.dimensions]
+            )
             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])
@@ -153,8 +155,9 @@ def add_gradient_orthogonal_constraints(
     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
+            grad, elements, inside = self.support.evaluate_shape_derivatives(
+                points[:, : self.dimensions]
+            )
             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)
@@ -255,10 +258,10 @@ def minimise_edge_jumps(
             if callable(vector_func):
                 norm = vector_func(cp[:, i, :])
             # evaluate the shape function for the edges for each neighbouring triangle
-            cp_Dt, cp_tri1 = self.support.evaluate_shape_derivatives(
+            cp_Dt, cp_tri1, _inside1 = self.support.evaluate_shape_derivatives(
                 cp[:, i, :], elements=self.support.shared_element_relationships[:, 0]
             )
-            cp_Dn, cp_tri2 = self.support.evaluate_shape_derivatives(
+            cp_Dn, cp_tri2, _inside2 = self.support.evaluate_shape_derivatives(
                 cp[:, i, :], elements=self.support.shared_element_relationships[:, 1]
             )
             # constraint for each cp is triangle - neighbour create a Nx12 matrix