Source code for ytree.frontends.consistent_trees.arbor

ConsistentTreesArbor 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 numpy as np
import operator
import os

from yt.funcs import \

from ytree.data_structures.arbor import \

from ytree.frontends.consistent_trees.fields import \
from import \
    ConsistentTreesDataFile, \
    ConsistentTreesTreeFieldIO, \
from ytree.frontends.consistent_trees.utilities import \
from ytree.frontends.rockstar.arbor import \

from ytree.utilities.exceptions import \
from import \

[docs]class ConsistentTreesArbor(SegmentedArbor): """ Arbors loaded from consistent-trees tree_*.dat files. """ _field_info_class = ConsistentTreesFieldInfo _tree_field_io_class = ConsistentTreesTreeFieldIO _default_dtype = np.float32 _node_io_attrs = ('_fi', '_si', '_ei') def _get_data_files(self): self.data_files = [ConsistentTreesDataFile(self.filename)] def _parse_parameter_file(self, filename=None, ntrees_in_file=True): if filename is None: filename = self.filename fi = parse_ctrees_header( self, filename, ntrees_in_file=ntrees_in_file) self.field_list = list(fi.keys()) self.field_info.update(fi) def _plant_trees(self): if self.is_planted or self._size == 0: return lkey = len("tree ")+1 block_size = 4096 data_file = self.data_files[0], 2) file_size = data_file.fh.tell() pbar = get_pbar("Loading tree roots", file_size) offset = self._hoffset itree = 0 nblocks = np.ceil(float(file_size-self._hoffset) / block_size).astype(np.int64) for ib in range(nblocks): my_block = min(block_size, file_size - offset) if my_block <= 0: break buff = lihash = -1 for ih in range(buff.count("#")): ihash = buff.find("#", lihash+1) inl = buff.find("\n", ihash+1) if inl < 0: buff += data_file.fh.readline() inl = len(buff) uid = int(buff[ihash+lkey:inl]) self._node_info['uid'][itree] = uid lihash = ihash self._node_info['_si'][itree] = offset + inl + 1 self._node_info['_fi'][itree] = 0 if itree > 0: self._node_info['_ei'][itree-1] = offset + ihash - 1 itree += 1 offset = data_file.fh.tell() pbar.update(offset) self._node_info['_ei'][-1] = offset data_file.close() pbar.finish() @classmethod def _is_valid(self, *args, **kwargs): """ File should end in .dat and have a line in the header with the string, "Consistent Trees". """ fn = args[0] if not fn.endswith(".dat"): return False with open(fn, "r") as f: valid = False while True: line = f.readline() if line is None or not line.startswith("#"): break if "Consistent Trees" in line: valid = True break if not valid: return False return True
[docs]class ConsistentTreesGroupArbor(ConsistentTreesArbor): """ Arbors loaded from consistent-trees locations.dat files. """ def _get_data_files(self): pass def _parse_parameter_file(self): f = open(self.filename, 'r') f.readline() self._hoffset = f.tell() line = f.readline() if not line: raise ArborDataFileEmpty(self.filename) fn = os.path.join(, line.split()[3]) super()._parse_parameter_file(filename=fn, ntrees_in_file=False) def _plant_trees(self): if self.is_planted: return f = open(self.filename, 'r') ldata = list(map( lambda x: [int(x[0]), int(x[1]), int(x[2]), x[3], len(x[0])], [line.split() for line, _ in f_text_block(f, pbar_string='Reading locations')] )) f.close() self._size = len(ldata) # It's faster to create and sort arrays and then sort ldata # for some reason. dfns = np.unique([datum[3] for datum in ldata]) dfns.sort() fids = np.array([datum[1] for datum in ldata]) fids.sort() ufids = np.unique(fids) ufids.sort() # Some data files may be empty and so unlisted. # Make sure file ids and names line up. data_files = [None]*(ufids.max()+1) for i,fid in enumerate(ufids): data_files[fid] = dfns[i] self.data_files = \ [None if fn is None else ConsistentTreesDataFile(os.path.join(, fn)) for fn in data_files] ldata.sort(key=operator.itemgetter(1, 2)) pbar = get_pbar("Loading tree roots", self._size) # Set end offsets for each tree. # We don't get them from the location file. lkey = len("tree ")+3 # length of the separation line between trees same_file = np.diff(fids, append=fids[-1]+1) == 0 for i, tdata in enumerate(ldata): self._node_info['uid'][i] = tdata[0] self._node_info['_fi'][i] = tdata[1] self._node_info['_si'][i] = tdata[2] # Get end index from next tree. if same_file[i]: self._node_info['_ei'][i] = ldata[i+1][2] - lkey - tdata[4] pbar.update(i+1) pbar.finish() # Get end index for last trees in files. for i in np.where(~same_file)[0]: data_file = self.data_files[fids[i]], 2) self._node_info['_ei'][i] = data_file.fh.tell() data_file.close() @classmethod def _is_valid(self, *args, **kwargs): """ File should end in .dat and have a line in the header with the string, "Consistent Trees". """ fn = args[0] if not os.path.basename(fn) == 'locations.dat': return False with open(fn, "r") as f: valid = False while True: line = f.readline() if line is None or not line.startswith("#"): break if "TreeRootID FileID Offset Filename" in line: valid = True if not valid: return False return True
[docs]class ConsistentTreesHlistArbor(RockstarArbor): """ Class for Arbors created from consistent-trees hlist_*.list files. This is a hybrid type with multiple catalog files like the rockstar frontend, but with headers structured like consistent-trees. """ _has_uids = True _field_info_class = ConsistentTreesFieldInfo _data_file_class = ConsistentTreesHlistDataFile def _parse_parameter_file(self): ConsistentTreesArbor._parse_parameter_file( self, ntrees_in_file=False) def _get_data_files(self): """ Get all out_*.list files and sort them in reverse order. """ prefix = os.path.join(os.path.dirname(self.filename), "hlist_") suffix = ".list" my_files = glob.glob(f"{prefix}*{suffix}") # sort by catalog number my_files.sort( key=lambda x: self._get_file_index(x, prefix, suffix), reverse=True) self.data_files = \ [self._data_file_class(f, self) for f in my_files] def _get_file_index(self, f, prefix, suffix): return float(f[f.find(prefix)+len(prefix):f.rfind(suffix)]) @classmethod def _is_valid(self, *args, **kwargs): """ File should end in .list. """ fn = args[0] if not os.path.basename(fn).startswith("hlist") or \ not fn.endswith(".list"): return False return True