Source code for yt.frontends.stream.data_structures

"""
Data structures for Streaming, in-memory datasets



"""

#-----------------------------------------------------------------------------
# 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 os
import time
import weakref
import numpy as np
import uuid
from itertools import \
    chain, \
    product, \
    repeat
from collections import defaultdict

from numbers import Number as numeric_type

from yt.funcs import \
    iterable, \
    ensure_list, \
    issue_deprecation_warning
from yt.utilities.io_handler import io_registry
from yt.data_objects.field_data import \
    YTFieldData
from yt.data_objects.particle_unions import \
    ParticleUnion
from yt.data_objects.grid_patch import \
    AMRGridPatch
from yt.data_objects.static_output import \
    ParticleFile
from yt.geometry.geometry_handler import \
    YTDataChunk
from yt.geometry.grid_geometry_handler import \
    GridIndex
from yt.data_objects.octree_subset import \
    OctreeSubset
from yt.geometry.oct_geometry_handler import \
    OctreeIndex
from yt.geometry.particle_geometry_handler import \
    ParticleIndex
from yt.geometry.oct_container import \
    OctreeContainer
from yt.geometry.unstructured_mesh_handler import \
    UnstructuredIndex
from yt.data_objects.static_output import \
    Dataset
from yt.utilities.logger import ytLogger as mylog
from yt.utilities.lib.misc_utilities import \
    get_box_grids_level
from yt.geometry.grid_container import \
    GridTree, \
    MatchPointsToGrids
from yt.utilities.decompose import \
    decompose_array, get_psize
from yt.utilities.exceptions import \
    YTIllDefinedAMR, \
    YTInconsistentGridFieldShape, \
    YTInconsistentParticleFieldShape, \
    YTInconsistentGridFieldShapeGridDims
from yt.units.yt_array import \
    YTQuantity, \
    uconcatenate
from yt.utilities.flagging_methods import \
    FlaggingGrid
from yt.data_objects.unstructured_mesh import \
    SemiStructuredMesh, \
    UnstructuredMesh
from yt.extern.six import string_types
from .fields import \
    StreamFieldInfo
from yt.frontends.exodus_ii.util import \
    get_num_pseudo_dims
from yt.data_objects.unions import MeshUnion


