Add homogeneous attribute

swiftsimio/masks.py

@@ -15,7 +15,7 @@
 from swiftsimio.objects import InvalidSnapshot
 
 from swiftsimio.accelerated import ranges_from_array
-from typing import Dict
+from typing import Dict, List
 
 
 class SWIFTMask(object):
@@ -476,21 +476,50 @@ def constrain_index(self, index: int):
         """
         Constrain the mask to a single row.
 
-        Intended for use with SOAP catalogues, mask to read only a single row.
+        Intended for use with halo catalogues, mask to read only a single row.
 
         Parameters
         ----------
         index : int
             The index of the row to select.
         """
-        if not self.metadata.filetype == "SOAP":
-            warnings.warn("Not masking a SOAP catalogue, nothing constrained.")
+        if not self.metadata.homogeneous_arrays:
+            warnings.warn(
+                "Different datasets correspond to different objects, nothing constrained."
+            )
             return
         for group_name in self.metadata.present_group_names:
             setattr(self, group_name, np.array([[index, index + 1]]))
             setattr(self, f"{group_name}_size", 1)
         return
 
+    def constrain_indices(self, indices: List):
+        """
+        Constrain the mask to a list of rows.
+
+        Parameters
+        ----------
+        indices : List
+            An list of the indices of the rows to mask.
+        """
+        if not self.metadata.homogeneous_arrays:
+            warnings.warn(
+                "Different datasets correspond to different objects, nothing constrained."
+            )
+            return
+        if self.spatial_only:
+            warnings.warn(
+                "You cannot constrain a mask if spatial_only=True \n"
+                "Please re-initialise the SWIFTMask object with spatial_only=False"
+                "Nothing constrained."
+            )
+            return
+
+        self._shared[...] = False
+        self._shared[indices] = True
+        self._shared_size = len(indices)
+        return
+
     def get_masked_counts_offsets(
         self
     ) -> tuple[dict[str, np.array], dict[str, np.array]]:

swiftsimio/metadata/objects.py

@@ -45,6 +45,10 @@ class SWIFTMetadata(ABC):
     # (as is the case in SOAP) or whether each type (e.g. Gas, Dark Matter, etc.)
     # has its own top-level cell grid counts.
     shared_cell_counts: str | None = None
+    # Whether all the arrays in this files have the same length and order (as is
+    # the case for SOAP, all arrays correspond to subhalos) or whether there are
+    # multiple types (e.g. Gas, Dark Matter, etc.)
+    homogeneous_arrays: bool = False
 
     @abstractmethod
     def __init__(self, filename, units: "SWIFTUnits"):
@@ -1223,6 +1227,8 @@ class SWIFTFOFMetadata(SWIFTMetadata):
     class.
     """
 
+    homogeneous_arrays: bool = True
+
     def __init__(self, filename: str, units: SWIFTUnits):
         self.filename = filename
         self.units = units
@@ -1265,6 +1271,7 @@ class SWIFTSOAPMetadata(SWIFTMetadata):
 
     masking_valid: bool = True
     shared_cell_counts: str = "Subhalos"
+    homogeneous_arrays: bool = True
 
     def __init__(self, filename: str, units: SWIFTUnits):
         self.filename = filename