diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index d3c67611b..c0e65cf82 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/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: diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 01e1a0582..21d190348 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/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, diff --git a/tests/test_1058_dropped_meshvariable_aux_layout.py b/tests/test_1058_dropped_meshvariable_aux_layout.py new file mode 100644 index 000000000..cdc9d3e6e --- /dev/null +++ b/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)