Skip to content

Commit

Permalink
Merge pull request #3254 from cphyc/adaptahop-support-longint
Browse files Browse the repository at this point in the history 
ENH: AdaptaHOP more flexible support
  • Loading branch information
@matthewturk
matthewturk authored Nov 1, 2021
2 parents 991fe07 + 8c28313 commit b7dfe4a
Show file tree
Hide file tree
Showing 6 changed files with 267 additions and 68 deletions.
1 change: 1 addition & 0 deletions tests/tests.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'
53 changes: 50 additions & 3 deletions yt/frontends/adaptahop/data_structures.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import os
import re
import stat
from itertools import product

import numpy as np

Expand All @@ -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


Expand Down Expand Up @@ -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,
Expand All @@ -76,6 +79,8 @@ def __init__(
)
self.parent_ds = parent_ds

self._guess_headers_from_file(filename)

super().__init__(
filename,
dataset_type,
Expand All @@ -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
Expand Down
199 changes: 166 additions & 33 deletions yt/frontends/adaptahop/definitions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
]
)
12 changes: 10 additions & 2 deletions yt/frontends/adaptahop/fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,23 @@ 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")),
("virial_temperature", ("K", [], "Halo Virial Temperature")),
("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):
Expand Down
Loading

0 comments on commit b7dfe4a

Please sign in to comment.