Source code for yt.frontends.enzo.io

"""
Enzo-specific IO functions



"""

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

from yt.utilities.io_handler import \
    BaseIOHandler
from yt.utilities.logger import ytLogger as mylog
from yt.extern.six import b, iteritems
from yt.utilities.on_demand_imports import _h5py as h5py
from yt.geometry.selection_routines import GridSelector
import numpy as np


_convert_mass = ("particle_mass","mass")

_particle_position_names = {}

[docs]class IOHandlerPackedHDF5(BaseIOHandler): _dataset_type = "enzo_packed_3d" _base = slice(None) _field_dtype = "float64" def _read_field_names(self, grid): if grid.filename is None: return [] f = h5py.File(grid.filename, "r") try: group = f["/Grid%08i" % grid.id] except KeyError: group = f fields = [] dtypes = set([]) add_io = "io" in grid.ds.particle_types for name, v in iteritems(group): # NOTE: This won't work with 1D datasets or references. # For all versions of Enzo I know about, we can assume all floats # are of the same size. So, let's grab one. if not hasattr(v, "shape") or v.dtype == "O": continue elif len(v.dims) == 1: if grid.ds.dimensionality == 1: fields.append( ("enzo", str(name)) ) elif add_io: fields.append( ("io", str(name)) ) else: fields.append( ("enzo", str(name)) ) dtypes.add(v.dtype) if len(dtypes) == 1: # Now, if everything we saw was the same dtype, we can go ahead and # set it here. We do this because it is a HUGE savings for 32 bit # floats, since our numpy copying/casting is way faster than # h5py's, for some reason I don't understand. This does *not* need # to be correct -- it will get fixed later -- it just needs to be # okay for now. self._field_dtype = list(dtypes)[0] f.close() return fields @property def _read_exception(self): return (KeyError,) def _read_particle_coords(self, chunks, ptf): for rv in self._read_particle_fields(chunks, ptf, None): yield rv def _read_particle_fields(self, chunks, ptf, selector): chunks = list(chunks) for chunk in chunks: # These should be organized by grid filename f = None for g in chunk.objs: if g.filename is None: continue if f is None: #print "Opening (read) %s" % g.filename f = h5py.File(g.filename, "r") nap = sum(g.NumberOfActiveParticles.values()) if g.NumberOfParticles == 0 and nap == 0: continue ds = f.get("/Grid%08i" % g.id) for ptype, field_list in sorted(ptf.items()): if ptype != "io": if g.NumberOfActiveParticles[ptype] == 0: continue pds = ds.get("Particles/%s" % ptype) else: pds = ds pn = _particle_position_names.get(ptype, r"particle_position_%s") x, y, z = (np.asarray(pds.get(pn % ax).value, dtype="=f8") for ax in 'xyz') if selector is None: # This only ever happens if the call is made from # _read_particle_coords. yield ptype, (x, y, z) continue mask = selector.select_points(x, y, z, 0.0) if mask is None: continue for field in field_list: data = np.asarray(pds.get(field).value, "=f8") if field in _convert_mass: data *= g.dds.prod(dtype="f8") yield (ptype, field), data[mask] if f: f.close()
[docs] def io_iter(self, chunks, fields): h5_dtype = self._field_dtype for chunk in chunks: fid = None filename = -1 for obj in chunk.objs: if obj.filename is None: continue if obj.filename != filename: # Note one really important thing here: even if we do # implement LRU caching in the _read_obj_field function, # we'll still be doing file opening and whatnot. This is a # problem, but one we can return to. if fid is not None: fid.close() fid = h5py.h5f.open(b(obj.filename), h5py.h5f.ACC_RDONLY) filename = obj.filename for field in fields: nodal_flag = self.ds.field_info[field].nodal_flag dims = obj.ActiveDimensions[::-1] + nodal_flag[::-1] data = np.empty(dims, dtype=h5_dtype) yield field, obj, self._read_obj_field( obj, field, (fid, data)) if fid is not None: fid.close()
def _read_obj_field(self, obj, field, fid_data): if fid_data is None: fid_data = (None, None) fid, data = fid_data if fid is None: close = True fid = h5py.h5f.open(b(obj.filename), h5py.h5f.ACC_RDONLY) else: close = False if data is None: data = np.empty(obj.ActiveDimensions[::-1], dtype=self._field_dtype) ftype, fname = field try: node = "/Grid%08i/%s" % (obj.id, fname) dg = h5py.h5d.open(fid, b(node)) except KeyError: if fname == "Dark_Matter_Density": data[:] = 0 return data.T raise dg.read(h5py.h5s.ALL, h5py.h5s.ALL, data) # I don't know why, but on some installations of h5py this works, but # on others, nope. Doesn't seem to be a version thing. #dg.close() if close: fid.close() return data.T
[docs]class IOHandlerPackedHDF5GhostZones(IOHandlerPackedHDF5): _dataset_type = "enzo_packed_3d_gz" def __init__(self, *args, **kwargs): super(IOHandlerPackedHDF5GhostZones, self).__init__(*args, **kwargs) NGZ = self.ds.parameters.get("NumberOfGhostZones", 3) self._base = (slice(NGZ, -NGZ), slice(NGZ, -NGZ), slice(NGZ, -NGZ)) def _read_obj_field(self, *args, **kwargs): return super(IOHandlerPackedHDF5GhostZones, self)._read_obj_field( *args, **kwargs)[self._base]
[docs]class IOHandlerInMemory(BaseIOHandler): _dataset_type = "enzo_inline" def __init__(self, ds, ghost_zones=3): self.ds = ds import enzo self.enzo = enzo self.grids_in_memory = enzo.grid_data self.old_grids_in_memory = enzo.old_grid_data self.my_slice = (slice(ghost_zones,-ghost_zones), slice(ghost_zones,-ghost_zones), slice(ghost_zones,-ghost_zones)) BaseIOHandler.__init__(self, ds) def _read_field_names(self, grid): fields = [] add_io = "io" in grid.ds.particle_types for name, v in self.grids_in_memory[grid.id].items(): # NOTE: This won't work with 1D datasets or references. if not hasattr(v, "shape") or v.dtype == "O": continue elif v.ndim == 1: if grid.ds.dimensionality == 1: fields.append( ("enzo", str(name)) ) elif add_io: fields.append( ("io", str(name)) ) else: fields.append( ("enzo", str(name)) ) return fields def _read_fluid_selection(self, chunks, selector, fields, size): rv = {} # Now we have to do something unpleasant chunks = list(chunks) if isinstance(selector, GridSelector): if not (len(chunks) == len(chunks[0].objs) == 1): raise RuntimeError g = chunks[0].objs[0] for ftype, fname in fields: rv[(ftype, fname)] = \ self.grids_in_memory[g.id][fname].swapaxes(0, 2) return rv if size is None: size = sum((g.count(selector) for chunk in chunks for g in chunk.objs)) for field in fields: ftype, fname = field fsize = size rv[field] = np.empty(fsize, dtype="float64") ng = sum(len(c.objs) for c in chunks) mylog.debug("Reading %s cells of %s fields in %s grids", size, [f2 for f1, f2 in fields], ng) ind = 0 for chunk in chunks: for g in chunk.objs: # We want a *hard error* here. #if g.id not in self.grids_in_memory: continue for field in fields: ftype, fname = field data_view = self.grids_in_memory[g.id][fname][self.my_slice].swapaxes(0,2) nd = g.select(selector, data_view, rv[field], ind) ind += nd assert(ind == fsize) return rv def _read_particle_coords(self, chunks, ptf): chunks = list(chunks) for chunk in chunks: # These should be organized by grid filename for g in chunk.objs: if g.id not in self.grids_in_memory: continue nap = sum(g.NumberOfActiveParticles.values()) if g.NumberOfParticles == 0 and nap == 0: continue for ptype, field_list in sorted(ptf.items()): x, y, z = self.grids_in_memory[g.id]['particle_position_x'], \ self.grids_in_memory[g.id]['particle_position_y'], \ self.grids_in_memory[g.id]['particle_position_z'] yield ptype, (x, y, z) def _read_particle_fields(self, chunks, ptf, selector): chunks = list(chunks) for chunk in chunks: # These should be organized by grid filename for g in chunk.objs: if g.id not in self.grids_in_memory: continue nap = sum(g.NumberOfActiveParticles.values()) if g.NumberOfParticles == 0 and nap == 0: continue for ptype, field_list in sorted(ptf.items()): x, y, z = self.grids_in_memory[g.id]['particle_position_x'], \ self.grids_in_memory[g.id]['particle_position_y'], \ self.grids_in_memory[g.id]['particle_position_z'] mask = selector.select_points(x, y, z, 0.0) if mask is None: continue for field in field_list: data = self.grids_in_memory[g.id][field] if field in _convert_mass: data = data * g.dds.prod(dtype="f8") yield (ptype, field), data[mask]
[docs]class IOHandlerPacked2D(IOHandlerPackedHDF5): _dataset_type = "enzo_packed_2d" _particle_reader = False def _read_data_set(self, grid, field): f = h5py.File(grid.filename, "r") ds = f["/Grid%08i/%s" % (grid.id, field)][:] f.close() return ds.transpose()[:,:,None]
[docs] def modify(self, field): pass
def _read_fluid_selection(self, chunks, selector, fields, size): rv = {} # Now we have to do something unpleasant chunks = list(chunks) if isinstance(selector, GridSelector): if not (len(chunks) == len(chunks[0].objs) == 1): raise RuntimeError g = chunks[0].objs[0] f = h5py.File(g.filename, 'r') gds = f.get("/Grid%08i" % g.id) for ftype, fname in fields: rv[(ftype, fname)] = np.atleast_3d( gds.get(fname).value.transpose()) f.close() return rv if size is None: size = sum((g.count(selector) for chunk in chunks for g in chunk.objs)) for field in fields: ftype, fname = field fsize = size rv[field] = np.empty(fsize, dtype="float64") ng = sum(len(c.objs) for c in chunks) mylog.debug("Reading %s cells of %s fields in %s grids", size, [f2 for f1, f2 in fields], ng) ind = 0 for chunk in chunks: f = None for g in chunk.objs: if f is None: #print "Opening (count) %s" % g.filename f = h5py.File(g.filename, "r") gds = f.get("/Grid%08i" % g.id) if gds is None: gds = f for field in fields: ftype, fname = field ds = np.atleast_3d(gds.get(fname).value.transpose()) nd = g.select(selector, ds, rv[field], ind) # caches ind += nd f.close() return rv
[docs]class IOHandlerPacked1D(IOHandlerPackedHDF5): _dataset_type = "enzo_packed_1d" _particle_reader = False def _read_data_set(self, grid, field): f = h5py.File(grid.filename, "r") ds = f["/Grid%08i/%s" % (grid.id, field)][:] f.close() return ds.transpose()[:,None,None]
[docs] def modify(self, field): pass