Share the proxy transfer operator across a swarm's variables; measure the SWARM-15 seam

docs/developer/subsystems/interpolation.md

@@ -330,12 +330,37 @@ particles and there is no halo exchange (SWARM-15, see
 `docs/developer/design/SWARM_MODERNIZATION_DESIGN_2026-07.md` §4). A proxy node
 near a partition seam therefore gathers from a one-sided neighbourhood.
 
-Linear exactness improves this but does not fix it. A linear-exact stencil
-reproduces a linear field exactly from *any* neighbourhood, one-sided or not,
-so for linear fields the seam error is zero and the proxy is np-independent
-(pinned by `tests/parallel/test_0776_linear_rbf_proxy_parallel.py`). For a
-field with curvature a one-sided stencil still differs from a centred one, so
-np-dependence remains. The halo exchange in SWARM-15 is still the real fix.
+Linear exactness improves this but does not fix it, and the size of what is
+left has now been measured — SWARM-15's migration plan asks for exactly this
+as its first step.
+
+Max relative proxy error against the analytic field, same mesh and same
+particles at each np (2D, `cellSize` 1/24, `fill_param` 4):
+
+| field | scheme | np=1 | np=2 | np=4 |
+|---|---|---|---|---|
+| linear | `order=0` | 2.0e-3 | 2.4e-3 | 4.5e-3 |
+| linear | `order=1` | **5.1e-16** | **5.1e-16** | **5.5e-16** |
+| curved | `order=0` | 9.5e-3 | 1.1e-2 | 1.7e-2 |
+| curved | `order=1` | 1.5e-4 | 1.9e-4 | 2.3e-4 |
+
+Two things to read off. A linear field is **exactly** np-independent under
+`order=1` — round-off at every np, so the seam contributes nothing (pinned by
+`tests/parallel/test_0776_linear_rbf_proxy_parallel.py`). For a field with
+curvature the np-dependence persists in both schemes, roughly doubling from
+np=1 to np=4 — but the *absolute* error at np=4 is 73x smaller under `order=1`
+(2.3e-4 against 1.7e-2).
+
+So `order=1` reduces the seam's magnitude by about two orders of magnitude
+without removing the np-dependence. The halo exchange in SWARM-15 is still the
+real fix.
+
+```{note}
+Node counts grow slightly with np (728 / 758 / 790 here) because shared
+partition-boundary nodes are counted on more than one rank, so the np>1
+maxima are taken over marginally more node instances. The effect is small
+next to the trends above but it is not zero.
+```
 
 ## Related
 

src/underworld3/swarm.py

@@ -1200,10 +1200,21 @@ def _rbf_to_meshVar(self, meshVar, nnn=None, verbose=False, order=1,
                     stacklevel=2,
                 )
             Values = current_values
-        else:
+        elif monotone:
+            # The limiter is data-dependent, so it cannot ride on a cached
+            # geometry-only operator; take the direct path.
             Values = self.rbf_interpolate(
                 new_coords, verbose=verbose, nnn=nnn, order=order, monotone=monotone
             )
+        else:
+            raw_data = self.unpack_raw_data_from_petsc(squeeze=False)
+            resolved_nnn, resolved_order = self._resolve_stencil(
+                nnn, order, raw_data.shape[0]
+            )
+            operator = self.swarm._proxy_interpolation_operator(
+                meshVar, resolved_nnn, 2, resolved_order
+            )
+            Values = operator @ raw_data
 
         meshVar.data[...] = Values[...]
 
@@ -1445,6 +1456,21 @@ def _object_viewer(self):
         display(self.data),
         return
 
