Replace read-everywhere checkpoint reads with swarm-routed reads

MeshVariable.read_timestep and SwarmVariable.read_timestep both used to
have every rank h5py-open the saved file and build a rank-local KDTree
over the *full* saved field. PR underworldcode#146 measured 3.92 TB resident at 1152
ranks for a 1/128 spherical case driven by exactly this pattern.

Both methods now use a transient swarm to route saved (coord, value)
pairs to the rank that owns each location. Per-rank memory is bounded
by file_size / n_ranks instead of file_size per rank.

Two pieces:

  * Swarm._route_by_nearest_centroid() (new, swarm.py) — deterministic
    centroid-distance routing. Bypasses Swarm.migrate's points_in_domain
    test, which can return True on multiple ranks for vertex DOFs sitting
    on a partition boundary (owner + ghost). With this routing rule the
    destination is a pure function of the coordinate, so a saved point
    and a query at the same coord always land on the same rank.

  * MeshVariable.read_timestep — round-trip pattern with two transient
    swarms. Source swarm carries (coord, saved_value); rank 0 reads the
    file once and ships chunks via the centroid router. Query swarm
    carries each rank's live DOF coordinates, also routed by centroid.
    The interpolation runs rank-local against the landed source data,
    then results migrate back to the live DOF's home rank via the bare
    DMSwarm migrate idiom (rank-stamped, no validation).

  * SwarmVariable.read_timestep — same recipe, single direction (the
    live swarm's particles already have a known home; the source data is
    routed to where the live particles landed, then a rank-local KDTree
    fills in the values).

The serial round-trip tests (tests/test_0003_save_load.py) pass exactly
as before. New parallel test ptest_0762_read_timestep_swarm_routed.py
exercises the MeshVariable parallel round-trip on 2 ranks; the matching
SwarmVariable parallel test is skipped because Swarm.write_timestep
itself hangs in parallel — a pre-existing issue unrelated to read.

Underworld development team with AI support from Claude Code

Author: lmoresi

src/underworld3/discretisation/discretisation_mesh_variables.py

@@ -70,6 +70,8 @@ def wrapper(final):
     return wrapper
 
 
+
+
 class _BaseMeshVariable(Stateful, uw_object):
     """
     The MeshVariable class generates a variable supported by a finite element mesh and the
@@ -1155,90 +1157,145 @@ def read_timestep(
     ):
         """
         Read a mesh variable from an arbitrary vertex-based checkpoint file
-        and reconstruct/interpolate the data field accordingly. The data sizes / meshes can be
-        different and will be matched using a kd-tree / inverse-distance weighting
-        to the new mesh.
+        and reconstruct/interpolate the data field accordingly. The saved
+        mesh and the live mesh may have different sizes/decompositions; the
+        values are matched by nearest-neighbour kd-tree interpolation to
+        the live mesh nodes.
 
-        """
+        Parallel-safe and memory-bounded. Two transient swarms route the
+        work without ever holding the full file on more than one rank:
+
+          1. **Source swarm** — rank 0 reads the file; saved
+             ``(coord, value)`` pairs migrate to whichever rank owns the
+             centroid-domain of each location.
 
-        # Fix this to match the write_timestep function
+          2. **Query swarm** — each rank inserts *its own* live DOF
+             coordinates. They migrate using the same centroid logic, so
+             a live DOF and a saved point at the same coordinate land on
+             the same rank regardless of how PETSc partitioned the DM.
+             Each rank then runs a rank-local KDTree against the saved
+             data it received, and the interpolated values migrate back to
+             the live DOF's home rank.
 
-        # mesh.write_timestep( "test", meshUpdates=False, meshVars=[X], outputPath="", index=0)
-        # swarm.write_timestep("test", "swarm", swarmVars=[var], outputPath="", index=0)
+        Per-rank memory is bounded by ``file_size / n_ranks`` rather than
+        ``file_size`` per rank.
+        """
 
         output_base_name = os.path.join(outputPath, data_filename)
         data_file = output_base_name + f".mesh.{data_name}.{index:05}.h5"
 
-        # check if data_file exists
-        if os.path.isfile(os.path.abspath(data_file)):
-            pass
-        else:
+        if not os.path.isfile(os.path.abspath(data_file)):
             raise RuntimeError(f"{os.path.abspath(data_file)} does not exist")
 
         import h5py
         import numpy as np
 
