"""
LHaloTreeArbor 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 numpy as np
# import weakref
from ytree.arbor.io import \
FieldIO, \
TreeFieldIO
# class LHaloTreeFileID(object):
# r"""Handle cached file ID.
# Args:
# filename (str): Full path to file that should be opened for reading.
# Attributes:
# fd (file): Open file descriptor.
# """
# def __init__(self, filename):
# self._finalizer = None
# self.fd = None
# self.update(filename)
# def close(self):
# r"""Close the file descriptor."""
# self._finalizer()
# def update(self, filename):
# r"""Update the file descriptor and finalizer.
# Args:
# filename (str): Full path to file that should be opened for reading.
# """
# if (self.fd is None) or (self.fd.name != filename):
# if self.fd is not None:
# self.close()
# self.fd = open(filename, 'rb')
# self._finalizer = weakref.finalize(self, self.fd.close)
[docs]class LHaloTreeTreeFieldIO(TreeFieldIO):
def _read_fields(self, root_node, fields, dtypes=None,
f=None, root_only=False):
"""
Read fields from disk for a single tree.
Here we accept a list of fields and return a dictionary of NumPy
arrays for each field.
dtypes will be an optional dictionary of type for each field
f will optionally be the already-opened file handle.
If root_only is true, we only want the field value for the root
of the tree.
Below is the example for ctrees.
"""
if dtypes is None:
dtypes = {}
lht = root_node._lht
# This stores the file ID in a class that handles clean up via
# weakref.finalize. The cached file object is checked to see if it
# access the correct file for this node. If not, the cached file
# object is closed and replaced with a file object for the correct
# file.
# This is only necessary if LHaloTreeReader is using np.fromfile (it
# currently uses np.memmap).
# Get cached fid.
# if f is None:
# if not hasattr(self, '_cached_fid'):
# self._cached_fd = LHaloTreeFileID(lht.filename)
# self._cached_fd.update(lht.filename)
# f = self._cached_fd.fd
# Don't read all data if only uid/desc_uid requested
data = None
for field in fields:
if field not in ['uid', 'desc_uid']:
# if root_only:
# data = lht.read_single_root(root_node._index_in_lht, fd=f)
# else:
data = lht.read_single_tree(root_node._index_in_lht, fd=f)
break
if data is None:
data = dict()
if root_only:
halonum = 0
else:
halonum = None
tot_idx = lht.get_total_index(root_node._index_in_lht, halonum)
data['uid'] = lht.all_uids[tot_idx]
data['desc_uid'] = lht.all_desc_uids[tot_idx]
field_data = data
# apply field units
fi = self.arbor.field_info
for field in fields:
units = fi[field].get("units", "")
if units != "":
field_data[field] = \
self.arbor.arr(field_data[field], units)
return field_data
[docs]class LHaloTreeRootFieldIO(FieldIO):
def _read_fields(self, storage_object, fields, dtypes=None):
r"""Add root fields in bulk to same time."""
if dtypes is None:
dtypes = {}
fi = self.arbor.field_info
# Consolidate root field data
field_data = {}
for field in fields:
dtype = dtypes.get(field, None)
units = fi[field].get("units", "")
field_data[field] = np.empty(self.arbor.trees.size, dtype=dtype)
if units:
field_data[field] = self.arbor.arr(field_data[field], units)
ntrees_prev = 0
for lht in self.arbor._lhtfiles:
ntrees = lht.ntrees
for field in fields:
field_data[field][ntrees_prev:(ntrees_prev + ntrees)] = \
lht._root_data[field]
ntrees_prev += ntrees
return field_data