Pack and index auxiliary fields by DM field, not by position in mesh.vars

src/underworld3/discretisation/discretisation_mesh.py

@@ -3896,13 +3896,21 @@ def update_lvec(self, swarm_sync=True):
             # The field decomposition seems to fail if coarse DMs are present
             names, isets, dms = self.dm.createFieldDecomposition()
 
-            # traverse subdms, taking user generated data in the subdm
-            # local vec, pushing it into a global sub vec
-            for var, subiset, subdm in zip(self.vars.values(), isets, dms):
-                # var.vec lazily creates the PETSc local vector on first access
-                lvec = var.vec
+            # Traverse the DM's fields BY NAME. `self.vars` holds its
+            # variables weakly, so a dropped-and-collected variable leaves
+            # a field behind in the DM; a positional zip would then pack
+            # every later variable into the wrong field (measured: the
+            # cell-size field landing in a P2 slot as garbage, NaN
+            # residuals in a solver that reads it). An orphaned field is
+            # zeroed so nothing stale can reach a kernel.
+            for name, subiset, subdm in zip(names, isets, dms):
+                var = self.vars.get(name)
                 subvec = a_global.getSubVector(subiset)
-                subdm.localToGlobal(lvec, subvec, addv=False)
+                if var is None:
+                    subvec.set(0.0)
+                else:
+                    # var.vec lazily creates the PETSc local vector on first access
+                    subdm.localToGlobal(var.vec, subvec, addv=False)
                 a_global.restoreSubVector(subiset, subvec)
 
             for iset in isets:

src/underworld3/utilities/_jitextension.py