-        # Keep vector available for future access
-        pass
-
-        ## Sub functions that are used to read / interpolate the mesh.
-        def field_from_checkpoint(
-            data_file=None,
-            data_name=None,
-        ):
-            """Read the mesh data as a swarm-like value"""
+        n_components = self.shape[1]
+        dim = self.mesh.dim
+
+        # ---- Phase 1: source swarm carries saved (coord, value) pairs ----
+        source_swarm = uw.swarm.Swarm(self.mesh)
+        saved = uw.swarm.SwarmVariable(
+            "_read_timestep_saved",
+            source_swarm,
+            vtype=uw.VarType.MATRIX,
+            size=(1, n_components),
+            dtype=float,
+            _proxy=False,
+            varsymbol=r"\cal{S}",
+        )
 
-            if verbose and uw.mpi.rank == 0:
+        if uw.mpi.rank == 0:
+            if verbose:
                 print(f"Reading data file {data_file}", flush=True)
+            with h5py.File(data_file, "r") as h5f:
+                X_src = h5f["fields"]["coordinates"][()].reshape(-1, dim)
+                D_src = h5f["fields"][data_name][()].reshape(-1, n_components)
+        else:
+            X_src = np.empty((0, dim), dtype=np.double)
+            D_src = np.empty((0, n_components), dtype=np.double)
+
+        src_size_before = max(source_swarm.dm.getLocalSize(), 0)
+        source_swarm.add_particles_with_global_coordinates(X_src, migrate=False)
+        source_swarm._invalidate_canonical_data()
+        saved.array[src_size_before:, 0, :] = D_src[:, :]
+        # Deterministic centroid-distance routing: nearest rank-centroid
+        # owns the point. Both swarms (source + query below) use the same
+        # rule, so a saved point at coord X and a live-DOF query at the
+        # same X always land on the same rank — exact match restored.
+        # ``Swarm.migrate``'s ``points_in_domain`` test isn't enough on its
+        # own: at partition boundaries it can return True on multiple ranks
+        # (vertex DOFs are shared) and source/query end up apart.
+        source_swarm._route_by_nearest_centroid()
+
+        landed_X = source_swarm._particle_coordinates.array[...].reshape(-1, dim)
+        landed_D = saved.array[:, 0, :]
+
+        # ---- Phase 2: query swarm round-trips live DOFs to source rank ----
+        query_coords = self.coords
+        if hasattr(query_coords, "magnitude"):
+            query_coords = query_coords.magnitude
+        n_query_local = query_coords.shape[0]
+        original_index = np.arange(n_query_local).reshape(-1, 1, 1)
+
+        query_swarm = uw.swarm.Swarm(self.mesh)
+        origin_rank = uw.swarm.SwarmVariable(
+            "rank", query_swarm,
+            vtype=uw.VarType.SCALAR, dtype=int, _proxy=False,
+            varsymbol=r"\cal{R}_o",
+        )
+        origin_index_var = uw.swarm.SwarmVariable(
+            "index", query_swarm,
+            vtype=uw.VarType.SCALAR, dtype=int, _proxy=False,
+            varsymbol=r"\cal{I}",
+        )
+        result = uw.swarm.SwarmVariable(
+            "_read_timestep_result", query_swarm,
+            vtype=uw.VarType.MATRIX, size=(1, n_components),
+            dtype=float, _proxy=False, varsymbol=r"\cal{D}",
+        )
 
-            h5f = h5py.File(data_file)
-            D = h5f["fields"][data_name][()].reshape(-1, self.shape[1])
-            X = h5f["fields"]["coordinates"][()].reshape(-1, self.mesh.dim)
-
-            h5f.close()
-
-            if len(D.shape) == 1:
-                D = D.reshape(-1, 1)
-
-            return X, D
-
-        def map_to_vertex_values(X, D, nnn=4, p=2, verbose=False):
-            # Map from "swarm" of points to nodal points
-            # This is a permutation if we building on the checkpointed
-            # mesh file
-
-            mesh_kdt = uw.kdtree.KDTree(X)
-
-            # Strip pint units from query coords — the KDTree was built
-            # from plain HDF5 floats (same physical units, no metadata).
-            query_coords = self.coords
-            if hasattr(query_coords, "magnitude"):
-                query_coords = query_coords.magnitude
-
-            return mesh_kdt.rbf_interpolator_local(query_coords, D, nnn, p, verbose)
-
-        def values_to_mesh_var(mesh_variable, Values):
-            mesh = mesh_variable.mesh
-
-            # This should be trivial but there may be problems if
-            # the kdtree does not have enough neighbours to allocate
-            # values for every point. We handle that here.
-
-            mesh_variable.data[...] = Values[...]
+        q_size_before = max(query_swarm.dm.getLocalSize(), 0)
+        query_swarm.add_particles_with_global_coordinates(query_coords, migrate=False)
+        query_swarm._invalidate_canonical_data()
+        origin_rank.array[q_size_before:, 0, 0] = uw.mpi.rank
+        origin_index_var.array[q_size_before:, 0, 0] = original_index[:, 0, 0]
 
