Merge pull request #171 from kyleaoman/intersect_spatial_masks

Intersect spatial masks

docs/source/masking/index.rst

@@ -68,6 +68,24 @@ relatively small overhead when reading the data. This allows you to use snapshot
 that are much larger than the available memory on your machine and process them
 with ease.
 
+It is also possible to build up a region with a more complicated geometry by
+making repeated calls to :meth:`~swiftsimio.reader.SWIFTMask.constrain_spatial`
+and setting the optional argument `intersect=True`. By default any existing
+selection of cells would be overwritten; this option adds any additional cells
+that need to be selected for the new region to the existing selection instead.
+For instance, to add the diagonally opposed octant to the selection made above
+(and so obtain a region shaped like two cubes with a single corner touching):
+
+.. code-block:: python
+
+   additional_region = [[0.5 * b, 1.0 * b] for b in boxsize]
+   mask.constrain_spatial(additional_region, intersect=True)
+
+In the first call to :meth:`~swiftsimio.reader.SWIFTMask.constrain_spatial` the
+`intersect` argument can be set to `True` or left `False` (the default): since
+no mask yet exists both give the same result.
+
+
 Full mask
 ---------
 

swiftsimio/masks.py

@@ -8,7 +8,7 @@
 
 import numpy as np
 
-from swiftsimio import metadata, SWIFTMetadata, SWIFTUnits
+from swiftsimio import SWIFTMetadata
 
 from swiftsimio.accelerated import ranges_from_array
 from typing import Dict
@@ -316,7 +316,7 @@ def _update_spatial_mask(self, restrict, ptype: str, cell_mask: np.array):
 
         return
 
-    def constrain_spatial(self, restrict):
+    def constrain_spatial(self, restrict, intersect: bool = False):
         """
         Uses the cell metadata to create a spatial mask.
 
@@ -339,13 +339,25 @@ def constrain_spatial(self, restrict):
             These values must have units associated with them. It is also acceptable
             to have a row as None to not restrict in this direction.
 
+        intersect : bool
+            If `True`, intersect the spatial mask with any existing spatial mask to
+            select two (or more) regions with repeated calls to `constrain_spatial`.
+            By default (`False`) any existing mask is overwritten.
+
         See Also
         -------
 
         constrain_mask : method to further refine mask
         """
 
-        self.cell_mask = self._generate_cell_mask(restrict)
+        if hasattr(self, "cell_mask") and intersect:
+            # we already have a mask and are in intersect mode
+            self.cell_mask = np.logical_or(
+                self.cell_mask, self._generate_cell_mask(restrict)
+            )
+        else:
+            # we just make a new mask
+            self.cell_mask = self._generate_cell_mask(restrict)
 
         for ptype in self.metadata.present_particle_names:
             self._update_spatial_mask(restrict, ptype, self.cell_mask)

tests/test_mask.py

@@ -4,6 +4,7 @@
 
 from tests.helper import requires
 from swiftsimio import load, mask
+import numpy as np
 
 from unyt import unyt_array as array, dimensionless
 
@@ -90,3 +91,26 @@ def test_region_mask_not_modified(filename):
     this_mask._generate_cell_mask(read)
 
     assert read == read_constant
+
+
+@requires("cosmological_volume.hdf5")
+def test_region_mask_intersection(filename):
+    """
+    Tests that the intersection of two spatial mask regions includes the same cells as two
+    separate masks of the same two regions.
+    """
+
+    mask_1 = mask(filename, spatial_only=True)
+    mask_2 = mask(filename, spatial_only=True)
+    mask_intersect = mask(filename, spatial_only=True)
+    bs = mask_intersect.metadata.boxsize
+    region_1 = [[0 * b, 0.1 * b] for b in bs]
+    region_2 = [[0.6 * b, 0.7 * b] for b in bs]
+    mask_1.constrain_spatial(region_1)
+    mask_2.constrain_spatial(region_2)
+    # the intersect=True flag is optional on the first call:
+    mask_intersect.constrain_spatial(region_1, intersect=True)
+    mask_intersect.constrain_spatial(region_2, intersect=True)
+    assert (
+        np.logical_or(mask_1.cell_mask, mask_2.cell_mask) == mask_intersect.cell_mask
+    ).all()