Source code for yt.frontends.athena_pp.data_structures

"""
Data structures for Athena.



"""

#-----------------------------------------------------------------------------
# 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 weakref

from yt.funcs import \
    mylog, get_pbar, \
    ensure_tuple
from yt.data_objects.grid_patch import \
    AMRGridPatch
from yt.geometry.grid_geometry_handler import \
    GridIndex
from yt.data_objects.static_output import \
    Dataset
from yt.geometry.geometry_handler import \
    YTDataChunk
from yt.utilities.file_handler import \
    HDF5FileHandler
from yt.geometry.unstructured_mesh_handler import \
    UnstructuredIndex
from yt.data_objects.unstructured_mesh import \
    SemiStructuredMesh
from itertools import chain, product
from .fields import AthenaPPFieldInfo

geom_map = {"cartesian": "cartesian",
            "cylindrical": "cylindrical",
            "spherical_polar": "spherical",
            "minkowski": "cartesian",
            "tilted": "cartesian",
            "sinusoidal": "cartesian",
            "schwarzschild": "spherical",
            "kerr-schild": "spherical"}

_cis = np.fromiter(chain.from_iterable(product([0,1], [0,1], [0,1])),
                   dtype=np.int64, count = 8*3)
_cis.shape = (8, 3)

[docs]class AthenaPPLogarithmicMesh(SemiStructuredMesh): _index_offset = 0 def __init__(self, mesh_id, filename, connectivity_indices, connectivity_coords, index, blocks, dims): super(AthenaPPLogarithmicMesh, self).__init__(mesh_id, filename, connectivity_indices, connectivity_coords, index) self.mesh_blocks = blocks self.mesh_dims = dims
[docs]class AthenaPPLogarithmicIndex(UnstructuredIndex): def __init__(self, ds, dataset_type = 'athena_pp'): self._handle = ds._handle super(AthenaPPLogarithmicIndex, self).__init__(ds, dataset_type) self.index_filename = self.dataset.filename self.directory = os.path.dirname(self.dataset.filename) self.dataset_type = dataset_type def _initialize_mesh(self): mylog.debug("Setting up meshes.") num_blocks = self._handle.attrs["NumMeshBlocks"] log_loc = self._handle['LogicalLocations'] levels = self._handle["Levels"] x1f = self._handle["x1f"] x2f = self._handle["x2f"] x3f = self._handle["x3f"] nbx, nby, nbz = tuple(np.max(log_loc, axis=0)+1) nlevel = self._handle.attrs["MaxLevel"]+1 nb = np.array([nbx, nby, nbz], dtype='int') self.mesh_factors = np.ones(3, dtype='int')*((nb > 1).astype("int")+1) block_grid = -np.ones((nbx,nby,nbz,nlevel), dtype=np.int) block_grid[log_loc[:,0],log_loc[:,1],log_loc[:,2],levels[:]] = np.arange(num_blocks) block_list = np.arange(num_blocks, dtype='int64') bc = [] for i in range(num_blocks): if block_list[i] >= 0: ii, jj, kk = log_loc[i] neigh = block_grid[ii:ii+2,jj:jj+2,kk:kk+2,levels[i]] if np.all(neigh > -1): loc_ids = neigh.transpose().flatten() bc.append(loc_ids) block_list[loc_ids] = -1 else: bc.append(np.array(i)) block_list[i] = -1 num_meshes = len(bc) self.meshes = [] pbar = get_pbar("Constructing meshes", num_meshes) for i in range(num_meshes): ob = bc[i][0] x = x1f[ob,:] y = x2f[ob,:] z = x3f[ob,:] if nbx > 1: x = np.concatenate([x, x1f[bc[i][1],1:]]) if nby > 1: y = np.concatenate([y, x2f[bc[i][2],1:]]) if nbz > 1: z = np.concatenate([z, x3f[bc[i][4],1:]]) nxm = x.size nym = y.size nzm = z.size coords = np.zeros((nxm, nym, nzm, 3), dtype="float64", order="C") coords[:,:,:,0] = x[:,None,None] coords[:,:,:,1] = y[None,:,None] coords[:,:,:,2] = z[None,None,:] coords.shape = (nxm * nym * nzm, 3) cycle = np.rollaxis(np.indices((nxm-1,nym-1,nzm-1)), 0, 4) cycle.shape = ((nxm-1)*(nym-1)*(nzm-1), 3) off = _cis + cycle[:, np.newaxis] connectivity = ((off[:,:,0] * nym) + off[:,:,1]) * nzm + off[:,:,2] mesh = AthenaPPLogarithmicMesh(i, self.index_filename, connectivity, coords, self, bc[i], np.array([nxm-1, nym-1, nzm-1])) self.meshes.append(mesh) pbar.update(i) pbar.finish() mylog.debug("Done setting up meshes.") def _detect_output_fields(self): self.field_list = [("athena_pp", k) for k in self.ds._field_map] def _chunk_io(self, dobj, cache = True, local_only = False): gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) for subset in gobjs: yield YTDataChunk(dobj, "io", [subset], self._count_selection(dobj, [subset]), cache = cache)
[docs]class AthenaPPGrid(AMRGridPatch): _id_offset = 0 def __init__(self, id, index, level): AMRGridPatch.__init__(self, id, filename = index.index_filename, index = index) self.Parent = None self.Children = [] self.Level = level def _setup_dx(self): # So first we figure out what the index is. We don't assume # that dx=dy=dz , at least here. We probably do elsewhere. id = self.id - self._id_offset LE, RE = self.index.grid_left_edge[id,:], \ self.index.grid_right_edge[id,:] self.dds = self.ds.arr((RE-LE)/self.ActiveDimensions, "code_length") if self.ds.dimensionality < 2: self.dds[1] = 1.0 if self.ds.dimensionality < 3: self.dds[2] = 1.0 self.field_data['dx'], self.field_data['dy'], self.field_data['dz'] = self.dds def __repr__(self): return "AthenaPPGrid_%04i (%s)" % (self.id, self.ActiveDimensions)
[docs]class AthenaPPHierarchy(GridIndex): grid = AthenaPPGrid _dataset_type='athena_pp' _data_file = None def __init__(self, ds, dataset_type='athena_pp'): self.dataset = weakref.proxy(ds) self.directory = os.path.dirname(self.dataset.filename) self.dataset_type = dataset_type # for now, the index file is the dataset! self.index_filename = self.dataset.filename self._handle = ds._handle GridIndex.__init__(self, ds, dataset_type) def _detect_output_fields(self): self.field_list = [("athena_pp", k) for k in self.dataset._field_map] def _count_grids(self): self.num_grids = self._handle.attrs["NumMeshBlocks"] def _parse_index(self): num_grids = self._handle.attrs["NumMeshBlocks"] self.grid_left_edge = np.zeros((num_grids, 3), dtype='float64') self.grid_right_edge = np.zeros((num_grids, 3), dtype='float64') self.grid_dimensions = np.zeros((num_grids, 3), dtype='int32') for i in range(num_grids): x = self._handle["x1f"][i,:] y = self._handle["x2f"][i,:] z = self._handle["x3f"][i,:] self.grid_left_edge[i] = np.array([x[0], y[0], z[0]], dtype='float64') self.grid_right_edge[i] = np.array([x[-1], y[-1], z[-1]], dtype='float64') self.grid_dimensions[i] = self._handle.attrs["MeshBlockSize"] levels = self._handle["Levels"][:] self.grid_left_edge = self.ds.arr(self.grid_left_edge, "code_length") self.grid_right_edge = self.ds.arr(self.grid_right_edge, "code_length") self.grids = np.empty(self.num_grids, dtype='object') for i in range(num_grids): self.grids[i] = self.grid(i, self, levels[i]) if self.dataset.dimensionality <= 2: self.grid_right_edge[:,2] = self.dataset.domain_right_edge[2] if self.dataset.dimensionality == 1: self.grid_right_edge[:,1:] = self.dataset.domain_right_edge[1:] self.grid_particle_count = np.zeros([self.num_grids, 1], dtype='int64') def _populate_grid_objects(self): for g in self.grids: g._prepare_grid() g._setup_dx() self.max_level = self._handle.attrs["MaxLevel"] def _chunk_io(self, dobj, cache = True, local_only = False): gobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) for subset in gobjs: yield YTDataChunk(dobj, "io", [subset], self._count_selection(dobj, [subset]), cache = cache)
[docs]class AthenaPPDataset(Dataset): _field_info_class = AthenaPPFieldInfo _dataset_type = "athena_pp" def __init__(self, filename, dataset_type='athena_pp', storage_filename=None, parameters=None, units_override=None, unit_system="code"): self.fluid_types += ("athena_pp",) if parameters is None: parameters = {} self.specified_parameters = parameters if units_override is None: units_override = {} self._handle = HDF5FileHandler(filename) xrat = self._handle.attrs["RootGridX1"][2] yrat = self._handle.attrs["RootGridX2"][2] zrat = self._handle.attrs["RootGridX3"][2] if xrat != 1.0 or yrat != 1.0 or zrat != 1.0: self._index_class = AthenaPPLogarithmicIndex self.logarithmic = True else: self._index_class = AthenaPPHierarchy self.logarithmic = False Dataset.__init__(self, filename, dataset_type, units_override=units_override, unit_system=unit_system) self.filename = filename if storage_filename is None: storage_filename = '%s.yt' % filename.split('/')[-1] self.storage_filename = storage_filename self.backup_filename = self.filename[:-4] + "_backup.gdf" def _set_code_unit_attributes(self): """ Generates the conversion to various physical _units based on the parameter file """ if "length_unit" not in self.units_override: self.no_cgs_equiv_length = True for unit, cgs in [("length", "cm"), ("time", "s"), ("mass", "g"), ("temperature", "K")]: # We set these to cgs for now, but they may have been overriden if getattr(self, unit+'_unit', None) is not None: continue mylog.warning("Assuming 1.0 = 1.0 %s", cgs) setattr(self, "%s_unit" % unit, self.quan(1.0, cgs)) self.magnetic_unit = np.sqrt(4*np.pi * self.mass_unit / (self.time_unit**2 * self.length_unit)) self.magnetic_unit.convert_to_units("gauss") self.velocity_unit = self.length_unit / self.time_unit def _parse_parameter_file(self): xmin, xmax = self._handle.attrs["RootGridX1"][:2] ymin, ymax = self._handle.attrs["RootGridX2"][:2] zmin, zmax = self._handle.attrs["RootGridX3"][:2] self.domain_left_edge = np.array([xmin, ymin, zmin], dtype='float64') self.domain_right_edge = np.array([xmax, ymax, zmax], dtype='float64') self.geometry = geom_map[self._handle.attrs["Coordinates"].decode('utf-8')] self.domain_width = self.domain_right_edge-self.domain_left_edge self.domain_dimensions = self._handle.attrs["RootGridSize"] self._field_map = {} k = 0 for (i, dname), num_var in zip(enumerate(self._handle.attrs["DatasetNames"]), self._handle.attrs["NumVariables"]): for j in range(num_var): fname = self._handle.attrs["VariableNames"][k].decode("ascii","ignore") self._field_map[fname] = (dname.decode("ascii","ignore"), j) k += 1 self.refine_by = 2 dimensionality = 3 if self.domain_dimensions[2] == 1: dimensionality = 2 if self.domain_dimensions[1] == 1: dimensionality = 1 self.dimensionality = dimensionality self.current_time = self._handle.attrs["Time"] self.unique_identifier = self.parameter_filename.__hash__() self.cosmological_simulation = False self.num_ghost_zones = 0 self.field_ordering = 'fortran' self.boundary_conditions = [1]*6 if 'periodicity' in self.specified_parameters: self.periodicity = ensure_tuple(self.specified_parameters['periodicity']) else: self.periodicity = (True,True,True,) if 'gamma' in self.specified_parameters: self.gamma = float(self.specified_parameters['gamma']) else: self.gamma = 5./3. self.current_redshift = self.omega_lambda = self.omega_matter = \ self.hubble_constant = self.cosmological_simulation = 0.0 self.parameters['Time'] = self.current_time # Hardcode time conversion for now. self.parameters["HydroMethod"] = 0 # Hardcode for now until field staggering is supported. if "gamma" in self.specified_parameters: self.parameters["Gamma"] = self.specified_parameters["gamma"] else: self.parameters["Gamma"] = 5./3. @classmethod def _is_valid(self, *args, **kwargs): try: if args[0].endswith('athdf'): return True except: pass return False @property def _skip_cache(self): return True def __repr__(self): return self.basename.rsplit(".", 1)[0]