-            return
+        # Forward: live DOF coords go to the rank whose centroid is
+        # closest — same deterministic rule as the source swarm above.
+        query_swarm._route_by_nearest_centroid()
 
-        ## Read file information
+        local_query = query_swarm._particle_coordinates.array[...].reshape(-1, dim)
 
-        X, D = field_from_checkpoint(
-            data_file,
-            data_name,
-        )
+        if landed_X.shape[0] > 0 and local_query.shape[0] > 0:
+            kdt = uw.kdtree.KDTree(landed_X)
+            # ``nnn=1`` — exact match for round-trip reads, sensible
+            # nearest-neighbour fallback for cross-mesh reads.
+            result.array[:, 0, :] = kdt.rbf_interpolator_local(
+                local_query, landed_D, 1, 2, verbose
+            )
+        elif local_query.shape[0] > 0:
+            # No saved data landed on this rank — leave query payload zero
+            # and warn; callers can detect via a missing-rank pattern.
+            if verbose:
+                print(
+                    f"[rank {uw.mpi.rank}] read_timestep: no saved points landed; "
+                    f"queries from this rank will receive zeros",
+                    flush=True,
+                )
 
-        remapped_D = map_to_vertex_values(X, D)
+        # Reverse: stamp the destination rank from origin_rank and use the
+        # bare DMSwarm migrate to send each query particle back home.
+        query_swarm._rank_var.array[...] = origin_rank.array[...]
+        query_swarm.dm.migrate(remove_sent_points=True)
+        uw.mpi.barrier()
+        query_swarm._invalidate_canonical_data()
 
-        # This is empty at the moment
-        values_to_mesh_var(self, remapped_D)
+        # Reorder by original_index and write into self.data
+        idx = origin_index_var.array[:, 0, 0]
+        out = np.zeros((n_query_local, n_components), dtype=np.double)
+        out[idx, :] = result.array[:, 0, :]
+        self.data[...] = out
 
         return
 

src/underworld3/swarm.py

