From 3292d563bb0e73ef73bb485a15f33315f82c5d90 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 27 Aug 2024 17:38:51 +0100 Subject: [PATCH] Implement an option to mask a SOAP catalogue down to a single row. --- swiftsimio/masks.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index 5e6acce1..dde09d3f 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -3,6 +3,8 @@ snapshots. """ +import warnings + import unyt import h5py @@ -409,6 +411,25 @@ def convert_masks_to_ranges(self): return + def constrain_index(self, index: int): + """ + Constrain the mask to a single row. + + Intended for use with SOAP catalogues, mask to read only a single row. + + Parameters + ---------- + index : int + The index of the row to select. + """ + if not self.metadata.filetype == "SOAP": + warnings.warn("Not masking a SOAP catalogue, nothing constrained.") + return + for group_name in self.metadata.present_group_names: + setattr(self, group_name, np.array([[index, index + 1]])) + setattr(self, f"{group_name}_size", 1) + return + def get_masked_counts_offsets(self) -> (Dict[str, np.array], Dict[str, np.array]): """ Returns the particle counts and offsets in cells selected by the mask