Source code for ytree.frontends.treefarm.io

"""
TreeFarmArbor io classes 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 h5py
import numpy as np

from ytree.data_structures.io import \
    CatalogDataFile

[docs] class TreeFarmDataFile(CatalogDataFile): def open(self): self.fh = h5py.File(self.filename, "r") def _parse_header(self): self.open() fh = self.fh self.redshift = fh.attrs["current_redshift"] self.nhalos = fh.attrs["num_halos"] # Files with no halos won't have the units. # Keep trying until we get one. if not hasattr(self.arbor, "field_list"): self._setup_field_info(fh) self.close() def _setup_field_info(self, fh): fields = list(fh.keys()) fi = {} for field in fields: if fh[field].size == 0: # Zero-sized arrays won't have units, so don't bother. return units = fh[field].attrs["units"] if isinstance(units, bytes): units = units.decode("utf") fi[field] = {"source": "file", "units": units} fields.append("redshift") fi["redshift"] = {"source": "header", "units": ""} self.arbor.field_list = fields self.arbor.field_info.update(fi) def _read_data_default(self, rfields, dtypes): field_data = {} if not rfields: return field_data self.open() fh = self.fh field_data = dict((field, fh[field][()]) for field in rfields) self.close() for field in rfields: dtype = dtypes[field] field_data[field] = field_data[field].astype(dtype) return field_data def _read_data_select(self, rfields, tree_nodes, dtypes): field_data = {} if not rfields: return field_data file_ids = np.array([node._fi for node in tree_nodes]) self.open() fh = self.fh for field in rfields: field_data[field] = fh[field][()][file_ids] self.close() for field in rfields: dtype = dtypes[field] field_data[field] = field_data[field].astype(dtype) return field_data