Allow SOAP to run without parallel h5py

README.md

@@ -3,19 +3,23 @@
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.08252/status.svg)](https://doi.org/10.21105/joss.08252) 
 
 This repository contains programs which can be used to compute
-properties of halos in spherical apertures in [SWIFT](https://swift.strw.leidenuniv.nl/) snapshots.
+properties of halos in spherical apertures in
+[SWIFT](https://swift.strw.leidenuniv.nl/) snapshots.
 The resulting output halo catalogues can be read using the
 [swiftsimio](https://swiftsimio.readthedocs.io/en/latest/)
 python package.
 
-Please cite SOAP using the [JOSS paper](https://ui.adsabs.harvard.edu/abs/2025JOSS...10.8252M)
+Please cite SOAP using the
+[JOSS paper](https://ui.adsabs.harvard.edu/abs/2025JOSS...10.8252M).
 
 ## Installation
 
 The code is written in python and uses mpi4py for parallelism.
-IO is carried out in parallel, and so [parallel h5py](https://docs.h5py.org/en/stable/mpi.html) is required. SOAP and it's dependencies can also be 
+IO is also intended to run in parallel, and so
+[parallel h5py](https://docs.h5py.org/en/stable/mpi.html) is recommended.
+SOAP and its dependencies can be
 installed directly using the command
-`pip install git+https://github.com/SWIFTSIM/SOAP.git@soap_runtime`
+`pip install git+https://github.com/SWIFTSIM/SOAP.git`
 but this may install a serial version of h5py. Therefore the following
 steps are recommended for install
 ```

SOAP/core/chunk_tasks.py

@@ -11,6 +11,7 @@
 from . import shared_mesh
 from . import result_set
 from . import memory_use
+from . import parallel_io
 from .dataset_names import mass_dataset, ptypes_for_so_masses
 from .halo_tasks import process_halos
 from .mask_cells import mask_cells
@@ -388,10 +389,9 @@ def message(m):
         comm_have_results = comm.Split(colour, comm_rank)
         if len(results) > 0:
             filename = scratch_file_format % {"file_nr": self.chunk_nr}
-            with h5py.File(
-                filename, "w", driver="mpio", comm=comm_have_results
-            ) as outfile:
-                results.collective_write(outfile, comm_have_results)
+            outfile = parallel_io.open_collective(filename, "w", comm_have_results)
+            results.collective_write(outfile, comm_have_results)
+            parallel_io.close_collective(outfile, comm_have_results)
         comm_have_results.Free()
         comm.barrier()
 

SOAP/core/combine_chunks.py

@@ -14,7 +14,7 @@
 from SOAP.catalogue_readers import read_hbtplus
 from SOAP.property_calculation.subhalo_rank import compute_subhalo_rank
 from SOAP.property_table import PropertyTable
-from . import lustre, swift_units
+from . import lustre, parallel_io, swift_units
 from .mpi_timer import MPITimer
 
 
@@ -364,8 +364,8 @@ def combine_chunks(
             outfile.close()
     comm_world.barrier()
 
-    # Reopen the output file in parallel mode
-    outfile = h5py.File(output_file, "r+", driver="mpio", comm=comm_world)
+    # Reopen the output file for writing by all ranks.
+    outfile = parallel_io.open_collective(output_file, "r+", comm_world)
     props_kept = {}
 
     with MPITimer("Writing output properties", comm_world):
@@ -392,7 +392,7 @@ def combine_chunks(
 
             # Write these properties to the output file
             for name in names:
-                phdf5.collective_write(
+                parallel_io.collective_write(
                     outfile, name, data[name], create_dataset=False, comm=comm_world
                 )
 
@@ -453,7 +453,7 @@ def combine_chunks(
         if not physical:
             soap_com_unit = soap_com_unit * cellgrid.get_unit("a") ** a_exponent
         fof_com = (fof_com * fof_com_unit).to(soap_com_unit)
-        phdf5.collective_write(
+        parallel_io.collective_write(
             outfile,
             "InputHalos/FOF/Centres",
             fof_com,
@@ -472,7 +472,7 @@ def combine_chunks(
         if not physical:
             soap_mass_unit = soap_mass_unit * cellgrid.get_unit("a") ** a_exponent
         fof_mass = (fof_mass * fof_mass_unit).to(soap_mass_unit)
-        phdf5.collective_write(
+        parallel_io.collective_write(
             outfile,
             "InputHalos/FOF/Masses",
             fof_mass,
@@ -484,7 +484,7 @@ def combine_chunks(
         fof_size[keep] = psort.fetch_elements(
             fof_file.read("Groups/Sizes"), indices, comm=comm_world
         )
-        phdf5.collective_write(
+        parallel_io.collective_write(
             outfile,
             "InputHalos/FOF/Sizes",
             fof_size,
@@ -518,7 +518,7 @@ def combine_chunks(
             if not physical:
                 soap_radii_unit = soap_radii_unit * cellgrid.get_unit("a") ** a_exponent
             fof_radii = (fof_radii * fof_com_unit).to(soap_radii_unit)
-            phdf5.collective_write(
+            parallel_io.collective_write(
                 outfile,
                 "InputHalos/FOF/Radii",
                 fof_radii,
@@ -558,7 +558,7 @@ def combine_chunks(
                 host_halo_index = -1 * np.ones(sat_mask.shape[0], dtype=np.int64)
                 host_halo_index[has_host_mask] = indices
 
-            phdf5.collective_write(
+            parallel_io.collective_write(
                 outfile,
                 "SOAP/HostHaloIndex",
                 host_halo_index,
@@ -587,7 +587,7 @@ def combine_chunks(
             subhalo_rank = compute_subhalo_rank(
                 host_id, props_kept["BoundSubhalo/TotalMass"], comm_world
             )
-            phdf5.collective_write(
+            parallel_io.collective_write(
                 outfile,
                 "SOAP/SubhaloRankByBoundMass",
                 subhalo_rank,
@@ -657,7 +657,7 @@ def combine_chunks(
                 assert n_keep[i_bin] <= np.sum(mask)
                 keep_idx = np.random.choice(idx, size=n_keep[i_bin], replace=False)
                 reduced_snapshot[keep_idx] = 1
-            phdf5.collective_write(
+            parallel_io.collective_write(
                 outfile,
                 "SOAP/IncludedInReducedSnapshot",
                 reduced_snapshot,
@@ -721,7 +721,7 @@ def combine_chunks(
                     track_id, prev_track_id, comm=comm_world
                 )
 
-            phdf5.collective_write(
+            parallel_io.collective_write(
                 outfile,
                 f"SOAP/{name}Index",
                 prev_index,
@@ -730,4 +730,4 @@ def combine_chunks(
             )
 
     # Done.
-    outfile.close()
+    parallel_io.close_collective(outfile, comm_world)

SOAP/core/parallel_io.py

@@ -0,0 +1,49 @@
+#!/bin/env python
+
+import h5py
+import virgo.mpi.parallel_hdf5 as phdf5
+
+# True if h5py was built without MPI support, in which case we can't use the
+# mpio driver and have to fall back to writing output files on a single rank.
+SERIAL_HDF5 = phdf5.SERIAL_HDF5
+
+
+def open_collective(filename, mode, comm):
+    """
+    Open a file which all ranks in comm will access.
+
+    With parallel HDF5 this returns a file handle opened in MPI mode on every
+    rank. Without it, files opened for reading are opened independently on
+    each rank, and files being written are opened on rank 0 only with the
+    other ranks getting None. Any code which touches the returned object
+    directly (creating groups, writing attributes) must check for None.
+    """
+    if SERIAL_HDF5:
+        if mode == "r":
+            return h5py.File(filename, "r")
+        if comm.Get_rank() == 0:
+            return h5py.File(filename, mode)
+        return None
+    return h5py.File(filename, mode, driver="mpio", comm=comm)
+
+
+def close_collective(outfile, comm):
+    """
+    Close a file opened with open_collective().
+    """
+    comm.barrier()
+    if outfile is not None:
+        outfile.close()
+    comm.barrier()
+
+
+def collective_write(group, name, data, comm, **kwargs):
+    """
+    Write a dataset by concatenating the contributions from all ranks in comm
+    along the first axis.
+
+    All ranks in comm must call this, including those with no data to write.
+    """
+    if SERIAL_HDF5:
+        return phdf5.serial_collective_write(group, name, data, comm, **kwargs)
+    return phdf5.collective_write(group, name, data, comm, **kwargs)

SOAP/core/result_set.py

@@ -6,9 +6,9 @@
 import numpy as np
 from mpi4py import MPI
 import unyt
-import virgo.mpi.parallel_hdf5 as phdf5
 import virgo.mpi.parallel_sort as psort
 
+from . import parallel_io
 from . import swift_units
 
 
@@ -243,6 +243,9 @@ def find_groups_to_create(paths):
     def collective_write(self, outfile, comm):
         """
         Write the results to a file in collective mode
+
+        outfile is None on ranks which don't have the file open, which is the
+        case for ranks other than rank 0 if we don't have parallel HDF5.
         """
 
         # Ensure arrays are exactly the right size
@@ -253,24 +256,26 @@ def collective_write(self, outfile, comm):
 
         # Ensure any HDF5 groups we need exist
         group_names = comm.bcast(self.find_groups_to_create(names))
-        for group_name in group_names:
-            outfile.create_group(group_name)
+        if outfile is not None:
+            for group_name in group_names:
+                outfile.create_group(group_name)
 
         # Loop over output arrays
         for name in names:
 
             # Write this array
             data, description, physical, a_exponent = self.result_arrays[name]
-            phdf5.collective_write(outfile, name, data, comm)
+            parallel_io.collective_write(outfile, name, data, comm)
 
             # Attach units metadata and description
-            if hasattr(data, "units"):
-                attrs = swift_units.attributes_from_units(
-                    data.units, physical, a_exponent
-                )
-                for attr_name, attr_value in attrs.items():
-                    outfile[name].attrs[attr_name] = attr_value
-            outfile[name].attrs["Description"] = description
+            if outfile is not None:
+                if hasattr(data, "units"):
+                    attrs = swift_units.attributes_from_units(
+                        data.units, physical, a_exponent
+                    )
+                    for attr_name, attr_value in attrs.items():
+                        outfile[name].attrs[attr_name] = attr_value
+                outfile[name].attrs["Description"] = description
 
     def get_metadata(self, comm):
         """

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "scipy",
     "matplotlib",
     "psutil",
-    "virgodc",
+    "virgodc>=1.0.5",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -6,6 +6,6 @@ astropy>=6
 scipy
 matplotlib
 psutil
-virgodc>=1.0.3
+virgodc>=1.0.5
 numba
 pytest-mpi

tests/test_subhalo_rank.py

@@ -1,11 +1,11 @@
 #!/bin/env python
 
 import numpy as np
-import h5py
 import pytest
 from mpi4py import MPI
 import virgo.mpi.parallel_hdf5 as phdf5
 
+from SOAP.core import parallel_io
 from SOAP.property_calculation.subhalo_rank import compute_subhalo_rank
 
 import helpers
@@ -19,7 +19,7 @@
 def test_subhalo_rank(filename):
 
     # Read HBT halos from a small DMO run
-    with h5py.File(filename, "r", driver="mpio", comm=comm) as file:
+    with parallel_io.open_collective(filename, "r", comm) as file:
         sub = phdf5.collective_read(file["Subhalos"], comm=comm)
     if comm_rank == 0:
         print("Read subhalos")