Use cell bounding box metadata

docs/source/masking/index.rst

@@ -10,14 +10,13 @@ with this. This functionality is provided through the :mod:`swiftsimio.masks`
 sub-module but is available easily through the :meth:`swiftsimio.mask`
 top-level function.
 
-This functionality is used heavily in our `VELOCIraptor integration library`_
-for only reading data that is near bound objects.
+This functionality is used heavily in `swiftgalaxy`_.
 
 There are two types of mask, with the default only allowing spatial masking.
 Full masks require significantly more memory overhead and are generally much
 slower than the spatial only mask.
 
-.. _`VELOCIraptor integration library`: https://github.com/swiftsim/velociraptor-python
+.. _`swiftgalaxy`: https://github.com/SWIFTSIM/swiftgalaxy
 
 Spatial-only masking
 --------------------
@@ -39,7 +38,7 @@ Example
 ^^^^^^^
 
 In this example we will use the :obj:`swiftsimio.masks.SWIFTMask` object
-to load the bottom left 'half' corner of the box.
+to load the the octant of the box closes to the origin.
 
 .. code-block:: python
 
@@ -69,8 +68,8 @@ 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
+making repeated calls to :meth:`~swiftsimio.masks.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
@@ -81,10 +80,143 @@ For instance, to add the diagonally opposed octant to the selection made above
    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
+In the first call to :meth:`~swiftsimio.masks.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.
 
+Periodic boundaries
+^^^^^^^^^^^^^^^^^^^
+
+The mask region is aware of the periodic box boundaries. Let's take for example a
+region shaped like a "slab" in the :math:`x-y` plane with :math:`|z|<0.1L_\mathrm{box}`.
+One way to write this is by thinking of the :math:`z<0` part as
+lying at the upper edge of the box:
+
+.. code-block:: python
+
+   mask = sw.mask(filename)
+   mask.constrain_spatial(
+       [
+           None,
+           None,
+           [0.0 * mask.metadata.boxsize[2], 0.1 * mask.metadata.boxsize[2]],
+       ]
+   )
+   mask.constrain_spatial(
+       [
+           None,
+           None,
+           [0.9 * mask.metadata.boxsize[2], 1.0 * mask.metadata.boxsize[2]],
+       ],
+       intersect=True,
+   )
+
+This is a bit inconvenient though since the region is actually contiguous if we
+account for the periodic boundary. :meth:`~swiftsimio.masks.SWIFTMask.constrain_spatial` allows us
+to select a region straddling the periodic boundary, for example this is an
+equivalent selection:
+
+.. code-block:: python
+
+   mask = sw.mask(filename)
+   mask.constrain_spatial(
+       [
+           None,
+	   None,
+	   [-0.1 * mask.metadata.boxsize[2], 0.1 * mask.metadata.boxsize[2]],
+       ]
+   )
+
+Note that masking never result in periodic copies of particles, nor does it shift
+particle coordinates to match the region defined; particle coordinates always
+lie in the range :math:`[0, L_\mathrm{box}]`. For example reading
+a region that extends beyond the box in all directions produces exactly one copy
+of every particle and is equivalent to providing no spatial mask:
+
+.. code-block:: python
+
+   mask = sw.mask(filename)
+   mask.constrain_spatial(
+       [[-0.1 * lbox, 1.1 * lbox] for lbox in mask.metadata.boxsize]
+   )
+
+Remember to wrap the coordinates yourself if relevant! Alternatively, the
+`swiftgalaxy`_ package offers support for coordinate transformations including
+periodic boundaries.
+
+Another equivalent region for the :math:`|z|<0.1L_\mathrm{box}` slab can be written
+by setting the lower bound to a greater value than the upper bound, the code will
+interpret this as a request to start at the lower bound, wrap through the upper
+periodic boundary and continue until the (numerically lower value of) the upper
+bound is reached:
+
+.. code-block:: python
+
+   mask = sw.mask(filename)
+   mask.constrain_spatial(
+       [
+           None,
+	   None,
+	   [0.9 * mask.metadata.boxsize[2], 0.1 * mask.metadata.boxsize[2]],
+       ]
+   )
+
+The coordinates defining the region must always be in the interval
+:math:`[-0.5L_\mathrm{box}, 1.5L_\mathrm{box}]`. This allows enough flexibility to
+define all possible regions.
+
+Implementation details
+^^^^^^^^^^^^^^^^^^^^^^
+
+SWIFT snapshots group particles according to the cell that they occupy so that
+particles belonging to a cell are stored contiguously. The cells form a regular grid
+covering the simulation domain. However, SWIFT does not guarantee that all particles
+that belong to a cell are within the boundaries of a cell at the time when a snapshot
+is produced (particles are moved between cells at intervals, but may drift outside of
+their current cell before being re-assigned). Snapshots contain metadata defining
+the "bounding box" of each cell that contains all particles assigned to it at the
+time that the snapshot was written. :mod:`swiftsimio` uses this information when
+deciding what cells to read, so you may find that the "extra" particles read in
+outside of the explicitly asked for have an irregular boundary with cuboid protrusions
+or indentations. This is normal: the cells read in are exactly those needed to
+guarantee that all particles in the specified region of interest are captured. It is
+therefore advantageous to make the region as small and tightly fit to the analysis
+task as possible - in particular, trying to align it with the cell boundaries will
+typically result in an I/O overhead as neighbouring cells with particles that have
+drifted into the region are read in. Unless these particles are actually needed, it
+is actually better for performance to *avoid* the cell boundaries when defining the
+region.
+
+Older SWIFT snapshots lack the metadata to know exactly how far particles have
+drifted out of their cells. In ``v10.2.0`` or newer, if :mod:`swiftsimio` does not
+find this metadata, it will pad the region (by 0.1 times the cell length by default),
+and issue a `UserWarning` indicating this.
+
+.. warning::
+
+   In the worst case that the region consists of one cell and the padding extends to all
+   neighbouring cells, this can result in up to a factor of :math:`3^3=27` additional
+   I/O overhead. Older :mod:`swiftsimio` versions instead risk missing particles near
+   the region boundary.
+
+In the unlikely case that particles drift more than 0.1 times
+the cell length away from their "home" cell and the cell bounding-box metadata is not
+present, some particles can be missed when applying a spatial mask. The padding of
+the region can be extended or switched off with the ``safe_padding`` parameter:
+
+.. code-block:: python
+
+   mask = sw.mask(filename)
+   lbox = mask.metadata.boxsize
+   mask.constrain_spatial(
+       [[0.4 * lbox, 0.6 * lbox] for lbox in mask.metadata.boxsize],
+       safe_padding=False,  # padding switched off
+   )
+   mask.constrain_spatial(
+       [[0.4 * lbox, 0.6 * lbox] for lbox in mask.metadata.boxsize],
+       safe_padding=1.0,  # pad more, by 1.0 instead of 0.1 cell lengths
+   )
+
 
 Full mask
 ---------
