From 790c97e11af3332921fa1468983ea1563ce9e2f6 Mon Sep 17 00:00:00 2001 From: robjmcgibbon Date: Thu, 19 Sep 2024 10:57:42 +0100 Subject: [PATCH 1/2] Add homogeneous attribute --- swiftsimio/masks.py | 8 +++++--- swiftsimio/metadata/objects.py | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index b3f406ca..4b1938dc 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -416,15 +416,17 @@ 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. + Intended for use with halo 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.") + if not self.metadata.homogeneous_arrays: + warnings.warn( + "Different datasets correspond to different objects, nothing constrained." + ) return for group_name in self.metadata.present_group_names: setattr(self, group_name, np.array([[index, index + 1]])) diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index e3bd3b00..b7aad62d 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -45,6 +45,10 @@ class SWIFTMetadata(ABC): # (as is the case in SOAP) or whether each type (e.g. Gas, Dark Matter, etc.) # has its own top-level cell grid counts. shared_cell_counts: str | None = None + # Whether all the arrays in this files have the same length and order (as is + # the case for SOAP, all arrays correspond to subhalos) or whether there are + # multiple types (e.g. Gas, Dark Matter, etc.) + homogeneous_arrays: bool = False @abstractmethod def __init__(self, filename, units: "SWIFTUnits"): @@ -1223,6 +1227,8 @@ class SWIFTFOFMetadata(SWIFTMetadata): class. """ + homogeneous_arrays: bool = True + def __init__(self, filename: str, units: SWIFTUnits): self.filename = filename self.units = units @@ -1265,6 +1271,7 @@ class SWIFTSOAPMetadata(SWIFTMetadata): masking_valid: bool = True shared_cell_counts: str = "Subhalos" + homogeneous_arrays: bool = True def __init__(self, filename: str, units: SWIFTUnits): self.filename = filename From 8dbf349bf5ce6e4630a797bfc1572e3035093539 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Sep 2024 21:12:57 +0100 Subject: [PATCH 2/2] Add a function to mask to a list of rows in a halo catalogue. --- swiftsimio/masks.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index 4b1938dc..a39d03a8 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -15,7 +15,7 @@ from swiftsimio.objects import InvalidSnapshot from swiftsimio.accelerated import ranges_from_array -from typing import Dict +from typing import Dict, List class SWIFTMask(object): @@ -433,6 +433,33 @@ def constrain_index(self, index: int): setattr(self, f"{group_name}_size", 1) return + def constrain_indices(self, indices: List): + """ + Constrain the mask to a list of rows. + + Parameters + ---------- + indices : List + An list of the indices of the rows to mask. + """ + if not self.metadata.homogeneous_arrays: + warnings.warn( + "Different datasets correspond to different objects, nothing constrained." + ) + return + if self.spatial_only: + warnings.warn( + "You cannot constrain a mask if spatial_only=True \n" + "Please re-initialise the SWIFTMask object with spatial_only=False" + "Nothing constrained." + ) + return + + self._shared[...] = False + self._shared[indices] = True + self._shared_size = len(indices) + 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