Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Unpack the boxsize as a cosmo_array. #214

Merged
merged 3 commits into from
Dec 5, 2024
Merged

Unpack the boxsize as a cosmo_array. #214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7e14dd7
Unpack the boxsize as a cosmo_array.
kyleaoman Nov 30, 2024
f11b395
Fix a cosmo_factor missing warning in test suite.
kyleaoman Nov 30, 2024
a190c39
Run formatter.
kyleaoman Nov 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions swiftsimio/metadata/metadata/metadata_fields.py
View file
Open in desktop
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
147 changes: 81 additions & 66 deletions swiftsimio/metadata/objects.py
View file
Open in desktop
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,6 +109,60 @@ 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,
Expand All @@ -118,17 +171,32 @@ def postprocess_header(self):
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 +247,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 +296,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 +312,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 +340,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 +516,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 @@ -886,7 +901,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 +913,7 @@ def metadata_discriminator(filename: str, units: SWIFTUnits) -> "SWIFTMetadata":
units : SWIFTUnits
The units object associated with the file


Returns
-------

Expand Down
4 changes: 3 additions & 1 deletion tests/test_soap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this needed, does masses2 not already support comparison?

Copy link
Member Author

@kyleaoman kyleaoman Dec 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lower and upper are unyt_quantitys, not cosmo_arrays. This is correct given the type hints of constrain_mask used further up in the test, but correctly triggers a warning on comparison with masses2. Perhaps constrain_mask should be modified to expect cosmo_arrays. It would still work if given unyt_quantitys in that case but would raise warnings on receiving unyt_quantitys (which is what should happen in that case, I guess). If that happens then the change to custom_mask should be reverted (it will emit a warning in that case so shouldn't go unnoticed).

) & (unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) <= upper)

assert len(masses2[custom_mask]) == len(masses)

Expand Down
Loading