def test_sample():
grid = {}
dims = np.array([64, 64, 64], dtype="int32")
inds = np.indices(dims)
grid["x"] = inds[0] + 0.5
grid["y"] = inds[1] + 0.5
grid["z"] = inds[2] + 0.5
num_particles = np.int64(1000)
xp = np.random.uniform(low=1.0, high=63.0, size=num_particles)
yp = np.random.uniform(low=1.0, high=63.0, size=num_particles)
zp = np.random.uniform(low=1.0, high=63.0, size=num_particles)
xfield = np.zeros((num_particles))
yfield = np.zeros((num_particles))
zfield = np.zeros((num_particles))
dx = 1.0
le = np.zeros((3))
CICSample_3(xp, yp, zp, xfield, num_particles, grid["x"], le, dims, dx)
CICSample_3(xp, yp, zp, yfield, num_particles, grid["y"], le, dims, dx)
CICSample_3(xp, yp, zp, zfield, num_particles, grid["z"], le, dims, dx)
assert_allclose(xp, xfield)
assert_allclose(yp, yfield)
assert_allclose(zp, zfield)
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)
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()
def _get_data(self, fields):
"""
Get a list of fields to include in the trajectory collection.
The trajectory collection itself is a dict of 2D numpy arrays,
with shape (num_indices, num_steps)
"""
missing_fields = [
field for field in fields if field not in self.field_data
]
if not missing_fields:
return
if self.suppress_logging:
old_level = int(ytcfg.get("yt", "loglevel"))
mylog.setLevel(40)
ds_first = self.data_series[0]
dd_first = ds_first.all_data()
fds = {}
new_particle_fields = []
for field in missing_fields:
fds[field] = dd_first._determine_fields(field)[0]
if field not in self.particle_fields:
if self.data_series[0]._get_field_info(
*fds[field]).particle_type:
self.particle_fields.append(field)
new_particle_fields.append(field)
grid_fields = [
field for field in missing_fields
if field not in self.particle_fields
]
step = int(0)
pbar = get_pbar(
f"Generating [{', '.join(missing_fields)}] fields in trajectories",
self.num_steps,
)
my_storage = {}
for i, (sto,
ds) in enumerate(self.data_series.piter(storage=my_storage)):
mask = self.masks[i]
sort = self.sorts[i]
pfield = {}
if new_particle_fields: # there's at least one particle field
dd = ds.all_data()
for field in new_particle_fields:
# This is easy... just get the particle fields
pfield[field] = dd[fds[field]].d[mask][sort]
if grid_fields:
# This is hard... must loop over grids
for field in grid_fields:
pfield[field] = np.zeros(self.num_indices)
x = self["particle_position_x"][:, step].d
y = self["particle_position_y"][:, step].d
z = self["particle_position_z"][:, step].d
particle_grids, particle_grid_inds = ds.index._find_points(
x, y, z)
# This will fail for non-grid index objects
for grid in particle_grids:
cube = grid.retrieve_ghost_zones(1, grid_fields)
for field in grid_fields:
CICSample_3(
x,
y,
z,
pfield[field],
self.num_indices,
cube[fds[field]],
np.array(grid.LeftEdge).astype(np.float64),
np.array(grid.ActiveDimensions).astype(np.int32),
grid.dds[0],
)
sto.result_id = ds.parameter_filename
sto.result = (self.array_indices[i], pfield)
pbar.update(step)
step += 1
pbar.finish()
output_field = np.empty((self.num_indices, self.num_steps))
output_field.fill(np.nan)
for field in missing_fields:
fd = fds[field]
for i, (_fn, (indices,
pfield)) in enumerate(sorted(my_storage.items())):
output_field[indices, i] = pfield[field]
self.field_data[field] = array_like_field(dd_first,
output_field.copy(), fd)
if self.suppress_logging:
mylog.setLevel(old_level)