From 7e14dd759b02eb7442ffdab000d2e1ee283c742c Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 30 Nov 2024 22:14:16 +0800 Subject: [PATCH 1/3] Unpack the boxsize as a cosmo_array. --- .../metadata/metadata/metadata_fields.py | 32 ++++ swiftsimio/metadata/objects.py | 169 ++++++++++-------- 2 files changed, 127 insertions(+), 74 deletions(-) diff --git a/swiftsimio/metadata/metadata/metadata_fields.py b/swiftsimio/metadata/metadata/metadata_fields.py index 182a3907..75e12a5e 100644 --- a/swiftsimio/metadata/metadata/metadata_fields.py +++ b/swiftsimio/metadata/metadata/metadata_fields.py @@ -2,6 +2,8 @@ Contains the description of the metadata fields in the SWIFT snapshots. """ +from ..objects import cosmo_factor, a + metadata_fields_to_read = { "Code": "code", "Cosmology": "cosmology_raw", @@ -54,6 +56,36 @@ def generate_units_header_unpack_arrays(m, l, t, I, T) -> dict: return units +def generate_cosmo_args_header_unpack_arrays(scale_factor) -> dict: + """ + Generates arguments for `cosmo_array` so that relevant header + items can be initialized as `cosmo_array`s. + + Parameters + ---------- + scale_factor : float + The scale factor. + + Returns + ------- + out : dict + A dictionary containing the `cosmo_array` arguments corresponding to + header items, omitting any that should not be `cosmo_array`s. + """ + + # Do not include those items that do not have units (and therefore + # should not be cosmo_array'd). + cosmo_args = { + "boxsize": dict( + cosmo_factor=cosmo_factor(a**1, scale_factor=scale_factor), + comoving=True, # if it's not, then a=1 and it doesn't matter + valid_transform=True, + ) + } + + return cosmo_args + + # These are the same as above, but they are extracted as real python strings # instead of bytecode characters. header_unpack_string = { diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index db0cdcec..14a7385b 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -6,7 +6,6 @@ object notation (e.g. PartType0/Coordinates -> gas.coordinates). """ - import numpy as np import unyt @@ -110,25 +109,100 @@ def postprocess_header(self): Some minor postprocessing on the header to local variables. """ + # We need the scale factor to initialize `cosmo_array`s, so start with the float + # items including the scale factor. + # These must be unpacked as they are stored as length-1 arrays + + header_unpack_float_units = ( + metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) + ) + for field, names in metadata.metadata_fields.header_unpack_single_float.items(): + try: + if isinstance(names, list): + # Sometimes we store a list in case we have multiple names, for + # example Redshift -> metadata.redshift AND metadata.z. Can't just do + # the iteration because we may loop over the letters in the string. + for variable in names: + if variable in header_unpack_float_units.keys(): + # We have an associated unit! + unit = header_unpack_float_units[variable] + setattr( + self, + variable, + unyt.unyt_quantity(self.header[field][0], units=unit), + ) + else: + # No unit + setattr(self, variable, self.header[field][0]) + else: + # We can just check for the unit and set the attribute + variable = names + if variable in header_unpack_float_units.keys(): + # We have an associated unit! + unit = header_unpack_float_units[variable] + setattr( + self, + variable, + unyt.unyt_quantity(self.header[field][0], units=unit), + ) + else: + # No unit + setattr(self, variable, self.header[field][0]) + except KeyError: + # Must not be present, just skip it + continue + # need the scale factor first for cosmology on other header attributes + try: + self.a = self.scale_factor + except AttributeError: + # These must always be present for the initialisation of cosmology properties + self.a = 1.0 + self.scale_factor = 1.0 + # These are just read straight in to variables - header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, + header_unpack_arrays_units = ( + metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) + ) + header_unpack_arrays_cosmo_args = ( + metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor + ) ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): try: if name in header_unpack_arrays_units.keys(): - setattr( - self, - name, - unyt.unyt_array( - self.header[field], units=header_unpack_arrays_units[name] - ), - ) + if name in header_unpack_arrays_cosmo_args.keys(): + setattr( + self, + name, + cosmo_array( + self.header[field], + units=header_unpack_arrays_units[name], + **header_unpack_arrays_cosmo_args[name], + ), + ) + else: + setattr( + self, + name, + unyt.unyt_array( + self.header[field], + units=header_unpack_arrays_units[name], + ), + ) # This is required or we automatically get everything in CGS! getattr(self, name).convert_to_units( header_unpack_arrays_units[name] @@ -179,52 +253,6 @@ def postprocess_header(self): # Must not be present, just skip it setattr(self, name, "") - # These must be unpacked as they are stored as length-1 arrays - - header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) - - for field, names in metadata.metadata_fields.header_unpack_single_float.items(): - try: - if isinstance(names, list): - # Sometimes we store a list in case we have multiple names, for example - # Redshift -> metadata.redshift AND metadata.z. Can't just do the iteration - # because we may loop over the letters in the string. - for variable in names: - if variable in header_unpack_float_units.keys(): - # We have an associated unit! - unit = header_unpack_float_units[variable] - setattr( - self, - variable, - unyt.unyt_quantity(self.header[field][0], units=unit), - ) - else: - # No unit - setattr(self, variable, self.header[field][0]) - else: - # We can just check for the unit and set the attribute - variable = names - if variable in header_unpack_float_units.keys(): - # We have an associated unit! - unit = header_unpack_float_units[variable] - setattr( - self, - variable, - unyt.unyt_quantity(self.header[field][0], units=unit), - ) - else: - # No unit - setattr(self, variable, self.header[field][0]) - except KeyError: - # Must not be present, just skip it - continue - # These are special cases, sorry! # Date and time of snapshot dump try: @@ -274,7 +302,7 @@ def postprocess_header(self): self.reduced_lightspeed = None # Store these separately as self.n_gas = number of gas particles for example - for (part_number, (_, part_name)) in enumerate( + for part_number, (_, part_name) in enumerate( metadata.particle_types.particle_name_underscores.items() ): try: @@ -290,13 +318,6 @@ def postprocess_header(self): # We can set a default and print a message whenever we require this value self.gas_gamma = None - try: - self.a = self.scale_factor - except AttributeError: - # These must always be present for the initialisation of cosmology properties - self.a = 1.0 - self.scale_factor = 1.0 - return def extract_cosmology(self): @@ -325,7 +346,7 @@ def present_groups(self) -> list[str]: A property giving the present particle groups in the file to be unpackaged into top-level properties. For instance, in a regular snapshot, this would be ["PartType0", "PartType1", "PartType4", ...]. In SOAP, this would be - ["SO/200_crit", "SO/200_mean", ...], i.e. one per aperture. + ["SO/200_crit", "SO/200_mean", ...], i.e. one per aperture. """ raise NotImplementedError @@ -501,7 +522,7 @@ def __init__( Parameters ---------- - group: str + group: str the name of the group in the hdf5 file group_name : str the corresponding group name for swiftsimio @@ -589,7 +610,7 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit ** unit_exponent + units *= unit**unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... @@ -687,7 +708,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a ** cosmo_exponent + a_factor_this_dataset = a**cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor) @@ -886,7 +907,7 @@ def __del__(self): def metadata_discriminator(filename: str, units: SWIFTUnits) -> "SWIFTMetadata": """ - Discriminates between the different types of metadata objects read from SWIFT-compatible + Discriminates between the different types of metadata objects read from SWIFT-compatile files. Parameters @@ -898,7 +919,7 @@ def metadata_discriminator(filename: str, units: SWIFTUnits) -> "SWIFTMetadata": units : SWIFTUnits The units object associated with the file - + Returns ------- From f11b39501c355dcdbdc89b1a1dd327597b2ef14b Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 30 Nov 2024 22:27:52 +0800 Subject: [PATCH 2/3] Fix a cosmo_factor missing warning in test suite. --- tests/test_soap.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_soap.py b/tests/test_soap.py index ad8274ea..66ec93fa 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -60,7 +60,9 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): masses2 = data2.spherical_overdensity_200_mean.total_mass # Manually mask - custom_mask = (masses2 >= lower) & (masses2 <= upper) + custom_mask = ( + unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) >= lower + ) & (unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) <= upper) assert len(masses2[custom_mask]) == len(masses) From a190c3912783d3d2a7b545485808b6d7675a1e4c Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 30 Nov 2024 22:37:41 +0800 Subject: [PATCH 3/3] Run formatter. --- .../metadata/metadata/metadata_fields.py | 2 +- swiftsimio/metadata/objects.py | 38 ++++++++----------- 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/swiftsimio/metadata/metadata/metadata_fields.py b/swiftsimio/metadata/metadata/metadata_fields.py index 75e12a5e..a032e039 100644 --- a/swiftsimio/metadata/metadata/metadata_fields.py +++ b/swiftsimio/metadata/metadata/metadata_fields.py @@ -77,7 +77,7 @@ def generate_cosmo_args_header_unpack_arrays(scale_factor) -> dict: # should not be cosmo_array'd). cosmo_args = { "boxsize": dict( - cosmo_factor=cosmo_factor(a**1, scale_factor=scale_factor), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=scale_factor), comoving=True, # if it's not, then a=1 and it doesn't matter valid_transform=True, ) diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index 14a7385b..411deeb7 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -113,14 +113,12 @@ def postprocess_header(self): # items including the scale factor. # These must be unpacked as they are stored as length-1 arrays - header_unpack_float_units = ( - metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) for field, names in metadata.metadata_fields.header_unpack_single_float.items(): try: @@ -166,19 +164,15 @@ def postprocess_header(self): self.scale_factor = 1.0 # These are just read straight in to variables - header_unpack_arrays_units = ( - metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) - header_unpack_arrays_cosmo_args = ( - metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( - self.scale_factor - ) + header_unpack_arrays_cosmo_args = metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): @@ -610,7 +604,7 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit**unit_exponent + units *= unit ** unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... @@ -708,7 +702,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a**cosmo_exponent + a_factor_this_dataset = a ** cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor)