Swarm.repopulate and population control; the cells fit guards against ill-conditioned cells

Swarm.repopulate(min_per_cell, max_per_cell, values, order) takes the
per-cell census and refills a starved cell from its own lattice at the
points farthest from the particles present. A new particle takes the
bounded Shepard reconstruction of every swarm variable from its nearest
neighbours (order=1 for linear-exact; measured extrapolating to 100 on a
field bounded by 1 in a starved corner), or a supplied value.
swarm.population_control = dict(...) runs it at the end of every
advection(), before the next fit. Collective (the domain test reduces).

The cells fit routes a cell whose Gram matrix condition number exceeds
cond_max (1e6) to the patch fit whatever its count, and a flat patch keeps
its mean: particles the advection clamps onto a wall slide along it as a
line, the P2 fit of a line is singular (condition 1e300 at 92 particles),
and that garbage grew by 1e12 in ten steps through the PIC read-back.

Rotating Gaussian, PIC, ten particles per cell, C = 0.25: population
control takes the L2 error from 1.6e-2 to 8.8e-3, level with the
integration-point history; the untapered box with wall in- and outflow
gives the same answer clamped or with exiting particles deleted. The cap
(max_per_cell, closest-pair removal) costs accuracy and is off by default.

Tests: test_0068 (refill exact for linear fields with order=1, value
override and cap, population control under rotation), a collinear-cell
test in test_0067; the history-before-advection test made rank-safe.

Underworld development team with AI support from Claude Code

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MSGAFeA7qYXgkuw9ud8F2G

Author: lmoresi

docs/developer/subsystems/integration-point-variables.md

