Source code for ytree.frontends.lhalotree_hdf5.fields

"""
LHaloTreeHDF5Arbor fields



"""

#-----------------------------------------------------------------------------
# Copyright (c) ytree development team. All rights reserved.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import re

from ytree.data_structures.fields import \
    FieldInfoContainer

m_unit = "1e10 * Msun/h"
p_unit = "Mpc/h"
r_unit = "kpc/h"
v_unit = "km/s"
j_unit = "kpc * km/s"

def _redshift(field, data):
    isn = data["SnapNum"].astype(int)
    return data.arbor.arr(data.arbor._redshifts[isn], "")

[docs] class LHaloTreeHDF5FieldInfo(FieldInfoContainer): alias_fields = ( ("position_x", "SubhaloPos_0", p_unit), ("position_y", "SubhaloPos_1", p_unit), ("position_z", "SubhaloPos_2", p_unit), ("velocity_x", "SubhaloVel_0", v_unit), ("velocity_y", "SubhaloVel_1", v_unit), ("velocity_z", "SubhaloVel_2", v_unit), ("angular_momentum_x", "SubhaloSpin_0", j_unit), ("angular_momentum_y", "SubhaloSpin_1", j_unit), ("angular_momentum_z", "SubhaloSpin_2", j_unit), ("velocity_dispersion", "SubhaloVelDisp", v_unit), ) known_fields = ( ("Group_M_Crit200", m_unit), ("Group_M_Mean200", m_unit), ("Group_M_TopHat200", m_unit), ("SubhaloHalfmassRad", r_unit), ("SubhaloHalfmassRadType", r_unit), ("SubhaloMassInRadType", m_unit), ("SubhaloMassType", m_unit), ("SubhaloPos", r_unit), ("SubhaloSpin", "kpc/h * km/s"), ("SubhaloVMax", v_unit), ("SubhaloVel", v_unit), ("SubhaloVelDisp", v_unit), ) def setup_known_fields(self): """ Add units for all <fieldname>_<number> fields as well. """ kfields = dict(self.known_fields) freg = re.compile(r"(^.+)_\d+$") for field in self: if self[field].get("units") is not None: continue if field in kfields: self[field]["units"] = kfields[field] continue fs = freg.search(field) if fs and fs.groups()[0] in kfields: self[field]["units"] = kfields[fs.groups()[0]] def setup_derived_fields(self): self.arbor.add_derived_field( "redshift", _redshift, units="", force_add=False) # this will add it to field list for saving self.arbor.field_list.append("redshift") super().setup_derived_fields()