def _fill_cache(self, ptype, index=0, offset=None): """Fills the particle position cache for the ``ptype``. Parameters ---------- ptype : str The on-disk name of the particle species index : int, optional offset : int, optional """ if str((ptype, index, offset)) not in self._cached_ptype: self._cached_ptype = str((ptype, index, offset)) pds = self._handle[self.base_path + self.particles_path + "/" + ptype] axes = list(pds["position"].keys()) if offset is None: if is_const_component(pds["position/" + axes[0]]): offset = pds["position/" + axes[0]].attrs["shape"] else: offset = pds["position/" + axes[0]].len() self.cache = np.empty((3, offset), dtype=np.float64) for i in np.arange(3): ax = "xyz"[i] if ax in axes: np.add( get_component(pds, "position/" + ax, index, offset), get_component(pds, "positionOffset/" + ax, index, offset), self.cache[i], ) else: # Pad accordingly with zeros to make 1D/2D datasets compatible # These have to be the same shape as the existing axes since that # equals the number of particles self.cache[i] = np.zeros(offset)
def _parse_index(self): """Fills each grid with appropriate properties (extent, dimensions, ...) This calculates the properties of every OpenPMDGrid based on the total number of grids in the simulation. The domain is divided into ``self.num_grids`` (roughly) equally sized chunks along the x-axis. ``grid_levels`` is always equal to 0 since we only have one level of refinement in openPMD. Notes ----- ``self.grid_dimensions`` is rounded to the nearest integer. Grid edges are calculated from this dimension. Grids with dimensions [0, 0, 0] are particle only. The others do not have any particles affiliated with them. """ f = self.dataset._handle bp = self.dataset.base_path pp = self.dataset.particles_path self.grid_levels.flat[:] = 0 self.grids = np.empty(self.num_grids, dtype="object") grid_index_total = 0 # Mesh grids for mesh in set(self.meshshapes.values()): (shape, spacing, offset, unit_si) = mesh shape = np.asarray(shape) spacing = np.asarray(spacing) offset = np.asarray(offset) # Total dimension of this grid domain_dimension = np.asarray(shape, dtype=np.int32) domain_dimension = np.append( domain_dimension, np.ones(3 - len(domain_dimension)) ) # Number of grids of this shape num_grids = min(shape[0], int(np.ceil(reduce(mul, shape) * self.vpg ** -1))) gle = offset * unit_si # self.dataset.domain_left_edge gre = ( domain_dimension[: spacing.size] * unit_si * spacing + gle ) # self.dataset.domain_right_edge gle = np.append(gle, np.zeros(3 - len(gle))) gre = np.append(gre, np.ones(3 - len(gre))) grid_dim_offset = np.linspace( 0, domain_dimension[0], num_grids + 1, dtype=np.int32 ) grid_edge_offset = ( grid_dim_offset * np.float(domain_dimension[0]) ** -1 * (gre[0] - gle[0]) + gle[0] ) mesh_names = [] for (mname, mdata) in self.meshshapes.items(): if mesh == mdata: mesh_names.append(str(mname)) prev = 0 for grid in np.arange(num_grids): self.grid_dimensions[grid_index_total] = domain_dimension self.grid_dimensions[grid_index_total][0] = ( grid_dim_offset[grid + 1] - grid_dim_offset[grid] ) self.grid_left_edge[grid_index_total] = gle self.grid_left_edge[grid_index_total][0] = grid_edge_offset[grid] self.grid_right_edge[grid_index_total] = gre self.grid_right_edge[grid_index_total][0] = grid_edge_offset[grid + 1] self.grid_particle_count[grid_index_total] = 0 self.grids[grid_index_total] = self.grid( grid_index_total, self, 0, fi=prev, fo=self.grid_dimensions[grid_index_total][0], ft=mesh_names, ) prev += self.grid_dimensions[grid_index_total][0] grid_index_total += 1 handled_ptypes = [] # Particle grids for (species, count) in self.numparts.items(): if "#" in species: # This is a particlePatch spec = species.split("#") patch = f[bp + pp + "/" + spec[0] + "/particlePatches"] domain_dimension = np.ones(3, dtype=np.int32) for (ind, axis) in enumerate(list(patch["extent"].keys())): domain_dimension[ind] = patch["extent/" + axis][()][int(spec[1])] num_grids = int(np.ceil(count * self.vpg ** -1)) gle = [] for axis in patch["offset"].keys(): gle.append( get_component(patch, "offset/" + axis, int(spec[1]), 1)[0] ) gle = np.asarray(gle) gle = np.append(gle, np.zeros(3 - len(gle))) gre = [] for axis in patch["extent"].keys(): gre.append( get_component(patch, "extent/" + axis, int(spec[1]), 1)[0] ) gre = np.asarray(gre) gre = np.append(gre, np.ones(3 - len(gre))) np.add(gle, gre, gre) npo = patch["numParticlesOffset"][()].item(int(spec[1])) particle_count = np.linspace( npo, npo + count, num_grids + 1, dtype=np.int32 ) particle_names = [str(spec[0])] elif str(species) not in handled_ptypes: domain_dimension = self.dataset.domain_dimensions num_grids = int(np.ceil(count * self.vpg ** -1)) gle = self.dataset.domain_left_edge gre = self.dataset.domain_right_edge particle_count = np.linspace(0, count, num_grids + 1, dtype=np.int32) particle_names = [] for (pname, size) in self.numparts.items(): if size == count: # Since this is not part of a particlePatch, we can include multiple same-sized ptypes particle_names.append(str(pname)) handled_ptypes.append(str(pname)) else: # A grid with this exact particle count has already been created continue for grid in np.arange(num_grids): self.grid_dimensions[grid_index_total] = domain_dimension self.grid_left_edge[grid_index_total] = gle self.grid_right_edge[grid_index_total] = gre self.grid_particle_count[grid_index_total] = ( particle_count[grid + 1] - particle_count[grid] ) * len(particle_names) self.grids[grid_index_total] = self.grid( grid_index_total, self, 0, pi=particle_count[grid], po=particle_count[grid + 1] - particle_count[grid], pt=particle_names, ) grid_index_total += 1
def _read_particle_selection(self, chunks, selector, fields): """Read particle fields for particle species masked by a selection. Parameters ---------- chunks A list of chunks A chunk is a list of grids selector A region (inside your domain) specifying which parts of the field you want to read. See [1] and [2] fields : array_like Tuples (ptype, pfield) representing a field Returns ------- dict keys are tuples (ptype, pfield) representing a field values are (N,) ndarrays with data from that field """ f = self._handle bp = self.base_path pp = self.particles_path ds = f[bp + pp] unions = self.ds.particle_unions chunks = list(chunks) # chunks is a generator rv = {} ind = {} particle_count = {} ptf = defaultdict(list) # ParticleTypes&Fields rfm = defaultdict(list) # RequestFieldMapping for (ptype, pname) in fields: pfield = (ptype, pname) # Overestimate the size of all pfields so they include all particles # and shrink it later particle_count[pfield] = 0 if ptype in unions: for pt in unions[ptype]: particle_count[pfield] += self.ds.particle_type_counts[pt] ptf[pt].append(pname) rfm[pt, pname].append(pfield) else: particle_count[pfield] = self.ds.particle_type_counts[ptype] ptf[ptype].append(pname) rfm[pfield].append(pfield) rv[pfield] = np.empty((particle_count[pfield],), dtype=np.float64) ind[pfield] = 0 for ptype in ptf: for chunk in chunks: for grid in chunk.objs: if str(ptype) == "io": species = list(ds.keys())[0] else: species = ptype if species not in grid.ptypes: continue # read particle coords into cache self._fill_cache(species, grid.pindex, grid.poffset) mask = selector.select_points( self.cache[0], self.cache[1], self.cache[2], 0.0 ) if mask is None: continue pds = ds[species] for field in ptf[ptype]: component = "/".join(field.split("_")[1:]) component = component.replace("positionCoarse", "position") component = component.replace("-", "_") data = get_component(pds, component, grid.pindex, grid.poffset)[ mask ] for request_field in rfm[(ptype, field)]: rv[request_field][ ind[request_field] : ind[request_field] + data.shape[0] ] = data ind[request_field] += data.shape[0] for field in fields: rv[field] = rv[field][: ind[field]] return rv
def _read_fluid_selection(self, chunks, selector, fields, size): """Reads given fields masked by a given selection. Parameters ---------- chunks A list of chunks A chunk is a list of grids selector A region (inside your domain) specifying which parts of the field you want to read. See [1] and [2] fields : array_like Tuples (fname, ftype) representing a field size : int Size of the data to read Returns ------- dict keys are tuples (ftype, fname) representing a field values are flat (``size``,) ndarrays with data from that field """ f = self._handle bp = self.base_path mp = self.meshes_path ds = f[bp + mp] chunks = list(chunks) rv = {} ind = {} if isinstance(selector, GridSelector): if not (len(chunks) == len(chunks[0].objs) == 1): raise RuntimeError if size is None: size = sum((g.count(selector) for chunk in chunks for g in chunk.objs)) for field in fields: rv[field] = np.empty(size, dtype=np.float64) ind[field] = 0 for (ftype, fname) in fields: field = (ftype, fname) for chunk in chunks: for grid in chunk.objs: mask = grid._get_selector_mask(selector) if mask is None: continue component = fname.replace("_", "/").replace("-", "_") if component.split("/")[0] not in grid.ftypes: data = np.full(grid.ActiveDimensions, 0, dtype=np.float64) else: data = get_component(ds, component, grid.findex, grid.foffset) # The following is a modified AMRGridPatch.select(...) data.shape = ( mask.shape ) # Workaround - casts a 2D (x,y) array to 3D (x,y,1) count = grid.count(selector) rv[field][ind[field] : ind[field] + count] = data[mask] ind[field] += count for field in fields: rv[field] = rv[field][: ind[field]] rv[field].flatten() return rv