+    def _resolve_stencil(self, nnn, order, n_particles):
+        """Stencil size and reproduction order this rank can actually support.
+
+        A rank holding fewer particles than the affine tail needs cannot
+        support a linear fit, so it degrades to inverse distance rather than
+        failing the whole refresh. Shared by the direct and the cached-operator
+        paths so they cannot disagree about what they asked for.
+        """
+        if nnn is None:
+            nnn = 2 * (self.swarm.mesh.dim + 1)
+        nnn = min(nnn, n_particles)
+        if order == 1 and nnn < self.swarm.mesh.dim + 2:
+            order = 0
+        return nnn, order
+
     def rbf_interpolate(self, new_coords, verbose=False, nnn=None, order=1,
                         monotone=False):
         """
@@ -1508,17 +1534,7 @@ def rbf_interpolate(self, new_coords, verbose=False, nnn=None, order=1,
                 )
             return np.zeros((new_coords.shape[0], data_size[1]))
 
-        if nnn is None:
-            nnn = 2 * (self.swarm.mesh.dim + 1)
-
-        if nnn > data_size[0]:
-            nnn = data_size[0]
-
-        # A rank holding fewer particles than the affine tail needs cannot
-        # support a linear fit at all. Inverse distance still gives a sensible
-        # answer there, so degrade rather than fail the whole refresh.
-        if order == 1 and nnn < self.swarm.mesh.dim + 2:
-            order = 0
+        nnn, order = self._resolve_stencil(nnn, order, data_size[0])
 
         # Use direct PETSc access to avoid callback circular dependency
         D = raw_data.copy()
@@ -2839,6 +2855,11 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True):
 
         self._X0_uninitialised = True
         self._index = None
+        # Particle -> proxy-node transfer operators, keyed by geometry and
+        # stencil and shared by every proxied variable of this swarm. Entries
+        # carry the kd-tree they were built from, so they self-invalidate.
+        self._proxy_interpolation_cache = {}
+        self._proxy_cache_mesh_version = None
         self._migration_disabled = False
 
         # Deterministic (SPMD-consistent) creation index — used to order
@@ -3041,6 +3062,54 @@ def _sync_before_assembly(self):
                 var._proxy_stale = True  # align ranks before the collective refresh
                 var._update_proxy_if_stale()
 
+    def _proxy_interpolation_operator(self, meshVar, nnn, p, order):
+        """Sparse particle -> proxy-node transfer, shared across variables.
+
+        The weights depend only on geometry, so every proxied variable whose
+        proxy has the same degree and continuity on the same mesh needs the
+        SAME operator. Measured: a refresh is ~75% weight solve, and the cost
+        of refreshing K proxied variables on one swarm scales linearly with K
+        (4 variables cost 3.95x one in 2D, 4.08x in 3D) because each solves
+        for identical weights independently. Building the operator once per
+        (geometry, stencil) collapses that to one solve plus K sparse
+        products.
+
+        Validity is tied to the kd-tree *instance* rather than to a flag, so
+        the cache cannot outlive the particle positions it was built from:
+        ``migrate()`` drops ``_kdtree``, the next lookup sees a different
+        object and rebuilds. A stale entry therefore keeps its old tree alive
+        until it is replaced -- one tree per distinct key, which is bounded by
+        the number of proxy discretisations in use.
+        """
+        # Two independent things can invalidate an operator, and each is
+        # handled where it can be detected structurally rather than by a flag
+        # someone has to remember to set:
+        #
+        #   mesh geometry  -- a deform or adapt bumps _mesh_version. The whole
+        #                     cache is dropped, because every entry was built
+        #                     against the old node positions. Keying on the
+        #                     version instead would keep the dead entries
+        #                     forever, one set per mesh generation.
+        #   particle motion -- migrate() replaces the kd-tree, so an entry that
+        #                     does not carry the current tree is stale.
+        version = self.mesh._mesh_version
+        if self._proxy_cache_mesh_version != version:
+            self._proxy_interpolation_cache.clear()
+            self._proxy_cache_mesh_version = version
+
+        kdtree = self._get_kdtree()
+        key = (meshVar.degree, meshVar.continuous, nnn, p, order)
+
+        cached = self._proxy_interpolation_cache.get(key)
+        if cached is not None and cached[0] is kdtree:
+            return cached[1]
+
+        operator = kdtree.interpolation_matrix(
+            meshVar.coords_nd, nnn=nnn, p=p, order=order
+        )
+        self._proxy_interpolation_cache[key] = (kdtree, operator)
+        return operator
+
     def _get_kdtree(self):
         """
         Return a cached KDTree for the swarm particle coordinates.

tests/test_0117_swarm_proxy_weight_cache.py

