Skip to content

Commit

Permalink
Unpack the boxsize as a cosmo_array.
Browse files Browse the repository at this point in the history
  • Loading branch information
@kyleaoman
kyleaoman committed Nov 30, 2024
1 parent 18d246b commit 7e14dd7
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 74 deletions.
32 changes: 32 additions & 0 deletions swiftsimio/metadata/metadata/metadata_fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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 = {
Expand Down
169 changes: 95 additions & 74 deletions swiftsimio/metadata/objects.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
object notation (e.g. PartType0/Coordinates -> gas.coordinates).
"""


import numpy as np
import unyt

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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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...
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -898,7 +919,7 @@ def metadata_discriminator(filename: str, units: SWIFTUnits) -> "SWIFTMetadata":
units : SWIFTUnits
The units object associated with the file
Returns
-------
Expand Down

0 comments on commit 7e14dd7

Please sign in to comment.