particle_generator.py 文件源码

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, string_types) else key
                if field not in self.default_fields:
                    self.particles[:,self.field_list.index(field)] = value[idxs]