Skip to content

Commit

Permalink
Add new SWIFT attributes
Browse files Browse the repository at this point in the history
  • Loading branch information
@robjmcgibbon
robjmcgibbon committed Jul 16, 2024
1 parent 4ac1cba commit 8b94a0b
Show file tree
Hide file tree
Showing 2 changed files with 122 additions and 8 deletions.
56 changes: 50 additions & 6 deletions swiftsimio/objects.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -718,6 +724,7 @@ def __new__(
name=None,
cosmo_factor=None,
comoving=True,
valid_transform=True,
compression=None,
):
"""
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -932,19 +952,23 @@ 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()

return copied_data

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):
"""
Expand All @@ -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.
Expand All @@ -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
-------
Expand All @@ -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.
Expand All @@ -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
--------
Expand All @@ -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
Expand Down
74 changes: 72 additions & 2 deletions swiftsimio/reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -601,14 +601,16 @@ def load_groups(self):
metadata.<group_name>_properties
This contains six arrays,
This contains eight arrays,
metadata.<type>_properties.field_names
metadata.<type>_properties.field_paths
metadata.<type>_properties.field_units
metadata.<type>_properties.field_cosmologies
metadata.<type>_properties.field_descriptions
metadata.<type>_properties.field_compressions
metadata.<type>_properties.field_physicals
metadata.<type>_properties.field_valid_transforms
As well as some more information about the group.
"""
Expand Down Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
):
"""
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -1234,6 +1290,8 @@ def getter(self):
cosmo_factor=cosmo_factor,
name=description,
compression=compression,
comoving=not physical,
valid_transform=valid_transform,
),
)
else:
Expand All @@ -1250,6 +1308,8 @@ def getter(self):
cosmo_factor=cosmo_factor,
name=description,
compression=compression,
comoving=not physical,
valid_transform=valid_transform,
),
)
except KeyError:
Expand Down Expand Up @@ -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
Expand All @@ -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,
)
Expand All @@ -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:
Expand All @@ -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),
Expand All @@ -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),
Expand Down

0 comments on commit 8b94a0b

Please sign in to comment.