diff --git a/docs/developer/subsystems/interpolation.md b/docs/developer/subsystems/interpolation.md index 261fb962..679eea66 100644 --- a/docs/developer/subsystems/interpolation.md +++ b/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 diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 06e3ae35..86fb29aa 100644 --- a/src/underworld3/swarm.py +++ b/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. diff --git a/tests/test_0117_swarm_proxy_weight_cache.py b/tests/test_0117_swarm_proxy_weight_cache.py new file mode 100644 index 00000000..3adef467 --- /dev/null +++ b/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