@@ -1910,37 +1910,66 @@ def read_timestep(
         else:
             raise RuntimeError(f"{os.path.abspath(filename)} does not exist")
 
-        ### open up file with coords on all procs and open up data on all procs. May be problematic for large problems.
-        with (
-            h5py.File(f"{filename}", "r") as h5f_data,
-            h5py.File(f"{swarmFilename}", "r") as h5f_swarm,
-        ):
-
-            # with self.swarm.access(self):
-            var_dtype = self.dtype
-            file_dtype = h5f_data["data"][:].dtype
-            file_length = h5f_data["data"][:].shape[0]
-
-            if var_dtype != file_dtype:
-                if comm.rank == 0:
+        # Memory-bounded parallel read: rank 0 streams coords + values from
+        # disk into a transient routing swarm; ``swarm.migrate`` ships each
+        # (coord, value) pair to the rank that owns its location; each rank
+        # then runs a small rank-local KDTree to fill the live swarm's
+        # particles. Per-rank memory scales as ``file_size / n_ranks``.
+
+        n_components = self.num_components
+        dim = self.swarm.mesh.dim
+
+        if uw.mpi.rank == 0:
+            with (
+                h5py.File(f"{filename}", "r") as h5f_data,
+                h5py.File(f"{swarmFilename}", "r") as h5f_swarm,
+            ):
+                file_dtype = h5f_data["data"].dtype
+                if self.dtype != file_dtype:
                     warnings.warn(
-                        f"{os.path.basename(filename)} dtype ({file_dtype}) does not match {self.name} swarm variable dtype ({var_dtype}) which may result in a loss of data.",
+                        f"{os.path.basename(filename)} dtype ({file_dtype}) "
+                        f"does not match {self.name} swarm variable dtype "
+                        f"({self.dtype}) which may result in a loss of data.",
                         stacklevel=2,
                     )
+                X_chunk = h5f_swarm["coordinates"][()].reshape(-1, dim)
+                D_chunk = h5f_data["data"][()].reshape(-1, n_components)
+        else:
+            X_chunk = np.empty((0, dim), dtype=np.double)
+            D_chunk = np.empty((0, n_components), dtype=np.double)
+
+        tmp_swarm = uw.swarm.Swarm(self.swarm.mesh)
+        saved = SwarmVariable(
+            "_read_timestep_saved",
+            tmp_swarm,
+            vtype=uw.VarType.MATRIX,
+            size=(1, n_components),
+            dtype=float,
+            _proxy=False,
+            varsymbol=r"\cal{S}",
+        )
 
-            # First work out which are local points and ignore the rest
-            # This might help speed up the load by dropping lots of particles
-
-            all_coords = h5f_swarm["coordinates"][()]
-            all_data = h5f_data["data"][()]
+        size_before = max(tmp_swarm.dm.getLocalSize(), 0)
+        tmp_swarm.add_particles_with_global_coordinates(X_chunk, migrate=False)
+        tmp_swarm._invalidate_canonical_data()
+        saved.array[size_before:, 0, :] = D_chunk[:, :]
 
-            local_coords = all_coords  # [local]
-            local_data = all_data  # [local]
+        # Deterministic centroid-distance routing — see Swarm._route_by_nearest_centroid.
+        tmp_swarm._route_by_nearest_centroid()
 
-            kdt = uw.kdtree.KDTree(local_coords)
+        landed_X = tmp_swarm._particle_coordinates.array[...].reshape(-1, dim)
+        landed_D = saved.array[:, 0, :]
 
+        if landed_X.shape[0] == 0:
+            warnings.warn(
+                f"[rank {uw.mpi.rank}] read_timestep: no saved swarm points "
+                f"landed locally; '{self.name}' on this rank will not be updated",
+                stacklevel=2,
+            )
+        else:
+            kdt = uw.kdtree.KDTree(landed_X)
             self.array[:, 0, :] = kdt.rbf_interpolator_local(
-                self.swarm._particle_coordinates.data, local_data, nnn=1
+                self.swarm._particle_coordinates.data, landed_D, nnn=1
             )
 
         return
@@ -2557,6 +2586,36 @@ def _invalidate_canonical_data(self):
             if hasattr(var, "_canonical_data"):
                 var._canonical_data = None
 
+    def _route_by_nearest_centroid(self):
+        """Migrate every particle to the rank whose domain-centroid is closest.
+
+        This is a deterministic alternative to :meth:`migrate`: the destination
+        is a pure function of the coordinate, computed identically on every
+        rank. Two swarms migrated this way are guaranteed to place equal
+        coordinates on the same rank — which the standard ``migrate`` does not
+        guarantee at partition boundaries (vertex DOFs sitting on a shared face
+        can return ``True`` from ``points_in_domain`` on multiple ranks).
+
+        Used by checkpoint readers and any consumer that needs source data and
+        query coordinates to converge on the same rank without relying on
+        PETSc's DOF distribution.
+        """
+        centroids = self.mesh._get_domain_centroids()
+        centroid_kdt = uw.kdtree.KDTree(centroids)
+
+        coords = self.dm.getField("DMSwarmPIC_coor").reshape(-1, self.dim).copy()
+        self.dm.restoreField("DMSwarmPIC_coor")
+
+        if coords.shape[0] > 0:
+            _, owner_rank = centroid_kdt.query(coords, k=1, sqr_dists=False)
+            rank_arr = self.dm.getField("DMSwarm_rank")
+            rank_arr[:, 0] = owner_rank.astype(rank_arr.dtype, copy=False)
+            self.dm.restoreField("DMSwarm_rank")
+
+        self.dm.migrate(remove_sent_points=True)
+        uw.mpi.barrier()
+        self._invalidate_canonical_data()
+
     @property
     def mesh(self):
         """The mesh this swarm operates on"""