Source code for yt.frontends.flash.io

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

import numpy as np
from itertools import groupby

from yt.utilities.io_handler import \
    BaseIOHandler
from yt.geometry.selection_routines import AlwaysSelector
from yt.utilities.lib.geometry_utils import \
    compute_morton

# http://stackoverflow.com/questions/2361945/detecting-consecutive-integers-in-a-list
[docs]def particle_sequences(grids): g_iter = sorted(grids, key = lambda g: g.id) for k, g in groupby(enumerate(g_iter), lambda i_x:i_x[0]-i_x[1].id): seq = list(v[1] for v in g) yield seq[0], seq[-1]
[docs]def grid_sequences(grids): g_iter = sorted(grids, key = lambda g: g.id) for k, g in groupby(enumerate(g_iter), lambda i_x1:i_x1[0]-i_x1[1].id): seq = list(v[1] for v in g) yield seq
[docs]def determine_particle_fields(handle): try: particle_fields = [s[0].decode("ascii","ignore").strip() for s in handle["/particle names"][:]] _particle_fields = dict([("particle_" + s, i) for i, s in enumerate(particle_fields)]) except KeyError: _particle_fields = {} return _particle_fields
[docs]class IOHandlerFLASH(BaseIOHandler): _particle_reader = False _dataset_type = "flash_hdf5" def __init__(self, ds): super(IOHandlerFLASH, self).__init__(ds) # Now we cache the particle fields self._handle = ds._handle self._particle_handle = ds._particle_handle self._particle_fields = determine_particle_fields(self._particle_handle) def _read_particles(self, fields_to_read, type, args, grid_list, count_list, conv_factors): pass
[docs] def io_iter(self, chunks, fields): f = self._handle for chunk in chunks: for field in fields: # Note that we *prefer* to iterate over the fields on the # outside; here, though, we're iterating over them on the # inside because we may exhaust our chunks. ftype, fname = field ds = f["/%s" % fname] for gs in grid_sequences(chunk.objs): start = gs[0].id - gs[0]._id_offset end = gs[-1].id - gs[-1]._id_offset + 1 data = ds[start:end,:,:,:] for i, g in enumerate(gs): yield field, g, self._read_obj_field(g, field, (data, i))
def _read_particle_coords(self, chunks, ptf): chunks = list(chunks) f_part = self._particle_handle p_ind = self.ds.index._particle_indices px, py, pz = (self._particle_fields["particle_pos%s" % ax] for ax in 'xyz') p_fields = f_part["/tracer particles"] assert(len(ptf) == 1) ptype = list(ptf.keys())[0] for chunk in chunks: start = end = None for g1, g2 in particle_sequences(chunk.objs): start = p_ind[g1.id - g1._id_offset] end = p_ind[g2.id - g2._id_offset + 1] x = np.asarray(p_fields[start:end, px], dtype="=f8") y = np.asarray(p_fields[start:end, py], dtype="=f8") z = np.asarray(p_fields[start:end, pz], dtype="=f8") yield ptype, (x, y, z) def _read_particle_fields(self, chunks, ptf, selector): chunks = list(chunks) f_part = self._particle_handle p_ind = self.ds.index._particle_indices px, py, pz = (self._particle_fields["particle_pos%s" % ax] for ax in 'xyz') p_fields = f_part["/tracer particles"] assert(len(ptf) == 1) ptype = list(ptf.keys())[0] field_list = ptf[ptype] for chunk in chunks: for g1, g2 in particle_sequences(chunk.objs): start = p_ind[g1.id - g1._id_offset] end = p_ind[g2.id - g2._id_offset + 1] x = np.asarray(p_fields[start:end, px], dtype="=f8") y = np.asarray(p_fields[start:end, py], dtype="=f8") z = np.asarray(p_fields[start:end, pz], dtype="=f8") mask = selector.select_points(x, y, z, 0.0) if mask is None: continue for field in field_list: fi = self._particle_fields[field] data = p_fields[start:end, fi] yield (ptype, field), data[mask] def _read_obj_field(self, obj, field, ds_offset = None): if ds_offset is None: ds_offset = (None, -1) ds, offset = ds_offset # our context here includes datasets and whatnot that are opened in the # hdf5 file if ds is None: ds = self._handle["/%s" % field[1]] if offset == -1: data = ds[obj.id - obj._id_offset, :,:,:].transpose() else: data = ds[offset, :,:,:].transpose() return data def _read_chunk_data(self, chunk, fields): f = self._handle rv = {} for g in chunk.objs: rv[g.id] = {} # Split into particles and non-particles fluid_fields, particle_fields = [], [] for ftype, fname in fields: if ftype in self.ds.particle_types: particle_fields.append((ftype, fname)) else: fluid_fields.append((ftype, fname)) if len(particle_fields) > 0: selector = AlwaysSelector(self.ds) rv.update(self._read_particle_selection( [chunk], selector, particle_fields)) if len(fluid_fields) == 0: return rv for field in fluid_fields: ftype, fname = field ds = f["/%s" % fname] for gs in grid_sequences(chunk.objs): start = gs[0].id - gs[0]._id_offset end = gs[-1].id - gs[-1]._id_offset + 1 data = ds[start:end,:,:,:].transpose() for i, g in enumerate(gs): rv[g.id][field] = np.asarray(data[...,i], "=f8") return rv
[docs]class IOHandlerFLASHParticle(BaseIOHandler): _particle_reader = True _dataset_type = "flash_particle_hdf5" def __init__(self, ds): super(IOHandlerFLASHParticle, self).__init__(ds) # Now we cache the particle fields self._handle = ds._handle self._particle_fields = determine_particle_fields(self._handle) self._position_fields = [self._particle_fields["particle_pos%s" % ax] for ax in 'xyz'] self._chunksize = 32**3 def _read_fluid_selection(self, chunks, selector, fields, size): raise NotImplementedError def _read_particle_coords(self, chunks, ptf): chunks = list(chunks) data_files = set([]) assert(len(ptf) == 1) for chunk in chunks: for obj in chunk.objs: data_files.update(obj.data_files) px, py, pz = self._position_fields p_fields = self._handle["/tracer particles"] assert(len(data_files) == 1) for data_file in sorted(data_files): pcount = self._count_particles(data_file)["io"] for ptype, field_list in sorted(ptf.items()): total = 0 while total < pcount: count = min(self._chunksize, pcount - total) x = np.asarray(p_fields[total:total+count, px], dtype="=f8") y = np.asarray(p_fields[total:total+count, py], dtype="=f8") z = np.asarray(p_fields[total:total+count, pz], dtype="=f8") total += count yield ptype, (x, y, z) def _read_particle_fields(self, chunks, ptf, selector): chunks = list(chunks) data_files = set([]) assert(len(ptf) == 1) for chunk in chunks: for obj in chunk.objs: data_files.update(obj.data_files) px, py, pz = self._position_fields p_fields = self._handle["/tracer particles"] assert(len(data_files) == 1) for data_file in sorted(data_files): pcount = self._count_particles(data_file)["io"] for ptype, field_list in sorted(ptf.items()): total = 0 while total < pcount: count = min(self._chunksize, pcount - total) x = np.asarray(p_fields[total:total+count, px], dtype="=f8") y = np.asarray(p_fields[total:total+count, py], dtype="=f8") z = np.asarray(p_fields[total:total+count, pz], dtype="=f8") total += count mask = selector.select_points(x, y, z, 0.0) del x, y, z if mask is None: continue for field in field_list: fi = self._particle_fields[field] data = p_fields[total-count:total, fi] yield (ptype, field), data[mask] def _initialize_index(self, data_file, regions): p_fields = self._handle["/tracer particles"] px, py, pz = self._position_fields pcount = self._count_particles(data_file)["io"] morton = np.empty(pcount, dtype='uint64') ind = 0 while ind < pcount: npart = min(self._chunksize, pcount - ind) pos = np.empty((npart, 3), dtype="=f8") pos[:,0] = p_fields[ind:ind+npart, px] pos[:,1] = p_fields[ind:ind+npart, py] pos[:,2] = p_fields[ind:ind+npart, pz] regions.add_data_file(pos, data_file.file_id) morton[ind:ind+npart] = \ compute_morton(pos[:,0], pos[:,1], pos[:,2], data_file.ds.domain_left_edge, data_file.ds.domain_right_edge) ind += self._chunksize return morton def _count_particles(self, data_file): pcount = {"io": self._handle["/localnp"][:].sum()} return pcount def _identify_fields(self, data_file): fields = [("io", field) for field in self._particle_fields] return fields, {}