Source code for yt.frontends.artio.data_structures

"""
ARTIO-specific data structures




"""

#-----------------------------------------------------------------------------
# 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 os
import stat
import weakref

from collections import defaultdict

from yt.frontends.artio._artio_caller import \
    artio_is_valid, \
    artio_fileset, \
    ARTIOSFCRangeHandler
from yt.frontends.artio import _artio_caller
from yt.frontends.artio.fields import \
    ARTIOFieldInfo

from yt.funcs import \
    mylog, \
    setdefaultattr
from yt.geometry.geometry_handler import \
    Index, \
    YTDataChunk
import yt.geometry.particle_deposit as particle_deposit
from yt.data_objects.static_output import \
    Dataset
from yt.data_objects.octree_subset import \
    OctreeSubset
from yt.data_objects.field_data import \
    YTFieldData
from yt.utilities.exceptions import \
    YTParticleDepositionNotImplemented

[docs]class ARTIOOctreeSubset(OctreeSubset): _domain_offset = 0 domain_id = -1 _con_args = ("base_region", "sfc_start", "sfc_end", "oct_handler", "ds") _type_name = 'octree_subset' _num_zones = 2 def __init__(self, base_region, sfc_start, sfc_end, oct_handler, ds): self.field_data = YTFieldData() self.field_parameters = {} self.sfc_start = sfc_start self.sfc_end = sfc_end self.oct_handler = oct_handler self.ds = ds self._last_mask = None self._last_selector_id = None self._current_particle_type = 'all' self._current_fluid_type = self.ds.default_fluid_type self.base_region = base_region self.base_selector = base_region.selector @property def min_ind(self): return self.sfc_start @property def max_ind(self): return self.sfc_end
[docs] def fill(self, fields, selector): if len(fields) == 0: return [] handle = self.oct_handler.artio_handle field_indices = [handle.parameters["grid_variable_labels"].index(f) for (ft, f) in fields] cell_count = selector.count_oct_cells( self.oct_handler, self.domain_id) self.data_size = cell_count levels, cell_inds, file_inds = self.oct_handler.file_index_octs( selector, self.domain_id, cell_count) domain_counts = self.oct_handler.domain_count(selector) tr = [np.zeros(cell_count, dtype="float64") for field in fields] self.oct_handler.fill_sfc(levels, cell_inds, file_inds, domain_counts, field_indices, tr) tr = dict((field, v) for field, v in zip(fields, tr)) return tr
[docs] def fill_particles(self, fields): if len(fields) == 0: return {} ptype_indices = self.ds.particle_types art_fields = [(ptype_indices.index(ptype), fname) for ptype, fname in fields] species_data = self.oct_handler.fill_sfc_particles(art_fields) tr = defaultdict(dict) # Now we need to sum things up and then fill for s, f in fields: count = 0 dt = "float64" # default i = ptype_indices.index(s) # No vector fields in ARTIO count += species_data[i, f].size dt = species_data[i, f].dtype tr[s][f] = np.zeros(count, dtype=dt) cp = 0 v = species_data.pop((i, f)) tr[s][f][cp:cp+v.size] = v cp += v.size return tr
# We create something of a fake octree here. This is primarily to enable us to # reuse code for things like __getitem__ and the like. We will also create a # new oct_handler type that is functionally equivalent, except that it will # only manage the root mesh.
[docs]class ARTIORootMeshSubset(ARTIOOctreeSubset): _num_zones = 1 _type_name = 'sfc_subset' _selector_module = _artio_caller domain_id = -1
[docs] def fill(self, fields, selector): # We know how big these will be. if len(fields) == 0: return [] handle = self.ds._handle field_indices = [handle.parameters["grid_variable_labels"].index(f) for (ft, f) in fields] tr = self.oct_handler.fill_sfc(selector, field_indices) self.data_size = tr[0].size tr = dict((field, v) for field, v in zip(fields, tr)) return tr
[docs] def deposit(self, positions, fields = None, method = None, kernel_name = 'cubic'): # Here we perform our particle deposition. if fields is None: fields = [] cls = getattr(particle_deposit, "deposit_%s" % method, None) if cls is None: raise YTParticleDepositionNotImplemented(method) nz = self.nz nvals = (nz, nz, nz, self.ires.size) # We allocate number of zones, not number of octs op = cls(nvals, kernel_name) op.initialize() mylog.debug("Depositing %s (%s^3) particles into %s Root Mesh", positions.shape[0], positions.shape[0]**0.3333333, nvals[-1]) pos = np.array(positions, dtype="float64") f64 = [np.array(f, dtype="float64") for f in fields] self.oct_handler.deposit(op, self.base_selector, pos, f64) vals = op.finalize() if vals is None: return return np.asfortranarray(vals)
[docs]class ARTIOIndex(Index): def __init__(self, ds, dataset_type='artio'): self.dataset_type = dataset_type self.dataset = weakref.proxy(ds) # for now, the index file is the dataset! self.index_filename = self.dataset.parameter_filename self.directory = os.path.dirname(self.index_filename) self.max_level = ds.max_level self.range_handlers = {} self.float_type = np.float64 super(ARTIOIndex, self).__init__(ds, dataset_type) @property def max_range(self): return self.dataset.max_range def _setup_geometry(self): mylog.debug("Initializing Geometry Handler empty for now.")
[docs] def get_smallest_dx(self): """ Returns (in code units) the smallest cell size in the simulation. """ return (self.dataset.domain_width / (self.dataset.domain_dimensions * 2**(self.max_level))).min()
def _get_particle_type_counts(self): # this could be done in the artio C interface without creating temporary # arrays but I don't want to touch that code # if a future brave soul wants to try, take a look at # `read_sfc_particles` in _artio_caller.pyx result = {} ad = self.ds.all_data() for ptype in self.ds.particle_types_raw: result[ptype] = ad[ptype, 'PID'].size return result
[docs] def convert(self, unit): return self.dataset.conversion_factors[unit]
[docs] def find_max(self, field, finest_levels=3): """ Returns (value, center) of location of maximum for a given field. """ if (field, finest_levels) in self._max_locations: return self._max_locations[(field, finest_levels)] mv, pos = self.find_max_cell_location(field, finest_levels) self._max_locations[(field, finest_levels)] = (mv, pos) return mv, pos
[docs] def find_max_cell_location(self, field, finest_levels=3): source = self.all_data() if finest_levels is not False: source.min_level = self.max_level - finest_levels mylog.debug("Searching for maximum value of %s", field) max_val, mx, my, mz = \ source.quantities["MaxLocation"](field) mylog.info("Max Value is %0.5e at %0.16f %0.16f %0.16f", max_val, mx, my, mz) self.ds.parameters["Max%sValue" % (field)] = max_val self.ds.parameters["Max%sPos" % (field)] = "%s" % ((mx, my, mz),) return max_val, np.array((mx, my, mz), dtype='float64')
def _detect_output_fields(self): self.fluid_field_list = self._detect_fluid_fields() self.particle_field_list = self._detect_particle_fields() self.field_list = self.fluid_field_list + self.particle_field_list mylog.debug("Detected fields: %s", (self.field_list,)) def _detect_fluid_fields(self): return [("artio", f) for f in self.ds.artio_parameters["grid_variable_labels"]] def _detect_particle_fields(self): fields = set() for i, ptype in enumerate(self.ds.particle_types): if ptype == "all": break # This will always be after all intrinsic for fname in self.ds.particle_variables[i]: fields.add((ptype, fname)) return list(fields) def _identify_base_chunk(self, dobj): if getattr(dobj, "_chunk_info", None) is None: try: all_data = all(dobj.left_edge == self.ds.domain_left_edge) and\ all(dobj.right_edge == self.ds.domain_right_edge) except: all_data = False base_region = getattr(dobj, "base_region", dobj) sfc_start = getattr(dobj, "sfc_start", None) sfc_end = getattr(dobj, "sfc_end", None) nz = getattr(dobj, "_num_zones", 0) if all_data: mylog.debug("Selecting entire artio domain") list_sfc_ranges = self.ds._handle.root_sfc_ranges_all( max_range_size = self.max_range) elif sfc_start is not None and sfc_end is not None: mylog.debug("Restricting to %s .. %s", sfc_start, sfc_end) list_sfc_ranges = [(sfc_start, sfc_end)] else: mylog.debug("Running selector on artio base grid") list_sfc_ranges = self.ds._handle.root_sfc_ranges( dobj.selector, max_range_size = self.max_range) ci = [] #v = np.array(list_sfc_ranges) #list_sfc_ranges = [ (v.min(), v.max()) ] for (start, end) in list_sfc_ranges: if (start, end) in self.range_handlers.keys(): range_handler = self.range_handlers[(start, end)] else: range_handler = ARTIOSFCRangeHandler( self.ds.domain_dimensions, self.ds.domain_left_edge, self.ds.domain_right_edge, self.ds._handle, start, end) range_handler.construct_mesh() self.range_handlers[(start, end)] = range_handler if nz != 2: ci.append(ARTIORootMeshSubset(base_region, start, end, range_handler.root_mesh_handler, self.ds)) if nz != 1 and range_handler.total_octs > 0: ci.append(ARTIOOctreeSubset(base_region, start, end, range_handler.octree_handler, self.ds)) dobj._chunk_info = ci if len(list_sfc_ranges) > 1: mylog.info("Created %d chunks for ARTIO" % len(list_sfc_ranges)) dobj._current_chunk = list(self._chunk_all(dobj))[0] def _data_size(self, dobj, dobjs): size = 0 for d in dobjs: size += d.data_size return size def _chunk_all(self, dobj): oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) yield YTDataChunk(dobj, "all", oobjs, None, cache = True) def _chunk_spatial(self, dobj, ngz, preload_fields = None): if ngz > 0: raise NotImplementedError sobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) for i,og in enumerate(sobjs): if ngz > 0: g = og.retrieve_ghost_zones(ngz, [], smoothed=True) else: g = og yield YTDataChunk(dobj, "spatial", [g], None, cache = True) def _chunk_io(self, dobj, cache = True, local_only = False): # _current_chunk is made from identify_base_chunk oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) for chunk in oobjs: yield YTDataChunk(dobj, "io", [chunk], None, cache = cache) def _read_fluid_fields(self, fields, dobj, chunk=None): if len(fields) == 0: return {}, [] if chunk is None: self._identify_base_chunk(dobj) fields_to_return = {} fields_to_read, fields_to_generate = self._split_fields(fields) if len(fields_to_read) == 0: return {}, fields_to_generate fields_to_return = self.io._read_fluid_selection( self._chunk_io(dobj), dobj.selector, fields_to_read) return fields_to_return, fields_to_generate
[docs]class ARTIODataset(Dataset): _handle = None _index_class = ARTIOIndex _field_info_class = ARTIOFieldInfo def __init__(self, filename, dataset_type='artio', storage_filename=None, max_range = 1024, units_override=None, unit_system="cgs"): from sys import version if self._handle is not None: return self.max_range = max_range self.fluid_types += ("artio",) self._filename = filename self._fileset_prefix = filename[:-4] if version < '3': self._handle = artio_fileset(self._fileset_prefix) else: self._handle = artio_fileset(bytes(self._fileset_prefix,'utf-8')) self.artio_parameters = self._handle.parameters # Here we want to initiate a traceback, if the reader is not built. Dataset.__init__(self, filename, dataset_type, units_override=units_override, unit_system=unit_system) self.storage_filename = storage_filename def _set_code_unit_attributes(self): setdefaultattr( self, 'mass_unit', self.quan(self.parameters["unit_m"], "g")) setdefaultattr( self, 'length_unit', self.quan(self.parameters["unit_l"], "cm")) setdefaultattr( self, 'time_unit', self.quan(self.parameters["unit_t"], "s")) setdefaultattr(self, 'velocity_unit', self.length_unit / self.time_unit) def _parse_parameter_file(self): # hard-coded -- not provided by headers self.dimensionality = 3 self.refine_by = 2 self.parameters["HydroMethod"] = 'artio' self.parameters["Time"] = 1. # default unit is 1... # read header self.unique_identifier = \ int(os.stat(self.parameter_filename)[stat.ST_CTIME]) self.num_grid = self._handle.num_grid self.domain_dimensions = np.ones(3, dtype='int32') * self.num_grid self.domain_left_edge = np.zeros(3, dtype="float64") self.domain_right_edge = np.ones(3, dtype='float64')*self.num_grid # TODO: detect if grid exists self.min_level = 0 # ART has min_level=0 self.max_level = self.artio_parameters["grid_max_level"][0] # TODO: detect if particles exist if self._handle.has_particles: self.num_species = self.artio_parameters["num_particle_species"][0] self.particle_variables = [["PID", "SPECIES"] for i in range(self.num_species)] self.particle_types_raw = \ self.artio_parameters["particle_species_labels"] self.particle_types = tuple(self.particle_types_raw) for species in range(self.num_species): # Mass would be best as a derived field, # but wouldn't detect under 'all' if self.artio_parameters["particle_species_labels"][species]\ == "N-BODY": self.particle_variables[species].append("MASS") if self.artio_parameters["num_primary_variables"][species] > 0: self.particle_variables[species].extend( self.artio_parameters[ "species_%02d_primary_variable_labels" % (species, )]) if self.artio_parameters["num_secondary_variables"][species] > 0: self.particle_variables[species].extend( self.artio_parameters[ "species_%02d_secondary_variable_labels" % (species, )]) else: self.num_species = 0 self.particle_variables = [] self.particle_types = () self.particle_types_raw = self.particle_types self.current_time = self.quan(self._handle.tphys_from_tcode(self.artio_parameters["tl"][0]),"yr") # detect cosmology if "abox" in self.artio_parameters: self.cosmological_simulation = True abox = self.artio_parameters["abox"][0] self.omega_lambda = self.artio_parameters["OmegaL"][0] self.omega_matter = self.artio_parameters["OmegaM"][0] self.hubble_constant = self.artio_parameters["hubble"][0] self.current_redshift = 1.0/self.artio_parameters["auni"][0] - 1.0 self.current_redshift_box = 1.0/abox - 1.0 self.parameters["initial_redshift"] =\ 1.0 / self.artio_parameters["auni_init"][0] - 1.0 self.parameters["CosmologyInitialRedshift"] =\ self.parameters["initial_redshift"] self.parameters['unit_m'] = self.artio_parameters["mass_unit"][0] self.parameters['unit_t'] =\ self.artio_parameters["time_unit"][0] * abox**2 self.parameters['unit_l'] =\ self.artio_parameters["length_unit"][0] * abox if self.artio_parameters["DeltaDC"][0] != 0: mylog.warn("DeltaDC != 0, which implies auni != abox. Be sure you understand which expansion parameter is appropriate for your use! (Gnedin, Kravtsov, & Rudd 2011)") else: self.cosmological_simulation = False self.parameters['unit_l'] = self.artio_parameters["length_unit"][0] self.parameters['unit_t'] = self.artio_parameters["time_unit"][0] self.parameters['unit_m'] = self.artio_parameters["mass_unit"][0] # hard coded assumption of 3D periodicity self.periodicity = (True, True, True) @classmethod def _is_valid(self, *args, **kwargs): from sys import version # a valid artio header file starts with a prefix and ends with .art if not args[0].endswith(".art"): return False if version < '3': return artio_is_valid(args[0][:-4]) else: return artio_is_valid(bytes(args[0][:-4],'utf-8'))