Merge pull request #77 from kyleaoman/maskable

Add ability to mask catalogue

Author: JBorrow

docs/source/basic_usage/index.rst

@@ -62,6 +62,11 @@ based on the unit metadata in the file.
 There is also full unit information available in the ``data.units`` object, with
 an ``astropy`` cosmology object provided as ``data.units.cosmology``.
 
+Subsets of the catalogue can be read by providing a mask, for instance to read only values for the 3rd object in the catalogue (indexed from 0):
+
+.. code-block:: python
+   masked_data = load("/path/to/catalogue.properties", mask=3)
+
 Creating your first plot
 ------------------------
 

docs/source/particles_files/index.rst

@@ -33,7 +33,8 @@ This returns two objects, ``particles``, and ``unbound_particles``,
 corresponding to both the bound and unbound component of your halo
 respectively. Each of these contains the information required to extract just
 those particles from a snapshot (this is made much easier by using the SWIFT
-integration, shown below).
+integration, shown below). Note that use of a masked ``catalogue`` is not
+supported.
 
 The instances of
 :class:`velociraptor.particles.particles.VelociraptorParticles` have several

velociraptor/__init__.py

@@ -37,6 +37,7 @@ def load(
     filename: str,
     disregard_units: bool = False,
     registration_file_path: Union[str, None] = None,
+    mask: slice = Ellipsis
 ) -> VelociraptorCatalogue:
     """
     Loads a velociraptor catalogue, producing a VelociraptorCatalogue
@@ -63,6 +64,11 @@ def load(
         additional properties with the catalogue. This is an
         advanced feature. See the documentation for more details.
 
+    mask: Union[None, NDArray[bool], int], optional
+        If a boolean array is provided, it is used to mask all
+        catalogue arrays. If an int is provided, catalogue arrays
+        are masked to the single corresponding element.
+
 
     Returns
     -------
@@ -72,7 +78,11 @@ def load(
         .properties file.
     """
 
-    catalogue = VelociraptorCatalogue(filename, disregard_units=disregard_units)
+    catalogue = VelociraptorCatalogue(
+        filename,
+        disregard_units=disregard_units,
+        mask=mask
+    )
 
     if registration_file_path is not None:
         catalogue.register_derived_quantities(registration_file_path)

velociraptor/catalogue/catalogue.py

@@ -88,7 +88,13 @@ def register_field_properties(self):
         return
 
 
-def generate_getter(filename, name: str, field: str, full_name: str, unit):
+def generate_getter(
+        filename,
+        name: str,
+        field: str,
+        full_name: str,
+        unit
+):
     """
     Generates a function that:
 
@@ -115,7 +121,8 @@ def getter(self):
         else:
             with h5py.File(filename, "r") as handle:
                 try:
-                    setattr(self, f"_{name}", unyt.unyt_array(handle[field][...], unit))
+                    mask = getattr(self, "mask")
+                    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:
@@ -161,6 +168,7 @@ def generate_sub_catalogue(
     registration_function: Callable,
     units: VelociraptorUnits,
     field_metadata: List[VelociraptorFieldMetadata],
+    mask: slice = Ellipsis
 ):
     """
     Generates a sub-catalogue object with the correct properties set.
@@ -190,7 +198,7 @@ def generate_sub_catalogue(
                 metadata.snake_case,
                 metadata.path,
                 metadata.name,
-                metadata.unit,
+                metadata.unit
             ),
             generate_setter(metadata.snake_case),
             generate_deleter(metadata.snake_case),
@@ -205,7 +213,7 @@ def generate_sub_catalogue(
     )
 
     # Finally, we can actually create an instance of our new class.
-    catalogue = ThisSubCatalogue(filename=filename)
+    catalogue = ThisSubCatalogue(filename=filename, mask=mask)
     catalogue.valid_sub_paths = valid_sub_paths
 
     return catalogue
@@ -224,8 +232,9 @@ class __VelociraptorSubCatalogue(object):
     # The valid paths contained within
     valid_sub_paths: List[str]
 
-    def __init__(self, filename):
+    def __init__(self, filename, mask=Ellipsis):
         self.filename = filename
+        self.mask = mask
 
         return
 
@@ -250,6 +259,7 @@ def __init__(
         filename: str,
         disregard_units: bool = False,
         extra_registration_functions: Union[None, Dict[str, Callable]] = None,
+        mask: slice = Ellipsis,
     ):
         """
         Initialise the velociraptor catalogue with all of the available
@@ -276,10 +286,16 @@ def __init__(
             should be a dictionary of strings pointing to callables, which
             conform to the registration function API. This is an advanced
             feature.
+
+        mask: slice, optional
+            If a boolean array is provided, it is used to mask all catalogue
+            arrays. If an int is provided, catalogue arrays are masked to the
+            single corresponding element. Default: Ellipsis (``...``).
         """
         self.filename = filename
         self.disregard_units = disregard_units
         self.extra_registration_functions = extra_registration_functions
+        self.mask = mask
 
         self.get_units()
         self.extract_properties_from_units()
@@ -296,11 +312,18 @@ def __str__(self):
         the memory location.
         """
 
-        return (
-            f"Velociraptor catalogue at {self.filename}. "
-            "Contains the following field collections: "
-            f"{', '.join(self.valid_field_metadata.keys())}"
-        )
+        if self.mask is Ellipsis:
+            return (
+                f"Velociraptor catalogue at {self.filename}. "
+                "Contains the following field collections: "
+                f"{', '.join(self.valid_field_metadata.keys())}"
+            )
+        else:
+            return (
+                f"Masked velociraptor catalogue at {self.filename}. "
+                "Contains the following field collections: "
+                f"{', '.join(self.valid_field_metadata.keys())}"
+            )
 
     def __repr__(self):
         return str(self)
@@ -388,6 +411,7 @@ def __create_sub_catalogues(self):
                     registration_function=self.registration_functions[attribute_name],
                     units=self.units,
                     field_metadata=field_metadata,
+                    mask=self.mask
                 ),
             )
 

velociraptor/particles/particles.py

@@ -57,6 +57,9 @@ def __init__(self, filename, catalogue: Union[VelociraptorCatalogue, None]):
         halo.
         """
 
+        if catalogue.mask is not Ellipsis:
+            raise ValueError('VelociraptorGroups not compatible with masked catalogue.')
+        
         self.filename = filename
         self.catalogue = catalogue