diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index a40468129..a2094c400 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -66,10 +66,6 @@ class SwarmType(Enum): DMSWARM_PIC = 1 -# We can grab this type from the PETSc module -# SwarmPICLayout has been moved to pic_swarm.py - - # Note - much of the setup is necessarily the same as the MeshVariable # and the duplication should be removed. @@ -106,8 +102,8 @@ class SwarmVariable(DimensionalityMixin, MathematicalMixin, Stateful, uw_object) varsymbol : str, optional LaTeX symbol for display. Defaults to ``name``. rebuild_on_cycle : bool, default=True - If True, rebuild the proxy when particles cycle through periodic - boundaries. Recommended for continuous fields. + No effect. Retained for backward compatibility with the removed + particle-recycling (streak swarm) feature. units : str or pint.Unit, optional Physical units for this variable (e.g., 'kelvin', 'Pa'). Requires reference quantities to be set on the model. @@ -335,7 +331,8 @@ def __init__( self._nn_proxy = _nn_proxy self._create_proxy_variable() - # recycle swarm + # Inert: kept for backward compatibility with the removed + # particle-recycling (streak swarm) feature. self._rebuild_on_cycle = rebuild_on_cycle self._register = _register @@ -683,49 +680,44 @@ def __setitem__(self, key, value): reshaped_data = modified_data.reshape(-1, self.parent.num_components) self.parent.data[:] = reshaped_data - # Forward common array methods + # Reduction methods follow the MeshVariable array-view contract: + # scalar (float) for single-component variables, per-component + # tuple for multi-component variables (LE-07 / BF-11). + # + # NOTE: these are simple arithmetic reductions over the particle + # values. Swarm particles are generally non-uniformly distributed + # in space, so mean()/std() only APPROXIMATE the spatial + # statistics — use mesh integrals of the proxy field for + # spatially-accurate statistics. + + def _per_component_reduction(self, reduction): + data = self._get_array_data() + if self.parent.num_components == 1: + return float(reduction(data)) + flat = np.asarray(data).reshape(data.shape[0], -1) + return tuple( + float(reduction(flat[:, i])) for i in range(self.parent.num_components) + ) + def max(self): - return self._get_array_data().max() + """Maximum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.max) def min(self): - return self._get_array_data().min() + """Minimum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.min) def mean(self): - """ - Compute mean value of swarm particles. - - ⚠️ WARNING: This computes a simple arithmetic mean of the particle values. - Since swarm particles are typically non-uniformly distributed in space, - this is an APPROXIMATION of the spatial mean. For accurate spatial - statistics, consider using integration via swarm proxy variables or - computing mesh integrals of the proxy field. - - Returns - ------- - float or tuple - Mean value (float for scalar variables, tuple for multi-component) - """ - return self._get_array_data().mean() + """Arithmetic particle mean (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.mean) def sum(self): - return self._get_array_data().sum() + """Sum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.sum) def std(self): - """ - Compute standard deviation of swarm particles. - - ⚠️ WARNING: This computes a simple numpy std of the particle values. - Since swarm particles are typically non-uniformly distributed in space, - this is an APPROXIMATION of the spatial standard deviation. For accurate - spatial statistics, consider using integration via swarm proxy variables - or computing mesh integrals of the proxy field. - - Returns - ------- - float or tuple - Standard deviation (float for scalar variables, tuple for multi-component) - """ - return self._get_array_data().std() + """Arithmetic particle standard deviation (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.std) @property def shape(self): @@ -898,49 +890,45 @@ def __setitem__(self, key, value): packed_data = self.parent._pack_array_to_data_format(modified_data) self.parent.data[:] = packed_data - # Forward common array methods + # Reduction methods follow the MeshVariable array-view contract: + # scalar (float) for single-component variables, per-component + # tuple for multi-component variables (LE-07 / BF-11). Components + # are ordered as in the flat canonical layout. + # + # NOTE: these are simple arithmetic reductions over the particle + # values. Swarm particles are generally non-uniformly distributed + # in space, so mean()/std() only APPROXIMATE the spatial + # statistics — use mesh integrals of the proxy field for + # spatially-accurate statistics. + + def _per_component_reduction(self, reduction): + data = self._get_array_data() + if self.parent.num_components == 1: + return float(reduction(data)) + flat = np.asarray(data).reshape(data.shape[0], -1) + return tuple( + float(reduction(flat[:, i])) for i in range(self.parent.num_components) + ) + def max(self): - return self._get_array_data().max() + """Maximum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.max) def min(self): - return self._get_array_data().min() + """Minimum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.min) def mean(self): - """ - Compute mean value of swarm particles. - - ⚠️ WARNING: This computes a simple arithmetic mean of the particle values. - Since swarm particles are typically non-uniformly distributed in space, - this is an APPROXIMATION of the spatial mean. For accurate spatial - statistics, consider using integration via swarm proxy variables or - computing mesh integrals of the proxy field. - - Returns - ------- - float or tuple - Mean value (float for scalar variables, tuple for multi-component) - """ - return self._get_array_data().mean() + """Arithmetic particle mean (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.mean) def sum(self): - return self._get_array_data().sum() + """Sum (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.sum) def std(self): - """ - Compute standard deviation of swarm particles. - - ⚠️ WARNING: This computes a simple numpy std of the particle values. - Since swarm particles are typically non-uniformly distributed in space, - this is an APPROXIMATION of the spatial standard deviation. For accurate - spatial statistics, consider using integration via swarm proxy variables - or computing mesh integrals of the proxy field. - - Returns - ------- - float or tuple - Standard deviation (float for scalar variables, tuple for multi-component) - """ - return self._get_array_data().std() + """Arithmetic particle standard deviation (float for scalars, per-component tuple otherwise).""" + return self._per_component_reduction(np.std) @property def shape(self): @@ -1561,17 +1549,9 @@ def rbf_interpolate(self, new_coords, verbose=False, nnn=None): nnn = data_size[0] # Use direct PETSc access to avoid callback circular dependency - if self.swarm.recycle_rate > 1: - not_remeshed = self.swarm._remeshed.data[:, 0] != 0 - D = raw_data[not_remeshed].copy() - - kdt = uw.kdtree.KDTree(self.swarm._particle_coordinates.data[not_remeshed, :]) - values = kdt.rbf_interpolator_local(new_coords, D, nnn, 2, verbose) - else: - D = raw_data.copy() - # Use cached KDTree for standard swarms - kdt = self.swarm._get_kdtree() - values = kdt.rbf_interpolator_local(new_coords, D, nnn, 2, verbose) + D = raw_data.copy() + kdt = self.swarm._get_kdtree() + values = kdt.rbf_interpolator_local(new_coords, D, nnn, 2, verbose) return values @@ -2617,15 +2597,6 @@ def _update_proxy_variables(self): return -## Import PIC-related classes from separate module to maintain compatibility -# from .pic_swarm import PICSwarm, NodalPointPICSwarm, SwarmPICLayout - -## This should be the basic swarm, and we can then create a sub-class that will -## be a PIC swarm - -# PICSwarm and NodalPointPICSwarm classes have been moved to pic_swarm.py - - ## New - Basic Swarm (no PIC skillz) ## What is missing: ## - no celldm @@ -2652,9 +2623,10 @@ class Swarm(Stateful, uw_object): The mesh object that defines the computational domain for particle operations. Particles will be associated with this mesh for spatial queries and operations. recycle_rate : int, optional - Rate at which particles are recycled for streak management. If > 1, enables - streak particle functionality where particles are duplicated and tracked - across multiple cycles. Default is 0 (no recycling). + Particle recycling (streak swarms) is NOT implemented: values > 1 + raise ``NotImplementedError``. The parameter is retained so that + existing calls passing the default (0 or 1, meaning no recycling) + keep working. verbose : bool, optional Enable verbose output for debugging and monitoring particle operations. Default is False. @@ -2668,11 +2640,6 @@ class Swarm(Stateful, uw_object): >>> swarm = uw.swarm.Swarm(mesh=mesh) >>> swarm.populate(fill_param=2) - Create a streak swarm with recycling: - - >>> streak_swarm = uw.swarm.Swarm(mesh=mesh, recycle_rate=5) - >>> streak_swarm.populate(fill_param=1) - Add custom particle data: >>> temperature = swarm.add_variable("temperature", 1) @@ -2695,6 +2662,16 @@ class Swarm(Stateful, uw_object): @timing.routine_timer_decorator def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True): + # Particle recycling (streak swarms) was excised in 2026-07: the + # machinery had been broken (NameError) and untested for some time + # (audit finding SWARM-08). Refuse rather than crash later. + if recycle_rate > 1: + raise NotImplementedError( + "Particle recycling / streak swarms (recycle_rate > 1) are not " + "implemented. Construct the Swarm without recycle_rate and manage " + "particle re-seeding explicitly (e.g. add_particles_with_coordinates)." + ) + Swarm.instances += 1 self.verbose = verbose @@ -2750,9 +2727,9 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True): #### - # Is the swarm a streak-swarm ? + # Retained attribute: always 0 or 1 (no recycling) — see the + # NotImplementedError guard above. self.recycle_rate = recycle_rate - self.cycle = 0 # dictionary for variables # Using WeakValueDictionary to prevent circular references @@ -2791,21 +2768,6 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True): rebuild_on_cycle=False, ) - # This is for swarm streak management: - # add variable to hold swarm origins - - if self.recycle_rate > 1: - - self._remeshed = uw.swarm.SwarmVariable( - "DMSwarm_remeshed", - self, - 1, - dtype=int, - _register=True, - _proxy=False, - rebuild_on_cycle=False, - ) - self._X0_uninitialised = True self._index = None self._nnmapdict = {} @@ -3641,49 +3603,6 @@ def populate( self.dm.restoreField("DMSwarmPIC_coor") self.dm.restoreField("DMSwarm_rank") - if self.recycle_rate > 1: - with self.access(): - # This is a mesh-local quantity, so let's just - # store it on the mesh in an ad_hoc fashion for now - - self.mesh.particle_X_orig = self._particle_coordinates.data.copy() - - with self.access(): - swarm_orig_size = self.local_size - all_local_coords = np.vstack( - (self._particle_coordinates.data,) * (self.recycle_rate) - ) - - swarm_new_size = all_local_coords.shape[0] - - self.dm.addNPoints(swarm_new_size - swarm_orig_size) - - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.cdim)) - - # Compute perturbation - extract magnitude if coordinates have units - # numpy.array(..., dtype=float64) forces conversion to plain array - coord_data = np.array(all_local_coords, dtype=np.float64) - search_lengths = np.array(self.mesh._search_lengths[all_local_cells], dtype=np.float64) - - perturbation = ( - (0.33 / (1 + fill_param)) - * (np.random.random(size=coord_data.shape) - 0.5) - * 0.00001 - * search_lengths # typical cell size - ) - - # Add perturbation (coords array stores dimensionless values) - coords[...] = coord_data + perturbation - - self.dm.restoreField("DMSwarmPIC_coor") - - ## Now set the cycle values - - with self.access(self._remeshed): - for i in range(0, self.recycle_rate): - offset = swarm_orig_size * i - self._remeshed.data[offset::, 0] = i - # Invalidate cached data — the swarm was just given its particles. # Any canonical `.data` array created before populate() (legitimate: # variables must be created first) is sized for the empty swarm and @@ -3956,13 +3875,6 @@ def add_particles_with_coordinates(self, coordinatesArray) -> int: self.dm.restoreField("DMSwarm_rank") self.dm.restoreField("DMSwarmPIC_coor") - # Here we update the swarm cycle values as required - - if self.recycle_rate > 1: - with self.access(self._remeshed): - # self._Xorig.data[...] = coordinatesArray - self._remeshed.data[...] = 0 - self.dm.migrate(remove_sent_points=True) # Invalidate cached data — particle count changed after addNPoints + migrate @@ -4053,13 +3965,6 @@ def add_particles_with_global_coordinates( self.dm.restoreField("DMSwarm_rank") self.dm.restoreField("DMSwarmPIC_coor") - # Here we update the swarm cycle values as required - - if self.recycle_rate > 1: - with self.access(self._remeshed): - # self._Xorig.data[...] = globalCoordinatesArray - self._remeshed.data[...] = 0 - # Invalidate cached data — the particle count changed via addNPoints # (mirrors add_particles_with_coordinates). This must not be left to # migrate(): with migrate=False nothing else invalidates, and every @@ -4146,8 +4051,14 @@ def save( else: # Sequential fallback: rank 0 creates the file and writes its slab, # then each higher rank appends in turn. - - points_data_copy = self.points[:].copy() + # + # MODEL-UNIT coordinates, exactly like the parallel branch above: + # the deprecated `self.points` used here previously applied the + # model length scale, so sequential checkpoints differed from + # parallel ones by that factor and could not round-trip through + # read_timestep, which re-inserts raw model-unit coordinates + # (SWARM-19 / BF-17). + points_data_copy = self._particle_coordinates.data[:].copy() local_n = points_data_copy.shape[0] if comm.rank == 0: @@ -4520,7 +4431,7 @@ def snapshot_payload(self) -> dict: Captured: per-rank particle coordinates (from ``DMSwarmPIC_coor``) and every user swarm-variable's data array. PETSc-internal variables (``DMSwarmPIC_coor``, - ``DMSwarm_X0``, ``DMSwarm_remeshed``) are excluded — their + ``DMSwarm_X0``) are excluded — their contents either come from the captured coords or are regenerated on the next solve. """ @@ -4997,12 +4908,17 @@ def advection( # Use internal model-unit coordinates directly (no conversion needed) v_at_Vpts = np.zeros_like(self._particle_coordinates.data[...]) - # First evaluate the velocity at the particle locations - # (this is a local operation) + # First evaluate the velocity at the launch points. This must + # be a GLOBAL evaluation: no migration happens inside the + # substep loop (deferred migration is suspended above, so + # arrays keep a stable row order), which means from substep 2 + # onward a particle can sit outside this rank's domain — a + # rank-local evaluation silently extrapolates wrong values + # for it (SWARM-16 / BF-16). - v_at_Vpts[...] = uw.function.evaluate(V_fn_matrix, self._particle_coordinates.data)[ - :, 0, : - ] + v_at_Vpts[...] = uw.function.global_evaluate( + V_fn_matrix, self._particle_coordinates.data + )[:, 0, :] mid_pt_coords = ( self._particle_coordinates.data[...] @@ -5061,97 +4977,6 @@ def advection( ## End of substepping loop self._deferred_migration_suspended = False - ## Cycling of the swarm is a cheap and cheerful version of population control for particles. It turns the - ## swarm into a streak-swarm where particles are Lagrangian for a number of steps and then reset to their - ## original location. - - if self.recycle_rate > 1: - # Restore particles which have cycle == cycle rate (use >= just in case) - - # Remove remesh points and recreate a new set at the mesh-local - # locations that we already have stored. - - with self.access(self._particle_coordinates, self._remeshed): - remeshed = self._remeshed.data[:, 0] == 0 - # This is one way to do it ... we can do this better though - self.data[remeshed, 0] = 1.0e100 - - swarm_size = self.dm.getLocalSize() - - num_remeshed_points = self.mesh.particle_X_orig.shape[0] - - self.dm.addNPoints(num_remeshed_points) - - # Informational: remesh just re-injected particles. - self._population_generation += 1 - - ## cellid = self.dm.getField("DMSwarm_cellid") - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.cdim)) - rmsh = self.dm.getField("DMSwarm_remeshed") - - # print(f"cellid -> {cellid.shape}") - # print(f"particle coords -> {coords.shape}") - # print(f"remeshed points -> {num_remeshed_points}") - - # Compute perturbation - extract magnitude if coordinates have units - # numpy.array(..., dtype=float64) forces conversion to plain array - coord_data = np.array(self.mesh.particle_X_orig[:, :], dtype=np.float64) - radii_data = np.array(self.mesh._radii[cellid[swarm_size::]], dtype=np.float64) - - perturbation = 0.00001 * ( - (0.33 / (1 + self.fill_param)) - * (np.random.random(size=(num_remeshed_points, self.dim)) - 0.5) - * radii_data.reshape(-1, 1) - ) - - # Add perturbation (coords array stores dimensionless values) - coords[swarm_size::] = coord_data + perturbation - ## cellid[swarm_size::] = self.mesh.particle_CellID_orig[:, 0] - rmsh[swarm_size::] = 0 - - # self.dm.restoreField("DMSwarm_cellid") - self.dm.restoreField("DMSwarmPIC_coor") - self.dm.restoreField("DMSwarm_remeshed") - - # when we let this go, the particles may be re-distributed to - # other processors, and we will need to rebuild the remeshed - # array before trying to compute / assign values to variables - - for swarmVar in self.vars.values(): - if swarmVar._rebuild_on_cycle: - with self.access(swarmVar): - if swarmVar.dtype is int: - nnn = 1 - else: - nnn = self.mesh.dim + 1 # 3 for triangles, 4 for tets ... - - interpolated_values = ( - swarmVar.rbf_interpolate(self.mesh.particle_X_orig, nnn=nnn) - # swarmVar._meshVar.fn, self.mesh.particle_X_orig - # ) - ).astype(swarmVar.dtype) - - swarmVar.data[swarm_size::] = interpolated_values - - ## - ## Determine RANK - ## - - # Migrate will already have been called by the access manager. - # Maybe we should hash the local particle coords to make this - # a little more user-friendly - - # self.dm.migrate(remove_sent_points=True) - - with self.access(self._remeshed): - self._remeshed.data[...] = np.mod(self._remeshed.data[...] - 1, self.recycle_rate) - - self.cycle += 1 - - ## End of cycle_swarm loop - # - # - # Re-route particles to their owning ranks and remove any that # have genuinely left the domain. Use the default max_its so that # boundary particles whose owner is the 2nd/3rd-closest centroid @@ -5188,8 +5013,14 @@ def estimate_dt(self, V_fn): # Plain UWQuantity without units context - use magnitude vel = vel.magnitude - # Ensure vel is a plain numpy array + # Ensure vel is a plain numpy array in flat (n_particles, dim) form. + # evaluate() returns matrix-shaped (n, 1, dim) arrays; indexing that + # shape as vel[:, 1] hit the size-1 axis and the swallowed IndexError + # made estimate_dt() return None for every non-trivial velocity — + # silently disabling advection's step_limit substepping (BF-16). vel = np.asarray(vel) + if vel.ndim == 3: + vel = vel.reshape(vel.shape[0], -1) try: magvel_squared = vel[:, 0] ** 2 + vel[:, 1] ** 2 @@ -5199,6 +5030,9 @@ def estimate_dt(self, V_fn): max_magvel = math.sqrt(magvel_squared.max()) except (ValueError, IndexError): + # Sanctioned: a rank holding zero particles has an empty vel + # array (its .max() raises); it contributes zero to the + # global maximum below. max_magvel = 0.0 from mpi4py import MPI @@ -5220,6 +5054,12 @@ def estimate_dt(self, V_fn): class NodalPointSwarm(Swarm): r"""BASIC_Swarm with particles located at the coordinate points of a meshVariable + .. deprecated:: 2026-07 + ``NodalPointSwarm`` is deprecated and will be removed in the next + release cycle. The semi-Lagrangian history managers in + ``uw.systems.ddt`` no longer use it and there are no remaining + internal callers. + The swarmVariable `X0` is defined so that the particles can "snap back" to their original locations after they have been moved. @@ -5231,13 +5071,26 @@ def __init__( trackedVariable: uw.discretisation.MeshVariable, verbose=False, ): + import warnings + + warnings.warn( + "NodalPointSwarm is deprecated and will be removed in the next " + "release cycle. Use the semi-Lagrangian history managers in " + "uw.systems.ddt, or a plain Swarm populated at the variable's " + "coordinates.", + DeprecationWarning, + stacklevel=2, + ) + self.trackedVariable = trackedVariable self.swarmVariable = None mesh = trackedVariable.mesh - # Set up a standard swarm - super().__init__(mesh, verbose, clip_to_mesh=False) + # Keyword-explicit: Swarm.__init__ takes recycle_rate as its second + # positional parameter, so a positional `verbose` here used to land + # in recycle_rate and be silently discarded (SWARM-11). + super().__init__(mesh, verbose=verbose, clip_to_mesh=False) nswarm = self diff --git a/src/underworld3/swarms/pic_swarm.py b/src/underworld3/swarms/pic_swarm.py deleted file mode 100644 index 20a0462f2..000000000 --- a/src/underworld3/swarms/pic_swarm.py +++ /dev/null @@ -1,1534 +0,0 @@ -from posixpath import pardir -import petsc4py.PETSc as PETSc - -import numpy as np -import sympy -import h5py -import os -import warnings -from typing import Optional, Tuple - -import underworld3 as uw -from underworld3.utilities._api_tools import Stateful -from underworld3.utilities._api_tools import uw_object - -import underworld3.timing as timing - -comm = uw.mpi.comm - -from enum import Enum - -# Import necessary classes from swarm module -from .swarm import SwarmType, SwarmVariable - - -class SwarmPICLayout(Enum): - """ - Particle population fill type: - - SwarmPICLayout.REGULAR defines points on a regular ijk mesh. Supported by simplex cell types only. - SwarmPICLayout.GAUSS defines points using an npoint Gauss-Legendre tensor product quadrature rule. - SwarmPICLayout.SUBDIVISION defines points on the centroid of a sub-divided reference cell. - """ - - REGULAR = 0 - GAUSS = 1 - SUBDIVISION = 2 - - -class PICSwarm(Stateful, uw_object): - """ - Particle swarm implementation with automatic mesh-particle interactions. - - The `Swarm` class is Underworld's primary particle management system, built on PETSc's - DMSWARM_PIC type. It provides automatic particle migration, mesh-particle connectivity, - and streamlined particle operations for Lagrangian particle tracking and data storage. - - Differences from UW Swarm: - - **Mesh Integration**: Built-in particle-in-cell (PIC) connectivity with automatic cell tracking - - **Migration**: Uses the standard PETSc strategy for migration which depends on the DM type. This requires - calculation of cell-relationships each time the coordinates are updated and particles that are not found will - be deleted. - - - Parameters - ---------- - mesh : uw.discretisation.Mesh - The mesh object that defines the computational domain. Particles will be - automatically associated with mesh cells for efficient spatial operations. - recycle_rate : int, optional - Rate at which particles are recycled for streak management. If > 1, enables - streak particle functionality where particles are duplicated and tracked - across multiple cycles. Default is 0 (no recycling). - verbose : bool, optional - Enable verbose output for debugging and monitoring particle operations. - Default is False. - - Attributes - ---------- - mesh : uw.discretisation.Mesh - Reference to the associated mesh object. - dim : int - Spatial dimension of the mesh (2D or 3D). - cdim : int - Coordinate dimension of the mesh. - data : numpy.ndarray - Direct access to particle coordinate data. - particle_coordinates : SwarmVariable - SwarmVariable containing particle coordinate information (auto-created). - particle_cellid : SwarmVariable - SwarmVariable containing particle cell ID information (auto-created). - recycle_rate : int - Current recycle rate for streak management. - cycle : int - Current cycle number for streak particles. - - Methods - ------- - populate_petsc(fill_param=1) - Populate swarm using PETSc's built-in particle generation. - populate(fill_param=1, layout=SwarmPICLayout.GAUSS) - Populate the swarm with particles using specified layout. - add_particles_with_coordinates(coords) - Add new particles at specified coordinate locations. - add_variable(name, size, dtype=float) - Add a new variable to track additional particle properties. - save(filename, meshUnits=1.0, swarmUnits=1.0, units="dimensionless") - Save swarm data to file. - read_timestep(filename, step_name, outputPath="./output/") - Read swarm data from a specific timestep file. - advection(V_fn, delta_t, evalf=False, corrector=True, restore_points_func=None) - Advect particles using a velocity field with automatic migration. - estimate_dt(V_fn, dt_min=1.0e-15, dt_max=1.0) - Estimate appropriate timestep for particle advection. - - Examples - -------- - Create a standard swarm with automatic features: - - >>> import underworld3 as uw - >>> mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0,0), maxCoords=(1,1)) - >>> swarm = uw.swarm.Swarm(mesh=mesh) - >>> swarm.populate(fill_param=2, layout=uw.swarm.SwarmPICLayout.GAUSS) - - Access automatic coordinate and cell ID fields: - - >>> coords = swarm._particle_coordinates.data - >>> cell_ids = swarm.particle_cellid.data - - Create a streak swarm with recycling: - - >>> streak_swarm = uw.swarm.Swarm(mesh=mesh, recycle_rate=5) - >>> streak_swarm.populate(fill_param=1) - - Add custom particle data and perform advection: - - >>> temperature = swarm.add_variable("temperature", 1) - >>> velocity_field = mesh.add_variable("velocity", mesh.dim) - >>> # ... set up velocity field ... - >>> swarm.advection(velocity_field.sym, delta_t=0.01) # Automatic migration - - Notes - ----- - - Particle migration occurs automatically during advection operations - - Coordinate and cell ID fields are created and managed automatically at the - PETSc level - - - """ - - instances = 0 - - @timing.routine_timer_decorator - def __init__(self, mesh, recycle_rate=0, verbose=False): - PICSwarm.instances += 1 - - self.celldm = mesh.dm.clone() - - self.verbose = verbose - self._mesh = mesh - self.dim = mesh.dim - self.cdim = mesh.cdim - self.dm = PETSc.DMSwarm().create() - self.dm.setDimension(self.dim) - self.dm.setType(SwarmType.DMSWARM_PIC.value) - self.dm.setCellDM(self.celldm) - self._data = None - - # Is the swarm a streak-swarm ? - self.recycle_rate = recycle_rate - self.cycle = 0 - - # dictionary for variables - - # import weakref (not helpful as garbage collection does not remove the fields from the DM) - # self._vars = weakref.WeakValueDictionary() - self._vars = {} - - # add variable to handle particle coords - predefined by DMSwarm, expose to UW - self._coord_var = SwarmVariable( - "DMSwarmPIC_coor", - self, - self.cdim, - dtype=float, - _register=False, - _proxy=False, - rebuild_on_cycle=False, - ) - - # add variable to handle particle cell id - predefined by DMSwarm, expose to UW - self._cellid_var = SwarmVariable( - "DMSwarm_cellid", - self, - 1, - dtype=int, - _register=False, - _proxy=False, - rebuild_on_cycle=False, - ) - - # add variable to hold swarm coordinates during position updates - self._X0 = uw.swarm.SwarmVariable( - "DMSwarm_X0", - self, - self.cdim, - dtype=float, - _register=True, - _proxy=False, - rebuild_on_cycle=False, - ) - - # This is for swarm streak management: - # add variable to hold swarm origins - - if self.recycle_rate > 1: - # self._Xorig = uw.swarm.SwarmVariable( - # "DMSwarm_Xorig", - # self, - # self.cdim, - # dtype=float, - # _register=True, - # _proxy=False, - # rebuild_on_cycle=False, - # ) - - self._remeshed = uw.swarm.SwarmVariable( - "DMSwarm_remeshed", - self, - 1, - dtype=int, - _register=True, - _proxy=False, - rebuild_on_cycle=False, - ) - - self._X0_uninitialised = True - self._kdtree = None - self._nnmapdict = {} - - super().__init__() - - @property - def mesh(self): - """The mesh associated with this swarm. - - Returns - ------- - Mesh - The computational mesh containing the particles. - """ - return self._mesh - - @property - def data(self): - """Particle coordinate data array. - - Provides direct read/write access to particle positions. - - Returns - ------- - numpy.ndarray - Coordinate array of shape ``(n_particles, dim)``. - """ - return self._particle_coordinates.data - - @property - def particle_coordinates(self): - """SwarmVariable holding particle coordinates. - - Returns - ------- - SwarmVariable - The internal coordinate variable (use ``.data`` for array access). - """ - return self._coord_var - - @property - def particle_cellid(self): - """SwarmVariable holding particle cell IDs. - - Each particle is associated with a mesh cell for efficient - spatial operations. - - Returns - ------- - SwarmVariable - Cell ID variable (use ``.data`` for array access). - """ - return self._cellid_var - - @timing.routine_timer_decorator - def populate_petsc( - self, - fill_param: Optional[int] = 3, - layout: Optional[SwarmPICLayout] = None, - ): - """ - Populate the swarm with particles throughout the domain. - - When using SwarmPICLayout.REGULAR, `fill_param` defines the number of points in each spatial direction. - When using SwarmPICLayout.GAUSS, `fill_param` defines the number of quadrature points in each spatial direction. - When using SwarmPICLayout.SUBDIVISION, `fill_param` defines the number times the reference cell is sub-divided. - - Parameters - ---------- - fill_param: - Parameter determining the particle count per cell for the given layout. - layout: - Type of layout to use. Defaults to `SwarmPICLayout.REGULAR` for mesh objects with simplex - type cells, and `SwarmPICLayout.GAUSS` otherwise. - - """ - - self.fill_param = fill_param - - """ - Currently (2021.11.15) supported by PETSc release 3.16.x - - When using a DMPLEX the following case are supported: - (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle), - (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex, - (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri. - - So this means, simplex mesh in 3D only supports GAUSS - This is based - on the tensor product locations so it is not uniform in the cells. - """ - - if layout == None: - layout = SwarmPICLayout.GAUSS - - if not isinstance(layout, SwarmPICLayout): - raise ValueError("'layout' must be an instance of 'SwarmPICLayout'") - - self.layout = layout - self.dm.finalizeFieldRegister() - - ## Commenting this out for now. - ## Code seems to operate fine without it, and the - ## existing values are wrong. It should be something like - ## `(elend-elstart)*fill_param^dim` for quads, and around - ## half that for simplices, depending on layout. - # elstart,elend = self.mesh.dm.getHeightStratum(0) - # self.dm.setLocalSizes((elend-elstart) * fill_param, 0) - - self.dm.insertPointUsingCellDM(self.layout.value, fill_param) - return - - # - - @timing.routine_timer_decorator - def populate( - self, - fill_param: Optional[int] = 1, - ): - """ - Populate the swarm with particles throughout the domain. - - Parameters - ---------- - fill_param: - Parameter determining the particle count per cell (per dimension) - for the given layout, using the mesh degree. - - cell_search: - Use k-d tree to locate nearest cells (fails if this swarm is used to build a k-d tree) - - """ - - self.fill_param = fill_param - - newp_coords0 = self.mesh._get_coords_for_basis(fill_param, continuous=False) - newp_cells0 = self.mesh.get_closest_local_cells(newp_coords0) - - if np.any(newp_cells0 > self.mesh._centroids.shape[0]): - raise RuntimeError("Some new coordinates can't find a owning cell - Error") - - # valid = newp_cells0 != -1 - # newp_coords = newp_coords0[valid] - # newp_cells = newp_cells0[valid] - newp_coords = newp_coords0 - newp_cells = newp_cells0 - - self.dm.finalizeFieldRegister() - self.dm.addNPoints(newp_coords.shape[0] + 1) - - cellid = self.dm.getField("DMSwarm_cellid") - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.dim)) - - coords[...] = newp_coords[...] - cellid[:] = newp_cells[:] - - self.dm.restoreField("DMSwarmPIC_coor") - self.dm.restoreField("DMSwarm_cellid") - - ## Now make a series of copies to allow the swarm cycling to - ## work correctly (if required) - - # cellid = self.dm.getField("DMSwarm_cellid") - # lost = np.where(cellid == -1) - # print(f"{uw.mpi.rank} - lost particles: {lost[0].shape} out of {cellid.shape}", flush=True) - # self.dm.restoreField("DMSwarm_cellid") - - if self.recycle_rate > 1: - with self.access(): - # This is a mesh-local quantity, so let's just - # store it on the mesh in an ad_hoc fashion for now - - self.mesh.particle_X_orig = self._particle_coordinates.data.copy() - self.mesh.particle_CellID_orig = self._cellid_var.data.copy() - - with self.access(): - swarm_orig_size = self.local_size - all_local_coords = np.vstack( - (self._particle_coordinates.data,) * (self.recycle_rate) - ) - all_local_cells = np.vstack((self._cellid_var.data,) * (self.recycle_rate)) - - swarm_new_size = all_local_coords.shape[0] - - self.dm.addNPoints(swarm_new_size - swarm_orig_size) - - cellid = self.dm.getField("DMSwarm_cellid") - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.dim)) - - coords[...] = ( - all_local_coords[...] - + (0.33 / (1 + fill_param)) - * (np.random.random(size=all_local_coords.shape) - 0.5) - * 0.00001 - * self.mesh._search_lengths[all_local_cells] # typical cell size - ) - cellid[:] = all_local_cells[:, 0] - - self.dm.restoreField("DMSwarmPIC_coor") - self.dm.restoreField("DMSwarm_cellid") - - ## Now set the cycle values - - with self.access(self._remeshed): - for i in range(0, self.recycle_rate): - offset = swarm_orig_size * i - self._remeshed.data[offset::, 0] = i - - return - - @timing.routine_timer_decorator - def add_particles_with_coordinates(self, coordinatesArray) -> int: - """ - Add particles to the swarm using particle coordinates provided - using a numpy array. - - Note that particles with coordinates NOT local to the current processor will - be rejected / ignored. - - Either include an array with all coordinates to all processors - or an array with the local coordinates. - - Parameters - ---------- - coordinatesArray : numpy.ndarray - The numpy array containing the coordinate of the new particles. Array is - expected to take shape n*dim, where n is the number of new particles, and - dim is the dimensionality of the swarm's supporting mesh. - - Returns - -------- - npoints: int - The number of points added to the local section of the swarm. - """ - - if not isinstance(coordinatesArray, np.ndarray): - raise TypeError("'coordinateArray' must be provided as a numpy array") - if not len(coordinatesArray.shape) == 2: - raise ValueError("The 'coordinateArray' is expected to be two dimensional.") - if not coordinatesArray.shape[1] == self.mesh.dim: - #### petsc appears to ignore columns that are greater than the mesh dim, but still worth including - raise ValueError( - """The 'coordinateArray' must have shape n*dim, where 'n' is the - number of particles to add, and 'dim' is the dimensionality of - the supporting mesh ({}).""".format( - self.mesh.dim - ) - ) - - cells = self.mesh.get_closest_local_cells(coordinatesArray) - - valid_coordinates = coordinatesArray[cells != -1] - valid_cells = cells[cells != -1] - - npoints = len(valid_coordinates) - swarm_size = self.dm.getLocalSize() - - # -1 means no particles have been added yet - if swarm_size == -1: - swarm_size = 0 - npoints = npoints + 1 - - self.dm.finalizeFieldRegister() - self.dm.addNPoints(npoints=npoints) - - cellid = self.dm.getField("DMSwarm_cellid") - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.dim)) - - coords[swarm_size::, :] = valid_coordinates[:, :] - cellid[swarm_size::] = valid_cells[:] - - self.dm.restoreField("DMSwarmPIC_coor") - self.dm.restoreField("DMSwarm_cellid") - - # Here we update the swarm cycle values as required - - if self.recycle_rate > 1: - with self.access(self._remeshed): - # self._Xorig.data[...] = coordinatesArray - self._remeshed.data[...] = 0 - - self.dm.migrate(remove_sent_points=True) - - return npoints - - @timing.routine_timer_decorator - def save( - self, - filename: int, - compression: Optional[bool] = False, - compressionType: Optional[str] = "gzip", - force_sequential=False, - ): - """ - - Save the swarm coordinates to a h5 file. - - Parameters - ---------- - filename : - The filename of the swarm checkpoint file to save to disk. - compression : - Add compression to the h5 files (saves space but increases write times with increasing no. of processors) - compressionType : - Type of compression to use, 'gzip' and 'lzf' supported. 'gzip' is default. Compression also needs to be set to 'True'. - - - - """ - if h5py.h5.get_config().mpi == False and comm.size > 1 and comm.rank == 0: - warnings.warn( - "Collective IO not possible as h5py not available in parallel mode. Switching to sequential. This will be slow for models running on multiple processors", - stacklevel=2, - ) - if filename.endswith(".h5") == False: - raise RuntimeError("The filename must end with .h5") - if compression == True and comm.rank == 0: - warnings.warn("Compression may slow down write times", stacklevel=2) - - if h5py.h5.get_config().mpi == True and not force_sequential: - # It seems to be a bad idea to mix mpi barriers with the access - # context manager so the copy-free version of this seems to hang - # when there are many active cores. This is probably why the parallel - # h5py write hangs - - with self.access(): - data_copy = self.data[:].copy() - - with h5py.File(f"{filename[:-3]}.h5", "w", driver="mpio", comm=comm) as h5f: - if compression == True: - h5f.create_dataset( - "coordinates", - data=data_copy[:], - compression=compressionType, - ) - else: - h5f.create_dataset("coordinates", data=data_copy[:]) - - del data_copy - - else: - # It seems to be a bad idea to mix mpi barriers with the access - # context manager so the copy-free version of this seems to hang - # when there are many active cores - - with self.access(): - data_copy = self.data[:].copy() - - if comm.rank == 0: - with h5py.File(f"{filename[:-3]}.h5", "w") as h5f: - if compression == True: - h5f.create_dataset( - "coordinates", - data=data_copy, - chunks=True, - maxshape=(None, data_copy.shape[1]), - compression=compressionType, - ) - else: - h5f.create_dataset( - "coordinates", - data=data_copy, - chunks=True, - maxshape=(None, data_copy.shape[1]), - ) - - comm.barrier() - for i in range(1, comm.size): - if comm.rank == i: - with h5py.File(f"{filename[:-3]}.h5", "a") as h5f: - h5f["coordinates"].resize( - (h5f["coordinates"].shape[0] + data_copy.shape[0]), - axis=0, - ) - # passive swarm, zero local particles is not unusual - if data_copy.shape[0] > 0: - h5f["coordinates"][-data_copy.shape[0] :] = data_copy[:] - comm.barrier() - comm.barrier() - - del data_copy - - return - - @timing.routine_timer_decorator - def read_timestep( - self, - base_filename: str, - swarm_id: str, - index: int, - outputPath: Optional[str] = "", - ): - """Load particle coordinates from a saved timestep file. - - Reads an HDF5 file containing particle coordinates and adds - them to the swarm using :meth:`add_particles_with_coordinates`. - - Parameters - ---------- - base_filename : str - Base name for the output files. - swarm_id : str - Identifier for this swarm in the output. - index : int - Timestep index to read. - outputPath : str, optional - Directory containing the output files. - """ - output_base_name = os.path.join(outputPath, base_filename) - swarm_file = output_base_name + f".{swarm_id}.{index:05}.h5" - - ### open up file with coords on all procs - with h5py.File(f"{swarm_file}", "r") as h5f: - coordinates = h5f["coordinates"][:] - - #### utilises the UW function for adding a swarm by an array - self.add_particles_with_coordinates(coordinates) - - return - - @timing.routine_timer_decorator - def add_variable( - self, - name, - size=1, - dtype=float, - proxy_degree=2, - _nn_proxy=False, - ): - """Create a new SwarmVariable attached to this swarm. - - Parameters - ---------- - name : str - Name for the variable. - size : int, optional - Number of components per particle (default: 1 for scalar). - dtype : type, optional - Data type for the variable (default: float). - proxy_degree : int, optional - Polynomial degree for mesh proxy variable (default: 2). - _nn_proxy : bool, optional - Use nearest-neighbor proxy (internal use). - - Returns - ------- - SwarmVariable - New variable attached to this swarm. - - Examples - -------- - >>> material = swarm.add_variable("material", size=1, dtype=int) - >>> temperature = swarm.add_variable("T", size=1) - """ - return SwarmVariable( - name, - self, - size, - dtype=dtype, - proxy_degree=proxy_degree, - _nn_proxy=_nn_proxy, - ) - - @timing.routine_timer_decorator - def petsc_save_checkpoint( - self, - swarmName: str, - index: int, - outputPath: Optional[str] = "", - ): - """ - - Use PETSc to save the swarm and attached data to a .pbin and xdmf file. - - Parameters - ---------- - swarmName : - Name of the swarm to save. - index : - An index which might correspond to the timestep or output number (for example). - outputPath : - Path to save the data. If left empty it will save the data in the current working directory. - """ - - x_swarm_fname = f"{outputPath}{swarmName}_{index:05d}.xmf" - self.dm.viewXDMF(x_swarm_fname) - - @timing.routine_timer_decorator - def write_timestep( - self, - filename: str, - swarmname: str, - index: int, - swarmVars: Optional[list] = None, - outputPath: Optional[str] = "", - time: Optional[int] = None, - compression: Optional[bool] = False, - compressionType: Optional[str] = "gzip", - force_sequential: Optional[bool] = False, - ): - """ - - Save data to h5 and a corresponding xdmf for visualisation using h5py. - - Parameters - ---------- - swarmName : - Name of the swarm to save. - swarmVars : - List of swarm objects to save. - index : - An index which might correspond to the timestep or output number (for example). - outputPath : - Path to save the data. If left empty it will save the data in the current working directory. - time : - Attach the time to the generated xdmf. - compression : - Whether to compress the h5 files [bool]. - compressionType : - The type of compression to use. 'gzip' and 'lzf' are the supported types, with 'gzip' as the default. - """ - - # This will eliminate the issue of whether or not to put path separators in the - # outputPath. Also does the right thing if outputPath is "" - - output_base_name = os.path.join(outputPath, filename) + "." + swarmname - - # check the directory where we will write checkpoint - dir_path = os.path.dirname(output_base_name) # get directory - - # check if path exists - if os.path.exists(os.path.abspath(dir_path)): # easier to debug abs - pass - else: - raise RuntimeError(f"{os.path.abspath(dir_path)} does not exist") - - # check if we have write access - if os.access(os.path.abspath(dir_path), os.W_OK): - pass - else: - raise RuntimeError(f"No write access to {os.path.abspath(dir_path)}") - - # could also try to coerce this to be a list and raise if it fails (tuple, singleton ... ) - # also ... why the typechecking if this can still happen - - if swarmVars is not None and not isinstance(swarmVars, list): - raise RuntimeError("`swarmVars` does not appear to be a list.") - - else: - ### save the swarm particle location - self.save( - filename=f"{output_base_name}.{index:05d}.h5", - compression=compression, - compressionType=compressionType, - force_sequential=force_sequential, - ) - - #### Generate a h5 file for each field - if swarmVars != None: - for field in swarmVars: - field.save( - filename=f"{output_base_name}.{field.name}.{index:05d}.h5", - compression=compression, - compressionType=compressionType, - force_sequential=force_sequential, - ) - - if uw.mpi.rank == 0: - ### only need to combine the h5 files to a single xdmf on one proc - with open(f"{output_base_name}.{index:05d}.xdmf", "w") as xdmf: - # Write the XDMF header - xdmf.write('\n') - xdmf.write('\n') - xdmf.write("\n") - xdmf.write(f'\n') - - if time != None: - xdmf.write(f' \n") - xdmf.write("\n") - xdmf.write("\n") - - @property - def vars(self): - """Dictionary of SwarmVariables attached to this swarm. - - Returns - ------- - dict - Mapping from variable names to :class:`SwarmVariable` objects. - """ - return self._vars - - def access(self, *writeable_vars: SwarmVariable): - """ - This context manager makes the underlying swarm variables data available to - the user. The data should be accessed via the variables `data` handle. - - As default, all data is read-only. To enable writeable data, the user should - specify which variable they wish to modify. - - At the conclusion of the users context managed block, numerous further operations - will be automatically executed. This includes swarm parallel migration routines - where the swarm's `particle_coordinates` variable has been modified. The swarm - variable proxy mesh variables will also be updated for modifed swarm variables. - - Parameters - ---------- - writeable_vars - The variables for which data write access is required. - - Example - ------- - - >>> import underworld3 as uw - >>> someMesh = uw.discretisation.FeMesh_Cartesian() - >>> with someMesh._deform_mesh(): - ... someMesh.data[0] = [0.1,0.1] - >>> someMesh.data[0] - array([ 0.1, 0.1]) - """ - import time - - uw.timing._incrementDepth() - stime = time.time() - - deaccess_list = [] - for var in self._vars.values(): - # if already accessed within higher level context manager, continue. - if var._is_accessed == True: - continue - # set flag so variable status can be known elsewhere - var._is_accessed = True - # add to de-access list to rewind this later - deaccess_list.append(var) - # grab numpy object, setting read only if necessary - var._data = self.dm.getField(var.clean_name).reshape((-1, var.num_components)) - assert var._data is not None - if var not in writeable_vars: - var._old_data_flag = var._data.flags.writeable - var._data.flags.writeable = False - else: - # increment variable state - var._increment() - - # make view for each var component - if var._proxy: - for i in range(0, var.shape[0]): - for j in range(0, var.shape[1]): - var._data_container[i, j] = var._data_container[i, j]._replace( - data=var.data[:, var._data_layout(i, j)], - ) - - # if particles moving, update swarm state - if self._particle_coordinates in writeable_vars: - self._increment() - - # Create a class which specifies the required context - # manager hooks (`__enter__`, `__exit__`). - class exit_manager: - def __init__(self, swarm): - self.em_swarm = swarm - - def __enter__(self): - - pass - - def __exit__(self, *args): - - for var in self.em_swarm.vars.values(): - # only de-access variables we have set access for. - if var not in deaccess_list: - continue - # set this back, although possibly not required. - if var not in writeable_vars: - var._data.flags.writeable = var._old_data_flag - var._data = None - self.em_swarm.dm.restoreField(var.clean_name) - var._is_accessed = False - # do particle migration if coords changes - - if self.em_swarm._particle_coordinates in writeable_vars: - # let's use the mesh index to update the particles owning cells. - # note that the `petsc4py` interface is more convenient here as the - # `SwarmVariable.data` interface is controlled by the context manager - # that we are currently within, and it is therefore too easy to - # get things wrong that way. - - cellid = self.em_swarm.dm.getField("DMSwarm_cellid") - coords = self.em_swarm.dm.getField("DMSwarmPIC_coor").reshape( - (-1, self.em_swarm.dim) - ) - - cellid[:] = self.em_swarm.mesh.get_closest_cells(coords).reshape(-1) - - # num_lost = np.where(cellid == -1)[0].shape[0] - # print( - # f"{uw.mpi.rank} - EM 1: illegal_cells - {num_lost}", flush=True - # ) - - # if num_lost != 0: - # print("LOST: ", coords[np.where(cellid == -1)]) - - self.em_swarm.dm.restoreField("DMSwarmPIC_coor") - self.em_swarm.dm.restoreField("DMSwarm_cellid") - # now migrate. - - self.em_swarm.dm.migrate(remove_sent_points=True) - - # void these things too - self.em_swarm._kdtree = None - self.em_swarm._nnmapdict = {} - - # do var updates - for var in self.em_swarm.vars.values(): - # if swarm migrated, update all. - # if var updated, update var. - if (self.em_swarm._particle_coordinates in writeable_vars) or ( - var in writeable_vars - ): - var._update() - - if var._proxy: - for i in range(0, var.shape[0]): - for j in range(0, var.shape[1]): - # var._data_ij[i, j] = None - var._data_container[i, j] = var._data_container[i, j]._replace( - data=f"SwarmVariable[...].data is only available within mesh.access() context", - ) - - uw.timing._decrementDepth() - uw.timing.log_result(time.time() - stime, "Swarm.access", 1) - - return exit_manager(self) - - ## Better to have one master copy - this one is cut'n'pasted from - ## the MeshVariable class - - def _data_layout(self, i, j=None): - # mapping - - if self.vtype == uw.VarType.SCALAR: - return 0 - if self.vtype == uw.VarType.VECTOR: - if j is None: - return i - elif i == 0: - return j - else: - raise IndexError(f"Vectors have shape {self.mesh.dim} or {(1, self.mesh.dim)} ") - if self.vtype == uw.VarType.TENSOR: - if self.mesh.dim == 2: - return ((0, 1), (2, 3))[i][j] - else: - return ((0, 1, 2), (3, 4, 5), (6, 7, 8))[i][j] - - if self.vtype == uw.VarType.SYM_TENSOR: - if self.mesh.dim == 2: - return ((0, 2), (2, 1))[i][j] - else: - return ((0, 3, 4), (3, 1, 5), (4, 5, 2))[i][j] - - if self.vtype == uw.VarType.MATRIX: - return i + j * self.shape[0] - - def _get_kdtree(self): - """ - Return a cached KDTree for the swarm particle coordinates. - Invalidated automatically whenever particles migrate or positions change. - """ - if not hasattr(self, "_kdtree") or self._kdtree is None: - with self.access(): - self._kdtree = uw.kdtree.KDTree(self._coord_var.data) - - return self._kdtree - - @timing.routine_timer_decorator - def _get_map(self, var): - # generate tree if not avaiable - kd = self._get_kdtree() - - # get or generate map - meshvar_coords = var._meshVar.coords - # we can't use numpy arrays directly as keys in python dicts, so - # we'll use `xxhash` to generate a hash of array. - # this shouldn't be an issue performance wise but we should test to be - # sufficiently confident of this. - import xxhash - - h = xxhash.xxh64() - h.update(meshvar_coords) - digest = h.intdigest() - if digest not in self._nnmapdict: - # self._nnmapdict[digest] = self._kdtree.find_closest_point(meshvar_coords)[0] - self._nnmapdict[digest] = kd.query(meshvar_coords, k=1, sqr_dists=False)[1] - return self._nnmapdict[digest] - - @timing.routine_timer_decorator - def advection( - self, - V_fn, - delta_t, - order=2, - corrector=False, - restore_points_to_domain_func=None, - evalf=False, - step_limit=True, - ): - """Advect particles using a velocity field. - - Moves particles according to the velocity field using forward - Euler integration with automatic substepping based on CFL - conditions. Handles particle migration between MPI ranks. - - Parameters - ---------- - V_fn : sympy.Matrix or MeshVariable - Velocity field (vector expression). - delta_t : float - Total time to advect. - order : int, optional - Integration order (default: 2, not currently used). - corrector : bool, optional - Apply predictor-corrector scheme (default: False). - restore_points_to_domain_func : callable, optional - Function to restore particles that leave the domain. - evalf : bool, optional - Use numerical evaluation (True) or symbolic (False). - step_limit : bool, optional - Apply CFL-based substepping limit (default: True). - """ - dt_limit = self.estimate_dt(V_fn) - - if step_limit and dt_limit is not None: - substeps = int(max(1, round(abs(delta_t) / dt_limit))) - else: - substeps = 1 - - if uw.mpi.rank == 0 and self.verbose: - print(f"Substepping {substeps} / {abs(delta_t) / dt_limit}, {delta_t} ") - - # X0 holds the particle location at the start of advection - # This is needed because the particles may be migrated off-proc - # during timestepping. - - X0 = self._X0 - - V_fn_matrix = self.mesh.vector.to_matrix(V_fn) - - # Use current velocity to estimate where the particles would have - # landed in an implicit step. WE CANT DO THIS WITH SUB-STEPPING unless - # We have a lot more information about the previous launch point / timestep - # Also: how does this interact with the particle restoration function ? - - # if corrector == True and not self._X0_uninitialised: - # with self.access(self._particle_coordinates): - # v_at_Vpts = np.zeros_like(self.data) - - # if evalf: - # for d in range(self.dim): - # v_at_Vpts[:, d] = uw.function.evalf( - # V_fn_matrix[d], self.data - # ).reshape(-1) - # else: - # for d in range(self.dim): - # v_at_Vpts[:, d] = uw.function.evaluate( - # V_fn_matrix[d], self.data - # ).reshape(-1) - - # corrected_position = X0.data.copy() + delta_t * v_at_Vpts - # if restore_points_to_domain_func is not None: - # corrected_position = restore_points_to_domain_func( - # corrected_position - # ) - - # updated_current_coords = 0.5 * (corrected_position + self.data.copy()) - - # # validate_coords to ensure they live within the domain (or there will be trouble) - - # if restore_points_to_domain_func is not None: - # updated_current_coords = restore_points_to_domain_func( - # updated_current_coords - # ) - - # self.data[...] = updated_current_coords[...] - - # del updated_current_coords - # del v_at_Vpts - - print(f"{ uw.mpi.rank}: Peace", flush=True) - - # Wrap this whole thing in sub-stepping loop - for step in range(0, substeps): - - with self.access(X0): - X0.data[...] = self._particle_coordinates.data[...] - - # Mid point algorithm (2nd order) - - if order == 2: - with self.access(self._particle_coordinates): - v_at_Vpts = np.zeros_like(self._particle_coordinates.data) - - # if evalf: - # for d in range(self.dim): - # v_at_Vpts[:, d] = uw.function.evalf( - # V_fn_matrix[d], self._particle_coordinates.data - # ).reshape(-1) - # else: - for d in range(self.dim): - v_at_Vpts[:, d] = uw.function.evaluate( - V_fn_matrix[d], - self._particle_coordinates.data, - evalf=evalf, - ).reshape(-1) - - mid_pt_coords = ( - self._particle_coordinates.data[...] + 0.5 * delta_t * v_at_Vpts / substeps - ) - - # validate_coords to ensure they live within the domain (or there will be trouble) - - if restore_points_to_domain_func is not None: - mid_pt_coords = restore_points_to_domain_func(mid_pt_coords) - - self._particle_coordinates.data[...] = mid_pt_coords[...] - - del mid_pt_coords - - ## Let the swarm be updated, and then move the rest of the way - - v_at_Vpts = np.zeros_like(self.data) - - # if evalf: - # for d in range(self.dim): - # v_at_Vpts[:, d] = uw.function.evalf( - # V_fn_matrix[d], self._particle_coordinates.data - # ).reshape(-1) - # else: - # - for d in range(self.dim): - v_at_Vpts[:, d] = uw.function.evaluate( - V_fn_matrix[d], - self._particle_coordinates.data, - evalf=evalf, - ).reshape(-1) - - # if (uw.mpi.rank == 0): - # print("Re-launch from X0", flush=True) - - new_coords = X0.data[...] + delta_t * v_at_Vpts / substeps - - # validate_coords to ensure they live within the domain (or there will be trouble) - if restore_points_to_domain_func is not None: - new_coords = restore_points_to_domain_func(new_coords) - - self._particle_coordinates.data[...] = new_coords[...] - - del new_coords - del v_at_Vpts - - # forward Euler (1st order) - else: - with self.access(self._particle_coordinates): - v_at_Vpts = np.zeros_like(self.data) - - # if evalf: - # for d in range(self.dim): - # v_at_Vpts[:, d] = uw.function.evalf( - # V_fn_matrix[d], self.data - # ).reshape(-1) - # else: - for d in range(self.dim): - v_at_Vpts[:, d] = uw.function.evaluate( - V_fn_matrix[d], - self.data, - evalf=evalf, - ).reshape(-1) - - new_coords = self.data + delta_t * v_at_Vpts / substeps - - # validate_coords to ensure they live within the domain (or there will be trouble) - - if restore_points_to_domain_func is not None: - new_coords = restore_points_to_domain_func(new_coords) - - self.data[...] = new_coords[...].copy() - - ## End of substepping loop - - ## Cycling of the swarm is a cheap and cheerful version of population control for particles. It turns the - ## swarm into a streak-swarm where particles are Lagrangian for a number of steps and then reset to their - ## original location. - - if self.recycle_rate > 1: - # Restore particles which have cycle == cycle rate (use >= just in case) - - # Remove remesh points and recreate a new set at the mesh-local - # locations that we already have stored. - - with self.access(self._particle_coordinates, self._remeshed): - remeshed = self._remeshed.data[:, 0] == 0 - # This is one way to do it ... we can do this better though - self.data[remeshed, 0] = 1.0e100 - - swarm_size = self.dm.getLocalSize() - - num_remeshed_points = self.mesh.particle_X_orig.shape[0] - - self.dm.addNPoints(num_remeshed_points) - - cellid = self.dm.getField("DMSwarm_cellid") - coords = self.dm.getField("DMSwarmPIC_coor").reshape((-1, self.dim)) - rmsh = self.dm.getField("DMSwarm_remeshed") - - # print(f"cellid -> {cellid.shape}") - # print(f"particle coords -> {coords.shape}") - # print(f"remeshed points -> {num_remeshed_points}") - - perturbation = 0.00001 * ( - (0.33 / (1 + self.fill_param)) - * (np.random.random(size=(num_remeshed_points, self.dim)) - 0.5) - * self.mesh._radii[cellid[swarm_size::]].reshape(-1, 1) - ) - - coords[swarm_size::] = self.mesh.particle_X_orig[:, :] + perturbation - cellid[swarm_size::] = self.mesh.particle_CellID_orig[:, 0] - rmsh[swarm_size::] = 0 - - self.dm.restoreField("DMSwarm_cellid") - self.dm.restoreField("DMSwarmPIC_coor") - self.dm.restoreField("DMSwarm_remeshed") - - # when we let this go, the particles may be re-distributed to - # other processors, and we will need to rebuild the remeshed - # array before trying to compute / assign values to variables - - for swarmVar in self.vars.values(): - if swarmVar._rebuild_on_cycle: - with self.access(swarmVar): - if swarmVar.dtype is int: - nnn = 1 - else: - nnn = self.mesh.dim + 1 # 3 for triangles, 4 for tets ... - - interpolated_values = ( - swarmVar.rbf_interpolate(self.mesh.particle_X_orig, nnn=nnn) - # swarmVar._meshVar.fn, self.mesh.particle_X_orig - # ) - ).astype(swarmVar.dtype) - - swarmVar.data[swarm_size::] = interpolated_values - - self.dm.migrate(remove_sent_points=True) - - with self.access(self._remeshed): - self._remeshed.data[...] = np.mod(self._remeshed.data[...] - 1, self.recycle_rate) - - self.cycle += 1 - - return - - @timing.routine_timer_decorator - def estimate_dt(self, V_fn): - """ - Calculates an appropriate advective timestep for the given - mesh and velocity configuration. - """ - # we'll want to do this on an element by element basis - # for more general mesh - - # first let's extract a max global velocity magnitude - import math - import numpy as np - - with self.access(): - vel = uw.function.evaluate(V_fn, self._particle_coordinates.data, evalf=True) - - # If vel is unit-aware (UnitAwareArray), nondimensionalise it to get - # consistent nondimensional values that match mesh._radii - # Note: .magnitude returns physical units, which would be wrong here - if hasattr(vel, "units") and vel.units is not None: - vel = uw.non_dimensionalise(vel) - elif hasattr(vel, "magnitude"): - # Plain UWQuantity without units context - use magnitude - vel = vel.magnitude - - # Ensure vel is a plain numpy array - vel = np.asarray(vel) - - try: - magvel_squared = vel[:, 0] ** 2 + vel[:, 1] ** 2 - if self.mesh.dim == 3: - magvel_squared += vel[:, 2] ** 2 - - max_magvel = math.sqrt(magvel_squared.max()) - - except (ValueError, IndexError): - max_magvel = 0.0 - - from mpi4py import MPI - - max_magvel_glob = comm.allreduce(max_magvel, op=MPI.MAX) - - min_dx = self.mesh.get_min_radius() - - # The assumption should be that we cross one or two elements (2-4 radii), not more, - # in a single step (order 2, means one element per half-step or something - # that we can broadly interpret that way) - - if max_magvel_glob != 0.0: - return min_dx / max_magvel_glob - else: - return None - - -class NodalPointPICSwarm(PICSwarm): - r"""Swarm with particles located at the coordinate points of a meshVariable - - The swarmVariable `X0` is defined so that the particles can "snap back" to their original locations - after they have been moved. - - The purpose of this Swarm is to manage sample points for advection schemes based on upstream sampling - (method of characteristics etc)""" - - def __init__( - self, - trackedVariable: uw.discretisation.MeshVariable, - verbose=False, - ): - self.trackedVariable = trackedVariable - self.swarmVariable = None - - mesh = trackedVariable.mesh - - # Set up a standard swarm - super().__init__(mesh, verbose) - - nswarm = self - - meshVar_name = trackedVariable.clean_name - meshVar_symbol = trackedVariable.symbol - - ks = str(self.instance_number) - name = f"{meshVar_name}_star" - symbol = rf"{{ {meshVar_symbol} }}^{{ <*> }}" - - self.swarmVariable = uw.swarm.SwarmVariable( - name, - nswarm, - vtype=trackedVariable.vtype, - _proxy=False, - # proxy_degree=trackedVariable.degree, - # proxy_continuous=trackedVariable.continuous, - varsymbol=symbol, - ) - - # The launch point location - name = f"ns_X0_{ks}" - symbol = r"X0^{*^{{[" + ks + "]}}}" - nX0 = uw.swarm.SwarmVariable(name, nswarm, nswarm.dim, _proxy=False) - - # The launch point index - name = f"ns_I_{ks}" - symbol = r"I^{*^{{[" + ks + "]}}}" - nI0 = uw.swarm.SwarmVariable(name, nswarm, 1, dtype=int, _proxy=False) - - # The launch point processor rank - name = f"ns_R0_{ks}" - symbol = r"R0^{*^{{[" + ks + "]}}}" - nR0 = uw.swarm.SwarmVariable(name, nswarm, 1, dtype=int, _proxy=False) - - nswarm.dm.finalizeFieldRegister() - nswarm.dm.addNPoints( - trackedVariable.coords.shape[0] + 1 - ) # why + 1 ? That's the number of spots actually allocated - - cellid = nswarm.dm.getField("DMSwarm_cellid") - coords = nswarm.dm.getField("DMSwarmPIC_coor").reshape((-1, nswarm.dim)) - coords[...] = trackedVariable.coords[...] - cellid[:] = self.mesh.get_closest_local_cells(coords) - - # Move slightly within the chosen cell to avoid edge effects - centroid_coords = self.mesh._centroids[cellid] - - shift = 0.001 - coords[:, :] = (1.0 - shift) * coords[:, :] + shift * centroid_coords[:, :] - - nswarm.dm.restoreField("DMSwarmPIC_coor") - nswarm.dm.restoreField("DMSwarm_cellid") - - nswarm.dm.migrate(remove_sent_points=True) - - with nswarm.access(nX0, nI0): - nX0.data[:, :] = coords - nI0.data[:, 0] = range(0, coords.shape[0]) - - self._nswarm = nswarm - self._nX0 = nX0 - self._nI0 = nI0 - self._nR0 = nR0 - - return - - @timing.routine_timer_decorator - def advection( - self, - V_fn, - delta_t, - order=2, - corrector=False, - restore_points_to_domain_func=None, - evalf=False, - step_limit=True, - ): - """Advect nodal point particles and track back to mesh nodes. - - Extends parent advection with node-tracking logic. After - advection, particles are mapped back to their original mesh - nodes for semi-Lagrangian interpolation. - - Parameters - ---------- - V_fn : sympy.Matrix or MeshVariable - Velocity field (vector expression). - delta_t : float - Total time to advect. - order : int, optional - Integration order (default: 2). - corrector : bool, optional - Apply predictor-corrector scheme (default: False). - restore_points_to_domain_func : callable, optional - Function to restore particles that leave the domain. - evalf : bool, optional - Use numerical evaluation (True) or symbolic (False). - step_limit : bool, optional - Apply CFL-based substepping limit (default: True). - """ - with self.access(self._X0): - self._X0.data[...] = self._nX0.data[...] - - with self.access(self._nR0): - self._nR0.data[...] = uw.mpi.rank - - super().advection( - V_fn, - delta_t, - order, - corrector, - restore_points_to_domain_func, - evalf, - step_limit, - ) - - return diff --git a/tests/parallel/test_0766_swarm_substep_advection.py b/tests/parallel/test_0766_swarm_substep_advection.py new file mode 100644 index 000000000..e4ede6222 --- /dev/null +++ b/tests/parallel/test_0766_swarm_substep_advection.py @@ -0,0 +1,108 @@ +""" +Regression test for substep advection across partition boundaries (SWARM-16 / BF-16). + +``Swarm.advection()`` with ``substeps > 1`` evaluates the launch-point velocity +of every substep, but no migration happens inside the substep loop (deferred +migration is deliberately suspended there — see #313). From substep 2 onward a +particle may therefore sit outside its rank's domain when its launch-point +velocity is evaluated, and a rank-local evaluation returns extrapolated +(wrong) values for it. The launch-point evaluation must be a *global* +evaluation, like the midpoint evaluation already is. + +The test advects a ring of particles through a solid-body rotation stored in a +degree-1 MeshVariable (linear field, so FE interpolation is exact to roundoff) +and compares each particle's final position against an exact numpy replication +of the substepped midpoint scheme. Any wrong-rank evaluation after a particle +crosses the partition seam shows up as a position error far above roundoff. + +Run with: + mpirun -n 2 python -m pytest --with-mpi tests/parallel/test_0766_swarm_substep_advection.py + mpirun -n 4 python -m pytest --with-mpi tests/parallel/test_0766_swarm_substep_advection.py +""" + +import pytest +import numpy as np +import underworld3 as uw + +pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(300)] + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_substep_advection_across_partition_matches_scheme(): + mesh = uw.meshing.StructuredQuadBox(elementRes=(16, 16)) + + # Solid-body rotation about the domain centre, stored in a degree-1 + # variable: a linear field, represented exactly by the element basis, so + # a correctly-located evaluation reproduces it to machine precision. + V = uw.discretisation.MeshVariable("V_sub_adv", mesh, mesh.dim, degree=1) + V.array[:, 0, 0] = -(V.coords[:, 1] - 0.5) + V.array[:, 0, 1] = V.coords[:, 0] - 0.5 + + swarm = uw.swarm.Swarm(mesh) + pid = uw.swarm.SwarmVariable("pid_sub_adv", swarm, 1, dtype=int, _proxy=False) + + # Ring of particles that the rotation sweeps across the partition seam. + # The small angular offset keeps every launch point strictly inside one + # rank's cells (a point exactly on the seam is claimed by both ranks). + n_pts = 32 + theta = 0.03 + np.linspace(0.0, 2.0 * np.pi, n_pts, endpoint=False) + pts = np.column_stack([0.5 + 0.3 * np.cos(theta), 0.5 + 0.3 * np.sin(theta)]) + + # Same array on every rank; each rank inserts only the points it owns. + swarm.add_particles_with_coordinates(pts) + + # Tag each local particle with the index of its (exactly matching) + # launch coordinate so trajectories can be compared after migration. + local0 = swarm._particle_coordinates.data + d2 = ((local0[:, None, :] - pts[None, :, :]) ** 2).sum(axis=-1) + assert d2.size == 0 or d2.min(axis=1).max() < 1.0e-24 + pid.data[:, 0] = d2.argmin(axis=1) if d2.size else np.empty((0,), dtype=int) + + comm = uw.mpi.comm + n_global = comm.allreduce(local0.shape[0]) + assert n_global == n_pts + + # Force a fixed number of substeps through the step_limit machinery + # (substeps = round(|dt| / dt_limit) inside advection()). + n_sub = 6 + dt_limit = swarm.estimate_dt(V.sym) + delta_t = n_sub * dt_limit + + swarm.advection(V.sym, delta_t, order=2, step_limit=True) + + # Exact replication of the substepped midpoint scheme with the analytic + # (linear) velocity — what advection() must produce when every velocity + # evaluation is performed on the rank that owns the point. + def v_exact(x): + return np.column_stack([-(x[:, 1] - 0.5), x[:, 0] - 0.5]) + + expected = pts.copy() + dt_sub = delta_t / n_sub + for _ in range(n_sub): + mid = expected + 0.5 * dt_sub * v_exact(expected) + expected = expected + dt_sub * v_exact(mid) + + # Gather (pid, final position) from every rank and compare per particle. + # np.asarray: the NDArray_With_Callback wrapper cannot be pickled by + # allgather (closure-defined callback), so gather plain copies. + final_local = np.asarray(swarm._particle_coordinates.data).copy() + pid_local = np.asarray(pid.data[:, 0]).copy() + + all_final = np.concatenate(comm.allgather(final_local), axis=0) + all_pid = np.concatenate(comm.allgather(pid_local), axis=0) + + assert all_final.shape[0] == n_pts, ( + f"Lost particles during substepped advection: {all_final.shape[0]} of {n_pts}" + ) + + order = np.argsort(all_pid) + np.testing.assert_allclose( + all_final[order], + expected, + atol=1.0e-8, + err_msg="Substepped advection deviates from the exact midpoint scheme — " + "launch-point velocities are being evaluated on the wrong rank " + "after particles cross a partition boundary (SWARM-16).", + ) diff --git a/tests/test_0110_basic_swarm.py b/tests/test_0110_basic_swarm.py index a2e7f99b5..9d3d9b4de 100644 --- a/tests/test_0110_basic_swarm.py +++ b/tests/test_0110_basic_swarm.py @@ -157,3 +157,82 @@ def test_particle_clip_context_manager(setup_data): npts1 = swarm._particle_coordinates.data.shape[0] assert npts1 == 0 + + +@pytest.mark.tier_a +def test_recycle_rate_not_implemented(): + """recycle_rate > 1 (streak swarms) must refuse at construction. + + The recycling machinery was excised in 2026-07 (audit SWARM-08/SWARM-09, + remediation D4): it had been broken for some time (NameError in populate + and advection) and had zero tests. A clear NotImplementedError at + construction replaces a crash deep inside populate(). + """ + from underworld3 import swarm + from underworld3.meshing import UnstructuredSimplexBox + + mesh = UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0) + + with pytest.raises(NotImplementedError, match="recycle_rate"): + swarm.Swarm(mesh, recycle_rate=5) + + # The no-recycling defaults still construct and populate normally + s = swarm.Swarm(mesh, recycle_rate=0) + s.populate(fill_param=1) + assert s._particle_coordinates.data.shape[0] > 0 + + +@pytest.mark.tier_a +def test_nodal_point_swarm_deprecated_but_working(): + """NodalPointSwarm: deprecation warning + construction smoke test. + + The class is deprecated (audit SWARM-11, remediation D5) and will be + removed next release cycle; during the warning period it must still + construct correctly. This also pins the SWARM-11 positional-argument + fix: `verbose` used to be passed positionally into Swarm's + `recycle_rate` slot and silently discarded. + """ + import numpy as np + import underworld3 as uw + from underworld3.meshing import UnstructuredSimplexBox + + mesh = UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0) + T = uw.discretisation.MeshVariable("T_nps_smoke", mesh, 1) + + with pytest.warns(DeprecationWarning, match="NodalPointSwarm is deprecated"): + nswarm = uw.swarm.NodalPointSwarm(T, verbose=True) + + # verbose must reach Swarm.__init__ (not land in recycle_rate) + assert nswarm.verbose is True + assert nswarm.recycle_rate == 0 + + # Still functional during the warning period: one particle per node of + # the tracked variable, with launch points recorded for snap-back. + n_local = nswarm._particle_coordinates.data.shape[0] + n_global = uw.mpi.comm.allreduce(n_local) + assert n_global == T.coords.shape[0] + assert nswarm._nX0.data.shape == (n_local, mesh.dim) + + +@pytest.mark.tier_a +def test_estimate_dt_with_mesh_variable_velocity(): + """estimate_dt must return a positive limit for a non-trivial velocity. + + Regression (BF-16 collateral): evaluate() returns matrix-shaped + (n, 1, dim) arrays; estimate_dt indexed vel[:, 1] into the size-1 axis + and the swallowed IndexError made it return None for every velocity — + silently disabling advection's step_limit substepping. + """ + import numpy as np + import underworld3 as uw + + mesh = uw.meshing.StructuredQuadBox(elementRes=(8, 8)) + V = uw.discretisation.MeshVariable("V_dt_est", mesh, mesh.dim, degree=1) + V.array[:, 0, 0] = -(V.coords[:, 1] - 0.5) + V.array[:, 0, 1] = V.coords[:, 0] - 0.5 + + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=1) + + dt_limit = swarm.estimate_dt(V.sym) + assert dt_limit is not None and dt_limit > 0.0 diff --git a/tests/test_0114_swarm_save_coordinate_units.py b/tests/test_0114_swarm_save_coordinate_units.py new file mode 100644 index 000000000..2743a78f8 --- /dev/null +++ b/tests/test_0114_swarm_save_coordinate_units.py @@ -0,0 +1,104 @@ +""" +Regression test for Swarm.save() coordinate-system consistency (SWARM-19 / BF-17). + +``Swarm.save()`` has two IO branches: a parallel-HDF5 path and a sequential +fallback (``force_sequential=True`` or h5py built without MPI). The parallel +branch saved ``_particle_coordinates.data`` (model units) while the sequential +branch saved the deprecated ``self.points`` property, which multiplies by the +model length scale when coordinate scaling is active — so the two branches +produced checkpoints that differed by the length scale, and the sequential +files could not round-trip through ``read_timestep`` (which re-inserts raw +coordinates as model units). + +Both branches must write MODEL-UNIT coordinates (the convention the parallel +branch and ``read_timestep`` already used). +""" + +import os + +import h5py +import numpy as np +import pytest + +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +@pytest.fixture +def scaled_model(): + """A default model with a length scale, so coordinate scaling is active.""" + uw.reset_default_model() + model = uw.get_default_model() + model.set_reference_quantities( + domain_depth=uw.quantity(1000, "km"), + plate_velocity=uw.quantity(5, "cm/year"), + ) + yield model + uw.reset_default_model() + + +def _saved_coordinates(filename): + with h5py.File(filename, "r") as h5f: + return h5f["coordinates"][:] + + +def test_save_writes_model_units_in_both_io_branches(scaled_model, tmp_path): + mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4)) + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=1) + + # Coordinate scaling must actually be active for this test to bite: + # with a 1000 km depth reference the length scale is 1e6 (metres). + assert mesh.CoordinateSystem._scaled + assert mesh.CoordinateSystem._length_scale != 1.0 + + # np.asarray: drop the NDArray_With_Callback wrapper for plain numpy ops + model_coords = np.asarray(swarm._particle_coordinates.data).copy() + + f_seq = str(tmp_path / "swarm_seq.h5") + swarm.save(f_seq, force_sequential=True) + seq_coords = _saved_coordinates(f_seq) + np.testing.assert_allclose( + np.sort(seq_coords, axis=0), + np.sort(model_coords, axis=0), + rtol=1e-12, + err_msg="sequential save() branch must write model-unit coordinates " + "(SWARM-19: it used to write physically-scaled self.points)", + ) + + if h5py.h5.get_config().mpi: + f_par = str(tmp_path / "swarm_par.h5") + swarm.save(f_par) + par_coords = _saved_coordinates(f_par) + np.testing.assert_allclose( + np.sort(par_coords, axis=0), + np.sort(seq_coords, axis=0), + rtol=1e-12, + err_msg="parallel and sequential save() branches must write the " + "same coordinate system (SWARM-19)", + ) + + +def test_save_read_roundtrip_with_scaling(scaled_model, tmp_path): + """A sequential checkpoint must round-trip through read_timestep.""" + mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4)) + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=1) + original = np.sort(np.asarray(swarm._particle_coordinates.data), axis=0) + + # save() then read back through the read_timestep naming convention + base = str(tmp_path / "chk") + filename = base + ".swarm.00000.h5" + swarm.save(filename, force_sequential=True) + + swarm2 = uw.swarm.Swarm(mesh) + swarm2.dm.finalizeFieldRegister() + swarm2.read_timestep(os.path.basename(base), "swarm", 0, outputPath=str(tmp_path)) + + restored = np.sort(np.asarray(swarm2._particle_coordinates.data), axis=0) + assert restored.shape == original.shape, ( + "read_timestep dropped particles: coordinates were saved in the wrong " + "unit system (physically scaled, outside the model-unit mesh domain)" + ) + np.testing.assert_allclose(restored, original, rtol=1e-12) diff --git a/tests/test_0850_comprehensive_reduction_operations.py b/tests/test_0850_comprehensive_reduction_operations.py index ec1617357..5f1c96c68 100644 --- a/tests/test_0850_comprehensive_reduction_operations.py +++ b/tests/test_0850_comprehensive_reduction_operations.py @@ -10,13 +10,11 @@ - Global reductions on _BaseMeshVariable (using PETSc) - Global reductions on UnitAwareArray (using MPI) -STATUS (2025-11-15): -- SWARM TESTS PARTIALLY FIXED: Corrected variable ordering and populate() API -- REAL CODE BUG FOUND: SwarmVariable reductions return scalars, should return tuples - - MeshVariable.array.max() → (1.0, 2.0) for 2D vector ✓ - - SwarmVariable.array.max() → 2.0 for 2D vector ✗ (interface inconsistency) -- Tests are CORRECT - implementation is WRONG -- Marked swarm tests as skip until reduction interface bug is fixed +STATUS (2026-07): +- SwarmVariable array-view reductions now follow the MeshVariable contract: + float for single-component variables, per-component tuple for + multi-component variables (LE-07 / BF-11 remediation). The swarm tests + below are unskipped and act as the regression net for that contract. """ import numpy as np @@ -25,11 +23,10 @@ @pytest.mark.level_2 # Intermediate - array reductions -@pytest.mark.tier_c # Experimental - tests reveal reduction interface bug in SwarmVariable +@pytest.mark.tier_b # Regression net for the LE-07/BF-11 reduction-contract fix class TestSwarmArrayViewReductions: """Test reduction operations on swarm array views.""" - @pytest.mark.skip(reason="BUG: SwarmVariable reductions return scalars instead of tuples for vector variables. Fix SwarmVariable reduction interface to match MeshVariable, then remove skip.") def test_simple_swarm_array_view_reductions(self): """Test all reduction operations on SimpleSwarmArrayView.""" swarm = uw.swarm.Swarm(uw.meshing.StructuredQuadBox(elementRes=(5, 5))) @@ -69,14 +66,14 @@ def test_simple_swarm_array_view_reductions(self): assert len(sum_result) == 2 assert len(std_result) == 2 - @pytest.mark.skip(reason="BUG: SwarmVariable(4) requires explicit vtype parameter + reduction interface bug. Fix SwarmVariable API, then remove skip.") def test_tensor_swarm_array_view_reductions(self): """Test all reduction operations on TensorSwarmArrayView.""" swarm = uw.swarm.Swarm(uw.meshing.StructuredQuadBox(elementRes=(5, 5))) # Create tensor swarm variable BEFORE populating (CRITICAL!) - # Note: This currently fails - need vtype parameter for 4-component variables - tensor_var = uw.swarm.SwarmVariable("tensor", swarm, 4) # 2x2 tensor + # An explicit vtype is required for ambiguous component counts — + # shared, deliberate design (MeshVariable raises the same ValueError). + tensor_var = uw.swarm.SwarmVariable("tensor", swarm, 4, vtype=uw.VarType.TENSOR) # 2x2 tensor # NOW populate the swarm with specific coordinates coords = np.random.RandomState(0).random((100, 2)) @@ -310,10 +307,10 @@ def test_all_reduction_methods_exist(self): assert callable(getattr(var.array, method)), f"Method {method} not callable" -@pytest.mark.xfail(reason="std() method not yet implemented on MeshVariable") class TestStdMethodNewImplementations: """Specific tests for the newly added std() method.""" + @pytest.mark.xfail(reason="std() method not yet implemented on MeshVariable") def test_mesh_variable_std_new_method(self): """Test the newly added std() method on mesh variables.""" mesh = uw.meshing.StructuredQuadBox(elementRes=(5, 5)) @@ -332,7 +329,6 @@ def test_mesh_variable_std_new_method(self): # but std should be positive for non-constant data assert result >= 0 - @pytest.mark.skip(reason="BUG: SwarmVariable populate() ordering issue. Create variable before populate(), then remove skip.") def test_swarm_std_new_method(self): """Test the newly added std() method on swarm variables.""" swarm = uw.swarm.Swarm(uw.meshing.StructuredQuadBox(elementRes=(5, 5)))