diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 7e58b5d2..4d7222c0 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -116,6 +116,8 @@ def wrapped(self, *args, **kwargs): ret.cosmo_factor = self.cosmo_factor if hasattr(self, "comoving"): ret.comoving = self.comoving + if hasattr(self, "valid_transform"): + ret.valid_transform = self.valid_transform return ret return wrapped @@ -591,7 +593,7 @@ class cosmo_array(unyt_array): Cosmology array class. This inherits from the unyt.unyt_array, and adds - three variables: compression, cosmo_factor, and comoving. + four variables: compression, cosmo_factor, comoving, and valid_transform. Data is assumed to be comoving when passed to the object but you can override this by setting the latter flag to be False. @@ -615,8 +617,12 @@ class cosmo_array(unyt_array): String describing any compression that was applied to this array in the hdf5 file. + valid_transform: bool + if True then the array can be converted from physical to comoving units + """ + # TODO: _cosmo_factor_ufunc_registry = { add: _preserve_cosmo_factor, subtract: _preserve_cosmo_factor, @@ -718,6 +724,7 @@ def __new__( name=None, cosmo_factor=None, comoving=True, + valid_transform=True, compression=None, ): """ @@ -753,6 +760,8 @@ def __new__( coordinates comoving : bool flag to indicate whether using comoving coordinates + valid_transform : bool + flag to indicate whether this array can be converted to comoving compression : string description of the compression filters that were applied to that array in the hdf5 file @@ -800,6 +809,11 @@ def __new__( obj.cosmo_factor = cosmo_factor obj.comoving = comoving obj.compression = compression + obj.valid_transform = valid_transform + if not obj.valid_transform: + assert ( + not obj.comoving + ), "Cosmo arrays without a valid transform to comoving units must be physical" return obj @@ -810,6 +824,7 @@ def __array_finalize__(self, obj): self.cosmo_factor = getattr(obj, "cosmo_factor", None) self.comoving = getattr(obj, "comoving", True) self.compression = getattr(obj, "compression", None) + self.valid_transform = getattr(obj, "valid_transform", True) def __str__(self): if self.comoving: @@ -837,7 +852,9 @@ def __reduce__(self): """ np_ret = super(cosmo_array, self).__reduce__() obj_state = np_ret[2] - cosmo_state = (((self.cosmo_factor, self.comoving),) + obj_state[:],) + cosmo_state = ( + ((self.cosmo_factor, self.comoving, self.valid_transform),) + obj_state[:], + ) new_ret = np_ret[:2] + cosmo_state + np_ret[3:] return new_ret @@ -849,7 +866,7 @@ def __setstate__(self, state): state and pass the rest to unyt_array.__setstate__. """ super(cosmo_array, self).__setstate__(state[1:]) - self.cosmo_factor, self.comoving = state[0] + self.cosmo_factor, self.comoving, self.valid_transform = state[0] # Wrap functions that return copies of cosmo_arrays so that our # attributes get passed through: @@ -889,6 +906,9 @@ def convert_to_comoving(self) -> None: """ if self.comoving: return + if not self.valid_transform: + # TODO: Decide on error/warning + print("What are you doing???") else: # Best to just modify values as otherwise we're just going to have # to do a convert_to_units anyway. @@ -932,6 +952,9 @@ def to_comoving(self): cosmo_array copy of cosmo_array in comoving units """ + if not self.valid_transform: + # TODO: Decide on error/warning + print("What are you doing???") copied_data = self.in_units(self.units, cosmo_factor=self.cosmo_factor) copied_data.convert_to_comoving() @@ -939,12 +962,13 @@ def to_comoving(self): def compatible_with_comoving(self): """ + # TODO: Is this the same question as "can be converted to comoving?" Is this cosmo_array compatible with a comoving cosmo_array? This is the case if the cosmo_array is comoving, or if the scale factor exponent is 0 (cosmo_factor.a_factor() == 1) """ - return self.comoving or (self.cosmo_factor.a_factor == 1.0) + return self.valid_transform def compatible_with_physical(self): """ @@ -953,11 +977,18 @@ def compatible_with_physical(self): This is the case if the cosmo_array is physical, or if the scale factor exponent is 0 (cosmo_factor.a_factor == 1) """ + # TODO: Isn't this always true? We can convert it to physical if needed? return (not self.comoving) or (self.cosmo_factor.a_factor == 1.0) @classmethod def from_astropy( - cls, arr, unit_registry=None, comoving=True, cosmo_factor=None, compression=None + cls, + arr, + unit_registry=None, + comoving=True, + cosmo_factor=None, + compression=None, + valid_transform=True, ): """ Convert an AstroPy "Quantity" to a cosmo_array. @@ -977,6 +1008,8 @@ def from_astropy( compression : string String describing any compression that was applied to this array in the hdf5 file. + valid_transform : bool + flag to indicate whether this array can be converted to comoving Example ------- @@ -989,12 +1022,19 @@ def from_astropy( obj.comoving = comoving obj.cosmo_factor = cosmo_factor obj.compression = compression + obj.valid_trasform = valid_transform return obj @classmethod def from_pint( - cls, arr, unit_registry=None, comoving=True, cosmo_factor=None, compression=None + cls, + arr, + unit_registry=None, + comoving=True, + cosmo_factor=None, + compression=None, + valid_transform=True, ): """ Convert a Pint "Quantity" to a cosmo_array. @@ -1014,6 +1054,8 @@ def from_pint( compression : string String describing any compression that was applied to this array in the hdf5 file. + valid_transform : bool + flag to indicate whether this array can be converted to comoving Examples -------- @@ -1032,9 +1074,11 @@ def from_pint( obj.comoving = comoving obj.cosmo_factor = cosmo_factor obj.compression = compression + obj.valid_trasform = valid_transform return obj + # TODO: def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): cms = [ (hasattr(inp, "comoving"), getattr(inp, "comoving", None)) for inp in inputs diff --git a/swiftsimio/reader.py b/swiftsimio/reader.py index d57374b3..116052dc 100644 --- a/swiftsimio/reader.py +++ b/swiftsimio/reader.py @@ -601,7 +601,7 @@ def load_groups(self): metadata._properties - This contains six arrays, + This contains eight arrays, metadata._properties.field_names metadata._properties.field_paths @@ -609,6 +609,8 @@ def load_groups(self): metadata._properties.field_cosmologies metadata._properties.field_descriptions metadata._properties.field_compressions + metadata._properties.field_physicals + metadata._properties.field_valid_transforms As well as some more information about the group. """ @@ -851,6 +853,10 @@ class SWIFTGroupMetadata(object): Loads in compressions of the fields for each dataset. load_cosmology(self): Loads in the field cosmologies. + load_physical(self): + Loads in whether the field is saved as comoving or physical. + load_valid_transforms(self): + Loads in whether the field can be converted to comoving. load_named_columns(self): Loads the named column data for relevant fields. """ @@ -903,6 +909,8 @@ def load_metadata(self): self.load_field_descriptions() self.load_field_compressions() self.load_cosmology() + self.load_physical() + self.load_valid_transforms() self.load_named_columns() def load_field_names(self): @@ -1000,7 +1008,7 @@ def get_desc(dataset): return description + " Not masked." # TODO: Update for https://github.com/SWIFTSIM/SOAP/pull/81 - # TODO: Also need to add group mask data for that PR + # TODO: Also need to add metadata for each halo type for that PR # If is_masked is true then these other attributes should exist mask_datasets = dataset.attrs["Mask Datasets"] mask_threshold = dataset.attrs["Mask Threshold"] @@ -1068,6 +1076,44 @@ def get_cosmo(dataset): return + def load_physical(self): + """ + Loads in whether the field is saved as comoving or physical. + """ + + def get_physical(dataset): + try: + physical = dataset.attrs["Value stored as physical"][0] == 1 + except: + physical = False + return physical + + self.field_physicals = [ + get_physical(self.metadata.handle[x]) for x in self.field_paths + ] + + return + + def load_valid_transforms(self): + """ + Loads in whether the field can be converted to comoving. + """ + + def get_valid_transform(dataset): + try: + valid_transform = ( + dataset.attrs["Property can be converted to comoving"][0] == 1 + ) + except: + valid_transform = True + return valid_transform + + self.field_valid_transforms = [ + get_valid_transform(self.metadata.handle[x]) for x in self.field_paths + ] + + return + def load_named_columns(self): """ Loads the named column data for relevant fields. @@ -1118,6 +1164,8 @@ def generate_getter( cosmo_factor: cosmo_factor, description: str, compression: str, + physical: bool, + valid_transform: bool, columns: Union[None, np.lib.index_tricks.IndexExpression] = None, ): """ @@ -1163,6 +1211,14 @@ def generate_getter( String describing the lossy compression filters that were applied to the data (read from the HDF5 file). + physical: bool + Bool that describes whether the data in the file is stored in comoving + or physical units. + + valid_transform: bool + Bool that describes whether converting this field from physical to comoving + units is a valid operation. + columns: np.lib.index_tricks.IndexEpression, optional Index expression corresponding to which columns to read from the numpy array. If not provided, we read all columns and return an n-dimensional array. @@ -1234,6 +1290,8 @@ def getter(self): cosmo_factor=cosmo_factor, name=description, compression=compression, + comoving=not physical, + valid_transform=valid_transform, ), ) else: @@ -1250,6 +1308,8 @@ def getter(self): cosmo_factor=cosmo_factor, name=description, compression=compression, + comoving=not physical, + valid_transform=valid_transform, ), ) except KeyError: @@ -1480,6 +1540,8 @@ def generate_datasets(group_metadata: SWIFTGroupMetadata, mask): field_names = group_metadata.field_names field_cosmologies = group_metadata.field_cosmologies field_units = group_metadata.field_units + field_physicals = group_metadata.field_physicals + field_valid_transforms = group_metadata.field_valid_transforms field_descriptions = group_metadata.field_descriptions field_compressions = group_metadata.field_compressions field_named_columns = group_metadata.named_columns @@ -1489,6 +1551,8 @@ def generate_datasets(group_metadata: SWIFTGroupMetadata, mask): field_names, field_cosmologies, field_units, + field_physicals, + field_valid_transforms, field_descriptions, field_compressions, ) @@ -1505,6 +1569,8 @@ def generate_datasets(group_metadata: SWIFTGroupMetadata, mask): field_name, field_cosmology, field_unit, + field_physical, + field_valid_transform, field_description, field_compression, ) in dataset_iterator: @@ -1522,6 +1588,8 @@ def generate_datasets(group_metadata: SWIFTGroupMetadata, mask): cosmo_factor=field_cosmology, description=field_description, compression=field_compression, + physical=field_physical, + valid_transform=field_valid_transform, ), generate_setter(field_name), generate_deleter(field_name), @@ -1548,6 +1616,8 @@ def generate_datasets(group_metadata: SWIFTGroupMetadata, mask): cosmo_factor=field_cosmology, description=f"{field_description} [Column {index}, {column}]", compression=field_compression, + physical=field_physical, + valid_transform=field_valid_transform, columns=np.s_[index], ), generate_setter(column),