"""
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