Source code for ytree.frontends.lhalotree_hdf5.arbor

"""
LHaloTreeHDF5Arbor class and member functions



"""

#-----------------------------------------------------------------------------
# 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 glob
import h5py
import numpy as np
import re

from yt.funcs import \
    get_pbar

from ytree.data_structures.arbor import \
    SegmentedArbor

from ytree.frontends.lhalotree_hdf5.fields import \
    LHaloTreeHDF5FieldInfo
from ytree.frontends.lhalotree_hdf5.io import \
    LHaloTreeHDF5DataFile, \
    LHaloTreeHDF5TreeFieldIO

[docs]class LHaloTreeHDF5Arbor(SegmentedArbor): """ IllustrisTNG format. """ _suffix = ".hdf5" _data_file_class = LHaloTreeHDF5DataFile _field_info_class = LHaloTreeHDF5FieldInfo _tree_field_io_class = LHaloTreeHDF5TreeFieldIO
[docs] def __init__(self, filename, hubble_constant=1.0, box_size=None, omega_matter=None, omega_lambda=None): self.hubble_constant = hubble_constant self.omega_matter = omega_matter self.omega_lambda = omega_lambda if box_size is not None: self.box_size = self.quan(box_size, "Mpc/h") super().__init__(filename)
def _get_data_files(self): suffix = self._suffix reg = re.search(rf"^(.+\D)\d+{suffix}$", self.filename) if reg is None: raise RuntimeError( f"Cannot determine numbering system for {self.filename}.") prefix = reg.groups()[0] files = glob.glob(f"{prefix}*{self._suffix}") self.data_files = [self._data_file_class(f, self) for f in files] self._file_count = np.array([df._size for df in self.data_files]) self._size = self._file_count.sum() def _parse_parameter_file(self): f = h5py.File(self.parameter_filename, mode='r') g = f["Header"] ntrees = g.attrs["NtreesPerFile"] self._redshifts = g["Redshifts"][()] field_list = [] if ntrees == 0: field_list = [] return g = f["Tree0"] for d in g: dshape = g[d].shape if len(dshape) == 1: field_list.append(d) else: field_list.extend( [f"{d}_{i}" for i in range(dshape[1])]) self.field_list = field_list fi = dict((field, {}) for field in field_list) for field in ["uid", "desc_uid"]: self.field_list.append(field) fi[field] = {"units": "", "source": "arbor"} fi["desc_uid"]["dependencies"] = ["Descendant"] self.field_info.update(fi) def _plant_trees(self): if self.is_planted or self._size == 0: return c = 0 file_offsets = self._file_count.cumsum() - self._file_count pbar = get_pbar('Planting trees', self._size) for idf, data_file in enumerate(self.data_files): data_file.open() tree_size = data_file.fh["Header"]["TreeNHalos"][()] data_file.close() size = data_file._size istart = file_offsets[idf] iend = istart + size self._node_info['_fi'][istart:iend] = idf self._node_info['_si'][istart:iend] = np.arange(size) self._node_info['_tree_size'][istart:iend] = tree_size c += size pbar.update(c) pbar.finish() uids = self._node_info['_tree_size'] self._node_info['uid'] = uids.cumsum() - uids @classmethod def _is_valid(self, *args, **kwargs): """ Should be an hdf5 file with a few key attributes. """ fn = args[0] if not h5py.is_hdf5(fn): return False attrs = ["FirstSnapshotNr", "LastSnapshotNr", "SnapSkipFac", "NtreesPerFile", "NhalosPerFile", "ParticleMass"] groups = ["Redshifts", "TotNsubhalos", "TreeNHalos"] with h5py.File(fn, mode='r') as f: g = f["Header"] for attr in attrs: if attr not in g.attrs: return False for group in groups: if group not in g: return False return True