Merge pull request #69 from kyleaoman/split_mask_funcs

Split spatial masking and masking by group ID into separate functions.

Author: JBorrow

docs/source/particles_files/index.rst

@@ -86,6 +86,9 @@ follows:
    sphviewer = SPHViewerWrapper(data.gas)
    sphviewer.quickview(xsize=1024,ysize=1024,r="infinity")
 
+The spatial mask and extra mask can be obtained individually by using
+the functions :func:`velociraptor.swift.swift.generate_spatial_mask` and
+:func:`velociraptor.swift.swift.generate_bound_mask`, respectively.
 
 To see these functions in action, you can check out the examples available in
 ``examples/swift_integration*.py``` in the repository.

velociraptor/swift/swift.py

@@ -16,44 +16,28 @@
 from collections import namedtuple
 
 
-def to_swiftsimio_dataset(
-    particles: VelociraptorParticles,
-    snapshot_filename,
-    generate_extra_mask: bool = False,
-) -> Union[
-    swiftsimio.reader.SWIFTDataset, Tuple[swiftsimio.reader.SWIFTDataset, namedtuple]
-]:
+def generate_spatial_mask(
+        particles: VelociraptorParticles,
+        snapshot_filename
+) -> swiftsimio.mask:
     """
-    Loads a VelociraptorParticles instance for one halo into a 
-    `swiftsimio` masked dataset.
+    Determines the mask defining the spatial region containing
+    the particles of interest.
 
-    Initially, this uses `r_max` to perform a spatial mask, and
-    then returns the `swiftsimio` dataset and a secondary mask
-    that may be used to extract only the particles that are
-    part of the FoF group.
-
-    You will need to instantiate the VelociraptorParticles instance
+    This uses `r_max` to define the extent of the region, so you
+    will need to instantiate the VelociraptorParticles instance
     with an associated catalogue to use this feature, as it requires
     the knowledge of `r_max`.
-    
-    Takes three arguments:
+
+    Takes two arguments:
 
     + particles, the VelociraptorParticles instance,
     + snapshot_filename, the path to the associated SWIFT snapshot.
-    + generate_extra_mask, whether or not to generate the secondary 
-                           mask object that allows for the extraction
-                           of particles that are present only in the
-                           FoF group.
 
     It returns:
 
-    + data, the swiftsimio dataset
-    + mask, an object containing for all available datasets in the
-            swift dataset. The initial masking is performed on a
-            spatial only basis, and this is required to only extract
-            the particles in the FoF group as identified by 
-            velociraptor. This is only provided if generate_extra_mask
-            has a truthy value.
+    + mask, an object containing masks for all available datasets in the
+            swift dataset.
     """
 
     # First use the swiftsimio spatial masking to constrain our dataset
@@ -70,7 +54,8 @@ def to_swiftsimio_dataset(
             length_factor = 1.0
     except AttributeError:
         raise RuntimeError(
-            "Please use a particles instance with an associated halo catalogue."
+            "Please use a particles instance with an associated halo "
+            "catalogue."
         )
 
     spatial_mask = [
@@ -90,14 +75,27 @@ def to_swiftsimio_dataset(
 
     swift_mask.constrain_spatial(spatial_mask)
 
-    # TODO: Make spatial masking work
-    # swift_mask = None
+    return swift_mask
 
-    data = swiftsimio.load(snapshot_filename, mask=swift_mask)
 
-    if not generate_extra_mask:
-        return data
+def generate_bound_mask(
+        data: swiftsimio.reader.SWIFTDataset,
+        particles: VelociraptorParticles
+) -> namedtuple:
+    """
+    Determines the mask defining the particles bound to the object
+    of interest.
+
+    Takes two arguments:
+
+    + data, a SWIFTDataset, which may be spatially masked,
+    + particles, the VelociraptorParticles instance.
 
+    It returns:
+
+    + mask, an object containing masks for all available datasets in the
+            swift dataset.
+    """
     # Now we must generate the secondary mask, for all available
     # particle types.
 
@@ -112,7 +110,63 @@ def to_swiftsimio_dataset(
 
     # 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_particle_names
+    )
     mask = MaskTuple(**particle_name_masks)
 
+    return mask
+
+
+def to_swiftsimio_dataset(
+    particles: VelociraptorParticles,
+    snapshot_filename,
+    generate_extra_mask: bool = False,
+) -> Union[
+    swiftsimio.reader.SWIFTDataset,
+    Tuple[swiftsimio.reader.SWIFTDataset, namedtuple]
+]:
+    """
+    Loads a VelociraptorParticles instance for one halo into a
+    `swiftsimio` masked dataset.
+
+    Initially, this uses `r_max` to perform a spatial mask, and
+    then returns the `swiftsimio` dataset and a secondary mask
+    that may be used to extract only the particles that are
+    part of the FoF group.
+
+    You will need to instantiate the VelociraptorParticles instance
+    with an associated catalogue to use this feature, as it requires
+    the knowledge of `r_max`.
+
+    Takes three arguments:
+
+    + particles, the VelociraptorParticles instance,
+    + snapshot_filename, the path to the associated SWIFT snapshot.
+    + generate_extra_mask, whether or not to generate the secondary
+                           mask object that allows for the extraction
+                           of particles that are present only in the
+                           FoF group.
+
+    It returns:
+
+    + data, the swiftsimio dataset
+    + mask, an object containing masks for all available datasets in the
+            swift dataset. The initial masking is performed on a
+            spatial only basis, and this is required to only extract
+            the particles in the FoF group as identified by
+            velociraptor. This is only provided if generate_extra_mask
+            has a truthy value.
+    """
+
+    swift_mask = generate_spatial_mask(particles, snapshot_filename)
+
+    data = swiftsimio.load(snapshot_filename, mask=swift_mask)
+
+    if not generate_extra_mask:
+        return data
+
+    mask = generate_bound_mask(data, particles)
+
     return data, mask