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)