@@ -0,0 +1,224 @@
+"""Proxy transfer operators are shared across variables and self-invalidate.
+
+The `order=1` weights depend only on geometry — proxy node coordinates and
+particle positions — so every proxied variable whose proxy has the same degree
+and continuity on the same mesh needs the *same* operator. Building it once per
+(geometry, stencil) instead of once per variable is worth ~70% of the refresh
+cost on a swarm carrying four proxied variables.
+
+The correctness requirement is that this changes nothing about the values, and
+that a cached operator can never outlive the particle positions it was built
+from. Validity is tied to the kd-tree instance rather than to a flag, so
+`migrate()` — which drops the tree — invalidates the cache structurally.
+
+Known inherited limitation, deliberately not guarded here: the kd-tree is a
+no-copy view of the coordinate buffer, so writing coordinates in place without
+migrating leaves the tree stale (SWARM-02, `test_0113`). A cached operator is
+stale in exactly the same circumstances and no others.
+"""
+
+import numpy as np
+import pytest
+
+import underworld3 as uw
+from underworld3.meshing import UnstructuredSimplexBox
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_b]
+
+
+def _swarm_with(n_vars, proxy_degrees=None, cell_size=1.0 / 8.0):
+    mesh = UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=cell_size
+    )
+    degrees = proxy_degrees or [2] * n_vars
+    swarm = uw.swarm.Swarm(mesh)
+    variables = [
+        swarm.add_variable(name=f"v{i}", size=1, proxy_degree=d)
+        for i, d in enumerate(degrees)
+    ]
+    swarm.populate(fill_param=3)
+    rng = np.random.default_rng(5)
+    for var in variables:
+        var.data[:, 0] = rng.random(swarm.local_size)
+    return mesh, swarm, variables
+
+
+def test_cached_refresh_matches_the_uncached_result():
+    """The cache must be invisible in the values it produces."""
+    mesh, swarm, variables = _swarm_with(2)
+    var = variables[0]
+
+    swarm._proxy_interpolation_cache.clear()
+    var._rbf_to_meshVar(var._meshVar)
+    first = np.asarray(var._meshVar.data[:, 0]).copy()
+
+    # Second refresh hits the cache; a third with the cache cleared rebuilds.
+    var._rbf_to_meshVar(var._meshVar)
+    cached = np.asarray(var._meshVar.data[:, 0]).copy()
+
+    swarm._proxy_interpolation_cache.clear()
+    var._rbf_to_meshVar(var._meshVar)
+    rebuilt = np.asarray(var._meshVar.data[:, 0]).copy()
+
+    assert np.array_equal(cached, first)
+    assert np.array_equal(rebuilt, first)
+
+    del swarm
+    del mesh
+
+
+def test_operator_is_shared_between_variables_of_the_same_proxy_degree():
+    mesh, swarm, variables = _swarm_with(3)
+    swarm._proxy_interpolation_cache.clear()
+
+    for var in variables:
+        var._rbf_to_meshVar(var._meshVar)
+
+    assert len(swarm._proxy_interpolation_cache) == 1, (
+        "three variables with the same proxy degree should share one operator, "
+        f"got {len(swarm._proxy_interpolation_cache)} cache entries"
+    )
+
+    del swarm
+    del mesh
+
+
+def test_different_proxy_degrees_do_not_share_an_operator():
+    """Different degrees mean different node coordinates."""
+    mesh, swarm, variables = _swarm_with(2, proxy_degrees=[1, 2])
+    swarm._proxy_interpolation_cache.clear()
+
+    for var in variables:
+        var._rbf_to_meshVar(var._meshVar)
+
+    assert len(swarm._proxy_interpolation_cache) == 2
+
+    coords = [np.asarray(v._meshVar.coords_nd) for v in variables]
+    assert coords[0].shape != coords[1].shape
+
+    del swarm
+    del mesh
+
+
+def test_migrate_invalidates_the_cached_operator():
+    """A cached operator must not survive the particle positions it used."""
+    mesh, swarm, variables = _swarm_with(1)
+    var = variables[0]
+
+    var._rbf_to_meshVar(var._meshVar)
+    key = next(iter(swarm._proxy_interpolation_cache))
+    tree_before, operator_before = swarm._proxy_interpolation_cache[key]
+
+    coords = swarm._particle_coordinates.data.copy()
+    coords[:, 0] = np.clip(coords[:, 0] + 0.05, 0.001, 0.999)
+    swarm._particle_coordinates.data[...] = coords
+    swarm.migrate()
+
+    var._rbf_to_meshVar(var._meshVar)
+    tree_after, operator_after = swarm._proxy_interpolation_cache[key]
+
+    assert tree_after is not tree_before, "kd-tree survived a migrate"
+    assert operator_after is not operator_before, (
+        "the operator was reused across a migrate — it encodes stale particle "
+        "positions"
+    )
+
+    del swarm
+    del mesh
+
+
+def test_linear_field_still_exact_through_the_cached_path():
+    """The guarantee the operator exists to provide survives caching."""
+    mesh, swarm, variables = _swarm_with(2)
+    var = variables[0]
+    proxy = var._meshVar
+
+    particle_coords = swarm._particle_coordinates.data
+    var.data[:, 0] = 0.5 + particle_coords @ np.array([1.0, 2.0])
+
+    var._rbf_to_meshVar(proxy)          # miss
+    var._rbf_to_meshVar(proxy)          # hit
+
+    node_coords = np.asarray(proxy.coords_nd)
+    expected = 0.5 + node_coords @ np.array([1.0, 2.0])
+    error = np.abs(np.asarray(proxy.data[:, 0]) - expected).max()
+
+    assert error < 1.0e-12, f"cached path lost linear exactness: {error:.3e}"
+
+    del swarm
+    del mesh
+
+
+def test_monotone_refresh_bypasses_the_cache():
+    """The limiter is data-dependent, so it cannot use a geometry-only
+    operator; it must still produce a bounded, correct result."""
+    mesh, swarm, variables = _swarm_with(1)
+    var = variables[0]
+    proxy = var._meshVar
+
+    particle_coords = swarm._particle_coordinates.data
+    var.data[:, 0] = 0.5 + particle_coords @ np.array([1.0, 2.0])
+
+    swarm._proxy_interpolation_cache.clear()
+    var._rbf_to_meshVar(proxy, monotone=True)
+
+    assert len(swarm._proxy_interpolation_cache) == 0, (
+        "the monotone path populated the geometry-only cache"
+    )
+
+    node_coords = np.asarray(proxy.coords_nd)
+    expected = 0.5 + node_coords @ np.array([1.0, 2.0])
+    assert np.abs(np.asarray(proxy.data[:, 0]) - expected).max() < 1.0e-12
+
+    del swarm
+    del mesh
+
+
+def test_cache_does_not_grow_across_mesh_generations():
+    """A deform invalidates every entry, so the cache must not accumulate.
+
+    Keying on the mesh version instead of clearing would keep one dead entry
+    set per mesh generation — each holding a kd-tree, which is a no-copy view
+    of a particle coordinate buffer — for the life of the swarm. Measured
+    before the fix: two entries after a single deform.
+    """
+    mesh, swarm, variables = _swarm_with(1)
+    var = variables[0]
+
+    var.data[:, 0] = 1.0
+    var._rbf_to_meshVar(var._meshVar)
+    assert len(swarm._proxy_interpolation_cache) == 1
+
+    for scale in (1.05, 1.05, 1.05):
+        coords = np.asarray(mesh.X.coords).copy()
+        coords[:, 1] *= scale
+        mesh.deform(coords)
+        var._rbf_to_meshVar(var._meshVar)
+        assert len(swarm._proxy_interpolation_cache) == 1, (
+            "cache accumulated an entry per mesh generation: "
+            f"{len(swarm._proxy_interpolation_cache)} entries"
+        )
+
+    del swarm
+    del mesh
+
+
+def test_operator_is_rebuilt_after_a_mesh_deform():
+    """Correctness half of the above: the new operator is not the old one."""
+    mesh, swarm, variables = _swarm_with(1)
+    var = variables[0]
+
+    var.data[:, 0] = 1.0
+    var._rbf_to_meshVar(var._meshVar)
+    before = next(iter(swarm._proxy_interpolation_cache.values()))[1]
+
+    coords = np.asarray(mesh.X.coords).copy()
+    coords[:, 1] *= 1.1
+    mesh.deform(coords)
+    var._rbf_to_meshVar(var._meshVar)
+    after = next(iter(swarm._proxy_interpolation_cache.values()))[1]
+
+    assert after is not before, "operator survived a mesh deform"
+
+    del swarm
+    del mesh