@@ -331,6 +331,47 @@ size are `nmin` and `patch_nnn` on `CellPolynomialProjector.fit`.
 M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="cells", proxy_degree=2)
 ```
 
+### Repopulation: keeping every cell fit-able
+
+A flow that empties cells starves the fit, and the two particle read-back
+schemes both diverged on emptied corner cells before repopulation existed.
+`Swarm.repopulate()` takes the per-cell census (owning cells from the strict
+locator) and refills a starved cell from its own lattice, the points
+`populate` uses, choosing the lattice points farthest from the particles
+present. A new particle takes, for every swarm variable, the bounded Shepard
+reconstruction from its nearest neighbours (`order=1` for the linear-exact
+reconstruction; a starved cell is where neighbours are far, and the linear
+tail extrapolated to values of 100 on a field bounded by 1), or a supplied
+value (`values={var: constant or callable}`, an inflow datum for instance).
+`swarm.population_control = dict(...)` makes every `advection()` end with a
+repopulation, which is what the cells proxy wants: the refill runs before
+the next fit.
+
+```python
+swarm.population_control = dict()             # refill to the populate() density
+swarm.population_control = dict(min_per_cell=8, values={T: 0.0})
+```
+
+Count is not the whole criterion. Particles the advection clamps back onto a
+wall (`mesh.return_coords_to_bounds`) slide along it as a line, and the P2
+fit of a collinear set is singular whatever its count (measured: condition
+number 1e300 at 92 particles in a wall cell, garbage that grew by 1e12 in
+ten steps through the read-back). The fit therefore routes a cell whose Gram
+matrix has condition number above `cond_max` (1e6) to the patch fit, and a
+patch that is itself flat keeps only its mean. With that guard and
+population control the untapered rotating box, where every wall has an
+inflow and an outflow segment, runs to the same answer whether exiting
+particles are clamped or deleted (`mesh.return_coords_to_bounds = None`,
+the right setting for a true outflow, which also keeps the particle count
+from growing).
+
+Measured on the rotating Gaussian (h = 0.1, C = 0.25, 10 particles per cell,
+PIC, one revolution): population control takes the L2 error from 1.6e-2 to
+8.8e-3, level with the integration-point history at 9.1e-3, because no cell
+is ever left to the linear patch fit. A cap (`max_per_cell`) thins over-full
+cells by removing the particles closest to a neighbour; measured it costs
+accuracy (6.8e-2) and is off by default.
+
 ### Why a least-squares fit and not a conservative transfer
 
 The conservative particle-to-mesh transfer solves the rule mass matrix

src/underworld3/swarm.py

@@ -2995,6 +2995,10 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True):
         # a Lagrangian history registers its first sampling here so it sees
         # the field at the launch positions, not at the landing ones.
         self._pre_advection_hooks = []
+        # Population control: a dict of repopulate() keyword arguments (or
+        # None). When set, advection() ends with repopulate(**population_control)
+        # so no cell is left starved before the next fit of a cells proxy.
+        self.population_control = None
         self._index = None
         # Particle -> proxy-node transfer operators, keyed by geometry and
         # stencil and shared by every proxied variable of this swarm. Entries
@@ -4952,6 +4956,183 @@ def _data_layout(self, i, j=None):
         if self.vtype == uw.VarType.MATRIX:
             return i + j * self.shape[0]
 
+    @timing.routine_timer_decorator
+    @uw.collective_operation
+    def repopulate(
+        self,
+        min_per_cell=None,
+        max_per_cell=None,
+        values=None,
+        nnn=None,
+        order=0,
+        verbose=False,
+    ):
+        """Add particles to cells that hold too few, remove from cells that hold
+        too many, so every cell can support a well-posed fit of its particles.
+
+        The trigger is the per-cell census (owning cells from the strict
+        locator). A starved cell is filled from its own lattice, the points
+        ``populate`` uses (degree ``fill_param``, cell interior), choosing the
+        lattice points farthest from the particles already present. A new
+        particle takes, for every variable, the RBF reconstruction from the
+        nearest existing particles at its position: bounded Shepard weights by
+        default (``order=0``), since a starved cell is where the neighbours
+        are far and a linear-exact tail extrapolates (measured: values of 100
+        on a field bounded by 1 in the emptied corners of a rotating box);
+        ``order=1`` gives the linear-exact reconstruction. ``values`` overrides
+        a variable with a callable ``f(coords) -> (n, components)`` or a
+        constant, an inflow datum for instance. A cell above
+        ``max_per_cell`` loses its most redundant particles, those closest to
+        a neighbour in the same cell.
+
+        Rank-local placement (a cell is filled by the rank that owns it), but
+        collective: every rank must call it, the domain test reduces.
+
+        Parameters
+        ----------
+        min_per_cell : int, optional
+            Particles a cell must hold; default the lattice count of
+            ``fill_param`` (the density ``populate`` gave).
+        max_per_cell : int, optional
+            Cap above which particles are removed; default no removal.
+        values : dict, optional
+            ``{variable or name: callable or constant}`` for new particles.
+        nnn : int, optional
+            Neighbours in the RBF reconstruction (default ``2 (dim + 1)``).
+        order : {0, 1}, optional
+            RBF reconstruction order for new particles: 0 bounded (default),
+            1 linear-exact.
+
+        Returns
+        -------
+        (added, removed) : the counts on this rank.
+        """
+        mesh = self.mesh
+        dim = self.cdim
+        fill = getattr(self, "fill_param", None) or 1
+        lattice = np.asarray(mesh._get_coords_for_basis(fill, continuous=False))
+        c0, c1 = mesh.dm.getHeightStratum(0)
+        ncells = c1 - c0
+        n_lat = lattice.shape[0] // max(ncells, 1)
+        if min_per_cell is None:
+            min_per_cell = n_lat
+        if max_per_cell is not None:
+            min_per_cell = min(min_per_cell, max_per_cell)   # a cap below the lattice count wins
+
+        # Every rank must reach the (collective) domain test before any
+        # rank-local branch; the census itself is rank-local.
+        lat_owned = np.asarray(mesh.points_in_domain(lattice, strict_validation=True), dtype=bool)
+        self._flush_pending_petsc_sync()
+        X = np.array(self._particle_coordinates.data, copy=True) if self.local_size > 0 \
+            else np.zeros((0, dim))
+        cells = np.asarray(mesh._robust_owning_cells(X), dtype=np.int64) if X.shape[0] else np.zeros(0, np.int64)
+        npc = np.bincount(cells[cells >= 0], minlength=ncells)
+        lat_cells = np.asarray(mesh._robust_owning_cells(lattice), dtype=np.int64)
+        owned = np.zeros(ncells, dtype=bool)
+        owned[lat_cells[lat_owned & (lat_cells >= 0)]] = True
+
+        added = removed = 0
+
+        # ---- removal: the most redundant particles of over-full cells ----------
+        if max_per_cell is not None and X.shape[0] > 0:
+            drop = []
+            for c in np.nonzero(owned & (npc > max_per_cell))[0]:
+                idx = np.nonzero(cells == c)[0]
+                P = X[idx]
+                d = np.linalg.norm(P[:, None, :] - P[None, :, :], axis=2)
+                np.fill_diagonal(d, np.inf)
+                nearest = d.min(axis=1)
+                surplus = int(npc[c] - max_per_cell)
+                drop.extend(idx[np.argsort(nearest)[:surplus]].tolist())
+            if drop:
+                for index in sorted(drop, reverse=True):
+                    self.dm.removePointAtIndex(int(index))
+                removed = len(drop)
+                keep = np.ones(X.shape[0], dtype=bool)
+                keep[drop] = False
+                X, cells = X[keep], cells[keep]
+                npc = np.bincount(cells[cells >= 0], minlength=ncells)
+                self._invalidate_canonical_data()
+
+        # ---- addition: starved cells, lattice points farthest from particles -
+        need = np.where(owned, np.maximum(min_per_cell - npc, 0), 0)
+        new_coords = []
+        if need.sum() > 0:
+            cand_ok = lat_owned & (lat_cells >= 0) & (need[np.maximum(lat_cells, 0)] > 0)
+            cand = lattice[cand_ok]
+            cand_cells = lat_cells[cand_ok]
+            if X.shape[0] > 0:
+                dist, _ = uw.kdtree.KDTree(X).query(cand, k=1, sqr_dists=False)
+                dist = np.asarray(dist).reshape(-1)
+            else:
+                dist = np.zeros(cand.shape[0])
+            sort_idx = np.lexsort((-dist, cand_cells))       # by cell, farthest first
+            cand, cand_cells, dist = cand[sort_idx], cand_cells[sort_idx], dist[sort_idx]
+            # rank within cell
+            start = np.searchsorted(cand_cells, np.arange(ncells), side="left")
+            rank_in_cell = np.arange(cand.shape[0]) - start[cand_cells]
+            take = rank_in_cell < need[cand_cells]
+            new_coords = cand[take]
+
+        n_new = int(len(new_coords))
+        if n_new > 0:
+            n_old = max(self.dm.getLocalSize(), 0)
+            nnn = nnn or 2 * (dim + 1)
+            nnn = min(nnn, max(n_old, 1))
+            rbf_order = order if nnn >= dim + 2 else 0
+            operator = None
+            if n_old > 0:
+                operator = uw.kdtree.KDTree(X).interpolation_matrix(
+                    np.asarray(new_coords), nnn=nnn, p=2, order=rbf_order)
+            # raw values of every variable at the old particles, BEFORE the add
+            raw_old = {}
+            for name, var in self._vars.items():
+                if var is self._particle_coordinates or var.clean_name in (
+                        "DMSwarmPIC_coor", "DMSwarm_rank", "DMSwarm_X0"):
+                    continue
+                raw_old[name] = np.asarray(var.unpack_raw_data_from_petsc(squeeze=False)).reshape(n_old, -1)
+
+            self.dm.finalizeFieldRegister()
+            self.dm.addNPoints(n_new)
+            coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, dim))
+            coords[n_old:, :] = np.asarray(new_coords)
+            self.dm.restoreField("DMSwarmPIC_coor")
+            ranks = self.dm.getField("DMSwarm_rank")
+            ranks.reshape(-1)[n_old:] = uw.mpi.rank
+            self.dm.restoreField("DMSwarm_rank")
+            x0 = getattr(self, "_X0", None)
+            if x0 is not None:
+                f = self.dm.getField(x0.clean_name).reshape((-1, dim))
+                f[n_old:, :] = np.asarray(new_coords)
+                self.dm.restoreField(x0.clean_name)
+
+            values = values or {}
+            for name, var in self._vars.items():
+                if name not in raw_old:
+                    continue
+                spec = values.get(var, values.get(name, values.get(var.clean_name)))
+                ncomp = raw_old[name].shape[1] if n_old > 0 else var.num_components
+                if spec is not None:
+                    vals = spec(np.asarray(new_coords)) if callable(spec) else spec
+                    vals = np.broadcast_to(np.asarray(vals, dtype=float).reshape(n_new, -1) if np.ndim(vals) > 0 else vals, (n_new, ncomp))
+                elif operator is not None:
+                    vals = operator @ raw_old[name]
+                else:
+                    vals = np.zeros((n_new, ncomp))
+                f = self.dm.getField(var.clean_name).reshape((-1, ncomp))
+                f[n_old:, :] = np.asarray(vals).reshape(n_new, ncomp)
+                self.dm.restoreField(var.clean_name)
+            added = n_new
+            self._invalidate_canonical_data()
+
+        if added or removed:
+            self._population_generation += 1
+        if verbose:
+            print(f"repopulate: rank {uw.mpi.rank} added {added}, removed {removed} "
+                  f"(cells starved {int((need > 0).sum())})", flush=True)
+        return added, removed
+
+
     @timing.routine_timer_decorator
     def advection(
         self,
@@ -5182,6 +5363,9 @@ def advection(
             delete_lost_points=True,
         )
 
+        if self.population_control is not None:
+            self.repopulate(**self.population_control)
+
         return
 
     @timing.routine_timer_decorator

src/underworld3/utilities/cell_polynomial_projection.py

@@ -118,7 +118,7 @@ def locate(self, coords):
 
     # -- the fit ------------------------------------------------------------
 
-    def fit(self, coords, values, nmin=None, patch_nnn=None, old=None):
+    def fit(self, coords, values, nmin=None, patch_nnn=None, old=None, cond_max=1.0e6):
         """Fit every cell; returns nodal values shaped like ``meshVar.data``.
 
         Parameters
@@ -131,6 +131,12 @@ def fit(self, coords, values, nmin=None, patch_nnn=None, old=None):
         old : current proxy values, shaped like ``meshVar.data``; after the
             first fit a cell with no particles keeps them (on the first fit,
             or without ``old``, it takes the linear patch fit).
+        cond_max : a cell whose Gram matrix has a condition number above this
+            is treated as thin (patch fit) however many particles it holds.
+            Count is not enough: particles clamped onto a wall by the
+            advection lie on a line, and the P2 fit of a line is singular
+            (measured: condition 1e300 at 92 particles, garbage that grew
+            by 1e12 in ten steps through the read-back).
         """
         coords = np.asarray(coords, dtype=np.float64).reshape(-1, self.dim)
         values = np.asarray(values, dtype=np.float64).reshape(coords.shape[0], -1)
@@ -149,6 +155,15 @@ def fit(self, coords, values, nmin=None, patch_nnn=None, old=None):
         U = np.zeros((self.ncells, self.Nb, nc))
         nmin = nmin or self.Nb + 2
         dense = npc >= nmin
+        self.n_ill_conditioned = 0
+        if dense.any():
+            ev = np.linalg.eigvalsh(G[dense])
+            cond = ev[:, -1] / np.maximum(ev[:, 0], 1e-300)
+            ill = cond > cond_max
+            if ill.any():
+                self.n_ill_conditioned = int(ill.sum())
+                idx = np.nonzero(dense)[0][ill]
+                dense[idx] = False
         if dense.any():
             ridge = 1e-10 * np.trace(G[dense], axis1=1, axis2=2)[:, None, None] / self.Nb
             U[dense] = np.linalg.solve(G[dense] + ridge * np.eye(self.Nb)[None], R[dense])
@@ -175,6 +190,13 @@ def fit(self, coords, values, nmin=None, patch_nnn=None, old=None):
             Rt = np.einsum("cpa,cpk->cak", A, psi[idx])
             ridge = 1e-10 * np.trace(Gt, axis1=1, axis2=2)[:, None, None] / (self.dim + 1) + 1e-30
             coef = np.linalg.solve(Gt + ridge * np.eye(self.dim + 1)[None], Rt)      # (nthin, dim+1, nc)
+            # A patch whose particles are themselves (nearly) collinear cannot
+            # carry a gradient: keep only the constant term (the patch mean).
+            evt = np.linalg.eigvalsh(Gt)
+            flat = evt[:, -1] / np.maximum(evt[:, 0], 1e-300) > cond_max
+            if flat.any():
+                coef[flat, 1:, :] = 0.0
+                coef[flat, 0, :] = psi[idx[flat]].mean(axis=1)
             Adof = np.concatenate([np.ones((self.Nb, 1)), self.xi_dof], axis=1)       # (Nb, dim+1)
             U[thin] = np.einsum("ba,cak->cbk", Adof, coef)
 

tests/test_0067_integration_point_proxy.py

@@ -299,13 +299,42 @@ def test_lagrangian_swarm_history_is_sampled_before_the_first_move():
         swarm=swarm, psi_fn=T.sym, vtype=uw.VarType.SCALAR, degree=2, continuous=False,
         order=1, proxy_location="cells",
     )
+    X0 = uw.swarm.SwarmVariable("X0", swarm, 2)               # launch position, carried by the particle
     swarm.populate(fill_param=2)
     assert not lag._history_initialised
-    X_before = np.array(swarm._particle_coordinates.data, copy=True)
+    with uw.synchronised_array_update():
+        X0.data[...] = np.asarray(swarm._particle_coordinates.data)
     swarm.advection(sympy.Matrix([[0.1, 0.0]]), 0.5, order=2)   # every particle moves +0.05 in x
-    X_after = np.asarray(swarm._particle_coordinates.data)
     assert lag._history_initialised
+    # Particles may have changed rank: compare each against the launch
+    # position it carries, not against a rank-local array from before.
+    X_before = np.asarray(X0.data)
+    X_after = np.asarray(swarm._particle_coordinates.data)
     kept = np.abs(X_after[:, 0] - X_before[:, 0] - 0.05) < 1e-12  # particles not returned to bounds
     vals = np.asarray(lag.psi_star[0].data[:, 0])
     assert np.allclose(vals[kept], X_before[kept, 0], atol=1e-10)   # launch positions ...
     assert not np.allclose(vals[kept], X_after[kept, 0], atol=1e-3) # ... not landing positions
+
+
+def test_cells_proxy_collinear_particles_do_not_blow_up():
+    """Particles clamped onto a wall lie on a line; the P2 fit of a line is
+    singular. The fit routes such cells to the patch on the Gram condition
+    number, and a patch that is itself flat keeps only its mean."""
+    mesh = uw.meshing.UnstructuredSimplexBox(cellSize=0.1, qdegree=2)
+    swarm = uw.swarm.Swarm(mesh)
+    M = uw.swarm.SwarmVariable("M", swarm, 1, proxy_location="cells", proxy_degree=2)
+    swarm.populate(fill_param=3)
+    X = np.array(swarm._particle_coordinates.data, copy=True)
+    # Pile every particle of the bottom row of cells onto the wall y = 1e-9
+    bottom = X[:, 1] < 0.1
+    X[bottom, 1] = 1e-9
+    with uw.synchronised_array_update():
+        M.data[:, 0] = 1.0 + X[:, 0]          # values first: the move below migrates
+    with uw.synchronised_array_update():
+        swarm._particle_coordinates.data[...] = X
+    swarm.migrate()
+    M._update_proxy_if_stale()
+    pr = M._cell_projector
+    assert pr.n_ill_conditioned > 0
+    vals = np.asarray(M._meshVar.data[:, 0])
+    assert np.isfinite(vals).all() and vals.min() > 0.5 and vals.max() < 2.5, (vals.min(), vals.max())