Merge pull request #110 from kyleaoman/load_fof_catalogues

Changes for when swiftsimio merges SOAP support
SWIFTSIM · Oct 9, 2024 · f9a4090 · f9a4090
2 parents 23a99a6 + a12c933
commit f9a4090
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 34 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -3,7 +3,7 @@ requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
-packages = ["velociraptor"]
+packages = ["velociraptor", "velociraptor.catalogue", "velociraptor.fitting_formulae", "velociraptor.observations", "velociraptor.autoplotter", "velociraptor.particles", "velociraptor.swift", "velociraptor.tools"]
 
 [project]
 name = "velociraptor-python"

diff --git a/setup.py b/setup.py
diff --git a/velociraptor/catalogue/velociraptor_catalogue.py b/velociraptor/catalogue/velociraptor_catalogue.py
@@ -117,9 +117,30 @@ def getter(self):
             with h5py.File(filename, "r") as handle:
                 try:
                     mask = getattr(self, "mask")
-                    setattr(
-                        self, f"_{name}", unyt.unyt_array(handle[field][mask], unit)
-                    )
+                    if (
+                        np.ndim(mask) != 0
+                        and np.issubdtype(np.array(mask).dtype, np.integer)
+                        and not np.all(mask[:-1] < mask[1:])
+                    ):
+                        # We have a mask picking out items by index, and it's
+                        # not sorted. hdf5 demands that it be sorted.
+                        # We sort, read and then reverse the sort.
+                        sort_mask = np.argsort(mask)
+                        unsort_mask = np.argsort(sort_mask)
+                        setattr(
+                            self,
+                            f"_{name}",
+                            unyt.unyt_array(
+                                handle[field][mask[sort_mask]][unsort_mask],
+                                unit,
+                            ),
+                        )
+                    else:
+                        setattr(
+                            self,
+                            f"_{name}",
+                            unyt.unyt_array(handle[field][mask], unit),
+                        )
                     getattr(self, f"_{name}").name = full_name
                     getattr(self, f"_{name}").file = filename
                 except KeyError:

diff --git a/velociraptor/swift/swift.py b/velociraptor/swift/swift.py
@@ -5,7 +5,6 @@
 datasets in a computationally efficient way.
 """
 
-
 import swiftsimio
 import numpy as np
 
@@ -98,7 +97,7 @@ def generate_bound_mask(
 
     particle_name_masks = {}
 
-    for particle_name in data.metadata.present_particle_names:
+    for particle_name in data.metadata.present_group_names:
         # This will change if we ever take advantage of the
         # parttypes available through velociraptor.
         particle_name_masks[particle_name] = np.in1d(
@@ -107,7 +106,7 @@ def generate_bound_mask(
 
     # Finally we generate a named tuple with the correct fields and
     # fill it with the contents of our dictionary
-    MaskTuple = namedtuple("MaskCollection", data.metadata.present_particle_names)
+    MaskTuple = namedtuple("MaskCollection", data.metadata.present_group_names)
     mask = MaskTuple(**particle_name_masks)
 
     return mask