Source code for ytree.frontends.rockstar.io

"""
RockstarArbor 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

from ytree.data_structures.io import \
    CatalogDataFile
from ytree.utilities.io import \
    f_text_block

[docs]class RockstarDataFile(CatalogDataFile):
[docs] def __init__(self, filename, arbor): self.offsets = None super().__init__(filename, arbor)
def open(self): self.fh = open(self.filename, "r") def _parse_header(self): self.open() f = self.fh f.seek(0, 2) self.file_size = f.tell() f.seek(0) while True: line = f.readline() if line is None: self._hoffset = f.tell() break elif not line.startswith("#"): self._hoffset = f.tell() - len(line) break elif line.startswith("#a = "): self.scale_factor = float(line.split(" = ")[1]) self.close() def _read_data_default(self, rfields, dtypes): if not rfields: return {} fi = self.arbor.field_info field_data = \ self._create_field_arrays(rfields, dtypes) offsets = [] self.open() f = self.fh f.seek(self._hoffset) file_size = self.file_size - self._hoffset for line, offset in f_text_block(f, file_size=file_size): offsets.append(offset) sline = line.split() for field in rfields: field_data[field].append(sline[fi[field]["column"]]) self.close() for field in rfields: field_data[field] = \ np.array(field_data[field], dtype=dtypes[field]) if self.offsets is None: self.offsets = np.array(offsets) return field_data def _read_data_select(self, rfields, tree_nodes, dtypes): if not rfields: return {} fi = self.arbor.field_info nt = len(tree_nodes) field_data = \ self._create_field_arrays(rfields, dtypes, size=nt) self.open() f = self.fh for i in range(nt): f.seek(self.offsets[tree_nodes[i]._fi]) line = f.readline() sline = line.split() for field in rfields: dtype = dtypes[field] field_data[field][i] = dtype(sline[fi[field]["column"]]) self.close() return field_data