@@ -773,6 +773,24 @@ def getext(
 
 
 @timing.routine_timer_decorator
+def _aux_component_offsets(mesh):
+    """Component offset of every field of the mesh DM, keyed by field id.
+
+    Read from the DM itself, not from ``mesh.vars``: a MeshVariable that
+    was dropped and collected leaves its PETSc field in the DM (a DMPlex
+    cannot shed a field), and PETSc lays the auxiliary arrays out over
+    ALL fields in field order. The offsets therefore have to count the
+    orphaned fields too.
+    """
+    offsets = {}
+    total = 0
+    for field_id in range(mesh.dm.getNumFields()):
+        fe, _label = mesh.dm.getField(field_id)
+        offsets[field_id] = total
+        total += fe.getNumComponents()
+    return offsets
+
+
 def generate_c_source(
     name,
     mesh: underworld3.discretisation.Mesh,
@@ -822,7 +840,7 @@ def generate_c_source(
         count_bd_residual_sig, count_bd_jacobian_sig = callbacks.counts
 
     # `_ccode` patching
-    def ccode_patch_fns(varlist, prefix_str):
+    def ccode_patch_fns(varlist, prefix_str, component_offsets=None):
         """
         This function patches uw functions with the necessary ccode
         routines for the code printing.
@@ -848,11 +866,22 @@ def ccode_patch_fns(varlist, prefix_str):
             ordered according to their `field_id`.
         prefix_str: str
             The string prefix to write.
+        component_offsets: dict, optional
+            Component offset of every field in the DM, by ``field_id``
+            (see ``_aux_component_offsets``). When given, each variable
+            is patched from ITS OWN field's offset instead of a running
+            count over ``varlist``: a field whose Python variable has
+            been dropped stays in the DM and still occupies its slots,
+            so a running count would shift every later variable onto
+            the wrong data.
         """
         u_i = 0  # variable increment
         u_x_i = 0  # variable gradient increment
         lambdafunc = lambda self, printer: self._ccodestr
         for var in varlist:
+            if component_offsets is not None:
+                u_i = component_offsets[var.field_id]
+                u_x_i = u_i * mesh.cdim
             if var.vtype == VarType.SCALAR:
                 # monkey patch this guy into the function
                 type(var.fn)._ccodestr = f"{prefix_str}[{u_i}]"
@@ -898,7 +927,8 @@ def ccode_patch_fns(varlist, prefix_str):
     # is important, as the secondary call will overwrite
     # those patched in the first call.
 
-    ccode_patch_fns(_stable_sorted(mesh.vars.values()), "petsc_a")
+    ccode_patch_fns(_stable_sorted(mesh.vars.values()), "petsc_a",
+                    component_offsets=_aux_component_offsets(mesh))
     ccode_patch_fns(primary_field_list, "petsc_u")
 
     # Also patch `BaseScalar` types. Nothing fancy - patch the overall type,

tests/test_1058_dropped_meshvariable_aux_layout.py

@@ -0,0 +1,113 @@
+"""A dropped MeshVariable must not corrupt the auxiliary data of later solves.
+
+`mesh.vars` holds variables weakly, but a DMPlex cannot shed a field: a
+variable that is dropped and garbage-collected leaves its PETSc field in
+the DM. Two places used to assume the registry and the DM field list line
+up by position:
+
+- `Mesh.update_lvec` zipped `mesh.vars.values()` against the DM's field
+  decomposition, so every later variable was packed into the wrong field
+  (the orphan's slot) and its own slot stayed at whatever it held;
+- the JIT's `petsc_a[]` offsets were a running count over the live
+  variables, skipping the orphan's components.
+
+Measured before the fix: a cell-size (P0) field landing in a P2 slot as
+garbage, NaN residuals (`DIVERGED_FUNCTION_NANORINF`) in one run and a
+subtly wrong answer in the next, depending on when the collector ran. The
+default Model holds the only strong reference to a variable (the mesh
+outlives the model it was created under), so `uw.reset_default_model()`,
+which the test suite runs between tests, releases every variable a script
+no longer names; the variable-statistics
+helpers also delete temporaries from the registry on purpose. The orphan is
+an ordinary state, not a misuse.
+
+Run: pixi run python -m pytest tests/test_1058_dropped_meshvariable_aux_layout.py -v
+"""
+import gc
+
+import numpy as np
+import pytest
+import sympy
+
+import underworld3 as uw
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_a]
+
+
+def _mesh():
+    return uw.meshing.UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.25, qdegree=3)
+
+
+def _poisson_with_field_coefficient(mesh, tag):
+    """A Poisson solve whose answer depends on an auxiliary field (the
+    diffusivity is a MeshVariable), so mis-packed aux data changes it."""
+    x, y = mesh.X
+    kappa = uw.discretisation.MeshVariable(f"kappa_{tag}", mesh, 1, degree=1)
+    kappa.array[:, 0, 0] = uw.function.evaluate(1.0 + 4.0 * x * y, kappa.coords).reshape(-1)
+    u = uw.discretisation.MeshVariable(f"u_{tag}", mesh, 1, degree=2)
+    poisson = uw.systems.Poisson(mesh, u)
+    poisson.constitutive_model = uw.constitutive_models.DiffusionModel
+    poisson.constitutive_model.Parameters.diffusivity = kappa.sym[0]
+    poisson.f = 1.0
+    for b in ("Left", "Right", "Top", "Bottom"):
+        poisson.add_dirichlet_bc(0.0, b)
+    poisson.solve()
+    return np.array(u.array), kappa, u
+
+
+def test_dropped_variable_leaves_an_orphaned_field():
+    """The premise: dropping a variable does not shrink the DM."""
+    mesh = _mesh()
+    n_fields = mesh.dm.getNumFields()
+    # The mesh keeps the model it was created under alive; a variable
+    # registers with the CURRENT default model, so a reset before and
+    # after creating it is what releases it (the suite's per-test reset).
+    uw.reset_default_model()
+    uw.discretisation.MeshVariable("temporary", mesh, 2, degree=2)
+    uw.reset_default_model()
+    gc.collect()
+    assert "temporary" not in mesh.vars
+    assert mesh.dm.getNumFields() == n_fields + 1
+
+
+def test_solve_after_a_dropped_variable_matches_a_clean_mesh():
+    reference, _k, _u = _poisson_with_field_coefficient(_mesh(), "ref")
+
+    mesh = _mesh()
+    uw.reset_default_model()
+    uw.discretisation.MeshVariable("dropped_vector", mesh, 2, degree=2)
+    uw.discretisation.MeshVariable("dropped_scalar", mesh, 1, degree=1)
+    uw.reset_default_model()
+    gc.collect()
+    assert mesh.dm.getNumFields() > len(mesh.vars)
+
+    answer, _k, _u = _poisson_with_field_coefficient(mesh, "orphan")
+    assert np.allclose(answer, reference, rtol=0, atol=1e-10)
+
+
+def test_packed_aux_vector_lands_in_the_named_fields():
+    mesh = _mesh()
+    uw.reset_default_model()
+    uw.discretisation.MeshVariable("dropped", mesh, 2, degree=1)
+    uw.reset_default_model()
+    gc.collect()
+    assert "dropped" not in mesh.vars
+    x, y = mesh.X
+    a = uw.discretisation.MeshVariable("a_live", mesh, 1, degree=1)
+    a.array[:, 0, 0] = uw.function.evaluate(x + 2 * y, a.coords).reshape(-1)
+
+    mesh.update_lvec()
+    names, isets, _dms = mesh.dm.createFieldDecomposition()
+    g = mesh.dm.getGlobalVec()
+    mesh.dm.localToGlobal(mesh.lvec, g)
+    packed = {}
+    for name, iset in zip(names, isets):
+        sub = g.getSubVector(iset)
+        packed[name] = (sub.min()[1], sub.max()[1])
+        g.restoreSubVector(iset, sub)
+    mesh.dm.restoreGlobalVec(g)
+
+    assert packed["dropped"] == (0.0, 0.0)
+    lo, hi = packed["a_live"]
+    assert lo == pytest.approx(0.0) and hi == pytest.approx(3.0)