def align_particles(self): tag_arr = self._data['tag'] if len(tag_arr) == 0: self.num_real_particles = 0 return num_particles = len(tag_arr) new_indices = array.empty(num_particles, dtype=np.int32, backend=self.backend) num_real_particles = array.empty(1, dtype=np.int32, backend=self.backend) align_particles_knl = self._get_align_kernel_without_strides() align_particles_knl(tag_arr=tag_arr, new_indices=new_indices, num_particles=num_particles, num_real_particles=num_real_particles) indices = {1: new_indices} for stride in set(self._particle_array.stride.values()): if stride > 1: indices[stride] = self._build_indices_with_strides( tag_arr, stride) self.align(indices) self.num_real_particles = int(num_real_particles.get())
def _make_strided_indices(self, indices_dict): '''Takes the indices in a dict assuming that the indices are for a stride of 1 and makes suitable indices for other possible stride values. ''' indices = indices_dict[1] n = indices.length if not self._strided_indices_knl: self._setup_strided_indices_kernel() for stride in set(self._particle_array.stride.values()): dest = array.empty(n*stride, dtype=np.int32, backend=self.backend) self._strided_indices_knl(indices, dest, stride) indices_dict[stride] = dest
def _build_indices_with_strides(self, tag_arr, stride): num_particles = len(tag_arr.dev) new_indices = array.empty(num_particles * stride, dtype=np.int32, backend=self.backend) align_particles_knl = self._get_align_kernel_with_strides() align_particles_knl(tag_arr=tag_arr, new_indices=new_indices, num_particles=num_particles, stride=stride) return new_indices
def test_copy_values(backend): check_import(backend) # Given dev_array = make_dev_array(backend) dev_array.fill(2) dest = array.empty(8, dtype=np.int32, backend=backend) indices = array.arange(0, 8, 1, dtype=np.int32, backend=backend) # When dev_array.copy_values(indices, dest) # Then assert np.all(dev_array[:len(indices)].get() == dest.get())
def _build_copy_indices_with_strides(self, indices): result = {1: indices} n_indices = len(indices) strides = set(self._particle_array.stride.values()) fill_indices_knl = self._get_copy_indices_kernel() for stride in strides: sz = n_indices * stride new_indices = array.empty(sz, dtype=np.uint32, backend=self.backend) fill_indices_knl(indices, new_indices, sz, stride) result[stride] = new_indices return result
def _remove_particles_bool(self, if_remove, align=True): """ Remove particle i if if_remove[i] is True """ num_indices = int(array.sum(if_remove, backend=self.backend)) if num_indices == 0: return num_particles = self.get_number_of_particles() new_indices = Array(np.uint32, n=(num_particles - num_indices), backend=self.backend) num_removed_particles = array.empty(1, dtype=np.int32, backend=self.backend) remove_knl, stride_knl = self._get_remove_particles_bool_kernels() remove_knl(if_remove=if_remove, new_indices=new_indices, num_removed_particles=num_removed_particles, num_particles=num_particles) new_num_particles = num_particles - int(num_removed_particles.get()) strides = set(self._particle_array.stride.values()) s_indices = {1: new_indices} for stride in strides: if stride == 1: continue size = new_num_particles * stride s_index = Array(np.uint32, n=size, backend=self.backend) stride_knl(new_indices, s_index, size, stride) s_indices[stride] = s_index for prop in self.properties: stride = self._particle_array.stride.get(prop, 1) s_index = s_indices[stride] self._data[prop].align(s_index) setattr(self, prop, self._data[prop]) if align: self.align_particles()