[docs]class StreamGrid(AMRGridPatch): """ Class representing a single In-memory Grid instance. """ __slots__ = ['proc_num'] _id_offset = 0 def __init__(self, id, index): """ Returns an instance of StreamGrid with *id*, associated with *filename* and *index*. """ #All of the field parameters will be passed to us as needed. AMRGridPatch.__init__(self, id, filename = None, index = index) self._children_ids = [] self._parent_id = -1 self.Level = -1
[docs] def set_filename(self, filename): pass
def __repr__(self): return "StreamGrid_%04i" % (self.id) @property def Parent(self): if self._parent_id == -1: return None return self.index.grids[self._parent_id - self._id_offset] @property def Children(self): return [self.index.grids[cid - self._id_offset] for cid in self._children_ids]
[docs]class StreamHandler(object): def __init__(self, left_edges, right_edges, dimensions, levels, parent_ids, particle_count, processor_ids, fields, field_units, code_units, io = None, particle_types = None, periodicity = (True, True, True)): if particle_types is None: particle_types = {} self.left_edges = np.array(left_edges) self.right_edges = np.array(right_edges) self.dimensions = dimensions self.levels = levels self.parent_ids = parent_ids self.particle_count = particle_count self.processor_ids = processor_ids self.num_grids = self.levels.size self.fields = fields self.field_units = field_units self.code_units = code_units self.io = io self.particle_types = particle_types self.periodicity = periodicity
[docs] def get_fields(self): return self.fields.all_fields
[docs] def get_particle_type(self, field): if field in self.particle_types: return self.particle_types[field] else: return False
[docs]class StreamHierarchy(GridIndex): grid = StreamGrid def __init__(self, ds, dataset_type = None): self.dataset_type = dataset_type self.float_type = 'float64' self.dataset = weakref.proxy(ds) # for _obtain_enzo self.stream_handler = ds.stream_handler self.float_type = "float64" self.directory = os.getcwd() GridIndex.__init__(self, ds, dataset_type) def _count_grids(self): self.num_grids = self.stream_handler.num_grids def _parse_index(self): self.grid_dimensions = self.stream_handler.dimensions self.grid_left_edge[:] = self.stream_handler.left_edges self.grid_right_edge[:] = self.stream_handler.right_edges self.grid_levels[:] = self.stream_handler.levels self.grid_procs = self.stream_handler.processor_ids self.grid_particle_count[:] = self.stream_handler.particle_count mylog.debug("Copying reverse tree") self.grids = [] # We enumerate, so it's 0-indexed id and 1-indexed pid for id in range(self.num_grids): self.grids.append(self.grid(id, self)) self.grids[id].Level = self.grid_levels[id, 0] parent_ids = self.stream_handler.parent_ids if parent_ids is not None: reverse_tree = self.stream_handler.parent_ids.tolist() # Initial setup: for gid, pid in enumerate(reverse_tree): if pid >= 0: self.grids[gid]._parent_id = pid self.grids[pid]._children_ids.append(self.grids[gid].id) else: mylog.debug("Reconstructing parent-child relationships") self._reconstruct_parent_child() self.max_level = self.grid_levels.max() mylog.debug("Preparing grids") temp_grids = np.empty(self.num_grids, dtype='object') for i, grid in enumerate(self.grids): if (i%1e4) == 0: mylog.debug("Prepared % 7i / % 7i grids", i, self.num_grids) grid.filename = None grid._prepare_grid() grid._setup_dx() grid.proc_num = self.grid_procs[i] temp_grids[i] = grid self.grids = temp_grids mylog.debug("Prepared") def _reconstruct_parent_child(self): mask = np.empty(len(self.grids), dtype='int32') mylog.debug("First pass; identifying child grids") for i, grid in enumerate(self.grids): get_box_grids_level(self.grid_left_edge[i,:], self.grid_right_edge[i,:], self.grid_levels[i] + 1, self.grid_left_edge, self.grid_right_edge, self.grid_levels, mask) ids = np.where(mask.astype("bool")) grid._children_ids = ids[0] # where is a tuple mylog.debug("Second pass; identifying parents") self.stream_handler.parent_ids = np.zeros( self.stream_handler.num_grids, "int64") - 1 for i, grid in enumerate(self.grids): # Second pass for child in grid.Children: child._parent_id = i # _id_offset = 0 self.stream_handler.parent_ids[child.id] = i def _initialize_grid_arrays(self): GridIndex._initialize_grid_arrays(self) self.grid_procs = np.zeros((self.num_grids,1),'int32') def _detect_output_fields(self): # NOTE: Because particle unions add to the actual field list, without # having the keys in the field list itself, we need to double check # here. fl = set(self.stream_handler.get_fields()) fl.update(set(getattr(self, "field_list", []))) self.field_list = list(fl) def _populate_grid_objects(self): for g in self.grids: g._setup_dx() self.max_level = self.grid_levels.max() def _setup_data_io(self): if self.stream_handler.io is not None: self.io = self.stream_handler.io else: self.io = io_registry[self.dataset_type](self.ds) def _reset_particle_count(self): self.grid_particle_count[:] = self.stream_handler.particle_count for i, grid in enumerate(self.grids): grid.NumberOfParticles = self.grid_particle_count[i, 0]
[docs] def update_data(self, data): """ Update the stream data with a new data dict. If fields already exist, they will be replaced, but if they do not, they will be added. Fields already in the stream but not part of the data dict will be left alone. """ particle_types = set_particle_types(data[0]) self.stream_handler.particle_types.update(particle_types) self.ds._find_particle_types() for i, grid in enumerate(self.grids): field_units, gdata, number_of_particles = process_data(data[i]) self.stream_handler.particle_count[i] = number_of_particles self.stream_handler.field_units.update(field_units) for field in gdata: if field in grid.field_data: grid.field_data.pop(field, None) self.stream_handler.fields[grid.id][field] = gdata[field] self._reset_particle_count() # We only want to create a superset of fields here. for field in self.ds.field_list: if field[0] == "all": self.ds.field_list.remove(field) self._detect_output_fields() self.ds.create_field_info() mylog.debug("Creating Particle Union 'all'") pu = ParticleUnion("all", list(self.ds.particle_types_raw)) self.ds.add_particle_union(pu) self.ds.particle_types = tuple(set(self.ds.particle_types))
[docs]class StreamDataset(Dataset): _index_class = StreamHierarchy _field_info_class = StreamFieldInfo _dataset_type = 'stream' def __init__(self, stream_handler, storage_filename=None, geometry="cartesian", unit_system="cgs"): self.fluid_types += ("stream",) self.geometry = geometry self.stream_handler = stream_handler self._find_particle_types() name = "InMemoryParameterFile_%s" % uuid.uuid4().hex from yt.data_objects.static_output import _cached_datasets _cached_datasets[name] = self Dataset.__init__(self, name, self._dataset_type, unit_system=unit_system) def _parse_parameter_file(self): self.basename = self.stream_handler.name self.parameters['CurrentTimeIdentifier'] = time.time() self.unique_identifier = self.parameters["CurrentTimeIdentifier"] self.domain_left_edge = self.stream_handler.domain_left_edge.copy() self.domain_right_edge = self.stream_handler.domain_right_edge.copy() self.refine_by = self.stream_handler.refine_by self.dimensionality = self.stream_handler.dimensionality self.periodicity = self.stream_handler.periodicity self.domain_dimensions = self.stream_handler.domain_dimensions self.current_time = self.stream_handler.simulation_time self.gamma = 5./3. self.parameters['EOSType'] = -1 self.parameters['CosmologyHubbleConstantNow'] = 1.0 self.parameters['CosmologyCurrentRedshift'] = 1.0 self.parameters['HydroMethod'] = -1 if self.stream_handler.cosmology_simulation: self.cosmological_simulation = 1 self.current_redshift = self.stream_handler.current_redshift self.omega_lambda = self.stream_handler.omega_lambda self.omega_matter = self.stream_handler.omega_matter self.hubble_constant = self.stream_handler.hubble_constant else: self.current_redshift = self.omega_lambda = self.omega_matter = \ self.hubble_constant = self.cosmological_simulation = 0.0 def _set_units(self): self.field_units = self.stream_handler.field_units def _set_code_unit_attributes(self): base_units = self.stream_handler.code_units attrs = ('length_unit', 'mass_unit', 'time_unit', 'velocity_unit', 'magnetic_unit') cgs_units = ('cm', 'g', 's', 'cm/s', 'gauss') for unit, attr, cgs_unit in zip(base_units, attrs, cgs_units): if isinstance(unit, string_types): uq = self.quan(1.0, unit) elif isinstance(unit, numeric_type): uq = self.quan(unit, cgs_unit) elif isinstance(unit, YTQuantity): uq = unit elif isinstance(unit, tuple): uq = self.quan(unit[0], unit[1]) else: raise RuntimeError("%s (%s) is invalid." % (attr, unit)) setattr(self, attr, uq) @classmethod def _is_valid(cls, *args, **kwargs): return False @property def _skip_cache(self): return True def _find_particle_types(self): particle_types = set([]) for k, v in self.stream_handler.particle_types.items(): if v: particle_types.add(k[0]) self.particle_types = tuple(particle_types) self.particle_types_raw = self.particle_types
[docs]class StreamDictFieldHandler(dict): _additional_fields = () @property def all_fields(self): self_fields = chain.from_iterable(s.keys() for s in self.values()) self_fields = list(set(self_fields)) fields = list(self._additional_fields) + self_fields fields = list(set(fields)) return fields
[docs]def set_particle_types(data): particle_types = {} for key in data.keys(): if key == "number_of_particles": continue if len(data[key].shape) == 1: particle_types[key] = True else: particle_types[key] = False return particle_types
[docs]def assign_particle_data(ds, pdata): """ Assign particle data to the grids using MatchPointsToGrids. This will overwrite any existing particle data, so be careful! """ for ptype in ds.particle_types_raw: check_fields = [(ptype, "particle_position_x"), (ptype, "particle_position")] if all(f not in pdata for f in check_fields): pdata_ftype = {} for f in [k for k in sorted(pdata)]: if not hasattr(pdata[f], "shape"): continue if f == 'number_of_particles': continue mylog.debug("Reassigning '%s' to ('%s','%s')", f, ptype, f) pdata_ftype[ptype, f] = pdata.pop(f) pdata_ftype.update(pdata) pdata = pdata_ftype # Note: what we need to do here is a bit tricky. Because occasionally this # gets called before we property handle the field detection, we cannot use # any information about the index. Fortunately for us, we can generate # most of the GridTree utilizing information we already have from the # stream handler. if len(ds.stream_handler.fields) > 1: pdata.pop("number_of_particles", None) num_grids = len(ds.stream_handler.fields) parent_ids = ds.stream_handler.parent_ids num_children = np.zeros(num_grids, dtype='int64') # We're going to do this the slow way mask = np.empty(num_grids, dtype="bool") for i in range(num_grids): np.equal(parent_ids, i, mask) num_children[i] = mask.sum() levels = ds.stream_handler.levels.astype("int64").ravel() grid_tree = GridTree(num_grids, ds.stream_handler.left_edges, ds.stream_handler.right_edges, ds.stream_handler.dimensions, ds.stream_handler.parent_ids, levels, num_children) grid_pdata = [] for i in range(num_grids): grid = {"number_of_particles": 0} grid_pdata.append(grid) for ptype in ds.particle_types_raw: if (ptype, "particle_position_x") in pdata: x, y, z = (pdata[ptype, "particle_position_%s" % ax] for ax in 'xyz') elif (ptype, "particle_position") in pdata: x, y, z = pdata[ptype, "particle_position"].T else: raise KeyError( "Cannot decompose particle data without position fields!") pts = MatchPointsToGrids(grid_tree, len(x), x, y, z) particle_grid_inds = pts.find_points_in_tree() idxs = np.argsort(particle_grid_inds) particle_grid_count = np.bincount(particle_grid_inds.astype("intp"), minlength=num_grids) particle_indices = np.zeros(num_grids + 1, dtype='int64') if num_grids > 1: np.add.accumulate(particle_grid_count.squeeze(), out=particle_indices[1:]) else: particle_indices[1] = particle_grid_count.squeeze() for i, pcount in enumerate(particle_grid_count): grid_pdata[i]["number_of_particles"] += pcount start = particle_indices[i] end = particle_indices[i+1] for key in pdata.keys(): if key[0] == ptype: grid_pdata[i][key] = pdata[key][idxs][start:end] else: grid_pdata = [pdata] for pd, gi in zip(grid_pdata, sorted(ds.stream_handler.fields)): ds.stream_handler.fields[gi].update(pd) ds.stream_handler.particle_types.update(set_particle_types(pd)) npart = ds.stream_handler.fields[gi].pop("number_of_particles", 0) ds.stream_handler.particle_count[gi] = npart
[docs]def process_data(data, grid_dims=None): new_data, field_units = {}, {} for field, val in data.items(): # val is a data array if isinstance(val, np.ndarray): # val is a YTArray if hasattr(val, "units"): field_units[field] = val.units new_data[field] = val.copy().d # val is a numpy array else: field_units[field] = "" new_data[field] = val.copy() # val is a tuple of (data, units) elif isinstance(val, tuple) and len(val) == 2: try: assert isinstance(field, (string_types, tuple)), \ "Field name is not a string!" assert isinstance(val[0], np.ndarray), \ "Field data is not an ndarray!" assert isinstance(val[1], string_types), \ "Unit specification is not a string!" field_units[field] = val[1] new_data[field] = val[0] except AssertionError as e: raise RuntimeError( "The data dict appears to be invalid.\n" + str(e)) # val is a list of data to be turned into an array elif iterable(val): field_units[field] = "" new_data[field] = np.asarray(val) else: raise RuntimeError("The data dict appears to be invalid. " "The data dictionary must map from field " "names to (numpy array, unit spec) tuples. ") data = new_data # At this point, we have arrays for all our fields new_data = {} for field in data: n_shape = len(data[field].shape) if isinstance(field, tuple): new_field = field elif n_shape in (1, 2): new_field = ("io", field) elif n_shape == 3: new_field = ("stream", field) else: raise RuntimeError new_data[new_field] = data[field] field_units[new_field] = field_units.pop(field) known_fields = StreamFieldInfo.known_particle_fields \ + StreamFieldInfo.known_other_fields # We do not want to override any of the known ones, if it's not # overridden here. if any(f[0] == new_field[1] for f in known_fields) and \ field_units[new_field] == "": field_units.pop(new_field) data = new_data # Sanity checking that all fields have the same dimensions. g_shapes = [] p_shapes = defaultdict(list) for field in data: f_shape = data[field].shape n_shape = len(f_shape) if n_shape in (1, 2): p_shapes[field[0]].append((field[1], f_shape[0])) elif n_shape == 3: g_shapes.append((field, f_shape)) if len(g_shapes) > 0: g_s = np.array([s[1] for s in g_shapes]) if not np.all(g_s == g_s[0]): raise YTInconsistentGridFieldShape(g_shapes) if grid_dims is not None: if not np.all(g_s == grid_dims): raise YTInconsistentGridFieldShapeGridDims(g_shapes, grid_dims) if len(p_shapes) > 0: for ptype, p_shape in p_shapes.items(): p_s = np.array([s[1] for s in p_shape]) if not np.all(p_s == p_s[0]): raise YTInconsistentParticleFieldShape(ptype, p_shape) # Now that we know the particle fields are consistent, determine the number # of particles. if len(p_shapes) > 0: number_of_particles = np.sum([s[0][1] for s in p_shapes.values()]) else: number_of_particles = 0 return field_units, data, number_of_particles
[docs]def load_uniform_grid(data, domain_dimensions, length_unit=None, bbox=None, nprocs=1, sim_time=0.0, mass_unit=None, time_unit=None, velocity_unit=None, magnetic_unit=None, periodicity=(True, True, True), geometry="cartesian", unit_system="cgs"): r"""Load a uniform grid of data into yt as a :class:`~yt.frontends.stream.data_structures.StreamHandler`. This should allow a uniform grid of data to be loaded directly into yt and analyzed as would any others. This comes with several caveats: * Units will be incorrect unless the unit system is explicitly specified. * Some functions may behave oddly, and parallelism will be disappointing or non-existent in most cases. * Particles may be difficult to integrate. Particle fields are detected as one-dimensional fields. Parameters ---------- data : dict This is a dict of numpy arrays or (numpy array, unit spec) tuples. The keys are the field names. domain_dimensions : array_like This is the domain dimensions of the grid length_unit : string Unit to use for lengths. Defaults to unitless. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units specified by length_unit. Defaults to a cubic unit-length domain. nprocs: integer, optional If greater than 1, will create this number of subarrays out of data sim_time : float, optional The simulation time in seconds mass_unit : string Unit to use for masses. Defaults to unitless. time_unit : string Unit to use for times. Defaults to unitless. velocity_unit : string Unit to use for velocities. Defaults to unitless. magnetic_unit : string Unit to use for magnetic fields. Defaults to unitless. periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis geometry : string or tuple "cartesian", "cylindrical", "polar", "spherical", "geographic" or "spectral_cube". Optionally, a tuple can be provided to specify the axis ordering -- for instance, to specify that the axis ordering should be z, x, y, this would be: ("cartesian", ("z", "x", "y")). The same can be done for other coordinates, for instance: ("spherical", ("theta", "phi", "r")). Examples -------- >>> bbox = np.array([[0., 1.0], [-1.5, 1.5], [1.0, 2.5]]) >>> arr = np.random.random((128, 128, 128)) >>> data = dict(density=arr) >>> ds = load_uniform_grid(data, arr.shape, length_unit='cm', ... bbox=bbox, nprocs=12) >>> dd = ds.all_data() >>> dd['density'] YTArray([ 0.87568064, 0.33686453, 0.70467189, ..., 0.70439916, 0.97506269, 0.03047113]) g/cm**3 >>> data = dict(density=(arr, 'kg/m**3')) >>> ds = load_uniform_grid(data, arr.shape, length_unit=3.03e24, ... bbox=bbox, nprocs=12) >>> dd = ds.all_data() >>> dd['density'] YTArray([ 8.75680644e-04, 3.36864527e-04, 7.04671886e-04, ..., 7.04399160e-04, 9.75062693e-04, 3.04711295e-05]) g/cm**3 """ domain_dimensions = np.array(domain_dimensions) if bbox is None: bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], 'float64') domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros(nprocs, dtype='int32').reshape((nprocs,1)) # If someone included this throw it away--old API if "number_of_particles" in data: issue_deprecation_warning("It is no longer necessary to include " "the number of particles in the data " "dict. The number of particles is " "determined from the sizes of the " "particle fields.") data.pop("number_of_particles") # First we fix our field names, apply units to data # and check for consistency of field shapes field_units, data, number_of_particles = process_data( data, grid_dims=tuple(domain_dimensions)) sfh = StreamDictFieldHandler() if number_of_particles > 0: particle_types = set_particle_types(data) # Used much further below. pdata = {"number_of_particles": number_of_particles} for key in list(data.keys()): if len(data[key].shape) == 1 or key[0] == 'io': if not isinstance(key, tuple): field = ("io", key) mylog.debug("Reassigning '%s' to '%s'", key, field) else: field = key sfh._additional_fields += (field,) pdata[field] = data.pop(key) else: particle_types = {} if nprocs > 1: temp = {} new_data = {} for key in data.keys(): psize = get_psize(np.array(data[key].shape), nprocs) grid_left_edges, grid_right_edges, shapes, slices = \ decompose_array(data[key].shape, psize, bbox) grid_dimensions = np.array([shape for shape in shapes], dtype="int32") temp[key] = [data[key][slice] for slice in slices] for gid in range(nprocs): new_data[gid] = {} for key in temp.keys(): new_data[gid].update({key:temp[key][gid]}) sfh.update(new_data) del new_data, temp else: sfh.update({0:data}) grid_left_edges = domain_left_edge grid_right_edges = domain_right_edge grid_dimensions = domain_dimensions.reshape(nprocs, 3).astype("int32") if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, -np.ones(nprocs, dtype='int64'), np.zeros(nprocs, dtype='int64').reshape(nprocs,1), # particle count np.zeros(nprocs).reshape((nprocs,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "UniformGridData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = 2 if np.all(domain_dimensions[1:] == 1): dimensionality = 1 elif domain_dimensions[2] == 1: dimensionality = 2 else: dimensionality = 3 handler.dimensionality = dimensionality handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamDataset(handler, geometry=geometry, unit_system=unit_system) # Now figure out where the particles go if number_of_particles > 0: # This will update the stream handler too assign_particle_data(sds, pdata) return sds
[docs]def load_amr_grids(grid_data, domain_dimensions, bbox=None, sim_time=0.0, length_unit=None, mass_unit=None, time_unit=None, velocity_unit=None, magnetic_unit=None, periodicity=(True, True, True), geometry="cartesian", refine_by=2, unit_system="cgs"): r"""Load a set of grids of data into yt as a :class:`~yt.frontends.stream.data_structures.StreamHandler`. This should allow a sequence of grids of varying resolution of data to be loaded directly into yt and analyzed as would any others. This comes with several caveats: * Units will be incorrect unless the unit system is explicitly specified. * Some functions may behave oddly, and parallelism will be disappointing or non-existent in most cases. * Particles may be difficult to integrate. * No consistency checks are performed on the index Parameters ---------- grid_data : list of dicts This is a list of dicts. Each dict must have entries "left_edge", "right_edge", "dimensions", "level", and then any remaining entries are assumed to be fields. Field entries must map to an NDArray. The grid_data may also include a particle count. If no particle count is supplied, the dataset is understood to contain no particles. The grid_data will be modified in place and can't be assumed to be static. domain_dimensions : array_like This is the domain dimensions of the grid length_unit : string or float Unit to use for lengths. Defaults to unitless. If set to be a string, the bbox dimensions are assumed to be in the corresponding units. If set to a float, the value is a assumed to be the conversion from bbox dimensions to centimeters. mass_unit : string or float Unit to use for masses. Defaults to unitless. time_unit : string or float Unit to use for times. Defaults to unitless. velocity_unit : string or float Unit to use for velocities. Defaults to unitless. magnetic_unit : string or float Unit to use for magnetic fields. Defaults to unitless. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units specified by length_unit. Defaults to a cubic unit-length domain. sim_time : float, optional The simulation time in seconds periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis geometry : string or tuple "cartesian", "cylindrical", "polar", "spherical", "geographic" or "spectral_cube". Optionally, a tuple can be provided to specify the axis ordering -- for instance, to specify that the axis ordering should be z, x, y, this would be: ("cartesian", ("z", "x", "y")). The same can be done for other coordinates, for instance: ("spherical", ("theta", "phi", "r")). refine_by : integer or list/array of integers. Specifies the refinement ratio between levels. Defaults to 2. This can be an array, in which case it specifies for each dimension. For instance, this can be used to say that some datasets have refinement of 1 in one dimension, indicating that they span the full range in that dimension. Examples -------- >>> grid_data = [ ... dict(left_edge = [0.0, 0.0, 0.0], ... right_edge = [1.0, 1.0, 1.], ... level = 0, ... dimensions = [32, 32, 32], ... number_of_particles = 0), ... dict(left_edge = [0.25, 0.25, 0.25], ... right_edge = [0.75, 0.75, 0.75], ... level = 1, ... dimensions = [32, 32, 32], ... number_of_particles = 0) ... ] ... >>> for g in grid_data: ... g["density"] = (np.random.random(g["dimensions"])*2**g["level"], "g/cm**3") ... >>> ds = load_amr_grids(grid_data, [32, 32, 32], length_unit=1.0) """ domain_dimensions = np.array(domain_dimensions) ngrids = len(grid_data) if bbox is None: bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], 'float64') domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros((ngrids, 1), dtype='int32') grid_left_edges = np.zeros((ngrids, 3), dtype="float64") grid_right_edges = np.zeros((ngrids, 3), dtype="float64") grid_dimensions = np.zeros((ngrids, 3), dtype="int32") number_of_particles = np.zeros((ngrids,1), dtype='int64') parent_ids = np.zeros(ngrids, dtype="int64") - 1 sfh = StreamDictFieldHandler() for i, g in enumerate(grid_data): grid_left_edges[i,:] = g.pop("left_edge") grid_right_edges[i,:] = g.pop("right_edge") grid_dimensions[i,:] = g.pop("dimensions") grid_levels[i,:] = g.pop("level") # If someone included this throw it away--old API if "number_of_particles" in g: issue_deprecation_warning("It is no longer necessary to include " "the number of particles in the data " "dict. The number of particles is " "determined from the sizes of the " "particle fields.") g.pop("number_of_particles") field_units, data, n_particles = process_data( g, grid_dims=tuple(grid_dimensions[i,:])) number_of_particles[i, :] = n_particles sfh[i] = data # We now reconstruct our parent ids, so that our particle assignment can # proceed. mask = np.empty(ngrids, dtype='int32') for gi in range(ngrids): get_box_grids_level(grid_left_edges[gi,:], grid_right_edges[gi,:], grid_levels[gi] + 1, grid_left_edges, grid_right_edges, grid_levels, mask) ids = np.where(mask.astype("bool")) for ci in ids: parent_ids[ci] = gi # Check if the grid structure is properly aligned (bug #1295) for lvl in range(grid_levels.min() + 1, grid_levels.max() + 1): idx = grid_levels.flatten() == lvl dims = domain_dimensions * refine_by ** (lvl - 1) for iax, ax in enumerate('xyz'): cell_edges = np.linspace(domain_left_edge[iax], domain_right_edge[iax], dims[iax], endpoint=False) if set(grid_left_edges[idx, iax]) - set(cell_edges): raise YTIllDefinedAMR(lvl, ax) if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' particle_types = {} for grid in sfh.values(): particle_types.update(set_particle_types(grid)) handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, parent_ids, number_of_particles, np.zeros(ngrids).reshape((ngrids,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "AMRGridData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = refine_by if np.all(domain_dimensions[1:] == 1): dimensionality = 1 elif domain_dimensions[2] == 1: dimensionality = 2 else: dimensionality = 3 handler.dimensionality = dimensionality handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamDataset(handler, geometry=geometry, unit_system=unit_system) return sds
[docs]def refine_amr(base_ds, refinement_criteria, fluid_operators, max_level, callback=None): r"""Given a base dataset, repeatedly apply refinement criteria and fluid operators until a maximum level is reached. Parameters ---------- base_ds : ~yt.data_objects.static_output.Dataset This is any static output. It can also be a stream static output, for instance as returned by load_uniform_data. refinement_critera : list of :class:`~yt.utilities.flagging_methods.FlaggingMethod` These criteria will be applied in sequence to identify cells that need to be refined. fluid_operators : list of :class:`~yt.utilities.initial_conditions.FluidOperator` These fluid operators will be applied in sequence to all resulting grids. max_level : int The maximum level to which the data will be refined callback : function, optional A function that will be called at the beginning of each refinement cycle, with the current dataset. Examples -------- >>> domain_dims = (32, 32, 32) >>> data = np.zeros(domain_dims) + 0.25 >>> fo = [ic.CoredSphere(0.05, 0.3, [0.7,0.4,0.75], {"Density": (0.25, 100.0)})] >>> rc = [fm.flagging_method_registry["overdensity"](8.0)] >>> ug = load_uniform_grid({'Density': data}, domain_dims, 1.0) >>> ds = refine_amr(ug, rc, fo, 5) """ # If we have particle data, set it aside for now number_of_particles = np.sum([grid.NumberOfParticles for grid in base_ds.index.grids]) if number_of_particles > 0: pdata = {} for field in base_ds.field_list: if not isinstance(field, tuple): field = ("unknown", field) fi = base_ds._get_field_info(*field) if fi.particle_type and field[0] in base_ds.particle_types_raw: pdata[field] = uconcatenate([grid[field] for grid in base_ds.index.grids]) pdata["number_of_particles"] = number_of_particles last_gc = base_ds.index.num_grids cur_gc = -1 ds = base_ds bbox = np.array([(ds.domain_left_edge[i], ds.domain_right_edge[i]) for i in range(3)]) while ds.index.max_level < max_level and last_gc != cur_gc: mylog.info("Refining another level. Current max level: %s", ds.index.max_level) last_gc = ds.index.grids.size for m in fluid_operators: m.apply(ds) if callback is not None: callback(ds) grid_data = [] for g in ds.index.grids: gd = dict(left_edge=g.LeftEdge, right_edge=g.RightEdge, level=g.Level, dimensions=g.ActiveDimensions) for field in ds.field_list: if not isinstance(field, tuple): field = ("unknown", field) fi = ds._get_field_info(*field) if not fi.particle_type: gd[field] = g[field] grid_data.append(gd) if g.Level < ds.index.max_level: continue fg = FlaggingGrid(g, refinement_criteria) nsg = fg.find_subgrids() for sg in nsg: LE = sg.left_index * g.dds + ds.domain_left_edge dims = sg.dimensions * ds.refine_by grid = ds.smoothed_covering_grid(g.Level + 1, LE, dims) gd = dict(left_edge=LE, right_edge=grid.right_edge, level=g.Level + 1, dimensions=dims) for field in ds.field_list: if not isinstance(field, tuple): field = ("unknown", field) fi = ds._get_field_info(*field) if not fi.particle_type: gd[field] = grid[field] grid_data.append(gd) ds = load_amr_grids(grid_data, ds.domain_dimensions, bbox=bbox) ds.particle_types_raw = base_ds.particle_types_raw ds.particle_types = ds.particle_types_raw # Now figure out where the particles go if number_of_particles > 0: # This will update the stream handler too assign_particle_data(ds, pdata) cur_gc = ds.index.num_grids return ds
[docs]class StreamParticleIndex(ParticleIndex): def __init__(self, ds, dataset_type = None): self.stream_handler = ds.stream_handler super(StreamParticleIndex, self).__init__(ds, dataset_type) def _setup_data_io(self): if self.stream_handler.io is not None: self.io = self.stream_handler.io else: self.io = io_registry[self.dataset_type](self.ds)
[docs]class StreamParticleFile(ParticleFile): pass
[docs]class StreamParticlesDataset(StreamDataset): _index_class = StreamParticleIndex _file_class = StreamParticleFile _field_info_class = StreamFieldInfo _dataset_type = "stream_particles" file_count = 1 filename_template = "stream_file" n_ref = 64 over_refine_factor = 1
[docs]def load_particles(data, length_unit = None, bbox=None, sim_time=0.0, mass_unit = None, time_unit = None, velocity_unit=None, magnetic_unit=None, periodicity=(True, True, True), n_ref = 64, over_refine_factor = 1, geometry = "cartesian", unit_system="cgs"): r"""Load a set of particles into yt as a :class:`~yt.frontends.stream.data_structures.StreamParticleHandler`. This should allow a collection of particle data to be loaded directly into yt and analyzed as would any others. This comes with several caveats: * There must be sufficient space in memory to contain both the particle data and the octree used to index the particles. * Parallelism will be disappointing or non-existent in most cases. This will initialize an Octree of data. Note that fluid fields will not work yet, or possibly ever. Parameters ---------- data : dict This is a dict of numpy arrays or (numpy array, unit name) tuples, where the keys are the field names. Particles positions must be named "particle_position_x", "particle_position_y", and "particle_position_z". length_unit : float Conversion factor from simulation length units to centimeters mass_unit : float Conversion factor from simulation mass units to grams time_unit : float Conversion factor from simulation time units to seconds velocity_unit : float Conversion factor from simulation velocity units to cm/s magnetic_unit : float Conversion factor from simulation magnetic units to gauss bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units of the length_unit sim_time : float, optional The simulation time in seconds periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis n_ref : int The number of particles that result in refining an oct used for indexing the particles. Examples -------- >>> pos = [np.random.random(128*128*128) for i in range(3)] >>> data = dict(particle_position_x = pos[0], ... particle_position_y = pos[1], ... particle_position_z = pos[2]) >>> bbox = np.array([[0., 1.0], [0.0, 1.0], [0.0, 1.0]]) >>> ds = load_particles(data, 3.08e24, bbox=bbox) """ domain_dimensions = np.ones(3, "int32") * (1 << over_refine_factor) nprocs = 1 if bbox is None: bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], 'float64') else: bbox = np.array(bbox) domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros(nprocs, dtype='int32').reshape((nprocs,1)) field_units, data, _ = process_data(data) sfh = StreamDictFieldHandler() pdata = {} for key in data.keys() : if not isinstance(key, tuple): field = ("io", key) mylog.debug("Reassigning '%s' to '%s'", key, field) else: field = key pdata[field] = data[key] sfh._additional_fields += (field,) data = pdata # Drop reference count particle_types = set_particle_types(data) sfh.update({'stream_file':data}) grid_left_edges = domain_left_edge grid_right_edges = domain_right_edge grid_dimensions = domain_dimensions.reshape(nprocs,3).astype("int32") if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' # I'm not sure we need any of this. handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, -np.ones(nprocs, dtype='int64'), np.zeros(nprocs, dtype='int64').reshape(nprocs,1), # Temporary np.zeros(nprocs).reshape((nprocs,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "ParticleData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = 2 handler.dimensionality = 3 handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamParticlesDataset(handler, geometry=geometry, unit_system=unit_system) sds.n_ref = n_ref sds.over_refine_factor = over_refine_factor return sds
_cis = np.fromiter(chain.from_iterable(product([0,1], [0,1], [0,1])), dtype=np.int64, count = 8*3) _cis.shape = (8, 3)
[docs]def hexahedral_connectivity(xgrid, ygrid, zgrid): r"""Define the cell coordinates and cell neighbors of a hexahedral mesh for a semistructured grid. Used to specify the connectivity and coordinates parameters used in :func:`~yt.frontends.stream.data_structures.load_hexahedral_mesh`. Parameters ---------- xgrid : array_like x-coordinates of boundaries of the hexahedral cells. Should be a one-dimensional array. ygrid : array_like y-coordinates of boundaries of the hexahedral cells. Should be a one-dimensional array. zgrid : array_like z-coordinates of boundaries of the hexahedral cells. Should be a one-dimensional array. Returns ------- coords : array_like The list of (x,y,z) coordinates of the vertices of the mesh. Is of size (M,3) where M is the number of vertices. connectivity : array_like For each hexahedron h in the mesh, gives the index of each of h's neighbors. Is of size (N,8), where N is the number of hexahedra. Examples -------- >>> xgrid = np.array([-1,-0.25,0,0.25,1]) >>> coords, conn = hexahedral_connectivity(xgrid,xgrid,xgrid) >>> coords array([[-1. , -1. , -1. ], [-1. , -1. , -0.25], [-1. , -1. , 0. ], ..., [ 1. , 1. , 0. ], [ 1. , 1. , 0.25], [ 1. , 1. , 1. ]]) >>> conn array([[ 0, 1, 5, 6, 25, 26, 30, 31], [ 1, 2, 6, 7, 26, 27, 31, 32], [ 2, 3, 7, 8, 27, 28, 32, 33], ..., [ 91, 92, 96, 97, 116, 117, 121, 122], [ 92, 93, 97, 98, 117, 118, 122, 123], [ 93, 94, 98, 99, 118, 119, 123, 124]]) """ nx = len(xgrid) ny = len(ygrid) nz = len(zgrid) coords = np.zeros((nx, ny, nz, 3), dtype="float64", order="C") coords[:,:,:,0] = xgrid[:,None,None] coords[:,:,:,1] = ygrid[None,:,None] coords[:,:,:,2] = zgrid[None,None,:] coords.shape = (nx * ny * nz, 3) cycle = np.rollaxis(np.indices((nx-1,ny-1,nz-1)), 0, 4) cycle.shape = ((nx-1)*(ny-1)*(nz-1), 3) off = _cis + cycle[:, np.newaxis] connectivity = ((off[:,:,0] * ny) + off[:,:,1]) * nz + off[:,:,2] return coords, connectivity
[docs]class StreamHexahedralMesh(SemiStructuredMesh): _connectivity_length = 8 _index_offset = 0
[docs]class StreamHexahedralHierarchy(UnstructuredIndex): def __init__(self, ds, dataset_type = None): self.stream_handler = ds.stream_handler super(StreamHexahedralHierarchy, self).__init__(ds, dataset_type) def _initialize_mesh(self): coords = self.stream_handler.fields.pop('coordinates') connec = self.stream_handler.fields.pop('connectivity') self.meshes = [StreamHexahedralMesh(0, self.index_filename, connec, coords, self)] def _setup_data_io(self): if self.stream_handler.io is not None: self.io = self.stream_handler.io else: self.io = io_registry[self.dataset_type](self.ds) def _detect_output_fields(self): self.field_list = list(set(self.stream_handler.get_fields()))
[docs]class StreamHexahedralDataset(StreamDataset): _index_class = StreamHexahedralHierarchy _field_info_class = StreamFieldInfo _dataset_type = "stream_hexahedral"
[docs]def load_hexahedral_mesh(data, connectivity, coordinates, length_unit = None, bbox=None, sim_time=0.0, mass_unit = None, time_unit = None, velocity_unit = None, magnetic_unit = None, periodicity=(True, True, True), geometry = "cartesian", unit_system="cgs"): r"""Load a hexahedral mesh of data into yt as a :class:`~yt.frontends.stream.data_structures.StreamHandler`. This should allow a semistructured grid of data to be loaded directly into yt and analyzed as would any others. This comes with several caveats: * Units will be incorrect unless the data has already been converted to cgs. * Some functions may behave oddly, and parallelism will be disappointing or non-existent in most cases. * Particles may be difficult to integrate. Particle fields are detected as one-dimensional fields. The number of particles is set by the "number_of_particles" key in data. Parameters ---------- data : dict This is a dict of numpy arrays, where the keys are the field names. There must only be one. Note that the data in the numpy arrays should define the cell-averaged value for of the quantity in in the hexahedral cell. connectivity : array_like This should be of size (N,8) where N is the number of zones. coordinates : array_like This should be of size (M,3) where M is the number of vertices indicated in the connectivity matrix. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units of the length unit. sim_time : float, optional The simulation time in seconds mass_unit : string Unit to use for masses. Defaults to unitless. time_unit : string Unit to use for times. Defaults to unitless. velocity_unit : string Unit to use for velocities. Defaults to unitless. magnetic_unit : string Unit to use for magnetic fields. Defaults to unitless. periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis geometry : string or tuple "cartesian", "cylindrical", "polar", "spherical", "geographic" or "spectral_cube". Optionally, a tuple can be provided to specify the axis ordering -- for instance, to specify that the axis ordering should be z, x, y, this would be: ("cartesian", ("z", "x", "y")). The same can be done for other coordinates, for instance: ("spherical", ("theta", "phi", "r")). """ domain_dimensions = np.ones(3, "int32") * 2 nprocs = 1 if bbox is None: bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], 'float64') domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros(nprocs, dtype='int32').reshape((nprocs,1)) field_units, data, _ = process_data(data) sfh = StreamDictFieldHandler() particle_types = set_particle_types(data) sfh.update({'connectivity': connectivity, 'coordinates': coordinates, 0: data}) # Simple check for axis length correctness if len(data) > 0: fn = list(sorted(data))[0] array_values = data[fn] if array_values.size != connectivity.shape[0]: mylog.error("Dimensions of array must be one fewer than the" + " coordinate set.") raise RuntimeError grid_left_edges = domain_left_edge grid_right_edges = domain_right_edge grid_dimensions = domain_dimensions.reshape(nprocs,3).astype("int32") if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' # I'm not sure we need any of this. handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, -np.ones(nprocs, dtype='int64'), np.zeros(nprocs, dtype='int64').reshape(nprocs,1), # Temporary np.zeros(nprocs).reshape((nprocs,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "HexahedralMeshData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = 2 handler.dimensionality = 3 handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamHexahedralDataset(handler, geometry=geometry, unit_system=unit_system) return sds
[docs]class StreamOctreeSubset(OctreeSubset): domain_id = 1 _domain_offset = 1 def __init__(self, base_region, ds, oct_handler, over_refine_factor = 1): self._num_zones = 1 << (over_refine_factor) self.field_data = YTFieldData() self.field_parameters = {} self.ds = ds self.oct_handler = oct_handler self._last_mask = None self._last_selector_id = None self._current_particle_type = 'io' self._current_fluid_type = self.ds.default_fluid_type self.base_region = base_region self.base_selector = base_region.selector
[docs] def fill(self, content, dest, selector, offset): # Here we get a copy of the file, which we skip through and read the # bits we want. oct_handler = self.oct_handler cell_count = selector.count_oct_cells(self.oct_handler, self.domain_id) levels, cell_inds, file_inds = self.oct_handler.file_index_octs( selector, self.domain_id, cell_count) levels[:] = 0 dest.update((field, np.empty(cell_count, dtype="float64")) for field in content) # Make references ... count = oct_handler.fill_level(0, levels, cell_inds, file_inds, dest, content, offset) return count
[docs]class StreamOctreeHandler(OctreeIndex): def __init__(self, ds, dataset_type = None): self.stream_handler = ds.stream_handler self.dataset_type = dataset_type super(StreamOctreeHandler, self).__init__(ds, dataset_type) def _setup_data_io(self): if self.stream_handler.io is not None: self.io = self.stream_handler.io else: self.io = io_registry[self.dataset_type](self.ds) def _initialize_oct_handler(self): header = dict(dims = [1, 1, 1], left_edge = self.ds.domain_left_edge, right_edge = self.ds.domain_right_edge, octree = self.ds.octree_mask, over_refine = self.ds.over_refine_factor, partial_coverage = self.ds.partial_coverage) self.oct_handler = OctreeContainer.load_octree(header) def _identify_base_chunk(self, dobj): if getattr(dobj, "_chunk_info", None) is None: base_region = getattr(dobj, "base_region", dobj) subset = [StreamOctreeSubset(base_region, self.dataset, self.oct_handler, self.ds.over_refine_factor)] dobj._chunk_info = subset dobj._current_chunk = list(self._chunk_all(dobj))[0] def _chunk_all(self, dobj): oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) yield YTDataChunk(dobj, "all", oobjs, None) def _chunk_spatial(self, dobj, ngz, sort = None, preload_fields = None): sobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) # We actually do not really use the data files except as input to the # ParticleOctreeSubset. # This is where we will perform cutting of the Octree and # load-balancing. That may require a specialized selector object to # cut based on some space-filling curve index. 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]) def _chunk_io(self, dobj, cache = True, local_only = False): oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info) for subset in oobjs: yield YTDataChunk(dobj, "io", [subset], None, cache = cache) def _setup_classes(self): dd = self._get_data_reader_dict() super(StreamOctreeHandler, self)._setup_classes(dd) def _detect_output_fields(self): # NOTE: Because particle unions add to the actual field list, without # having the keys in the field list itself, we need to double check # here. fl = set(self.stream_handler.get_fields()) fl.update(set(getattr(self, "field_list", []))) self.field_list = list(fl)
[docs]class StreamOctreeDataset(StreamDataset): _index_class = StreamOctreeHandler _field_info_class = StreamFieldInfo _dataset_type = "stream_octree"
[docs]def load_octree(octree_mask, data, bbox=None, sim_time=0.0, length_unit=None, mass_unit=None, time_unit=None, velocity_unit=None, magnetic_unit=None, periodicity=(True, True, True), over_refine_factor = 1, partial_coverage = 1, unit_system="cgs"): r"""Load an octree mask into yt. Octrees can be saved out by calling save_octree on an OctreeContainer. This enables them to be loaded back in. This will initialize an Octree of data. Note that fluid fields will not work yet, or possibly ever. Parameters ---------- octree_mask : np.ndarray[uint8_t] This is a depth-first refinement mask for an Octree. It should be of size n_octs * 8 (but see note about the root oct below), where each item is 1 for an oct-cell being refined and 0 for it not being refined. For over_refine_factors != 1, the children count will still be 8, so there will stil be n_octs * 8 entries. Note that if the root oct is not refined, there will be only one entry for the root, so the size of the mask will be (n_octs - 1)*8 + 1. data : dict A dictionary of 1D arrays. Note that these must of the size of the number of "False" values in the ``octree_mask``. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units of length sim_time : float, optional The simulation time in seconds length_unit : string Unit to use for lengths. Defaults to unitless. mass_unit : string Unit to use for masses. Defaults to unitless. time_unit : string Unit to use for times. Defaults to unitless. velocity_unit : string Unit to use for velocities. Defaults to unitless. magnetic_unit : string Unit to use for magnetic fields. Defaults to unitless. periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis partial_coverage : boolean Whether or not an oct can be refined cell-by-cell, or whether all 8 get refined. Example ------- >>> import yt >>> import numpy as np >>> oct_mask = [8, 0, 0, 0, 0, 8, 0, 8, ... 0, 0, 0, 0, 0, 0, 0, 0, ... 8, 0, 0, 0, 0, 0, 0, 0, ... 0] >>> >>> octree_mask = np.array(oct_mask, dtype=np.uint8) >>> quantities = {} >>> quantities['gas', 'density'] = np.random.random((22, 1), dtype='f8') >>> bbox = np.array([[-10., 10.], [-10., 10.], [-10., 10.]]) >>> >>> ds = yt.load_octree(octree_mask=octree_mask, ... data=quantities, ... bbox=bbox, ... over_refine_factor=0, ... partial_coverage=0) """ if not isinstance(octree_mask, np.ndarray) or octree_mask.dtype != np.uint8: raise TypeError("octree_mask should be a Numpy array with type uint8") nz = (1 << (over_refine_factor)) domain_dimensions = np.array([nz, nz, nz]) nprocs = 1 if bbox is None: bbox = np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], 'float64') domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros(nprocs, dtype='int32').reshape((nprocs,1)) field_units, data, _ = process_data(data) sfh = StreamDictFieldHandler() particle_types = set_particle_types(data) sfh.update({0:data}) grid_left_edges = domain_left_edge grid_right_edges = domain_right_edge grid_dimensions = domain_dimensions.reshape(nprocs,3).astype("int32") if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' # I'm not sure we need any of this. handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, -np.ones(nprocs, dtype='int64'), np.zeros(nprocs, dtype='int64').reshape(nprocs,1), # Temporary np.zeros(nprocs).reshape((nprocs,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "OctreeData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = 2 handler.dimensionality = 3 handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamOctreeDataset(handler, unit_system=unit_system) sds.octree_mask = octree_mask sds.partial_coverage = partial_coverage sds.over_refine_factor = over_refine_factor return sds
[docs]class StreamUnstructuredMesh(UnstructuredMesh): _index_offset = 0 def __init__(self, *args, **kwargs): super(StreamUnstructuredMesh, self).__init__(*args, **kwargs) self._connectivity_length = self.connectivity_indices.shape[1]
[docs]class StreamUnstructuredIndex(UnstructuredIndex): def __init__(self, ds, dataset_type = None): self.stream_handler = ds.stream_handler super(StreamUnstructuredIndex, self).__init__(ds, dataset_type) def _initialize_mesh(self): coords = ensure_list(self.stream_handler.fields.pop("coordinates")) connec = ensure_list(self.stream_handler.fields.pop("connectivity")) self.meshes = [StreamUnstructuredMesh( i, self.index_filename, c1, c2, self) for i, (c1, c2) in enumerate(zip(connec, repeat(coords[0])))] self.mesh_union = MeshUnion("mesh_union", self.meshes) def _setup_data_io(self): if self.stream_handler.io is not None: self.io = self.stream_handler.io else: self.io = io_registry[self.dataset_type](self.ds) def _detect_output_fields(self): self.field_list = list(set(self.stream_handler.get_fields())) fnames = list(set([fn for ft, fn in self.field_list])) self.field_list += [('all', fname) for fname in fnames]
[docs]class StreamUnstructuredMeshDataset(StreamDataset): _index_class = StreamUnstructuredIndex _field_info_class = StreamFieldInfo _dataset_type = "stream_unstructured" def _find_particle_types(self): pass
[docs]def load_unstructured_mesh(connectivity, coordinates, node_data=None, elem_data=None, length_unit=None, bbox=None, sim_time=0.0, mass_unit=None, time_unit=None, velocity_unit=None, magnetic_unit=None, periodicity=(False, False, False), geometry = "cartesian", unit_system="cgs"): r"""Load an unstructured mesh of data into yt as a :class:`~yt.frontends.stream.data_structures.StreamHandler`. This should allow an unstructured mesh data to be loaded directly into yt and analyzed as would any others. Not all functionality for visualization will be present, and some analysis functions may not yet have been implemented. Particle fields are detected as one-dimensional fields. The number of particles is set by the "number_of_particles" key in data. In the parameter descriptions below, a "vertex" is a 3D point in space, an "element" is a single polyhedron whose location is defined by a set of vertices, and a "mesh" is a set of polyhedral elements, each with the same number of vertices. Parameters ---------- connectivity : list of array_like or array_like This should either be a single 2D array or list of 2D arrays. If this is a list, each element in the list corresponds to the connectivity information for a distinct mesh. Each array can have different connectivity length and should be of shape (N,M) where N is the number of elements and M is the number of vertices per element. coordinates : array_like The 3D coordinates of mesh vertices. This should be of size (L, D) where L is the number of vertices and D is the number of coordinates per vertex (the spatial dimensions of the dataset). Currently this must be either 2 or 3. When loading more than one mesh, the data for each mesh should be concatenated into a single coordinates array. node_data : dict or list of dicts For a single mesh, a dict mapping field names to 2D numpy arrays, representing data defined at element vertices. For multiple meshes, this must be a list of dicts. elem_data : dict or list of dicts For a single mesh, a dict mapping field names to 1D numpy arrays, where each array has a length equal to the number of elements. The data must be defined at the center of each mesh element and there must be only one data value for each element. For multiple meshes, this must be a list of dicts, with one dict for each mesh. bbox : array_like (xdim:zdim, LE:RE), optional Size of computational domain in units of the length unit. sim_time : float, optional The simulation time in seconds mass_unit : string Unit to use for masses. Defaults to unitless. time_unit : string Unit to use for times. Defaults to unitless. velocity_unit : string Unit to use for velocities. Defaults to unitless. magnetic_unit : string Unit to use for magnetic fields. Defaults to unitless. periodicity : tuple of booleans Determines whether the data will be treated as periodic along each axis geometry : string or tuple "cartesian", "cylindrical", "polar", "spherical", "geographic" or "spectral_cube". Optionally, a tuple can be provided to specify the axis ordering -- for instance, to specify that the axis ordering should be z, x, y, this would be: ("cartesian", ("z", "x", "y")). The same can be done for other coordinates, for instance: ("spherical", ("theta", "phi", "r")). >>> # Coordinates for vertices of two tetrahedra >>> coordinates = np.array([[0.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.5, 1, 0.5], [0.5, 0.5, 0.0], [0.5, 0.5, 1.0]]) >>> # The indices in the coordinates array of mesh vertices. >>> # This mesh has two elements. >>> connectivity = np.array([[0, 1, 2, 4], [0, 1, 2, 3]]) >>> # Field data defined at the centers of the two mesh elements. >>> elem_data = { ... ('connect1', 'elem_field'): np.array([1, 2]) ... } >>> # Field data defined at node vertices >>> node_data = { ... ('connect1', 'node_field'): np.array([[0.0, 1.0, 2.0, 4.0], ... [0.0, 1.0, 2.0, 3.0]]) ... } >>> ds = yt.load_unstructured_mesh(connectivity, coordinates, ... elem_data=elem_data, ... node_data=node_data) """ dimensionality = coordinates.shape[1] domain_dimensions = np.ones(3, "int32") * 2 nprocs = 1 if elem_data is None and node_data is None: raise RuntimeError("No data supplied in load_unstructured_mesh.") if isinstance(connectivity, list): num_meshes = len(connectivity) else: num_meshes = 1 connectivity = ensure_list(connectivity) if elem_data is None: elem_data = [{} for i in range(num_meshes)] elem_data = ensure_list(elem_data) if node_data is None: node_data = [{} for i in range(num_meshes)] node_data = ensure_list(node_data) data = [{} for i in range(num_meshes)] for elem_dict, data_dict in zip(elem_data, data): for field, values in elem_dict.items(): data_dict[field] = values for node_dict, data_dict in zip(node_data, data): for field, values in node_dict.items(): data_dict[field] = values data = ensure_list(data) if bbox is None: bbox = [[coordinates[:,i].min() - 0.1 * abs(coordinates[:,i].min()), coordinates[:,i].max() + 0.1 * abs(coordinates[:,i].max())] for i in range(dimensionality)] if dimensionality < 3: bbox.append([0.0, 1.0]) if dimensionality < 2: bbox.append([0.0, 1.0]) # handle pseudo-dims here num_pseudo_dims = get_num_pseudo_dims(coordinates) dimensionality -= num_pseudo_dims for i in range(dimensionality, 3): bbox[i][0] = 0.0 bbox[i][1] = 1.0 bbox = np.array(bbox, dtype=np.float64) domain_left_edge = np.array(bbox[:, 0], 'float64') domain_right_edge = np.array(bbox[:, 1], 'float64') grid_levels = np.zeros(nprocs, dtype='int32').reshape((nprocs,1)) field_units = {} particle_types = {} sfh = StreamDictFieldHandler() sfh.update({'connectivity': connectivity, 'coordinates': coordinates}) for i, d in enumerate(data): _f_unit, _data, _ = process_data(d) field_units.update(_f_unit) sfh[i] = _data particle_types.update(set_particle_types(d)) # Simple check for axis length correctness if 0 and len(data) > 0: fn = list(sorted(data))[0] array_values = data[fn] if array_values.size != connectivity.shape[0]: mylog.error("Dimensions of array must be one fewer than the" + " coordinate set.") raise RuntimeError grid_left_edges = domain_left_edge grid_right_edges = domain_right_edge grid_dimensions = domain_dimensions.reshape(nprocs, 3).astype("int32") if length_unit is None: length_unit = 'code_length' if mass_unit is None: mass_unit = 'code_mass' if time_unit is None: time_unit = 'code_time' if velocity_unit is None: velocity_unit = 'code_velocity' if magnetic_unit is None: magnetic_unit = 'code_magnetic' # I'm not sure we need any of this. handler = StreamHandler( grid_left_edges, grid_right_edges, grid_dimensions, grid_levels, -np.ones(nprocs, dtype='int64'), np.zeros(nprocs, dtype='int64').reshape(nprocs,1), # Temporary np.zeros(nprocs).reshape((nprocs,1)), sfh, field_units, (length_unit, mass_unit, time_unit, velocity_unit, magnetic_unit), particle_types=particle_types, periodicity=periodicity ) handler.name = "UnstructuredMeshData" handler.domain_left_edge = domain_left_edge handler.domain_right_edge = domain_right_edge handler.refine_by = 2 handler.dimensionality = dimensionality handler.domain_dimensions = domain_dimensions handler.simulation_time = sim_time handler.cosmology_simulation = 0 sds = StreamUnstructuredMeshDataset(handler, geometry=geometry, unit_system=unit_system) fluid_types = ['all'] for i in range(1, num_meshes + 1): fluid_types += ['connect%d' % i] sds.fluid_types = tuple(fluid_types) def flatten(l): return [item for sublist in l for item in sublist] sds._node_fields = flatten([[f[1] for f in m] for m in node_data if m]) sds._elem_fields = flatten([[f[1] for f in m] for m in elem_data if m]) sds.default_field = [f for f in sds.field_list if f[0] == 'connect1'][-1] return sds