diff --git a/tests/tests.yaml b/tests/tests.yaml index 169f84c9bcb..7164b25b8cb 100644 --- a/tests/tests.yaml +++ b/tests/tests.yaml @@ -195,5 +195,6 @@ other_tests: - "--ignore-files=test_invalid_origin.py" - "--ignore-files=test_outputs_pytest\\.py" - "--exclude-test=yt.frontends.gdf.tests.test_outputs.TestGDF" + - "--exclude-test=yt.frontends.adaptahop.tests.test_outputs" cookbook: - 'doc/source/cookbook/tests/test_cookbook.py' diff --git a/yt/frontends/adaptahop/data_structures.py b/yt/frontends/adaptahop/data_structures.py index 6347b8b1613..4213f94e0fc 100644 --- a/yt/frontends/adaptahop/data_structures.py +++ b/yt/frontends/adaptahop/data_structures.py @@ -10,6 +10,7 @@ import os import re import stat +from itertools import product import numpy as np @@ -18,12 +19,12 @@ ) from yt.data_objects.static_output import Dataset from yt.frontends.halo_catalog.data_structures import HaloCatalogFile -from yt.funcs import setdefaultattr +from yt.funcs import mylog, setdefaultattr from yt.geometry.particle_geometry_handler import ParticleIndex from yt.units import Mpc from yt.utilities.cython_fortran_utils import FortranFile -from .definitions import HEADER_ATTRIBUTES +from .definitions import ADAPTAHOP_TEMPLATES, ATTR_T, HEADER_ATTRIBUTES from .fields import AdaptaHOPFieldInfo @@ -56,6 +57,8 @@ class AdaptaHOPDataset(Dataset): # AdaptaHOP internally assumes 1Mpc == 3.0824cm _code_length_to_Mpc = (1.0 * Mpc).to("cm").value / 3.08e24 + _header_attributes: ATTR_T = None + _halo_attributes: ATTR_T = None def __init__( self, @@ -76,6 +79,8 @@ def __init__( ) self.parent_ds = parent_ds + self._guess_headers_from_file(filename) + super().__init__( filename, dataset_type, @@ -89,9 +94,51 @@ def _set_code_unit_attributes(self): setdefaultattr(self, "velocity_unit", self.quan(1.0, "km / s")) setdefaultattr(self, "time_unit", self.length_unit / self.velocity_unit) + def _guess_headers_from_file(self, filename) -> ATTR_T: + with FortranFile(filename) as fpu: + ok = False + for dp, longint in product((True, False), (True, False)): + fpu.seek(0) + try: + header_attributes = HEADER_ATTRIBUTES(double=dp, longint=longint) + fpu.read_attrs(header_attributes) + ok = True + break + except (ValueError, OSError): + pass + + if not ok: + raise OSError("Could not read headers from file %s" % filename) + + istart = fpu.tell() + fpu.seek(0, 2) + iend = fpu.tell() + + # Try different templates + ok = False + for name, cls in ADAPTAHOP_TEMPLATES.items(): + fpu.seek(istart) + attributes = cls(longint, dp).HALO_ATTRIBUTES + mylog.debug("Trying %s(longint=%s, dp=%s)", name, longint, dp) + try: + # Try to read two halos to be sure + fpu.read_attrs(attributes) + if fpu.tell() < iend: + fpu.read_attrs(attributes) + ok = True + break + except (ValueError, OSError): + continue + + if not ok: + raise OSError("Could not guess fields from file %s" % filename) + + self._header_attributes = header_attributes + self._halo_attributes = attributes + def _parse_parameter_file(self): with FortranFile(self.parameter_filename) as fpu: - params = fpu.read_attrs(HEADER_ATTRIBUTES) + params = fpu.read_attrs(self._header_attributes) self.dimensionality = 3 self.unique_identifier = int(os.stat(self.parameter_filename)[stat.ST_CTIME]) # Domain related things diff --git a/yt/frontends/adaptahop/definitions.py b/yt/frontends/adaptahop/definitions.py index 732ebd4db14..824de1fe469 100644 --- a/yt/frontends/adaptahop/definitions.py +++ b/yt/frontends/adaptahop/definitions.py @@ -4,38 +4,171 @@ """ +import abc +from typing import Tuple, Union +from yt.funcs import mylog -HEADER_ATTRIBUTES = ( - ("npart", 1, "i"), - ("massp", 1, "f"), - ("aexp", 1, "f"), - ("omega_t", 1, "f"), - ("age", 1, "f"), - (("nhalos", "nsubs"), 2, "i"), -) - -HALO_ATTRIBUTES = ( - ("npart", 1, "i"), - ("particle_identities", -1, "i"), - ("particle_identifier", 1, "i"), - ("timestep", 1, "i"), - (("level", "host_id", "first_subhalo_id", "n_subhalos", "next_subhalo_id"), 5, "i"), - ("particle_mass", 1, "f"), - (("raw_position_x", "raw_position_y", "raw_position_z"), 3, "f"), - (("particle_velocity_x", "particle_velocity_y", "particle_velocity_z"), 3, "f"), - ( - ( - "particle_angular_momentum_x", - "particle_angular_momentum_y", - "particle_angular_momentum_z", - ), - 3, - "f", - ), - (("r", "a", "b", "c"), 4, "f"), - (("ek", "ep", "etot"), 3, "f"), - ("spin", 1, "f"), - (("virial_radius", "virial_mass", "virial_temperature", "virial_velocity"), 4, "f"), - (("rho0", "R_c"), 2, "f"), -) +ATTR_T = Tuple[Tuple[Union[Tuple[str, str], str], int, str]] + + +def HEADER_ATTRIBUTES(*, double: bool, longint: bool) -> ATTR_T: + int_type = "l" if longint else "i" + float_type = "d" if double else "f" + return ( + ("npart", 1, int_type), + ("massp", 1, float_type), + ("aexp", 1, float_type), + ("omega_t", 1, float_type), + ("age", 1, float_type), + (("nhalos", "nsubs"), 2, "i"), + ) + + +ADAPTAHOP_TEMPLATES = {} + + +class AdaptaHOPDefTemplate(abc.ABC): + templates = [] + + def __init_subclass__(cls, *args, **kwargs): + super().__init_subclass__(*args, **kwargs) + mylog.debug("Registering AdaptaHOP template class %s", cls.__name__) + ADAPTAHOP_TEMPLATES[cls.__name__] = cls + + def __init__(self, longint, double_precision): + self.longint = longint + self.double_precision = double_precision + + +class AdaptaHOPOld(AdaptaHOPDefTemplate): + @property + def HALO_ATTRIBUTES(self) -> ATTR_T: + int_type = "l" if self.longint else "i" + float_type = "d" if self.double_precision else "f" + return ( + ("npart", 1, int_type), + ("particle_identities", -1, int_type), + ("particle_identifier", 1, "i"), # this is the halo id, always an int32 + ("timestep", 1, "i"), + ( + ( + "level", + "host_id", + "first_subhalo_id", + "n_subhalos", + "next_subhalo_id", + ), + 5, + "i", + ), + ("particle_mass", 1, float_type), + (("raw_position_x", "raw_position_y", "raw_position_z"), 3, float_type), + ( + ("particle_velocity_x", "particle_velocity_y", "particle_velocity_z"), + 3, + float_type, + ), + ( + ( + "particle_angular_momentum_x", + "particle_angular_momentum_y", + "particle_angular_momentum_z", + ), + 3, + float_type, + ), + (("r", "a", "b", "c"), 4, float_type), + (("ek", "ep", "etot"), 3, float_type), + ("spin", 1, float_type), + ( + ( + "virial_radius", + "virial_mass", + "virial_temperature", + "virial_velocity", + ), + 4, + float_type, + ), + (("rho0", "R_c"), 2, float_type), + ) + + +class AdaptaHOPNewNoContam(AdaptaHOPDefTemplate): + @property + def HALO_ATTRIBUTES(self) -> ATTR_T: + int_type = "l" if self.longint else "i" + float_type = "d" if self.double_precision else "f" + return ( + ("npart", 1, int_type), + ("particle_identities", -1, int_type), + ("particle_identifier", 1, "i"), # this is the halo id, always an int32 + ("timestep", 1, "i"), + ( + ( + "level", + "host_id", + "first_subhalo_id", + "n_subhalos", + "next_subhalo_id", + ), + 5, + "i", + ), + ("particle_mass", 1, float_type), + ("npart_tot", 1, int_type), + ("particle_mass_tot", 1, float_type), + (("raw_position_x", "raw_position_y", "raw_position_z"), 3, float_type), + ( + ("particle_velocity_x", "particle_velocity_y", "particle_velocity_z"), + 3, + float_type, + ), + ( + ( + "particle_angular_momentum_x", + "particle_angular_momentum_y", + "particle_angular_momentum_z", + ), + 3, + float_type, + ), + (("r", "a", "b", "c"), 4, float_type), + (("ek", "ep", "etot"), 3, float_type), + ("spin", 1, float_type), + ("velocity_dispersion", 1, float_type), + ( + ( + "virial_radius", + "virial_mass", + "virial_temperature", + "virial_velocity", + ), + 4, + float_type, + ), + (("rmax", "vmax"), 2, float_type), + ("concentration", 1, float_type), + (("radius_200", "mass_200"), 2, float_type), + (("radius_50", "mass_50"), 2, float_type), + ("radius_profile", -1, float_type), + ("rho_profile", -1, float_type), + (("rho0", "R_c"), 2, float_type), + ) + + +class AdaptaHOPNewContam(AdaptaHOPNewNoContam): + @property + def HALO_ATTRIBUTES(self) -> ATTR_T: + attrs = list(super().HALO_ATTRIBUTES) + int_type = "l" if self.longint else "i" + float_type = "d" if self.double_precision else "f" + return tuple( + attrs + + [ + ("contaminated", 1, "i"), + (("m_contam", "mtot_contam"), 2, float_type), + (("n_contam", "ntot_contam"), 2, int_type), + ] + ) diff --git a/yt/frontends/adaptahop/fields.py b/yt/frontends/adaptahop/fields.py index 56e4280ada6..a6339df11db 100644 --- a/yt/frontends/adaptahop/fields.py +++ b/yt/frontends/adaptahop/fields.py @@ -39,7 +39,7 @@ class AdaptaHOPFieldInfo(FieldInfoContainer): ("ek", (e_units, [], "Halo Kinetic Energy")), ("ep", (e_units, [], "Halo Gravitational Energy")), ("ek", (e_units, [], "Halo Total Energy")), - ("spin", ("", [], "Halo Spin")), + ("spin", ("", [], "Halo Spin Parameter")), # Virial parameters ("virial_radius", (r_units, [], "Halo Virial Radius")), ("virial_mass", (m_units, [], "Halo Virial Mass")), @@ -47,7 +47,15 @@ class AdaptaHOPFieldInfo(FieldInfoContainer): ("virial_velocity", (v_units, [], "Halo Virial Velocity")), # NFW parameters ("rho0", (dens_units, [], "Halo NFW Density")), - ("R_c", (dens_units, [], "Halo NFW Scale Radius")), + ("R_c", (r_units, [], "Halo NFW Scale Radius")), + ("velocity_dispersion", ("km/s", [], "Velocity Dispersion")), + ("radius_200", (r_units, [], r"$R_\mathrm{200}$")), + ("radius_50", (r_units, [], r"$R_\mathrm{50}$")), + ("mass_200", (m_units, [], r"$M_\mathrm{200}$")), + ("mass_50", (m_units, [], r"$M_\mathrm{50}$")), + # Contamination + ("contaminated", ("", [], "Contaminated")), + ("m_contam", (m_units, [], "Contaminated Mass")), ) def setup_particle_fields(self, ptype): diff --git a/yt/frontends/adaptahop/io.py b/yt/frontends/adaptahop/io.py index df211314de4..1b07cd3d494 100644 --- a/yt/frontends/adaptahop/io.py +++ b/yt/frontends/adaptahop/io.py @@ -7,6 +7,7 @@ """ +from functools import partial from operator import attrgetter import numpy as np @@ -14,7 +15,7 @@ from yt.utilities.cython_fortran_utils import FortranFile from yt.utilities.io_handler import BaseIOHandler -from .definitions import HALO_ATTRIBUTES, HEADER_ATTRIBUTES +from .definitions import ATTR_T class IOHandlerAdaptaHOPBinary(BaseIOHandler): @@ -70,6 +71,10 @@ def iterate_over_attributes(attr_list): else: yield attr + halo_attributes = self.ds._halo_attributes + + attr_pos = partial(_find_attr_position, halo_attributes=halo_attributes) + for data_file in sorted(data_files, key=attrgetter("filename")): pcount = ( data_file.ds.parameters["nhalos"] + data_file.ds.parameters["nsubs"] @@ -77,16 +82,14 @@ def iterate_over_attributes(attr_list): if pcount == 0: continue ptype = "halos" - field_list0 = sorted(ptf[ptype], key=_find_attr_position) + field_list0 = sorted(ptf[ptype], key=attr_pos) field_list_pos = [f"raw_position_{k}" for k in "xyz"] - field_list = sorted( - set(field_list0 + field_list_pos), key=_find_attr_position - ) + field_list = sorted(set(field_list0 + field_list_pos), key=attr_pos) with FortranFile(self.ds.parameter_filename) as fpu: - params = fpu.read_attrs(HEADER_ATTRIBUTES) + params = fpu.read_attrs(self.ds._header_attributes) - todo = _todo_from_attributes(field_list) + todo = _todo_from_attributes(field_list, self.ds._halo_attributes) nhalos = params["nhalos"] + params["nsubs"] data = np.zeros((nhalos, len(field_list))) @@ -121,7 +124,8 @@ def _count_particles(self, data_file): def _identify_fields(self, data_file): fields = [] - for attr, _1, _2 in HALO_ATTRIBUTES: + + for attr, _1, _2 in self.ds._halo_attributes: if isinstance(attr, str): fields.append(("halos", attr)) else: @@ -138,7 +142,7 @@ def _get_particle_positions(self): return data with FortranFile(self.ds.parameter_filename) as fpu: - params = fpu.read_attrs(HEADER_ATTRIBUTES) + params = fpu.read_attrs(self.ds._header_attributes) todo = _todo_from_attributes( ( @@ -146,7 +150,8 @@ def _get_particle_positions(self): "raw_position_x", "raw_position_y", "raw_position_z", - ) + ), + self.ds._halo_attributes, ) nhalos = params["nhalos"] + params["nsubs"] @@ -180,7 +185,7 @@ def _get_particle_positions(self): def members(self, ihalo): offset = self._offsets[ihalo, 1] - todo = _todo_from_attributes(("particle_identities",)) + todo = _todo_from_attributes(("particle_identities",), self.ds._halo_attributes) with FortranFile(self.ds.parameter_filename) as fpu: fpu.seek(offset) if isinstance(todo[0], int): @@ -189,7 +194,7 @@ def members(self, ihalo): return members -def _todo_from_attributes(attributes): +def _todo_from_attributes(attributes: ATTR_T, halo_attributes: ATTR_T): # Helper function to generate a list of read-skip instructions given a list of # attributes. This is used to skip fields most of the fields when reading # the tree_brick files. @@ -198,7 +203,7 @@ def _todo_from_attributes(attributes): attributes = set(attributes) - for i, (attrs, l, k) in enumerate(HALO_ATTRIBUTES): + for i, (attrs, l, k) in enumerate(halo_attributes): if not isinstance(attrs, tuple): attrs_list = (attrs,) else: @@ -235,9 +240,9 @@ def _todo_from_attributes(attributes): return todo -def _find_attr_position(key): +def _find_attr_position(key, halo_attributes: ATTR_T) -> int: j = 0 - for attrs, *_ in HALO_ATTRIBUTES: + for attrs, *_ in halo_attributes: if not isinstance(attrs, tuple): attrs = (attrs,) for a in attrs: diff --git a/yt/frontends/adaptahop/tests/test_outputs.py b/yt/frontends/adaptahop/tests/test_outputs.py index 70c48458d0a..e0835c0d7aa 100644 --- a/yt/frontends/adaptahop/tests/test_outputs.py +++ b/yt/frontends/adaptahop/tests/test_outputs.py @@ -6,20 +6,24 @@ """ import numpy as np +import pytest from yt.frontends.adaptahop.data_structures import AdaptaHOPDataset from yt.testing import requires_file from yt.utilities.answer_testing.framework import data_dir_load r0 = "output_00080/info_00080.txt" -r1 = "output_00080_halos/tree_bricks080" +brick_old = "output_00080_halos/tree_bricks080" +brick_new = "output_00080_new_halos/tree_bricks080" @requires_file(r0) -@requires_file(r1) -def test_opening(): +@requires_file(brick_old) +@requires_file(brick_new) +@pytest.mark.parametrize("brick", [brick_old, brick_new]) +def test_opening(brick): ds_parent = data_dir_load(r0) - ds = data_dir_load(r1, kwargs=dict(parent_ds=ds_parent)) + ds = data_dir_load(brick, kwargs=dict(parent_ds=ds_parent)) ds.index @@ -27,12 +31,14 @@ def test_opening(): @requires_file(r0) -@requires_file(r1) -def test_field_access(): +@requires_file(brick_old) +@requires_file(brick_new) +@pytest.mark.parametrize("brick", [brick_old, brick_new]) +def test_field_access(brick): ds_parent = data_dir_load(r0) - ds = data_dir_load(r1, kwargs=dict(parent_ds=ds_parent)) + ds = data_dir_load(brick, kwargs=dict(parent_ds=ds_parent)) - skip_list = ("particle_identities", "mesh_id") + skip_list = ("particle_identities", "mesh_id", "radius_profile", "rho_profile") fields = [ (ptype, field) for (ptype, field) in ds.derived_field_list @@ -41,18 +47,17 @@ def test_field_access(): ad = ds.all_data() - def test_access(ptype, field): - ad[ptype, field] - for (ptype, field) in fields: - test_access(ptype, field) + ad[(ptype, field)] @requires_file(r0) -@requires_file(r1) -def test_get_halo(): +@requires_file(brick_old) +@requires_file(brick_new) +@pytest.mark.parametrize("brick", [brick_old, brick_new]) +def test_get_halo(brick): ds_parent = data_dir_load(r0) - ds = data_dir_load(r1, kwargs=dict(parent_ds=ds_parent)) + ds = data_dir_load(brick, kwargs=dict(parent_ds=ds_parent)) halo = ds.halo(1, ptype="io")