"""
TreeFarmArbor 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.
# -----------------------------------------------------------------------------
from collections import defaultdict
import glob
import h5py
import re
from unyt.unit_registry import UnitRegistry
from ytree.data_structures.arbor import CatalogArbor
from ytree.frontends.treefarm.fields import TreeFarmFieldInfo
from ytree.frontends.treefarm.io import TreeFarmDataFile
from ytree.utilities.io import _hdf5_yt_attr, parse_h5_attr
[docs]
class TreeFarmArbor(CatalogArbor):
"""
Class for Arbors created with :class:`~ytree.treefarm.TreeFarm`.
"""
_suffix = ".h5"
_field_info_class = TreeFarmFieldInfo
_data_file_class = TreeFarmDataFile
def _parse_parameter_file(self):
fh = h5py.File(self.filename, "r")
for attr in ["hubble_constant", "omega_matter", "omega_lambda"]:
setattr(self, attr, fh.attrs[attr])
my_ur = UnitRegistry.from_json(parse_h5_attr(fh, "unit_registry_json"))
right = _hdf5_yt_attr(fh, "domain_right_edge", unit_registry=my_ur)
left = _hdf5_yt_attr(fh, "domain_left_edge", unit_registry=my_ur)
# Drop the "cm" suffix because all lengths will
# be in comoving units.
self.box_size = self.quan((right - left)[0].to("Mpccm/h"), "Mpc/h")
fh.close()
def _get_data_files(self):
suffix = ".0" + 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}")
fids = defaultdict(list)
for my_file in files:
fid = int(my_file[len(prefix) : -len(suffix)])
fids[fid].append(my_file)
my_files = [fids[myfid] for myfid in sorted(fids.keys(), reverse=True)]
self.data_files = [
[self._data_file_class(f, self) for f in fl] for fl in my_files
]
@classmethod
def _is_valid(self, *args, **kwargs):
"""
File should end in .h5, be loadable as an hdf5 file,
and have a "data_type" attribute.
"""
fn = args[0]
if not fn.endswith(self._suffix):
return False
try:
with h5py.File(fn, "r") as f:
dtype = f.attrs.get("data_type")
if isinstance(dtype, bytes):
dtype = dtype.astype(str)
if dtype != "halo_catalog":
return False
except BaseException:
return False
return True