# Authors: Berk Geveci, Axel Huebl, Utkarsh Ayachit # from vtkmodules.util.vtkAlgorithm import VTKPythonAlgorithmBase from .. import print_error try: import numpy as np import openpmd_api as io _has_openpmd = True except ImportError as ie: print_error( "Missing required Python modules/packages. Algorithms in this module may " "not work as expected! \n {0}".format(ie)) _has_openpmd = False def createModifiedCallback(anobject): import weakref weakref_obj = weakref.ref(anobject) anobject = None def _markmodified(*args, **kwars): o = weakref_obj() if o is not None: o.Modified() return _markmodified class openPMDReader(VTKPythonAlgorithmBase): """A reader that reads openPMD format.""" def __init__(self): VTKPythonAlgorithmBase.__init__(self, nInputPorts=0, nOutputPorts=2, outputType='vtkPartitionedDataSet') self._filename = None self._timevalues = None self._series = None self._timemap = {} from vtkmodules.vtkCommonCore import vtkDataArraySelection self._arrayselection = vtkDataArraySelection() self._arrayselection.AddObserver("ModifiedEvent", createModifiedCallback(self)) self._speciesselection = vtkDataArraySelection() self._speciesselection.AddObserver("ModifiedEvent", createModifiedCallback(self)) self._particlearrayselection = vtkDataArraySelection() self._particlearrayselection.AddObserver("ModifiedEvent", createModifiedCallback(self)) def _get_update_time(self, outInfo): from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline executive = vtkStreamingDemandDrivenPipeline timevalues = self._timevalues if timevalues is None or len(timevalues) == 0: return None elif outInfo.Has(executive.UPDATE_TIME_STEP()) and len(timevalues) > 0: utime = outInfo.Get(executive.UPDATE_TIME_STEP()) dtime = timevalues[0] for atime in timevalues: if atime > utime: return dtime else: dtime = atime return dtime else: assert(len(timevalues) > 0) return timevalues[0] def _get_array_selection(self): return self._arrayselection def _get_particle_array_selection(self): return self._particlearrayselection def _get_species_selection(self): return self._speciesselection def SetFileName(self, name): """Specify filename for the file to read.""" if self._filename != name: self._filename = name self._timevalues = None if self._series: self._series = None self.Modified() def GetTimestepValues(self): return self._timevalues() def GetDataArraySelection(self): return self._get_array_selection() def GetSpeciesSelection(self): return self._get_species_selection() def GetParticleArraySelection(self): return self._get_particle_array_selection() def FillOutputPortInformation(self, port, info): from vtkmodules.vtkCommonDataModel import vtkDataObject if port == 0: info.Set(vtkDataObject.DATA_TYPE_NAME(), "vtkPartitionedDataSet") else: info.Set(vtkDataObject.DATA_TYPE_NAME(), "vtkPartitionedDataSetCollection") return 1 def RequestInformation(self, request, inInfoVec, outInfoVec): global _has_openpmd if not _has_openpmd: print_error("Required Python module 'openpmd_api' missing!") return 0 from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline, vtkAlgorithm executive = vtkStreamingDemandDrivenPipeline for i in (0,1): outInfo = outInfoVec.GetInformationObject(i) outInfo.Remove(executive.TIME_STEPS()) outInfo.Remove(executive.TIME_RANGE()) outInfo.Set(vtkAlgorithm.CAN_HANDLE_PIECE_REQUEST(), 1) # Why is this a string when it is None? if self._filename == 'None': return 1 mfile = open(self._filename, "r") pattern = mfile.readlines()[0][0:-1] del mfile import os if not self._series: self._series = io.Series(os.path.dirname(self._filename)+'/'+pattern, io.Access_Type.read_only) # This is how we get time values and arrays self._timemap = {} timevalues = [] arrays = set() particles = set() species = set() for idx, iteration in self._series.iterations.items(): # extract the time if callable(iteration.time): # prior to openPMD-api 0.13.0 time = iteration.time() * iteration.time_unit_SI() else: time = iteration.time * iteration.time_unit_SI timevalues.append(time) self._timemap[time] = idx arrays.update([ mesh_name for mesh_name, mesh in iteration.meshes.items()]) particles.update([ species_name + "_" + record_name for species_name, species in iteration.particles.items() for record_name, record in species.items() ]) species.update([ species_name for species_name, _ in iteration.particles.items() ]) for array in arrays: self._arrayselection.AddArray(array) for particle_array in particles: self._particlearrayselection.AddArray(particle_array) for species_name in species: self._speciesselection.AddArray(species_name) timesteps = list(self._series.iterations) self._timevalues = timevalues if len(timevalues) > 0: for i in (0,1): outInfo = outInfoVec.GetInformationObject(i) for t in timevalues: outInfo.Append(executive.TIME_STEPS(), t) outInfo.Append(executive.TIME_RANGE(), timevalues[0]) outInfo.Append(executive.TIME_RANGE(), timevalues[-1]) return 1 def _get_array_and_component(self, itr, name): for mesh_name, mesh in itr.meshes.items(): if mesh_name == name: return (mesh_name, None) for comp_name, _ in mesh.items(): if name == mesh_name + "_" + comp_name: return (mesh_name, comp_name) return (None, None) def _get_particle_array_and_component(self, itr, name): for species_name, species in itr.particles.items(): for record_name, record in species.items(): if name == species_name + "_" + record_name: return (species_name, record_name) return (None, None) def _load_array(self, var, chunk_offset, chunk_extent): arrays = [] for name, scalar in var.items(): comp = scalar.load_chunk(chunk_offset, chunk_extent) self._series.flush() comp = comp * scalar.unit_SI arrays.append(comp) ncomp = len(var) if ncomp > 1: flt = np.ravel(arrays, order='F') return flt.reshape((flt.shape[0]//ncomp, ncomp)) else: return arrays[0].flatten(order='F') def _find_array(self, itr, name): var = itr.meshes[name] theta_modes = None if var.geometry == io.Geometry.thetaMode: theta_modes = 3 # hardcoded, parse from geometry_parameters return (var, np.array(var.grid_spacing) * var.grid_unit_SI, np.array(var.grid_global_offset) * var.grid_unit_SI, theta_modes) def _get_num_particles(self, itr, species): sp = itr.particles[species] var = sp["position"] return var['x'].shape[0] def _load_particle_array(self, itr, species, name, start, end): sp = itr.particles[species] var = sp[name] arrays = [] for name, scalar in var.items(): comp = scalar.load_chunk([start], [end-start+1]) self._series.flush() comp = comp * scalar.unit_SI arrays.append(comp) ncomp = len(var) if ncomp > 1: flt = np.ravel(arrays, order='F') return flt.reshape((flt.shape[0]//ncomp, ncomp)) else: return arrays[0] def _load_particles(self, itr, species, start, end): sp = itr.particles[species] var = sp["position"] ovar = sp["positionOffset"] position_arrays = [] for name, scalar in var.items(): pos = scalar.load_chunk([start], [end-start+1]) self._series.flush() pos = pos * scalar.unit_SI off = ovar[name].load_chunk([start], [end-start+1]) self._series.flush() off = off * ovar[name].unit_SI position_arrays.append(pos + off) flt = np.ravel(position_arrays, order='F') num_components = len(var) # 1D, 2D and 3D positions flt = flt.reshape((flt.shape[0]//num_components, num_components)) # 1D and 2D particles: pad additional components with zero while flt.shape[1] < 3: flt = np.column_stack([flt, np.zeros_like(flt[:, 0])]) return flt def _load_species(self, itr, species, arrays, piece, npieces, ugrid): nparticles = self._get_num_particles(itr, species) nlocalparticles = nparticles // npieces start = nlocalparticles * piece end = start + nlocalparticles - 1 if piece == npieces - 1: end = nparticles - 1 pts = self._load_particles(itr, species, start, end) npts = pts.shape[0] ugrid.Points = pts for array in arrays: if array[1] == 'position' or array[1] == 'positionOffset': continue ugrid.PointData.append( self._load_particle_array(itr, array[0], array[1], start, end), array[1]) from vtkmodules.vtkCommonDataModel import vtkCellArray ca = vtkCellArray() if npts < np.iinfo(np.int32).max: dtype = np.int32 else: dtype = np.int64 offsets = np.linspace(0, npts, npts+1, dtype=dtype) cells = np.linspace(0, npts-1, npts, dtype=dtype) from vtkmodules.numpy_interface import dataset_adapter offsets = dataset_adapter.numpyTovtkDataArray(offsets) offsets2 = offsets.NewInstance() offsets2.DeepCopy(offsets) cells = dataset_adapter.numpyTovtkDataArray(cells) cells2 = cells.NewInstance() cells2.DeepCopy(cells) ca.SetData(offsets2, cells2) from vtkmodules.util import vtkConstants ugrid.VTKObject.SetCells(vtkConstants.VTK_VERTEX, ca) def _RequestFieldData(self, executive, output, outInfo): from vtkmodules.numpy_interface import dataset_adapter as dsa from vtkmodules.vtkCommonDataModel import vtkImageData from vtkmodules.vtkCommonExecutionModel import vtkExtentTranslator piece = outInfo.Get(executive.UPDATE_PIECE_NUMBER()) npieces = outInfo.Get(executive.UPDATE_NUMBER_OF_PIECES()) nghosts = outInfo.Get(executive.UPDATE_NUMBER_OF_GHOST_LEVELS()) et = vtkExtentTranslator() data_time = self._get_update_time(outInfo) idx = self._timemap[data_time] itr = self._series.iterations[idx] arrays = [] narrays = self._arrayselection.GetNumberOfArrays() for i in range(narrays): if self._arrayselection.GetArraySetting(i): name = self._arrayselection.GetArrayName(i) arrays.append((name, self._find_array(itr, name))) shp = None spacing = None theta_modes = None grid_offset = None for _, ary in arrays: var = ary[0] for name, scalar in var.items(): shape = scalar.shape break spc = list(ary[1]) if not spacing: spacing = spc elif spacing != spc: # all meshes need to have the same spacing return 0 offset = list(ary[2]) if not grid_offset: grid_offset = offset elif grid_offset != offset: # all meshes need to have the same spacing return 0 if not shp: shp = shape elif shape != shp: # all arrays needs to have the same shape return 0 if not theta_modes: theta_modes = ary[3] # fields/meshes: RZ if theta_modes and shp is not None: et.SetWholeExtent(0, shp[0]-1, 0, shp[1]-1, 0, shp[2]-1) et.SetSplitModeToZSlab() # note: Y and Z are both fine et.SetPiece(piece) et.SetNumberOfPieces(npieces) # et.SetGhostLevel(nghosts) et.PieceToExtentByPoints() ext = et.GetExtent() chunk_offset = [ext[0], ext[2], ext[4]] chunk_extent = [ext[1] - ext[0] + 1, ext[3] - ext[2] + 1, ext[5] - ext[4] + 1] data = [] nthetas = 100 # user parameter thetas = np.linspace(0., 2.*np.pi, nthetas) chunk_cyl_shape = (chunk_extent[1], chunk_extent[2], nthetas) # z, r, theta for name, var in arrays: cyl_values = np.zeros(chunk_cyl_shape) values = self._load_array(var[0], chunk_offset, chunk_extent) self._series.flush() print(chunk_cyl_shape) print(values.shape) print("+++++++++++") for ntheta in range(nthetas): cyl_values[:, :, ntheta] += values[0, :, :] data.append((name, cyl_values)) # add all other modes via loop # for m in range(theta_modes): cyl_spacing = [spacing[0], spacing[1], thetas[1]-thetas[0]] z_coord = np.linspace(0., cyl_spacing[0] * chunk_cyl_shape[0], chunk_cyl_shape[0]) r_coord = np.linspace(0., cyl_spacing[1] * chunk_cyl_shape[1], chunk_cyl_shape[1]) t_coord = thetas # to cartesian print(z_coord.shape, r_coord.shape, t_coord.shape) cyl_coords = np.meshgrid(r_coord, z_coord, t_coord) rs = cyl_coords[1] zs = cyl_coords[0] thetas = cyl_coords[2] y_coord = rs * np.sin(thetas) x_coord = rs * np.cos(thetas) z_coord = zs # mesh_pts = np.zeros((chunk_cyl_shape[0], chunk_cyl_shape[1], chunk_cyl_shape[2], 3)) # mesh_pts[:, :, :, 0] = z_coord img = vtkImageData() img.SetExtent( chunk_offset[1], chunk_offset[1] + chunk_cyl_shape[0] - 1, chunk_offset[2], chunk_offset[2] + chunk_cyl_shape[1] - 1, 0, nthetas-1) img.SetSpacing(cyl_spacing) imgw = dsa.WrapDataObject(img) output.SetPartition(0, img) for name, array in data: # print(array.shape) # print(array.transpose(2,1,0).flatten(order='C').shape) imgw.PointData.append(array.transpose(2,1,0).flatten(order='C'), name) # data = [] # for name, var in arrays: # unit_SI = var[0].unit_SI # data.append((name, unit_SI * var[0].load_chunk(chunk_offset, chunk_extent))) # self._series.flush() # fields/meshes: 1D-3D elif shp is not None: whole_extent = [] # interleave shape with zeros for s in shp: whole_extent.append(0) whole_extent.append(s-1) # 1D and 2D data: pad with 0, 0 for extra dimensions while len(whole_extent) < 6: whole_extent.append(0) whole_extent.append(0) et.SetWholeExtent(*whole_extent) et.SetPiece(piece) et.SetNumberOfPieces(npieces) et.SetGhostLevel(nghosts) et.PieceToExtent() ext = et.GetExtent() chunk_offset = [ext[0], ext[2], ext[4]] chunk_extent = [ext[1] - ext[0] + 1, ext[3] - ext[2] + 1, ext[5] - ext[4] + 1] # 1D and 2D data: remove extra dimensions for load del chunk_offset[len(shp):] del chunk_extent[len(shp):] data = [] for name, var in arrays: values = self._load_array(var[0], chunk_offset, chunk_extent) self._series.flush() data.append((name, values)) # 1D and 2D data: pad spacing with extra 1 and grid_offset with # extra 9 values until 3D i = iter(spacing) spacing = [next(i, 1) for _ in range(3)] i = iter(grid_offset) grid_offset = [next(i, 0) for _ in range(3)] img = vtkImageData() img.SetExtent(ext[0], ext[1], ext[2], ext[3], ext[4], ext[5]) img.SetSpacing(spacing) img.SetOrigin(grid_offset) et.SetGhostLevel(0) et.PieceToExtent() ext = et.GetExtent() ext = [ext[0], ext[1], ext[2], ext[3], ext[4], ext[5]] img.GenerateGhostArray(ext) imgw = dsa.WrapDataObject(img) output.SetPartition(0, img) for name, array in data: imgw.PointData.append(array, name) def _RequestParticleData(self, executive, poutput, outInfo): from vtkmodules.numpy_interface import dataset_adapter as dsa from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid, vtkPartitionedDataSet piece = outInfo.Get(executive.UPDATE_PIECE_NUMBER()) npieces = outInfo.Get(executive.UPDATE_NUMBER_OF_PIECES()) data_time = self._get_update_time(outInfo) idx = self._timemap[data_time] itr = self._series.iterations[idx] array_by_species = {} narrays = self._particlearrayselection.GetNumberOfArrays() for i in range(narrays): if self._particlearrayselection.GetArraySetting(i): name = self._particlearrayselection.GetArrayName(i) names = self._get_particle_array_and_component( itr, name) if names[0] and self._speciesselection.ArrayIsEnabled(names[0]): if not names[0] in array_by_species: array_by_species[names[0]] = [] array_by_species[names[0]].append(names) ids = 0 for species, arrays in array_by_species.items(): pds = vtkPartitionedDataSet() ugrid = vtkUnstructuredGrid() pds.SetPartition(0, ugrid) poutput.SetPartitionedDataSet(ids, pds) ids += 1 self._load_species( itr, species, arrays, piece, npieces, dsa.WrapDataObject(ugrid)) def RequestData(self, request, inInfoVec, outInfoVec): global _has_openpmd if not _has_openpmd: print_error("Required Python module 'openpmd_api' missing!") return 0 from vtkmodules.vtkCommonDataModel import vtkPartitionedDataSet, vtkPartitionedDataSetCollection from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline executive = vtkStreamingDemandDrivenPipeline numInfo = outInfoVec.GetNumberOfInformationObjects() for i in range(numInfo): outInfo = outInfoVec.GetInformationObject(i) if i == 0: output = vtkPartitionedDataSet.GetData(outInfoVec, 0) self._RequestFieldData(executive, output, outInfo) elif i == 1: poutput = vtkPartitionedDataSetCollection.GetData(outInfoVec, 1) self._RequestParticleData(executive, poutput, outInfo) else: print_error("numInfo number is wrong! " "It should be exactly 2, is=", numInfo) return 0 return 1