Implement lazy masks in the ToyHF.

Author: kyleaoman

swiftgalaxy/demo_data.py

@@ -7,13 +7,16 @@
 import numpy as np
 import unyt as u
 import subprocess
-from typing import Optional, Callable, Any, Union, Sequence
+from types import EllipsisType
+from typing import Optional, Callable, Any, Union, List, Sequence
+from numpy.typing import NDArray
 from astropy.cosmology import LambdaCDM
 from astropy import units as U
 from swiftsimio.objects import cosmo_array
 import swiftsimio
 from swiftsimio import Writer, SWIFTMask
-from swiftgalaxy import MaskCollection, SWIFTGalaxy
+from swiftgalaxy import SWIFTGalaxy
+from swiftgalaxy.masks import MaskCollection, LazyMask
 from swiftgalaxy.halo_catalogues import _HaloCatalogue
 from swiftsimio.units import cosmo_units
 
@@ -314,7 +317,7 @@ class ToyHF(_HaloCatalogue):
     def __init__(
         self,
         snapfile: Union[str, Path] = _toysnap_filename,
-        index: Union[int, Sequence[int]] = 0,
+        index: Union[int, List[int]] = 0,
     ) -> None:
         self.snapfile = snapfile
         if isinstance(index, Sequence):
@@ -394,23 +397,87 @@ def _generate_bound_only_mask(self, sg: SWIFTGalaxy) -> MaskCollection:
         out : :class:`~swiftgalaxy.masks.MaskCollection`
             The extra mask.
         """
-        # the two objects are in different cells, remember we're masking cell particles
-        if self.index == 0:
-            extra_mask = MaskCollection(
-                gas=np.s_[-_n_g_1:],
-                dark_matter=np.s_[-_n_dm_1:],
-                stars=np.s_[...],
-                black_holes=np.s_[...],
-            )
-        else:  # self.index == 1
-            extra_mask = MaskCollection(
-                gas=np.s_[-_n_g_2:],
-                dark_matter=np.s_[-_n_dm_2:],
-                stars=np.s_[...],
-                black_holes=np.s_[...],
-            )
 
-        return extra_mask
+        def generate_lazy_mask(group_name: str) -> LazyMask:
+            """
+            Generate a function that evaluates a mask for bound particles of a specified
+            particle type. The generated function should have one parameter, accepting a
+            boolean, that toggles masking the data loaded during the construction of the
+            mask on and off.
+
+            Parameters
+            ----------
+            group_name : :obj:`str`
+                The particle type to evaluate a mask for.
+
+            Returns
+            -------
+            out : Callable
+                The generated function that evaluates a mask.
+            """
+
+            def lazy_mask(
+                mask_loaded_data: bool = True,
+            ) -> Union[NDArray, slice, EllipsisType]:
+                """
+                "Evaluate" a mask that selects bound particles. In reality we know what
+                the mask is a priori. We pretend that we need to load the particle ids
+                so that we can test the behaviour of a dataset loaded while constructing
+                the mask.
+
+                This function must optionally mask the data (``particle_ids``) that it
+                has loaded.
+
+                Parameters
+                ----------
+                mask_loaded_data : :obj:`bool`, default ``True``
+                    If ``True``, data loaded while constructing the mask is masked
+                    during this function call. Set ``False`` when called from
+                    a :class:`~swiftgalaxy.iterator.SWIFTGalaxies` "server".
+
+                Returns
+                -------
+                out : :class:`~numpy.ndarray`, :obj:`slice` or :obj:`Ellipsis`
+                    The mask that selects bound particles.
+                """
+                getattr(
+                    getattr(sg, group_name)._particle_dataset,
+                    sg.id_particle_dataset_name,
+                )  # load the ids
+                assert isinstance(self._mask_index, int)  # placate mypy
+                mask = {
+                    "gas": (np.s_[-_n_g_1:], np.s_[-_n_g_2:])[self._mask_index],
+                    "dark_matter": (np.s_[-_n_dm_1:], np.s_[-_n_dm_2:])[
+                        self._mask_index
+                    ],
+                    "stars": np.s_[...],
+                    "black_holes": np.s_[...],
+                }[group_name]
+                if mask_loaded_data:
+                    # mask the particle_ids
+                    setattr(
+                        getattr(sg, group_name)._particle_dataset,
+                        f"_{sg.id_particle_dataset_name}",
+                        getattr(
+                            getattr(sg, group_name)._particle_dataset,
+                            f"_{sg.id_particle_dataset_name}",
+                        )[mask],
+                    )
+                assert (
+                    isinstance(mask, np.ndarray)
+                    or isinstance(mask, slice)
+                    or (mask is Ellipsis)
+                )  # placate mypy
+                return mask
+
+            return LazyMask(mask_function=lazy_mask)
+
+        return MaskCollection(
+            **{
+                group_name: generate_lazy_mask(group_name)
+                for group_name in sg.metadata.present_group_names
+            }
+        )
 
     @property
     def centre(self) -> cosmo_array:

swiftgalaxy/halo_catalogues.py

@@ -806,7 +806,7 @@ def lazy_mask(mask_loaded_data: bool = True) -> NDArray:
                 Evaluate a mask that selects bound particles by comparing the particle
                 group membership dataset ``group_nr_bound`` to the halo catalogue index.
 
-                This function must optionally mask the data (``group_nr_bound``) that is
+                This function must optionally mask the data (``group_nr_bound``) that it
                 has loaded.
 
                 Parameters
@@ -1178,7 +1178,7 @@ def lazy_mask(mask_loaded_data: bool = True) -> NDArray:
                 Evaluate a mask that selects bound particles by comparing the
                 ``particle_ids`` to the list of bound particle IDs.
 
-                This function must optionally mask the data (``particle_ids``) that is has
+                This function must optionally mask the data (``particle_ids``) that it has
                 loaded.
 
                 Parameters
@@ -1799,7 +1799,7 @@ def lazy_mask(mask_loaded_data: bool = True) -> Union[NDArray, slice]:
                 Evaluate a mask that selects bound particles by comparing the lists of
                 bound particle indices to the ranges read in the spatial mask.
 
-                This function must optionally mask the data that is has loaded, but it
+                This function must optionally mask the data that it has loaded, but it
                 loads nothing.
 
                 Parameters

swiftgalaxy/reader.py

@@ -844,6 +844,12 @@ def _mask_dataset(self, mask: LazyMask) -> None:
         particle_metadata = getattr(
             self._particle_dataset.metadata, f"{particle_name}_properties"
         )
+        # force old mask evaluation to ensure any data loaded during evaluation
+        # are in memory and have the old mask applied, if any:
+        old_mask = getattr(self._swiftgalaxy._extra_mask, particle_name)
+        if old_mask is not None:
+            old_mask._evaluate()
+        # apply the new mask to any data already in memory:
         for field_name in particle_metadata.field_names:
             if self._is_namedcolumns(field_name):
                 for named_column in getattr(self, field_name).named_columns:
@@ -861,6 +867,7 @@ def _mask_dataset(self, mask: LazyMask) -> None:
                         )
             elif getattr(self._particle_dataset, f"_{field_name}") is not None:
                 setattr(self, field_name, getattr(self, field_name)[mask.mask])
+        # also the derived coordinates, if any:
         self._mask_derived_coordinates(mask)
         if getattr(self._swiftgalaxy._extra_mask, particle_name) is None:
             setattr(self._swiftgalaxy._extra_mask, particle_name, mask)
@@ -875,7 +882,6 @@ def _mask_dataset(self, mask: LazyMask) -> None:
                         particle_name
                     ]
                 )
-            old_mask = getattr(self._swiftgalaxy._extra_mask, particle_name)
             # need to convert to an integer mask to combine
             # (boolean is insufficient in case of re-ordering masks)
             setattr(

tests/test_masking.py

@@ -35,6 +35,8 @@ def test_reordering_slice_mask(self, sg, particle_name, before_load):
         ids_before = getattr(sg, particle_name).particle_ids
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         sg.mask_particles(MaskCollection(**{particle_name: mask}))
         ids = getattr(sg, particle_name).particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -54,6 +56,8 @@ def test_reordering_int_mask(self, sg, particle_name, before_load):
         mask = mask[: mask.size // 2]
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         sg.mask_particles(MaskCollection(**{particle_name: mask}))
         ids = getattr(sg, particle_name).particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -69,6 +73,8 @@ def test_bool_mask(self, sg, particle_name, before_load):
         mask = np.random.rand(ids_before.size) > 0.5
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         sg.mask_particles(MaskCollection(**{particle_name: mask}))
         ids = getattr(sg, particle_name).particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -82,6 +88,8 @@ def test_namedcolumn_masked(self, sg, before_load):
         mask = np.random.rand(neutral_before.size) > 0.5
         if before_load:
             sg.gas.hydrogen_ionization_fractions._named_column_dataset._neutral = None
+            del sg._extra_mask.gas._mask
+            sg._extra_mask.gas._evaluated = False
         sg.mask_particles(MaskCollection(**{"gas": mask}))
         neutral = sg.gas.hydrogen_ionization_fractions.neutral
         assert_allclose_units(
@@ -135,6 +143,8 @@ def test_reordering_slice_mask(self, sg, particle_name, before_load):
         ids_before = getattr(sg, particle_name).particle_ids
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         masked_dataset = getattr(sg, particle_name)[mask]
         ids = masked_dataset.particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -154,6 +164,8 @@ def test_reordering_int_mask(self, sg, particle_name, before_load):
         mask = mask[: mask.size // 2]
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         masked_dataset = getattr(sg, particle_name)[mask]
         ids = masked_dataset.particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -169,6 +181,8 @@ def test_bool_mask(self, sg, particle_name, before_load):
         mask = np.random.rand(ids_before.size) > 0.5
         if before_load:
             getattr(sg, particle_name)._particle_dataset._particle_ids = None
+            del getattr(sg._extra_mask, particle_name)._mask
+            getattr(sg._extra_mask, particle_name)._evaluated = False
         masked_dataset = getattr(sg, particle_name)[mask]
         ids = masked_dataset.particle_ids
         assert_allclose_units(ids_before[mask], ids, rtol=0, atol=0)
@@ -184,6 +198,8 @@ def test_reordering_slice_mask(self, sg, before_load):
         fractions_before = sg.gas.hydrogen_ionization_fractions.neutral
         if before_load:
             sg.gas.hydrogen_ionization_fractions._neutral = None
+            del sg._extra_mask.gas._mask
+            sg._extra_mask.gas._evaluated = False
         masked_namedcolumnsdataset = sg.gas.hydrogen_ionization_fractions[mask]
         fractions = masked_namedcolumnsdataset.neutral
         assert_allclose_units(
@@ -204,6 +220,8 @@ def test_reordering_int_mask(self, sg, before_load):
         mask = mask[: mask.size // 2]
         if before_load:
             sg.gas.hydrogen_ionization_fractions._neutral = None
+            del sg._extra_mask.gas._mask
+            sg._extra_mask.gas._mask = False
         masked_namedcolumnsdataset = sg.gas.hydrogen_ionization_fractions[mask]
         fractions = masked_namedcolumnsdataset.neutral
         assert_allclose_units(
@@ -220,6 +238,8 @@ def test_bool_mask(self, sg, before_load):
         mask = np.random.rand(fractions_before.size) > 0.5
         if before_load:
             sg.gas.hydrogen_ionization_fractions._neutral = None
+            del sg._extra_mask.gas._mask
+            sg._extra_mask.gas._evaluated = False
         masked_namedcolumnsdataset = sg.gas.hydrogen_ionization_fractions[mask]
         fractions = masked_namedcolumnsdataset.neutral
         assert_allclose_units(