diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 57c20a563..6bb6b71b0 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -56,6 +56,17 @@ class SolverBaseClass(uw_object): self.petsc_options_prefix = self.name self.petsc_options = PETSc.Options(self.petsc_options_prefix) + self.dm = None + self.dm_hierarchy = [] + self.snes = None + self.petsc_fe_u = None + self.petsc_fe_p = None + + # Solution decomposition caches (IS sets) + self._velocity_is = None + self._pressure_is = None + self._subdict = {} + def _check_expression_meshes(self): """Check that all MeshVariable symbols in solver expressions belong to this solver's mesh. @@ -665,6 +676,20 @@ class SolverBaseClass(uw_object): if verbose and uw.mpi.rank == 0: print(f"Destroy solver DM hierarchy", flush=True) + # Clean up field decomposition cache (Stokes specific but safe to check here) + if hasattr(self, "_pressure_is") and self._pressure_is is not None: + self._pressure_is.destroy() + self._pressure_is = None + if hasattr(self, "_velocity_is") and self._velocity_is is not None: + self._velocity_is.destroy() + self._velocity_is = None + + if hasattr(self, "_subdict") and self._subdict: + for name, (is_set, subdm) in self._subdict.items(): + is_set.destroy() + subdm.destroy() + self._subdict = {} + if hasattr(self, "dm_hierarchy") and self.dm_hierarchy: # Destroys each level — including dm_hierarchy[-1] which # is self.dm — so no separate self.dm.destroy() needed. @@ -1695,7 +1720,6 @@ class SNES_Scalar(SolverBaseClass): print(f"Caution - the mesh quadrature ({mesh.qdegree})is lower") print(f"than {degree} which is required by the {self.name} solver") - self.dm_hierarchy = mesh.clone_dm_hierarchy() self.dm = self.dm_hierarchy[-1] @@ -3391,6 +3415,7 @@ class SNES_MultiComponent(SolverBaseClass): print(f"SNES_MultiComponent ({self.name}): Discretisation does not need to be rebuilt", flush=True) return + # Keep a note of the coordinates that we use for this setup self.mesh_dm_coordinate_hash = mesh_dm_coord_hash cdef PtrContainer ext = self.compiled_extensions @@ -5875,65 +5900,60 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): if verbose and uw.mpi.rank == 0: print(f"SNES Compute Boundary FEM Successfull", flush=True) - # get index set of pressure and velocity to separate solution from localvec # get local section local_section = self.dm.getLocalSection() - # Get the index sets for velocity and pressure fields - # Field numbers (adjust based on your setup) - velocity_field_num = 0 - pressure_field_num = 1 - - # Function to get index set for a field - def get_local_field_is(section, field, unconstrained=False): - """ - This function returns the index set of unconstrained points if True, or all points if False. - """ - pStart, pEnd = section.getChart() - indices = [] - for p in range(pStart, pEnd): - dof = section.getFieldDof(p, field) - if dof > 0: - offset = section.getFieldOffset(p, field) - if not unconstrained and self.Unknowns.p.continuous: - indices.append(offset) - else: - cind = section.getFieldConstraintIndices(p, field) - constrained = set(cind) if cind is not None else set() - for i in range(dof): - if i not in constrained: - index = offset + i - indices.append(index) - is_field = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF) - return is_field - - # Get index sets for pressure (both constrained and unconstrained points) - # we need indexset of pressure field to separate the solution from localvec. - # so we don't care whether a point is constrained by bc or not - pressure_is = get_local_field_is(local_section, pressure_field_num) - - # Get the total number of entries in the local vector - size = self.dm.getLocalVec().getLocalSize() - - # Create a list of all indices - all_indices = set(range(size)) - - # Get indices of the pressure field - pressure_indices = set(pressure_is.getIndices()) - - # Compute the complement for the velocity field - velocity_indices = sorted(list(all_indices - pressure_indices)) - - # Create the index set for velocity - velocity_is = PETSc.IS().createGeneral(velocity_indices, comm=PETSc.COMM_SELF) + # Get index sets for solution decomposition (velocity and pressure) + # We cache these to avoid expensive per-solve Python list allocations and IS creation + if not hasattr(self, "_pressure_is") or self._pressure_is is None: + + # Function to get index set for a field + def get_local_field_is(section, field, unconstrained=False): + """ + This function returns the index set of unconstrained points if True, or all points if False. + """ + pStart, pEnd = section.getChart() + indices = [] + for p in range(pStart, pEnd): + dof = section.getFieldDof(p, field) + if dof > 0: + offset = section.getFieldOffset(p, field) + if not unconstrained and self.Unknowns.p.continuous: + indices.append(offset) + else: + cind = section.getFieldConstraintIndices(p, field) + constrained = set(cind) if cind is not None else set() + for i in range(dof): + if i not in constrained: + index = offset + i + indices.append(index) + is_field = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF) + return is_field + + # Field numbers (adjust based on your setup) + velocity_field_num = 0 + pressure_field_num = 1 + + self._pressure_is = get_local_field_is(local_section, pressure_field_num) + + # Get indices for velocity (complement of pressure) + size = clvec.getLocalSize() + all_indices = set(range(size)) + pressure_indices = set(self._pressure_is.getIndices()) + velocity_indices = sorted(list(all_indices - pressure_indices)) + self._velocity_is = PETSc.IS().createGeneral(velocity_indices, comm=PETSc.COMM_SELF) # Copy solution back into pressure and velocity variables # with self.mesh.access(self.Unknowns.p, self.Unknowns.u): for name, var in self.fields.items(): if name=='velocity': - var.vec.array[:] = clvec.getSubVector(velocity_is).array[:] + subvec = clvec.getSubVector(self._velocity_is) + var.vec.array[:] = subvec.array[:] + clvec.restoreSubVector(self._velocity_is, subvec) elif name=='pressure': - var.vec.array[:] = clvec.getSubVector(pressure_is).array[:] + subvec = clvec.getSubVector(self._pressure_is) + var.vec.array[:] = subvec.array[:] + clvec.restoreSubVector(self._pressure_is, subvec) self.mesh._stale_lvec = True # Sync _gvec so downstream consumers (write, stats) see the result @@ -5943,7 +5963,8 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): if hasattr(target_var, "_canonical_data"): target_var._canonical_data = None - self.dm.restoreGlobalVec(clvec) + # Clean up local vectors + self.dm.restoreLocalVec(clvec) self.dm.restoreGlobalVec(gvec) self._warn_on_divergence() diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index bd699cc19..510787459 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/src/underworld3/discretisation/discretisation_mesh_variables.py @@ -1726,6 +1726,7 @@ def _sync_lvec_to_gvec(self): indexset, subdm = self.mesh.dm.createSubDM(self.field_id) subdm.localToGlobal(self._lvec, self._gvec, addv=False) indexset.destroy() + subdm.destroy() @property def old_data(self) -> numpy.ndarray: diff --git a/src/underworld3/function/dminterpolation_cache.py b/src/underworld3/function/dminterpolation_cache.py index 78a1019b0..e115b8b0b 100644 --- a/src/underworld3/function/dminterpolation_cache.py +++ b/src/underworld3/function/dminterpolation_cache.py @@ -35,12 +35,14 @@ def __init__(self, mesh, name: str = "default"): self.mesh = mesh self.name = name self._cache: Dict[Tuple[int, int], object] = {} # Stores CachedDMInterpolationInfo + self.max_entries = 10 # Statistics self._stats = { 'hits': 0, 'misses': 0, 'invalidations': 0, + 'evictions': 0, 'time_saved': 0.0, 'time_computing': 0.0, } @@ -77,7 +79,8 @@ def get_structure(self, coords: np.ndarray, dofcount: int): if key in self._cache: # CACHE HIT! self._stats['hits'] += 1 - cached_info = self._cache[key] + cached_info = self._cache.pop(key) + self._cache[key] = cached_info # LRU move to end return cached_info else: @@ -104,6 +107,13 @@ def store_structure(self, coords: np.ndarray, dofcount: int, cached_info): coords_hash = self._hash_coords(coords) key = (coords_hash, dofcount) + # LRU management: remove oldest entries if we exceed max_entries + if len(self._cache) >= self.max_entries: + # Remove oldest (first) item + oldest_key = next(iter(self._cache)) + del self._cache[oldest_key] + self._stats['evictions'] += 1 + self._cache[key] = cached_info # Python GC keeps it alive def _hash_coords(self, coords: np.ndarray) -> int: diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index f55f9089e..8b2af5d88 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -272,7 +272,7 @@ def f(self): @f.setter def f(self, value): """Set the source term (handles units and scaling).""" - self.is_setup = False + self._needs_function_rewire = True # Handle UWQuantity with units - enforce "units everywhere" principle if hasattr(value, "value") and hasattr(value, "units"): @@ -478,7 +478,7 @@ def f(self): @f.setter def f(self, value): """Set the source term.""" - self.is_setup = False + self._needs_function_rewire = True self._f = sympy.Matrix((value,)) @property @@ -627,7 +627,7 @@ def __init__( self.theta = theta self._delta_t = expression(R"\Delta t", 0, "Physically motivated timestep") self._storage = sympy.sympify(1) - self.is_setup = False + self._needs_function_rewire = True if DuDt is None: self.Unknowns.DuDt = Eulerian_DDt( @@ -667,7 +667,7 @@ def storage(self): @storage.setter def storage(self, value): - self.is_setup = False + self._needs_function_rewire = True self._storage = sympy.sympify(value) @property @@ -677,7 +677,7 @@ def delta_t(self): @delta_t.setter def delta_t(self, value): - self.is_setup = False + self._needs_function_rewire = True if hasattr(value, "value"): self._delta_t.sym = value.value elif hasattr(value, "magnitude"): @@ -798,7 +798,7 @@ def solve( self.delta_t = timestep if not self.constitutive_model._solver_is_setup: - self.is_setup = False + self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T if not self.is_setup: @@ -958,7 +958,7 @@ def water_content(self): @water_content.setter def water_content(self, value): - self.is_setup = False + self._needs_function_rewire = True self._water_content = sympy.sympify(value) if value is not None else None @property @@ -972,7 +972,7 @@ def capacity(self): @capacity.setter def capacity(self, value): - self.is_setup = False + self._needs_function_rewire = True self._capacity = sympy.sympify(value) @property @@ -1143,7 +1143,7 @@ def __init__( ) # this attrib records if we need to setup the problem (again) - self.is_setup = False + self._needs_function_rewire = True self._constitutive_model = None @@ -1185,7 +1185,7 @@ def set_jacobian_F1_source(self, F1_source, linesearch="cp"): Has no effect when ``F1_source is None``. """ self._F1_jacobian_source = F1_source - self.is_setup = False + self._needs_function_rewire = True if F1_source is not None and linesearch is not None: self.petsc_options["snes_linesearch_type"] = linesearch @@ -1237,7 +1237,9 @@ def solve( zero_init_guess: bool = True, timestep: float = None, _force_setup: bool = False, - verbose=False, + verbose: bool = False, + debug: bool = False, + debug_name: str = None, evalf=False, order=None, picard: int = 0, @@ -1297,19 +1299,19 @@ def solve( order = self._order if _force_setup: - self.is_setup = False + self._needs_function_rewire = True # Re-setup when effective_order changes (DDt history ramp-up) _current_eff_order = self.constitutive_model.effective_order if not hasattr(self, '_prev_effective_order'): self._prev_effective_order = None if _current_eff_order != self._prev_effective_order: - self.is_setup = False + self._needs_function_rewire = True self.constitutive_model._solver_is_setup = False self._prev_effective_order = _current_eff_order if not self.constitutive_model._solver_is_setup: - self.is_setup = False + self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T if not self.is_setup: @@ -1674,7 +1676,7 @@ def bodyforce(self): @bodyforce.setter def bodyforce(self, value): """Set the body force vector (e.g., gravity, buoyancy).""" - self.is_setup = False + self._needs_function_rewire = True if isinstance(value, uw.function.expressions.UWexpression): self._bodyforce.sym = value.sym else: @@ -1715,7 +1717,7 @@ def saddle_preconditioner(self): @saddle_preconditioner.setter def saddle_preconditioner(self, value): """Set the Schur complement preconditioner.""" - self.is_setup = False + self._needs_function_rewire = True symval = sympify(value) self._saddle_preconditioner = symval @@ -1756,7 +1758,7 @@ def penalty(self): @penalty.setter def penalty(self, value): """Set the augmented Lagrangian penalty parameter.""" - self.is_setup = False + self._needs_function_rewire = True self._penalty.sym = value # @property @@ -1991,7 +1993,7 @@ def __init__( verbose, ) - self.is_setup = False + self._needs_function_rewire = True self._smoothing = sympy.sympify(0) self._uw_weighting_function = sympy.sympify(1) self._uw_function = sympy.Matrix([0]) # Default: project zero @@ -2035,7 +2037,7 @@ def smoothing(self): @smoothing.setter def smoothing(self, smoothing_factor): """Set the smoothing regularization parameter.""" - self.is_setup = False + self._needs_function_rewire = True self._smoothing = sympify(smoothing_factor) @property @@ -2046,7 +2048,7 @@ def uw_weighting_function(self): @uw_weighting_function.setter def uw_weighting_function(self, user_uw_function): """Set the weighting function for the projection.""" - self.is_setup = False + self._needs_function_rewire = True self._uw_weighting_function = user_uw_function @@ -2114,7 +2116,7 @@ def __init__( verbose, ) - self.is_setup = False + self._needs_function_rewire = True self._smoothing = 0.0 self._penalty = 0.0 self._uw_weighting_function = 1.0 @@ -2171,7 +2173,7 @@ def smoothing(self): @smoothing.setter def smoothing(self, smoothing_factor): """Set the smoothing regularization parameter.""" - self.is_setup = False + self._needs_function_rewire = True self._smoothing = sympify(smoothing_factor) @property @@ -2182,7 +2184,7 @@ def penalty(self): @penalty.setter def penalty(self, value): """Set the divergence penalty parameter.""" - self.is_setup = False + self._needs_function_rewire = True symval = sympify(value) self._penalty = symval @@ -2194,7 +2196,7 @@ def uw_weighting_function(self): @uw_weighting_function.setter def uw_weighting_function(self, user_uw_function): """Set the weighting function for the projection.""" - self.is_setup = False + self._needs_function_rewire = True self._uw_weighting_function = user_uw_function @@ -2342,7 +2344,7 @@ def uw_scalar_function(self): @uw_scalar_function.setter def uw_scalar_function(self, user_uw_function): """Set the scalar component function for current tensor element.""" - self.is_setup = False + self._needs_function_rewire = True self._uw_scalar_function = user_uw_function @@ -2401,7 +2403,7 @@ def __init__( verbose=verbose, ) - self.is_setup = False + self._needs_function_rewire = True self._smoothing = sympify(0) self._uw_weighting_function = sympify(1) self._uw_function = sympy.zeros(1, self._n_components) @@ -2431,7 +2433,7 @@ def smoothing(self): @smoothing.setter def smoothing(self, value): - self.is_setup = False + self._needs_function_rewire = True self._smoothing = sympify(value) @property @@ -2441,7 +2443,7 @@ def uw_weighting_function(self): @uw_weighting_function.setter def uw_weighting_function(self, value): - self.is_setup = False + self._needs_function_rewire = True self._uw_weighting_function = value @@ -2555,7 +2557,7 @@ def __init__( # These are unique to the advection solver self._delta_t = expression(R"\Delta t", 0, "Physically motivated timestep") - self.is_setup = False + self._needs_function_rewire = True self.restore_points_to_domain_func = restore_points_func ### Setup the history terms ... This version should not build anything @@ -2670,7 +2672,7 @@ def f(self): @f.setter def f(self, value): """Set the volumetric source term.""" - self.is_setup = False + self._needs_function_rewire = True self._f = sympy.Matrix((value,)) @property @@ -2730,7 +2732,17 @@ def delta_t(self): @delta_t.setter def delta_t(self, value): """Set the timestep (handles unit conversion if provided).""" - self.is_setup = False + # Note: comparison must handle potential UWexpressions / UWQuantities + # Use .data or float() to get numeric values for stable comparison + try: + old_dt = float(self._delta_t.data) + new_dt = float(value.data) if hasattr(value, 'data') else float(value) + if np.isclose(old_dt, new_dt, rtol=1e-12, atol=1e-15): + return + except: + pass + + self._needs_function_rewire = True # Handle Pint Quantities with time dimensions if hasattr(value, "dimensionality"): @@ -2931,10 +2943,10 @@ def solve( self.delta_t = timestep # this will force an initialisation because the functions need to be updated if _force_setup: - self.is_setup = False + self._needs_function_rewire = True if not self.constitutive_model._solver_is_setup: - self.is_setup = False + self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T if not self.is_setup: @@ -3052,7 +3064,7 @@ def __init__( # These are unique to the advection solver self._delta_t = expression(R"\Delta t", 0, "Physically motivated timestep") - self.is_setup = False + self._needs_function_rewire = True ### Setup the history terms ... This version should not build anything ### by default - it's the template / skeleton @@ -3152,7 +3164,7 @@ def f(self): @f.setter def f(self, value): """Set the volumetric source term.""" - self.is_setup = False + self._needs_function_rewire = True self._f = sympy.Matrix((value,)) @property @@ -3163,7 +3175,17 @@ def delta_t(self): @delta_t.setter def delta_t(self, value): """Set the timestep (handles unit conversion if provided).""" - self.is_setup = False + # Note: comparison must handle potential UWexpressions / UWQuantities + # Use .data or float() to get numeric values for stable comparison + try: + old_dt = float(self._delta_t.data) + new_dt = float(value.data) if hasattr(value, 'data') else float(value) + if np.isclose(old_dt, new_dt, rtol=1e-12, atol=1e-15): + return + except: + pass + + self._needs_function_rewire = True # Handle Pint Quantities with time dimensions if hasattr(value, "dimensionality"): @@ -3307,10 +3329,10 @@ def solve( self.delta_t = timestep # this will force an initialisation because the functions need to be updated if _force_setup: - self.is_setup = False + self._needs_function_rewire = True if not self.constitutive_model._solver_is_setup: - self.is_setup = False + self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T # self._flux = self.constitutive_model.flux.T # self._flux_star = self._flux.copy() @@ -3450,7 +3472,7 @@ def __init__( # These are unique to the advection solver self._delta_t = expression(r"\Delta t", sympy.oo, "Navier-Stokes timestep") - self.is_setup = False + self._needs_function_rewire = True self._rho = expression(R"{\uprho}", rho, "Density") self._first_solve = True @@ -3594,7 +3616,7 @@ def delta_t(self): @delta_t.setter def delta_t(self, value): """Set the timestep value.""" - self.is_setup = False + self._needs_function_rewire = True self._delta_t.sym = value @property @@ -3605,7 +3627,7 @@ def rho(self): @rho.setter def rho(self, value): """Set the fluid density.""" - self.is_setup = False + self._needs_function_rewire = True self._rho.sym = value @property @@ -3616,7 +3638,7 @@ def f(self): @f.setter def f(self, value): """Set the volumetric source term.""" - self.is_setup = False + self._needs_function_rewire = True self._f = sympy.Matrix((value,)) @property @@ -3675,7 +3697,7 @@ def bodyforce(self): @bodyforce.setter def bodyforce(self, value): """Set the body force vector.""" - self.is_setup = False + self._needs_function_rewire = True self._bodyforce = self.mesh.vector.to_matrix(value) @property @@ -3686,7 +3708,7 @@ def saddle_preconditioner(self): @saddle_preconditioner.setter def saddle_preconditioner(self, value): """Set the Schur complement preconditioner.""" - self.is_setup = False + self._needs_function_rewire = True symval = sympify(value) self._saddle_preconditioner = symval @@ -3698,7 +3720,7 @@ def penalty(self): @penalty.setter def penalty(self, value): """Set the augmented Lagrangian penalty parameter.""" - self.is_setup = False + self._needs_function_rewire = True self._penalty.sym = value @timing.routine_timer_decorator @@ -3732,10 +3754,10 @@ def solve( self.delta_t = timestep # this will force an initialisation because the functions need to be updated if _force_setup: - self.is_setup = False + self._needs_function_rewire = True if not self.constitutive_model._solver_is_setup: - self.is_setup = False + self._needs_function_rewire = True self.DFDt.psi_fn = self.constitutive_model.flux.T if not self.is_setup: diff --git a/src/underworld3/utilities/_api_tools.py b/src/underworld3/utilities/_api_tools.py index d4649aaba..75a81d78d 100644 --- a/src/underworld3/utilities/_api_tools.py +++ b/src/underworld3/utilities/_api_tools.py @@ -97,14 +97,6 @@ def __get__(self, obj, objtype=None): def __set__(self, obj, value): """Set the value, with automatic unwrapping.""" - # Mark solver as needing setup when property changes - if hasattr(obj, "is_setup"): - obj.is_setup = False - - # Check None constraint - if value is None and not self.allow_none: - raise ValueError(f"Cannot set {self.attr_name[1:]} to None") - # Auto-unwrap objects with _sympify_() protocol if value is not None and hasattr(value, "_sympify_"): value = value._sympify_() @@ -117,6 +109,15 @@ def __set__(self, obj, value): if not isinstance(value, sympy.matrices.MatrixBase): value = sympy.Matrix([value]) + # Mark solver as needing setup ONLY when property actually changes + if hasattr(obj, "is_setup"): + old_value = getattr(obj, self.attr_name, None) + # Use 'is not' check first for speed + if old_value is not value: + # Check for mathematical inequality (handles SymPy expressions) + if old_value != value: + obj.is_setup = False + # Store the value setattr(obj, self.attr_name, value) diff --git a/tests/test_0006_memory_leak.py b/tests/test_0006_memory_leak.py new file mode 100644 index 000000000..ea294b855 --- /dev/null +++ b/tests/test_0006_memory_leak.py @@ -0,0 +1,98 @@ +import underworld3 as uw +import numpy as np +import os +import resource +import sympy +import gc +from mpi4py import MPI +import pytest + +def get_memory_usage(): + try: + if os.uname().sysname == 'Darwin': + rss = int(os.popen('ps -p %d -o rss=' % os.getpid()).read()) + return rss / 1024.0 # MiB + else: + with open('/proc/self/status') as f: + for line in f: + if line.startswith('VmRSS:'): + return int(line.split()[1]) / 1024.0 # MiB + except: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0 + +def test_stokes_advdiff_memory_leak(): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + + # Parameters + res = 32 + n_steps = 100 + warmup_step = 30 + + # Setup mesh + mesh = uw.meshing.StructuredQuadBox(elementRes=(res, res)) + mesh._dminterpolation_cache.max_entries = 10 # CAP CACHE + + # Variables + v = uw.discretisation.MeshVariable("u", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + + # Swarm for advection + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=2) + + # Stokes + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.viscosity = 1.0 + stokes.bodyforce = sympy.Matrix([0, T.sym[0]]) + stokes.add_dirichlet_bc([0.0, 0.0], "Top") + stokes.add_dirichlet_bc([0.0, 0.0], "Bottom") + stokes.add_dirichlet_bc([0.0, sympy.oo], "Left") + stokes.add_dirichlet_bc([0.0, sympy.oo], "Right") + + # AdvDiff + advdiff = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=v) + advdiff.constitutive_model = uw.constitutive_models.DiffusionModel + advdiff.constitutive_model.Parameters.diffusivity = 1.0 + advdiff.add_dirichlet_bc([0.0], "Top") + advdiff.add_dirichlet_bc([1.0], "Bottom") + + # Initial T + T.data[:] = 0.0 + with mesh.access(T): + T.data[:, 0] = 1.0 - mesh.data[:, 1] + + mem_warmup = 0.0 + + for i in range(n_steps): + stokes.solve() + dt = 0.001 + advdiff.solve(timestep=dt) + + with swarm.access(swarm): + swarm.data[...] = swarm.data[...] + 0.01 * np.random.rand(*swarm.data.shape) + swarm.data[...] = np.clip(swarm.data[...], 0, 1) + + v_at_swarm = uw.function.evaluate(v.sym, swarm.data) + + gc.collect() + + if i == warmup_step: + local_mem = get_memory_usage() + mem_warmup = comm.reduce(local_mem, op=MPI.SUM, root=0) + + # Final memory check + local_mem = get_memory_usage() + mem_final = comm.reduce(local_mem, op=MPI.SUM, root=0) + + if rank == 0: + growth = mem_final - mem_warmup + print(f"\nMemory growth from step {warmup_step} to {n_steps}: {growth:.2f} MiB") + + # We allow for some small growth (fragmentation, PETSc pool expansion) + # but catch the large unbounded leak (which was ~180 MiB per 100 steps). + # Stable growth should be < 20 MiB over the last 70 steps for this resolution. + assert growth < 30.0, f"Significant memory leak detected: {growth:.2f} MiB growth after warm-up."