ENH: AdaptaHOP more flexible support

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'

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

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),
+            ]
+        )

yt/frontends/adaptahop/fields.py

@@ -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):