Particle Container to Pure SoA Again (#4653)

* AMReX & pyAMReX: Latest `development`

More pure SoA and id handling goodness.

* Particle Container to Pure SoA Again

Transition to new, purely SoA particle containers.

This was originally merged in #3850 and reverted in #4652, since
we discovered issues loosing particles & laser particles on GPU.

* Modernize `idcpu` Treatment

- faster: less emitted operations, no jumps
- cheaper: less used registers
- safer: no read-before-write warnings
- cooler: no explanation needed

Author: ax3l

.github/workflows/cuda.yml

@@ -115,7 +115,7 @@ jobs:
         which nvcc || echo "nvcc not in PATH!"
 
         git clone https://github.com/AMReX-Codes/amrex.git ../amrex
-        cd ../amrex && git checkout --detach 24.02 && cd -
+        cd ../amrex && git checkout --detach 296ed40e16ae1877640f5b78e9162dbd4ba1c279 && cd -
         make COMP=gcc QED=FALSE USE_MPI=TRUE USE_GPU=TRUE USE_OMP=FALSE USE_PSATD=TRUE USE_CCACHE=TRUE -j 2
 
         ccache -s

Docs/source/developers/amrex_basics.rst

@@ -13,7 +13,7 @@ WarpX is built on the Adaptive Mesh Refinement (AMR) library `AMReX <https://git
 
 * ``amrex::MultiFab``: Collection of `FAB` (= ``FArrayBox``) on a single AMR level, distributed over MPI ranks. The concept of `ghost cells` is defined at the ``MultiFab`` level.
 
-* ``amrex::ParticleContainer``: A collection of particles, typically for particles of a physical species. Particles in a ``ParticleContainer`` are organized per ``Box``. Particles in a ``Box`` are organized per tile (this feature is off when running on GPU). Particles within a tile are stored in several structures, each being contiguous in memory: (i) an Array-Of-Struct (AoS) (often called `data`, they are the 3D position, the particle ID and the index of the CPU owning the particle), where the Struct is an ``amrex::Particle`` and (ii) Struct-Of-Arrays (SoA) for extra variables (often called ``attribs``, in WarpX they are the momentum, field on particle etc.).
+* ``amrex::ParticleContainer``: A collection of particles, typically for particles of a physical species. Particles in a ``ParticleContainer`` are organized per ``Box``. Particles in a ``Box`` are organized per tile (this feature is off when running on GPU). Particles within a tile are stored in several structures, each being contiguous in memory: (i) a Struct-of-Array (SoA) for ``amrex::ParticleReal`` data such as positions, weight, momentum, etc., (ii) a Struct-of-Array (SoA) for ``int`` data, such as ionization levels, and (iii) a Struct-of-Array (SoA) for a ``uint64_t`` unique identifier index per particle (containing a 40bit id and 24bit cpu sub-identifier as assigned at particle creation time). This id is also used to check if a particle is active/valid or marked for removal.
 
 The simulation domain is decomposed in several ``Box``, and each MPI rank owns (and performs operations on) the fields and particles defined on a few of these ``Box``, but has the metadata of all of them. For convenience, AMReX provides iterators, to easily iterate over all ``FArrayBox`` (or even tile-by-tile, optionally) in a ``MultiFab`` own by the MPI rank (``MFIter``), or over all particles in a ``ParticleContainer`` on a per-box basis (``ParIter``, or its derived class ``WarpXParIter``). These are respectively done in loops like:
 

Docs/source/developers/dimensionality.rst

@@ -49,12 +49,12 @@ WarpX axis labels     ``x, y, z``   ``x, z``    ``z``        ``x, z``
 --------------------  -----------  -----------  -----------  -----------
 *Particles*
 ------------------------------------------------------------------------
-AMReX AoS ``.pos()``  ``0, 1, 2``  ``0, 1``     ``0``        ``0, 1``
+AMReX ``.pos()``      ``0, 1, 2``  ``0, 1``     ``0``        ``0, 1``
 WarpX position names  ``x, y, z``  ``x, z``     ``z``        ``r, z``
 extra SoA attribute                                          ``theta``
 ====================  ===========  ===========  ===========  ===========
 
-Please see the following sections for particle AoS and SoA details.
+Please see the following sections for particle SoA details.
 
 Conventions
 -----------

Docs/source/developers/particles.rst

@@ -109,24 +109,26 @@ Particle attributes
 -------------------
 
 WarpX adds the following particle attributes by default to WarpX particles.
-These attributes are either stored in an Array-of-Struct (AoS) or Struct-of-Array (SoA) location of the AMReX particle containers.
+These attributes are stored in Struct-of-Array (SoA) locations of the AMReX particle containers: one SoA for ``amrex::ParticleReal`` attributes, one SoA for ``int`` attributes and one SoA for a ``uint64_t`` global particle index per particle.
 The data structures for those are either pre-described at compile-time (CT) or runtime (RT).
 
-====================  ================  ==================================  ===== ==== =====================
+====================  ================  ==================================  ===== ==== ======================
 Attribute name        ``int``/``real``  Description                         Where When Notes
-====================  ================  ==================================  ===== ==== =====================
-``position_x/y/z``    ``real``          Particle position.                  AoS   CT
-``cpu``               ``int``           CPU index where the particle        AoS   CT
+====================  ================  ==================================  ===== ==== ======================
+``position_x/y/z``    ``real``          Particle position.                  SoA   CT
+``weight``            ``real``          Particle position.                  SoA   CT
+``momentum_x/y/z``    ``real``          Particle position.                  SoA   CT
+``id``                ``amrex::Long``   CPU-local particle index            SoA   CT   First 40 bytes of
+                                        where the particle was created.                idcpu
+``cpu``               ``int``           CPU index where the particle        SoA   CT   Last 24 bytes of idcpu
                                         was created.
-``id``                ``int``           CPU-local particle index            AoS   CT
-                                        where the particle was created.
 ``ionizationLevel``   ``int``           Ion ionization level                SoA   RT   Added when ionization
                                                                                        physics is used.
 ``opticalDepthQSR``   ``real``          QED: optical depth of the Quantum-  SoA   RT   Added when PICSAR QED
                                         Synchrotron process                            physics is used.
 ``opticalDepthBW``    ``real``          QED: optical depth of the Breit-    SoA   RT   Added when PICSAR QED
                                         Wheeler process                                physics is used.
-====================  ================  ==================================  ===== ==== =====================
+====================  ================  ==================================  ===== ==== ======================
 
 WarpX allows extra runtime attributes to be added to particle containers (through ``AddRealComp("attrname")`` or ``AddIntComp("attrname")``).
 The attribute name can then be used to access the values of that attribute.

Docs/source/usage/workflows/python_extend.rst

@@ -190,6 +190,7 @@ Particles
    @callfromafterstep
    def my_after_step_callback():
        warpx = sim.extension.warpx
+       Config = sim.extension.Config
 
        # data access
        multi_pc = warpx.multi_particle_container()
@@ -200,27 +201,27 @@ Particles
        for lvl in range(pc.finest_level + 1):
            # get every local chunk of particles
            for pti in pc.iterator(pc, level=lvl):
-               # default layout: AoS with positions and cpuid
-               aos = pti.aos().to_numpy()
-
-               # additional compile-time and runtime attributes in SoA format
-               soa = pti.soa().to_numpy()
+               # compile-time and runtime attributes in SoA format
+               soa = pti.soa().to_cupy() if Config.have_gpu else \
+                     pti.soa().to_numpy()
 
                # notes:
                # Only the next lines are the "HOT LOOP" of the computation.
-               # For efficiency, use numpy array operation for speed on CPUs.
-               # For GPUs use .to_cupy() above and compute with cupy or numba.
+               # For speed, use array operation.
 
                # write to all particles in the chunk
-               # note: careful, if you change particle positions, you need to
+               # note: careful, if you change particle positions, you might need to
                #       redistribute particles before continuing the simulation step
-               # aos[()]["x"] = 0.30
-               # aos[()]["y"] = 0.35
-               # aos[()]["z"] = 0.40
+               soa.real[0][()] = 0.30  # x
+               soa.real[1][()] = 0.35  # y
+               soa.real[2][()] = 0.40  # z
 
-               for soa_real in soa.real:
+               # all other attributes: weight, momentum x, y, z, ...
+               for soa_real in soa.real[3:]:
                    soa_real[()] = 42.0
 
+               # by default empty unless ionization or QED physics is used
+               # or other runtime attributes were added manually
                for soa_int in soa.int:
                    soa_int[()] = 12
 
@@ -252,7 +253,8 @@ 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_real_arrays()``, ``get_particle_int_arrays()`` and
+``get_particle_idcpu_arrays()`` functions are 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``.

Examples/Tests/particle_data_python/PICMI_inputs_2d.py

@@ -153,10 +153,10 @@ def add_particles():
 ##########################
 
 assert (elec_wrapper.nps == 270 / (2 - args.unique))
-assert (elec_wrapper.particle_container.get_comp_index('w') == 0)
-assert (elec_wrapper.particle_container.get_comp_index('newPid') == 4)
+assert (elec_wrapper.particle_container.get_comp_index('w') == 2)
+assert (elec_wrapper.particle_container.get_comp_index('newPid') == 6)
 
-new_pid_vals = elec_wrapper.get_particle_arrays('newPid', 0)
+new_pid_vals = elec_wrapper.get_particle_real_arrays('newPid', 0)
 for vals in new_pid_vals:
     assert np.allclose(vals, 5)
 

Examples/Tests/particle_data_python/PICMI_inputs_prev_pos_2d.py

@@ -120,12 +120,12 @@
 elec_count = elec_wrapper.nps
 
 # check that the runtime attributes have the right indices
-assert (elec_wrapper.particle_container.get_comp_index('prev_x') == 4)
-assert (elec_wrapper.particle_container.get_comp_index('prev_z') == 5)
+assert (elec_wrapper.particle_container.get_comp_index('prev_x') == 6)
+assert (elec_wrapper.particle_container.get_comp_index('prev_z') == 7)
 
 # sanity check that the prev_z values are reasonable and
 # that the correct number of values are returned
-prev_z_vals = elec_wrapper.get_particle_arrays('prev_z', 0)
+prev_z_vals = elec_wrapper.get_particle_real_arrays('prev_z', 0)
 running_count = 0
 
 for z_vals in prev_z_vals:

Examples/Tests/restart/PICMI_inputs_runtime_component_analyze.py

@@ -158,10 +158,10 @@ def add_particles():
 ##########################
 
 assert electron_wrapper.nps == 90
-assert electron_wrapper.particle_container.get_comp_index("w") == 0
-assert electron_wrapper.particle_container.get_comp_index("newPid") == 4
+assert electron_wrapper.particle_container.get_comp_index("w") == 2
+assert electron_wrapper.particle_container.get_comp_index("newPid") == 6
 
-new_pid_vals = electron_wrapper.get_particle_arrays("newPid", 0)
+new_pid_vals = electron_wrapper.get_particle_real_arrays("newPid", 0)
 for vals in new_pid_vals:
     assert np.allclose(vals, 5)
 

Python/pywarpx/_libwarpx.py

@@ -5,8 +5,8 @@
 #
 # NOTE: We will reduce the libwarpx.py level of abstraction eventually!
 # Please add new functionality directly to pybind11-bound modules
-# and call them via sim.extension.libwarpx_so. ... and sim.extension.warpx.
-# ... from user code.
+# and call them via sim.extension.libwarpx_so. ... and sim.extension.Config and
+# sim.extension.warpx. ... from user code.
 #
 # Authors: Axel Huebl, Andrew Myers, David Grote, Remi Lehe, Weiqun Zhang
 #
@@ -108,6 +108,8 @@ def load_library(self):
                 from . import warpx_pybind_3d as cxx_3d
                 self.libwarpx_so = cxx_3d
                 self.dim = 3
+
+            self.Config = self.libwarpx_so.Config
         except ImportError:
             raise Exception(f"Dimensionality '{self.geometry_dim}' was not compiled in this Python install. Please recompile with -DWarpX_DIMS={_dims}")