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
