Source code for yt.frontends.sdf.data_structures

"""
Data structures for a generic SDF frontend




"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# 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 stat
import time
import os
import sys
import contextlib

from yt.utilities.logger import ytLogger as mylog
from yt.geometry.particle_geometry_handler import \
    ParticleIndex
from yt.data_objects.static_output import \
    ParticleDataset, \
    ParticleFile
from yt.funcs import \
    get_requests, \
    setdefaultattr
from .fields import \
    SDFFieldInfo
from yt.utilities.sdf import \
    SDFRead,\
    SDFIndex,\
    HTTPSDFRead

@contextlib.contextmanager
[docs]def safeopen(*args, **kwargs): if sys.version[0] != '3': kwargs.pop("encoding") with open(*args, **kwargs) as f: yield f
# currently specified by units_2HOT == 2 in header # in future will read directly from file units_2HOT_v2_length = 3.08567802e21 units_2HOT_v2_mass = 1.98892e43 units_2HOT_v2_time = 3.1558149984e16
[docs]class SDFFile(ParticleFile): pass
[docs]class SDFDataset(ParticleDataset): _index_class = ParticleIndex _file_class = SDFFile _field_info_class = SDFFieldInfo _particle_mass_name = None _particle_coordinates_name = None _particle_velocity_name = None _midx = None _skip_cache = True _subspace = False def __init__(self, filename, dataset_type="sdf_particles", n_ref=64, over_refine_factor=1, bounding_box=None, sdf_header=None, midx_filename=None, midx_header=None, midx_level=None, field_map=None, units_override=None, unit_system="cgs"): if bounding_box is not None: self._subspace = True bbox = np.array(bounding_box, dtype="float64") if bbox.shape == (2, 3): bbox = bbox.transpose() self.domain_left_edge = bbox[:,0] self.domain_right_edge = bbox[:,1] else: self.domain_left_edge = self.domain_right_edge = None self.sdf_header = sdf_header self.midx_filename = midx_filename self.midx_header = midx_header self.midx_level = midx_level if field_map is None: field_map = {} self._field_map = field_map prefix = '' if self.midx_filename is not None: prefix += 'midx_' if filename.startswith("http"): prefix += 'http_' dataset_type = prefix + 'sdf_particles' super(SDFDataset, self).__init__( filename, dataset_type=dataset_type, units_override=units_override, unit_system=unit_system, n_ref=n_ref, over_refine_factor=over_refine_factor) def _parse_parameter_file(self): if self.parameter_filename.startswith("http"): sdf_class = HTTPSDFRead else: sdf_class = SDFRead self.sdf_container = sdf_class(self.parameter_filename, header=self.sdf_header) # Reference self.parameters = self.sdf_container.parameters self.dimensionality = 3 self.refine_by = 2 try: self.unique_identifier = \ int(os.stat(self.parameter_filename)[stat.ST_CTIME]) except: self.unique_identifier = time.time() if self.domain_left_edge is None or self.domain_right_edge is None: R0 = self.parameters['R0'] if 'offset_center' in self.parameters and self.parameters['offset_center']: self.domain_left_edge = np.array([0, 0, 0], dtype=np.float64) self.domain_right_edge = np.array([ 2.0 * self.parameters.get("R%s" % ax, R0) for ax in 'xyz'], dtype=np.float64) else: self.domain_left_edge = np.array([ -self.parameters.get("R%s" % ax, R0) for ax in 'xyz'], dtype=np.float64) self.domain_right_edge = np.array([ +self.parameters.get("R%s" % ax, R0) for ax in 'xyz'], dtype=np.float64) self.domain_left_edge *= self.parameters.get("a", 1.0) self.domain_right_edge *= self.parameters.get("a", 1.0) nz = 1 << self.over_refine_factor self.domain_dimensions = np.ones(3, "int32") * nz if "do_periodic" in self.parameters and self.parameters["do_periodic"]: self.periodicity = (True, True, True) else: self.periodicity = (False, False, False) self.cosmological_simulation = 1 self.current_redshift = self.parameters.get("redshift", 0.0) self.omega_lambda = self.parameters["Omega0_lambda"] self.omega_matter = self.parameters["Omega0_m"] if "Omega0_fld" in self.parameters: self.omega_lambda += self.parameters["Omega0_fld"] if "Omega0_r" in self.parameters: # not correct, but most codes can't handle Omega0_r self.omega_matter += self.parameters["Omega0_r"] self.hubble_constant = self.parameters["h_100"] self.current_time = units_2HOT_v2_time * self.parameters.get("tpos", 0.0) mylog.info("Calculating time to be %0.3e seconds", self.current_time) self.filename_template = self.parameter_filename self.file_count = 1 @property def midx(self): if self._midx is None: if self.midx_filename is not None: if 'http' in self.midx_filename: sdf_class = HTTPSDFRead else: sdf_class = SDFRead indexdata = sdf_class(self.midx_filename, header=self.midx_header) self._midx = SDFIndex(self.sdf_container, indexdata, level=self.midx_level) else: raise RuntimeError("SDF index0 file not supplied in load.") return self._midx def _set_code_unit_attributes(self): setdefaultattr( self, 'length_unit', self.quan(1.0, self.parameters.get("length_unit", 'kpc'))) setdefaultattr( self, 'velocity_unit', self.quan(1.0, self.parameters.get("velocity_unit", 'kpc/Gyr'))) setdefaultattr( self, 'time_unit', self.quan(1.0, self.parameters.get("time_unit", 'Gyr'))) mass_unit = self.parameters.get("mass_unit", '1e10 Msun') if ' ' in mass_unit: factor, unit = mass_unit.split(' ') else: factor = 1.0 unit = mass_unit setdefaultattr(self, 'mass_unit', self.quan(float(factor), unit)) @classmethod def _is_valid(cls, *args, **kwargs): sdf_header = kwargs.get('sdf_header', args[0]) if sdf_header.startswith("http"): requests = get_requests() if requests is None: return False hreq = requests.get(sdf_header, stream=True) if hreq.status_code != 200: return False # Grab a whole 4k page. line = next(hreq.iter_content(4096)) elif os.path.isfile(sdf_header): with safeopen(sdf_header, "r", encoding = 'ISO-8859-1') as f: line = f.read(10).strip() else: return False return line.startswith("# SDF")