@@ -171,7 +303,7 @@ as follows
     sw.subset_writer.write_subset("test_subset.hdf5", mask)
 
 This will write a snapshot which contains the particles from the specified snapshot 
-whose *x*-coordinate is within the range [100, 1000] kpc. This function uses the 
+whose :math:`x`-coordinate is within the range [100, 1000] kpc. This function uses the 
 cell mask which encompases the specified spatial domain to successively read portions 
 of datasets from the input file and writes them to a new snapshot. 
 

pyproject.toml

@@ -23,7 +23,7 @@ packages = [
 
 [project]
 name = "swiftsimio"
-version="10.1.0"
+version="10.2.0"
 authors = [
     { name="Josh Borrow", email="josh@joshborrow.com" },
 ]

swiftsimio/__init__.py

@@ -1,3 +1,5 @@
+from typing import Optional as _Optional, Union as _Union
+
 from .reader import *
 from .snapshot_writer import SWIFTSnapshotWriter
 from .masks import SWIFTMask
@@ -17,7 +19,7 @@
 name = "swiftsimio"
 
 
-def validate_file(filename):
+def validate_file(filename: str):
     """
     Checks that the provided file is a SWIFT dataset.
 
@@ -47,7 +49,9 @@ def validate_file(filename):
     return True
 
 
-def mask(filename, spatial_only=True) -> SWIFTMask:
+def mask(
+    filename: str, spatial_only: bool = True, safe_padding: _Union[bool, float] = True
+) -> SWIFTMask:
     """
     Sets up a masking object for you to use with the correct units and
     metadata available.
@@ -61,6 +65,19 @@ def mask(filename, spatial_only=True) -> SWIFTMask:
         allow you to use masking on other variables (e.g. density).
         Defaults to True.
 
+    safe_padding : bool or float, optional
+        If snapshot does not specify bounding box of cell particles (MinPositions &
+        MaxPositions), pad the mask to gurantee that *all* particles in requested
+        spatial region(s) are selected. If the bounding box metadata is present, this
+        argument is ignored. The default (``True``) is to pad by one cell length.
+        Padding can be disabled (``False``) or set to a different fraction of the
+        cell length (e.g. ``0.2``). Only entire cells are loaded, but if the region
+        boundary is more than ``safe_padding`` from a cell boundary the neighbouring
+        cell is not read. Switching off can reduce I/O load by up to a factor of 10
+        in some cases (but a few particles in region could be missing). See
+        https://swiftsimio.readthedocs.io/en/latest/masking/index.html for further
+        details.
+
     Returns
     -------
     SWIFTMask
@@ -78,10 +95,12 @@ def mask(filename, spatial_only=True) -> SWIFTMask:
     units = SWIFTUnits(filename)
     metadata = metadata_discriminator(filename, units)
 
-    return SWIFTMask(metadata=metadata, spatial_only=spatial_only)
+    return SWIFTMask(
+        metadata=metadata, spatial_only=spatial_only, safe_padding=safe_padding
+    )
 
 
-def load(filename, mask=None) -> SWIFTDataset:
+def load(filename: str, mask: _Optional[SWIFTMask] = None) -> SWIFTDataset:
     """
     Loads the SWIFT dataset at filename.
 
@@ -96,7 +115,7 @@ def load(filename, mask=None) -> SWIFTDataset:
     return SWIFTDataset(filename, mask=mask)
 
 
-def load_statistics(filename) -> SWIFTStatisticsFile:
+def load_statistics(filename: str) -> SWIFTStatisticsFile:
     """
     Loads a SWIFT statistics file (``SFR.txt``, ``energy.txt``).
 

swiftsimio/conversions.py

@@ -94,7 +94,8 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology:
 
         try:
             Tcmb0 = cosmo["T_CMB_0 [K]"][0]
-        except (IndexError, KeyError, AttributeError):
+            assert Tcmb0 != 0
+        except (IndexError, KeyError, AttributeError, AssertionError):
             # expressions taken directly from astropy, since they do no longer
             # allow access to these attributes (since version 5.1+)
             critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value