diff --git a/Docs/source/usage/workflows/python_extend.rst b/Docs/source/usage/workflows/python_extend.rst index 6c9286c02ce..1bfa7e53b8a 100644 --- a/Docs/source/usage/workflows/python_extend.rst +++ b/Docs/source/usage/workflows/python_extend.rst @@ -252,7 +252,7 @@ Particles can be added to the simulation at specific positions and with specific .. autoclass:: pywarpx.particle_containers.ParticleContainerWrapper :members: -The ``get_particle_structs()`` and ``get_particle_arrays()`` functions are called +The ``get_particle_arrays()`` function is called by several utility functions of the form ``get_particle_{comp_name}`` where ``comp_name`` is one of ``x``, ``y``, ``z``, ``r``, ``theta``, ``id``, ``cpu``, ``weight``, ``ux``, ``uy`` or ``uz``. diff --git a/Python/pywarpx/particle_containers.py b/Python/pywarpx/particle_containers.py index 273a981f4bd..a2354c9b344 100644 --- a/Python/pywarpx/particle_containers.py +++ b/Python/pywarpx/particle_containers.py @@ -117,8 +117,10 @@ def add_particles(self, x=None, y=None, z=None, ux=None, uy=None, kwargs[key] = np.full(maxlen, val) # --- The number of built in attributes + # --- The positions + built_in_attrs = libwarpx.dim # --- The three velocities - built_in_attrs = 3 + built_in_attrs += 3 if libwarpx.geometry_dim == 'rz': # --- With RZ, there is also theta built_in_attrs += 1 @@ -188,56 +190,6 @@ def add_real_comp(self, pid_name, comm=True): self.particle_container.add_real_comp(pid_name, comm) - def get_particle_structs(self, level, copy_to_host=False): - ''' - This returns a list of numpy or cupy arrays containing the particle struct data - on each tile for this process. The particle data is represented as a structured - array and contains the particle 'x', 'y', 'z', and 'idcpu'. - - Unless copy_to_host is specified, the data for the structs are - not copied, but share the underlying memory buffer with WarpX. The - arrays are fully writeable. - - Note that cupy does not support structs: - https://github.com/cupy/cupy/issues/2031 - and will return arrays of binary blobs for the AoS (DP: ``"|V24"``). If copied - to host via copy_to_host, we correct for the right numpy AoS type. - - Parameters - ---------- - - level : int - The refinement level to reference (default=0) - - copy_to_host : bool - For GPU-enabled runs, one can either return the GPU - arrays (the default) or force a device-to-host copy. - - Returns - ------- - - List of arrays - The requested particle struct data - ''' - particle_data = [] - for pti in libwarpx.libwarpx_so.WarpXParIter(self.particle_container, level): - if copy_to_host: - particle_data.append(pti.aos().to_numpy(copy=True)) - else: - if libwarpx.amr.Config.have_gpu: - libwarpx.amr.Print( - "get_particle_structs: cupy does not yet support structs. " - "https://github.com/cupy/cupy/issues/2031" - "Did you mean copy_to_host=True?" - ) - xp, cupy_status = load_cupy() - if cupy_status is not None: - libwarpx.amr.Print(cupy_status) - aos_arr = xp.array(pti.aos(), copy=False) # void blobs for cupy - particle_data.append(aos_arr) - return particle_data - - def get_particle_arrays(self, comp_name, level, copy_to_host=False): ''' This returns a list of numpy or cupy arrays containing the particle array data @@ -284,6 +236,31 @@ def get_particle_arrays(self, comp_name, level, copy_to_host=False): return data_array + def get_particle_idcpu(self, level=0, copy_to_host=False): + ''' + Return a list of numpy or cupy arrays containing the particle 'idcpu' + numbers on each tile. + + Parameters + ---------- + + level : int + The refinement level to reference (default=0) + + copy_to_host : bool + For GPU-enabled runs, one can either return the GPU + arrays (the default) or force a device-to-host copy. + + Returns + ------- + + List of arrays + The requested particle idcpu + ''' + return self.get_particle_arrays('idcpu', level, copy_to_host=copy_to_host) + idcpu = property(get_particle_idcpu) + + def get_particle_id(self, level=0, copy_to_host=False): ''' Return a list of numpy or cupy arrays containing the particle 'id' @@ -305,8 +282,8 @@ def get_particle_id(self, level=0, copy_to_host=False): List of arrays The requested particle ids ''' - structs = self.get_particle_structs(level, copy_to_host) - return [libwarpx.amr.unpack_ids(struct['cpuid']) for struct in structs] + idcpu = self.get_particle_idcpu(level, copy_to_host) + return libwarpx.amr.unpack_ids(idcpu) def get_particle_cpu(self, level=0, copy_to_host=False): @@ -330,8 +307,8 @@ def get_particle_cpu(self, level=0, copy_to_host=False): List of arrays The requested particle cpus ''' - structs = self.get_particle_structs(level, copy_to_host) - return [libwarpx.amr.unpack_cpus(struct['cpuid']) for struct in structs] + idcpu = self.get_particle_idcpu(level, copy_to_host) + return libwarpx.amr.unpack_cpus(idcpu) def get_particle_x(self, level=0, copy_to_host=False): @@ -355,20 +332,7 @@ def get_particle_x(self, level=0, copy_to_host=False): List of arrays The requested particle x position ''' - if copy_to_host: - # Use the numpy version of cosine - xp = np - else: - xp, cupy_status = load_cupy() - - structs = self.get_particle_structs(level, copy_to_host) - if libwarpx.geometry_dim == '3d' or libwarpx.geometry_dim == '2d': - return [struct['x'] for struct in structs] - elif libwarpx.geometry_dim == 'rz': - theta = self.get_particle_theta(level, copy_to_host) - return [struct['x']*xp.cos(theta) for struct, theta in zip(structs, theta)] - elif libwarpx.geometry_dim == '1d': - raise Exception('get_particle_x: There is no x coordinate with 1D Cartesian') + return self.get_particle_arrays('x', level, copy_to_host=copy_to_host) xp = property(get_particle_x) @@ -393,20 +357,7 @@ def get_particle_y(self, level=0, copy_to_host=False): List of arrays The requested particle y position ''' - if copy_to_host: - # Use the numpy version of sine - xp = np - else: - xp, cupy_status = load_cupy() - - structs = self.get_particle_structs(level, copy_to_host) - if libwarpx.geometry_dim == '3d': - return [struct['y'] for struct in structs] - elif libwarpx.geometry_dim == 'rz': - theta = self.get_particle_theta(level, copy_to_host) - return [struct['x']*xp.sin(theta) for struct, theta in zip(structs, theta)] - elif libwarpx.geometry_dim == '1d' or libwarpx.geometry_dim == '2d': - raise Exception('get_particle_y: There is no y coordinate with 1D or 2D Cartesian') + return self.get_particle_arrays('y', level, copy_to_host=copy_to_host) yp = property(get_particle_y) @@ -433,11 +384,12 @@ def get_particle_r(self, level=0, copy_to_host=False): ''' xp, cupy_status = load_cupy() - structs = self.get_particle_structs(level, copy_to_host) if libwarpx.geometry_dim == 'rz': - return [struct['x'] for struct in structs] + return self.get_particle_x(level, copy_to_host) elif libwarpx.geometry_dim == '3d': - return [xp.sqrt(struct['x']**2 + struct['y']**2) for struct in structs] + x = self.get_particle_x(level, copy_to_host) + y = self.get_particle_y(level, copy_to_host) + return xp.sqrt(x**2 + y**2) elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == '1d': raise Exception('get_particle_r: There is no r coordinate with 1D or 2D Cartesian') rp = property(get_particle_r) @@ -469,8 +421,9 @@ def get_particle_theta(self, level=0, copy_to_host=False): if libwarpx.geometry_dim == 'rz': return self.get_particle_arrays('theta', level, copy_to_host) elif libwarpx.geometry_dim == '3d': - structs = self.get_particle_structs(level, copy_to_host) - return [xp.arctan2(struct['y'], struct['x']) for struct in structs] + x = self.get_particle_x(level, copy_to_host) + y = self.get_particle_y(level, copy_to_host) + return xp.arctan2(y, x) elif libwarpx.geometry_dim == '2d' or libwarpx.geometry_dim == '1d': raise Exception('get_particle_theta: There is no theta coordinate with 1D or 2D Cartesian') thetap = property(get_particle_theta) @@ -497,13 +450,7 @@ def get_particle_z(self, level=0, copy_to_host=False): List of arrays The requested particle z position ''' - structs = self.get_particle_structs(level, copy_to_host) - if libwarpx.geometry_dim == '3d': - return [struct['z'] for struct in structs] - elif libwarpx.geometry_dim == 'rz' or libwarpx.geometry_dim == '2d': - return [struct['y'] for struct in structs] - elif libwarpx.geometry_dim == '1d': - return [struct['x'] for struct in structs] + return self.get_particle_arrays('z', level, copy_to_host=copy_to_host) zp = property(get_particle_z) @@ -720,70 +667,6 @@ def get_particle_boundary_buffer_size(self, species_name, boundary, local=False) ) - def get_particle_boundary_buffer_structs( - self, species_name, boundary, level, copy_to_host=False - ): - ''' - This returns a list of numpy or cupy arrays containing the particle struct data - for a species that has been scraped by a specific simulation boundary. The - particle data is represented as a structured array and contains the - particle 'x', 'y', 'z', and 'idcpu'. - - Unless copy_to_host is specified, the data for the structs are - not copied, but share the underlying memory buffer with WarpX. The - arrays are fully writeable. - - Note that cupy does not support structs: - https://github.com/cupy/cupy/issues/2031 - and will return arrays of binary blobs for the AoS (DP: ``"|V24"``). If copied - to host via copy_to_host, we correct for the right numpy AoS type. - - Parameters - ---------- - - species_name : str - The species name that the data will be returned for - - boundary : str - The boundary from which to get the scraped particle data in the - form x/y/z_hi/lo or eb. - - level : int - The refinement level to reference (default=0) - - copy_to_host : bool - For GPU-enabled runs, one can either return the GPU - arrays (the default) or force a device-to-host copy. - - Returns - ------- - - List of arrays - The requested particle struct data - ''' - particle_container = self.particle_buffer.get_particle_container( - species_name, self._get_boundary_number(boundary) - ) - - particle_data = [] - for pti in libwarpx.libwarpx_so.BoundaryBufferParIter(particle_container, level): - if copy_to_host: - particle_data.append(pti.aos().to_numpy(copy=True)) - else: - if libwarpx.amr.Config.have_gpu: - libwarpx.amr.Print( - "get_particle_structs: cupy does not yet support structs. " - "https://github.com/cupy/cupy/issues/2031" - "Did you mean copy_to_host=True?" - ) - xp, cupy_status = load_cupy() - if cupy_status is not None: - libwarpx.amr.Print(cupy_status) - aos_arr = xp.array(pti.aos(), copy=False) # void blobs for cupy - particle_data.append(aos_arr) - return particle_data - - def get_particle_boundary_buffer(self, species_name, boundary, comp_name, level): ''' This returns a list of numpy or cupy arrays containing the particle array data diff --git a/Source/Python/Particles/ParticleBoundaryBuffer.cpp b/Source/Python/Particles/ParticleBoundaryBuffer.cpp index 2a35faece9b..b04ac75e600 100644 --- a/Source/Python/Particles/ParticleBoundaryBuffer.cpp +++ b/Source/Python/Particles/ParticleBoundaryBuffer.cpp @@ -10,13 +10,13 @@ namespace warpx { class BoundaryBufferParIter - : public amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator> + : public amrex::ParIterSoA { public: - using amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>::ParIter; + using amrex::ParIterSoA::ParIterSoA; BoundaryBufferParIter(ContainerType& pc, int level) : - amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>(pc, level) {} + amrex::ParIterSoA(pc, level) {} }; } @@ -24,9 +24,9 @@ void init_BoundaryBufferParIter (py::module& m) { py::class_< warpx::BoundaryBufferParIter, - amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator> + amrex::ParIterSoA >(m, "BoundaryBufferParIter") - .def(py::init::ContainerType&, int>(), + .def(py::init::ContainerType&, int>(), py::arg("particle_container"), py::arg("level") ) ; diff --git a/Source/Python/Particles/PinnedMemoryParticleContainer.cpp b/Source/Python/Particles/PinnedMemoryParticleContainer.cpp index 600d56a62c9..d4f6a422dbe 100644 --- a/Source/Python/Particles/PinnedMemoryParticleContainer.cpp +++ b/Source/Python/Particles/PinnedMemoryParticleContainer.cpp @@ -13,6 +13,6 @@ void init_PinnedMemoryParticleContainer (py::module& m) { py::class_< PinnedMemoryParticleContainer, - amrex::ParticleContainer<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator> + amrex::ParticleContainerPureSoA > pmpc (m, "PinnedMemoryParticleContainer"); } diff --git a/Source/Python/Particles/WarpXParticleContainer.cpp b/Source/Python/Particles/WarpXParticleContainer.cpp index 1473a750941..a295157fdae 100644 --- a/Source/Python/Particles/WarpXParticleContainer.cpp +++ b/Source/Python/Particles/WarpXParticleContainer.cpp @@ -12,11 +12,11 @@ void init_WarpXParIter (py::module& m) { py::class_< - WarpXParIter, amrex::ParIter<0,0,PIdx::nattribs> + WarpXParIter, amrex::ParIterSoA >(m, "WarpXParIter") - .def(py::init::ContainerType&, int>(), + .def(py::init::ContainerType&, int>(), py::arg("particle_container"), py::arg("level")) - .def(py::init::ContainerType&, int, amrex::MFItInfo&>(), + .def(py::init::ContainerType&, int, amrex::MFItInfo&>(), py::arg("particle_container"), py::arg("level"), py::arg("info")) ; @@ -26,11 +26,11 @@ void init_WarpXParticleContainer (py::module& m) { py::class_< WarpXParticleContainer, - amrex::ParticleContainer<0, 0, PIdx::nattribs, 0> + amrex::ParticleContainerPureSoA > wpc (m, "WarpXParticleContainer"); wpc .def("add_real_comp", - [](WarpXParticleContainer& pc, const std::string& name, bool const comm) { pc.AddRealComp(name, comm); }, + [](WarpXParticleContainer& pc, const std::string& name, bool comm) { pc.AddRealComp(name, comm); }, py::arg("name"), py::arg("comm") ) .def("add_n_particles",