import numpy as np
from yt.funcs import get_pbar
from yt.units._numpy_wrapper_functions import uconcatenate
from yt.utilities.lib.particle_mesh_operations import CICSample_3
[docs]
class ParticleGenerator:
def __init__(self, ds, num_particles, field_list, ptype="io"):
"""
Base class for generating particle fields which may be applied to
streams. Normally this would not be called directly, since it doesn't
really do anything except allocate memory. Takes a *ds* to serve as the
basis for determining grids, the number of particles *num_particles*,
a list of fields, *field_list*, and the particle type *ptype*, which
has a default value of "io".
"""
self.ds = ds
self.num_particles = num_particles
self.field_list = [
(ptype, fd) if isinstance(fd, str) else fd for fd in field_list
]
self.field_list.append((ptype, "particle_index"))
self.field_units = {
(ptype, f"particle_position_{ax}"): "code_length" for ax in "xyz"
}
self.field_units[ptype, "particle_index"] = ""
self.ptype = ptype
self._set_default_fields()
try:
self.posx_index = self.field_list.index(self.default_fields[0])
self.posy_index = self.field_list.index(self.default_fields[1])
self.posz_index = self.field_list.index(self.default_fields[2])
except Exception as e:
raise KeyError(
"You must specify position fields: "
+ " ".join(f"particle_position_{ax}" for ax in "xyz")
) from e
self.index_index = self.field_list.index((ptype, "particle_index"))
self.num_grids = self.ds.index.num_grids
self.NumberOfParticles = np.zeros(self.num_grids, dtype="int64")
self.ParticleGridIndices = np.zeros(self.num_grids + 1, dtype="int64")
self.num_fields = len(self.field_list)
self.particles = np.zeros(
(self.num_particles, self.num_fields), dtype="float64"
)
def _set_default_fields(self):
self.default_fields = [
(self.ptype, "particle_position_x"),
(self.ptype, "particle_position_y"),
(self.ptype, "particle_position_z"),
]
[docs]
def has_key(self, key):
"""
Check to see if *key* is in the particle field list.
"""
return key in self.field_list
[docs]
def keys(self):
"""
Return the list of particle fields.
"""
return self.field_list
def __getitem__(self, key):
"""
Get the field associated with key.
"""
return self.particles[:, self.field_list.index(key)]
def __setitem__(self, key, val):
"""
Sets a field to be some other value. Note that we assume
that the particles have been sorted by grid already, so
make sure the setting of the field is consistent with this.
"""
self.particles[:, self.field_list.index(key)] = val[:]
def __len__(self):
"""
The number of particles
"""
return self.num_particles
[docs]
def get_for_grid(self, grid):
"""
Return a dict containing all of the particle fields in the specified grid.
"""
ind = grid.id - grid._id_offset
start = self.ParticleGridIndices[ind]
end = self.ParticleGridIndices[ind + 1]
tr = {}
for field in self.field_list:
fi = self.field_list.index(field)
if field in self.field_units:
tr[field] = self.ds.arr(
self.particles[start:end, fi], self.field_units[field]
)
else:
tr[field] = self.particles[start:end, fi]
return tr
def _setup_particles(self, x, y, z, setup_fields=None):
"""
Assigns grids to particles and sets up particle positions. *setup_fields* is
a dict of fields other than the particle positions to set up.
"""
particle_grids, particle_grid_inds = self.ds.index._find_points(x, y, z)
idxs = np.argsort(particle_grid_inds)
self.particles[:, self.posx_index] = x[idxs]
self.particles[:, self.posy_index] = y[idxs]
self.particles[:, self.posz_index] = z[idxs]
self.NumberOfParticles = np.bincount(
particle_grid_inds.astype("intp"), minlength=self.num_grids
)
if self.num_grids > 1:
np.add.accumulate(
self.NumberOfParticles.squeeze(), out=self.ParticleGridIndices[1:]
)
else:
self.ParticleGridIndices[1] = self.NumberOfParticles.squeeze()
if setup_fields is not None:
for key, value in setup_fields.items():
field = (self.ptype, key) if isinstance(key, str) else key
if field not in self.default_fields:
self.particles[:, self.field_list.index(field)] = value[idxs]
[docs]
def assign_indices(self, function=None, **kwargs):
"""
Assign unique indices to the particles. The default is to just use
numpy.arange, but any function may be supplied with keyword arguments.
"""
if function is None:
self.particles[:, self.index_index] = np.arange(self.num_particles)
else:
self.particles[:, self.index_index] = function(**kwargs)
[docs]
def map_grid_fields_to_particles(self, mapping_dict):
r"""
For the fields in *mapping_dict*, map grid fields to the particles
using CIC sampling.
Examples
--------
>>> field_map = {
... "density": "particle_density",
... "temperature": "particle_temperature",
... }
>>> particles.map_grid_fields_to_particles(field_map)
"""
pbar = get_pbar("Mapping fields to particles", self.num_grids)
for i, grid in enumerate(self.ds.index.grids):
pbar.update(i + 1)
if self.NumberOfParticles[i] > 0:
start = self.ParticleGridIndices[i]
end = self.ParticleGridIndices[i + 1]
# Note we add one ghost zone to the grid!
cube = grid.retrieve_ghost_zones(1, list(mapping_dict.keys()))
le = np.array(grid.LeftEdge).astype(np.float64)
dims = np.array(grid.ActiveDimensions).astype(np.int32)
for gfield, pfield in mapping_dict.items():
self.field_units[pfield] = cube[gfield].units
field_index = self.field_list.index(pfield)
CICSample_3(
self.particles[start:end, self.posx_index],
self.particles[start:end, self.posy_index],
self.particles[start:end, self.posz_index],
self.particles[start:end, field_index],
np.int64(self.NumberOfParticles[i]),
cube[gfield],
le,
dims,
grid.dds[0],
)
pbar.finish()
[docs]
def apply_to_stream(self, overwrite=False, **kwargs):
"""
Apply the particles to a grid-based stream dataset. If particles
already exist, and overwrite=False, do not overwrite them, but add
the new ones to them.
"""
grid_data = []
for i, g in enumerate(self.ds.index.grids):
data = {}
number_of_particles = self.NumberOfParticles[i]
if not overwrite:
number_of_particles += g.NumberOfParticles
grid_particles = self.get_for_grid(g)
for field in self.field_list:
if number_of_particles > 0:
if (
g.NumberOfParticles > 0
and not overwrite
and field in self.ds.field_list
):
# We have particles in this grid, we're not
# overwriting them, and the field is in the field
# list already
data[field] = uconcatenate([g[field], grid_particles[field]])
else:
# Otherwise, simply add the field in
data[field] = grid_particles[field]
else:
# We don't have particles in this grid
data[field] = np.array([], dtype="float64")
grid_data.append(data)
self.ds.index.update_data(grid_data)
[docs]
class FromListParticleGenerator(ParticleGenerator):
def __init__(self, ds, num_particles, data, ptype="io"):
r"""
Generate particle fields from array-like lists contained in a dict.
Parameters
----------
ds : `Dataset`
The dataset which will serve as the base for these particles.
num_particles : int
The number of particles in the dict.
data : dict of NumPy arrays
The particle fields themselves.
ptype : string, optional
The particle type for these particle fields. Default: "io"
Examples
--------
>>> num_p = 100000
>>> posx = np.random.random(num_p)
>>> posy = np.random.random(num_p)
>>> posz = np.random.random(num_p)
>>> mass = np.ones(num_p)
>>> data = {
... "particle_position_x": posx,
... "particle_position_y": posy,
... "particle_position_z": posz,
... "particle_mass": mass,
... }
>>> particles = FromListParticleGenerator(ds, num_p, data)
"""
field_list = list(data.keys())
if "particle_position_x" in data:
x = data.pop("particle_position_x")
y = data.pop("particle_position_y")
z = data.pop("particle_position_z")
elif (ptype, "particle_position_x") in data:
x = data.pop((ptype, "particle_position_x"))
y = data.pop((ptype, "particle_position_y"))
z = data.pop((ptype, "particle_position_z"))
xcond = np.logical_or(x < ds.domain_left_edge[0], x >= ds.domain_right_edge[0])
ycond = np.logical_or(y < ds.domain_left_edge[1], y >= ds.domain_right_edge[1])
zcond = np.logical_or(z < ds.domain_left_edge[2], z >= ds.domain_right_edge[2])
cond = np.logical_or(xcond, ycond)
cond = np.logical_or(zcond, cond)
if np.any(cond):
raise ValueError("Some particles are outside of the domain!!!")
super().__init__(ds, num_particles, field_list, ptype=ptype)
self._setup_particles(x, y, z, setup_fields=data)
[docs]
class LatticeParticleGenerator(ParticleGenerator):
def __init__(
self,
ds,
particles_dims,
particles_left_edge,
particles_right_edge,
field_list,
ptype="io",
):
r"""
Generate particles in a lattice arrangement.
Parameters
----------
ds : `Dataset`
The dataset which will serve as the base for these particles.
particles_dims : int, array-like
The number of particles along each dimension
particles_left_edge : float, array-like
The 'left-most' starting positions of the lattice.
particles_right_edge : float, array-like
The 'right-most' ending positions of the lattice.
field_list : list of strings
A list of particle fields
ptype : string, optional
The particle type for these particle fields. Default: "io"
Examples
--------
>>> dims = (128, 128, 128)
>>> le = np.array([0.25, 0.25, 0.25])
>>> re = np.array([0.75, 0.75, 0.75])
>>> fields = [
... ("all", "particle_position_x"),
... ("all", "particle_position_y"),
... ("all", "particle_position_z"),
... ("all", "particle_density"),
... ("all", "particle_temperature"),
... ]
>>> particles = LatticeParticleGenerator(ds, dims, le, re, fields)
"""
num_x = particles_dims[0]
num_y = particles_dims[1]
num_z = particles_dims[2]
xmin = particles_left_edge[0]
ymin = particles_left_edge[1]
zmin = particles_left_edge[2]
xmax = particles_right_edge[0]
ymax = particles_right_edge[1]
zmax = particles_right_edge[2]
DLE = ds.domain_left_edge.in_units("code_length").d
DRE = ds.domain_right_edge.in_units("code_length").d
xcond = (xmin < DLE[0]) or (xmax >= DRE[0])
ycond = (ymin < DLE[1]) or (ymax >= DRE[1])
zcond = (zmin < DLE[2]) or (zmax >= DRE[2])
cond = xcond or ycond or zcond
if cond:
raise ValueError("Proposed bounds for particles are outside domain!!!")
super().__init__(ds, num_x * num_y * num_z, field_list, ptype=ptype)
dx = (xmax - xmin) / (num_x - 1)
dy = (ymax - ymin) / (num_y - 1)
dz = (zmax - zmin) / (num_z - 1)
inds = np.indices((num_x, num_y, num_z))
xpos = inds[0] * dx + xmin
ypos = inds[1] * dy + ymin
zpos = inds[2] * dz + zmin
self._setup_particles(xpos.flat[:], ypos.flat[:], zpos.flat[:])
[docs]
class WithDensityParticleGenerator(ParticleGenerator):
def __init__(
self,
ds,
data_source,
num_particles,
field_list,
density_field=("gas", "density"),
ptype="io",
):
r"""
Generate particles based on a density field.
Parameters
----------
ds : `Dataset`
The dataset which will serve as the base for these particles.
data_source :
`yt.data_objects.selection_objects.base_objects.YTSelectionContainer`
The data source containing the density field.
num_particles : int
The number of particles to be generated
field_list : list of strings
A list of particle fields
density_field : tuple, optional
A density field which will serve as the distribution function for the
particle positions. Theoretically, this could be any 'per-volume' field.
ptype : string, optional
The particle type for these particle fields. Default: "io"
Examples
--------
>>> sphere = ds.sphere(ds.domain_center, 0.5)
>>> num_p = 100000
>>> fields = [
... ("all", "particle_position_x"),
... ("all", "particle_position_y"),
... ("all", "particle_position_z"),
... ("all", "particle_density"),
... ("all", "particle_temperature"),
... ]
>>> particles = WithDensityParticleGenerator(
... ds, sphere, num_particles, fields, density_field="Dark_Matter_Density"
... )
"""
super().__init__(ds, num_particles, field_list, ptype=ptype)
num_cells = len(data_source["index", "x"].flat)
max_mass = (
data_source[density_field] * data_source["gas", "cell_volume"]
).max()
num_particles_left = num_particles
all_x = []
all_y = []
all_z = []
pbar = get_pbar("Generating Particles", num_particles)
tot_num_accepted = 0
rng = np.random.default_rng()
while num_particles_left > 0:
m = rng.uniform(high=1.01 * max_mass, size=num_particles_left)
idxs = rng.integers(low=0, high=num_cells, size=num_particles_left)
m_true = (
data_source[density_field] * data_source["gas", "cell_volume"]
).flat[idxs]
accept = m <= m_true
num_accepted = accept.sum()
accepted_idxs = idxs[accept]
xpos = (
data_source["index", "x"].flat[accepted_idxs]
+ rng.uniform(low=-0.5, high=0.5, size=num_accepted)
* data_source["index", "dx"].flat[accepted_idxs]
)
ypos = (
data_source["index", "y"].flat[accepted_idxs]
+ rng.uniform(low=-0.5, high=0.5, size=num_accepted)
* data_source["index", "dy"].flat[accepted_idxs]
)
zpos = (
data_source["index", "z"].flat[accepted_idxs]
+ rng.uniform(low=-0.5, high=0.5, size=num_accepted)
* data_source["index", "dz"].flat[accepted_idxs]
)
all_x.append(xpos)
all_y.append(ypos)
all_z.append(zpos)
num_particles_left -= num_accepted
tot_num_accepted += num_accepted
pbar.update(tot_num_accepted)
pbar.finish()
x = uconcatenate(all_x)
y = uconcatenate(all_y)
z = uconcatenate(all_z)
self._setup_particles(x, y, z)