Python: Pure SoA Particle

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``.

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

Source/Python/Particles/ParticleBoundaryBuffer.cpp

@@ -10,23 +10,23 @@
 
 namespace warpx {
     class BoundaryBufferParIter
-        : public amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>
+        : public amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>
     {
     public:
-        using amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>::ParIter;
+        using amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>::ParIterSoA;
 
         BoundaryBufferParIter(ContainerType& pc, int level) :
-            amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>(pc, level) {}
+            amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>(pc, level) {}
     };
 }
 
 void init_BoundaryBufferParIter (py::module& m)
 {
     py::class_<
         warpx::BoundaryBufferParIter,
-        amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>
+        amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>
     >(m, "BoundaryBufferParIter")
-        .def(py::init<amrex::ParIter<0,0,PIdx::nattribs,0,amrex::PinnedArenaAllocator>::ContainerType&, int>(),
+        .def(py::init<amrex::ParIterSoA<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>::ContainerType&, int>(),
             py::arg("particle_container"), py::arg("level")
         )
     ;

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<PIdx::nattribs, 0, amrex::PinnedArenaAllocator>
     > pmpc (m, "PinnedMemoryParticleContainer");
 }

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<PIdx::nattribs, 0>
     >(m, "WarpXParIter")
-        .def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int>(),
+        .def(py::init<amrex::ParIterSoA<PIdx::nattribs, 0>::ContainerType&, int>(),
             py::arg("particle_container"), py::arg("level"))
-        .def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int, amrex::MFItInfo&>(),
+        .def(py::init<amrex::ParIterSoA<PIdx::nattribs, 0>::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<PIdx::nattribs, 0>
     > 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",