"""
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 get_pbar
from ytree.data_structures.arbor import SegmentedArbor
from ytree.frontends.consistent_trees.fields import ConsistentTreesFieldInfo
from ytree.frontends.consistent_trees.io import (
ConsistentTreesDataFile,
ConsistentTreesTreeFieldIO,
ConsistentTreesHlistDataFile,
)
from ytree.frontends.consistent_trees.utilities import parse_ctrees_header
from ytree.frontends.rockstar.arbor import RockstarArbor
from ytree.utilities.exceptions import ArborDataFileEmpty
from ytree.utilities.io import f_text_block
[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]
data_file.open()
data_file.fh.seek(0, 2)
file_size = data_file.fh.tell()
pbar = get_pbar("Loading tree roots", file_size)
data_file.fh.seek(self._hoffset)
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 = data_file.fh.read(my_block)
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(self.directory, 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")
f.seek(self._hoffset)
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(self.directory, 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]]
data_file.open()
data_file.fh.seek(0, 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