Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.

Author: ax3l

cmake/dependencies/ABLASTR.cmake

@@ -163,18 +163,18 @@ set(ImpactX_openpmd_src ""
     "Local path to openPMD-api source directory (preferred if set)")
 
 # Git fetcher
-set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git"
+set(ImpactX_ablastr_repo "https://github.com/ax3l/WarpX.git"
     CACHE STRING
     "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
-set(ImpactX_ablastr_branch "23.05"
+set(ImpactX_ablastr_branch "topic-particle-soa"
     CACHE STRING
     "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")
 
 # AMReX is transitively pulled through ABLASTR
 set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
     CACHE STRING
     "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
-set(ImpactX_amrex_branch ""
+set(ImpactX_amrex_branch "a393d7ff7e320cefeeb55e31f1b0e0b5ac2d90ca"
     CACHE STRING
     "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")
 

cmake/dependencies/pyAMReX.cmake

@@ -76,7 +76,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
 set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
     CACHE STRING
     "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)")
-set(ImpactX_pyamrex_branch "23.05"
+set(ImpactX_pyamrex_branch "96743d0d5ebeea83b1ca04a4e293173123081756"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

examples/fodo/run_fodo_programmable.py

@@ -78,16 +78,16 @@ def my_drift(pge, pti, refpart):
 
     else:
         array = np.array
-    # access AoS data such as positions and cpu/id
-    aos = pti.aos()
-    aos_arr = array(aos, copy=False)
 
-    # access SoA data such as momentum
+    # access particle attributes
     soa = pti.soa()
     real_arrays = soa.GetRealData()
-    px = array(real_arrays[0], copy=False)
-    py = array(real_arrays[1], copy=False)
-    pt = array(real_arrays[2], copy=False)
+    x = array(real_arrays[0], copy=False)
+    y = array(real_arrays[1], copy=False)
+    t = array(real_arrays[2], copy=False)
+    px = array(real_arrays[3], copy=False)
+    py = array(real_arrays[4], copy=False)
+    pt = array(real_arrays[5], copy=False)
 
     # length of the current slice
     slice_ds = pge.ds / pge.nslice
@@ -97,9 +97,9 @@ def my_drift(pge, pti, refpart):
     betgam2 = pt_ref**2 - 1.0
 
     # advance position and momentum (drift)
-    aos_arr[:]["x"] += slice_ds * px[:]
-    aos_arr[:]["y"] += slice_ds * py[:]
-    aos_arr[:]["z"] += (slice_ds / betgam2) * pt[:]
+    x[:] += slice_ds * px[:]
+    y[:] += slice_ds * py[:]
+    t[:] += (slice_ds / betgam2) * pt[:]
 
 
 def my_ref_drift(pge, refpart):

src/initialization/InitMeshRefinement.cpp

@@ -201,6 +201,8 @@ namespace impactx
             amrex::Geometry g = Geom(lev);
             g.ProbDomain(rb);
             amrex::AmrMesh::SetGeometry(lev, g);
+
+            amrex::Print() << "New Geometry g=" << g << "\n";
         }
     }
 } // namespace impactx

src/initialization/InitParser.cpp

@@ -25,6 +25,11 @@ namespace impactx::initialization
         bool the_arena_is_managed = false; // AMReX' default: true
         pp_amrex.queryAdd("the_arena_is_managed", the_arena_is_managed);
 
+        // debugging
+        // https://amrex-codes.github.io/amrex/docs_html/Debugging.html#breaking-into-debuggers
+        pp_amrex.add("throw_exception", 1);
+        pp_amrex.add("signal_handling", 0);
+
         // Here we override the default tiling option for particles, which is always
         // "false" in AMReX, to "false" if compiling for GPU execution and "true"
         // if compiling for CPU.

src/particles/ImpactXParticleContainer.H

@@ -17,6 +17,7 @@
 #include <AMReX_MultiFab.H>
 #include <AMReX_ParIter.H>
 #include <AMReX_Particles.H>
+#include <AMReX_ParticleTile.H>
 
 #include <AMReX_IntVect.H>
 #include <AMReX_Vector.H>
@@ -28,43 +29,16 @@
 
 namespace impactx
 {
-    /** AMReX pre-defined Real attributes
-     *
-     * These are the AMReX pre-defined struct indexes for the Real attributes
-     * stored in an AoS in ImpactXParticleContainer. We document this here,
-     * because we change the meaning of these "positions" depending on the
-     * coordinate system we are currently in.
-     */
-    struct RealAoS
-    {
-        enum
-        {
-            x,  ///< position in x [m] (at fixed s OR fixed t)
-            y,  ///< position in y [m] (at fixed s OR fixed t)
-            t,  ///< c * time-of-flight [m] (at fixed s)
-            nattribs ///< the number of attributes above (always last)
-        };
-
-        // at fixed t, the third component represents the position z
-        enum {
-            z = t  ///< position in z [m] (at fixed t)
-        };
-
-        //! named labels for fixed s
-        static constexpr auto names_s = { "position_x", "position_y", "position_t" };
-        //! named labels for fixed t
-        static constexpr auto names_t = { "position_x", "position_y", "position_z" };
-        static_assert(names_s.size() == nattribs);
-        static_assert(names_t.size() == nattribs);
-    };
-
-    /** This struct indexes the additional Real attributes
+    /** This struct indexes the Real attributes
      *  stored in an SoA in ImpactXParticleContainer
      */
     struct RealSoA
     {
         enum
         {
+            x,   ///< position in x [m] (at fixed s or t)
+            y,   ///< position in y [m] (at fixed s or t)
+            t,   ///< time-of-flight ct [m] (at fixed s)
             px,  ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t)
             py,  ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed s or t)
             pt,  ///< energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s)
@@ -73,75 +47,85 @@ namespace impactx
             nattribs ///< the number of attributes above (always last)
         };
 
-        // at fixed t, the third component represents the momentum in z
+        // at fixed t, the third component represents the position z, the 6th component represents the momentum in z
         enum {
+            z = t,   ///< position in z [m] (at fixed t)
             pz = pt  ///< momentum in z, scaled by the magnitude of the reference momentum [unitless] (at fixed t)
         };
 
         //! named labels for fixed s
-        static constexpr auto names_s = { "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" };
+        static constexpr auto names_s = { "position_x", "position_y", "position_t", "momentum_x", "momentum_y", "momentum_t", "qm", "weighting" };
         //! named labels for fixed t
-        static constexpr auto names_t = { "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" };
+        static constexpr auto names_t = { "position_x", "position_y", "position_z", "momentum_x", "momentum_y", "momentum_z", "qm", "weighting" };
         static_assert(names_s.size() == nattribs);
         static_assert(names_t.size() == nattribs);
     };
 
-    /** This struct indexes the additional Integer attributes
+    /** This struct indexes the Integer attributes
      *  stored in an SoA in ImpactXParticleContainer
      */
     struct IntSoA
     {
         enum
         {
-            nattribs ///< the number of particles above (always last)
+            id,      ///< unique particle id on the rank of particle creation
+            cpu,     ///< rank the particle was created on
+            nattribs ///< the number of attributes above (always last)
         };
+
+        //! named labels for fixed s
+        static constexpr auto names_s = { "id", "cpu" };
+        //! named labels for fixed t
+        static constexpr auto names_t = { "id", "cpu" };
+        static_assert(names_s.size() == nattribs);
+        static_assert(names_t.size() == nattribs);
     };
 
     /** AMReX iterator for particle boxes
      *
      * We subclass here to change the default threading strategy, which is
      * `static` in AMReX, to `dynamic` in ImpactX.
      */
-    class ParIter
-        : public amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+    class ParIterSoA
+        : public amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
-        using amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParIter;
+        using amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>::ParIterSoA;
 
-        ParIter (ContainerType& pc, int level);
+        ParIterSoA (ContainerType& pc, int level);
 
-        ParIter (ContainerType& pc, int level, amrex::MFItInfo& info);
+        ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
     };
 
     /** Const AMReX iterator for particle boxes - data is read only.
      *
      * We subclass here to change the default threading strategy, which is
      * `static` in AMReX, to `dynamic` in ImpactX.
      */
-    class ParConstIter
-        : public amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+    class ParConstIterSoA
+        : public amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
-        using amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParConstIter;
+        using amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>::ParConstIterSoA;
 
-        ParConstIter (ContainerType& pc, int level);
+        ParConstIterSoA (ContainerType& pc, int level);
 
-        ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info);
+        ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info);
     };
 
     /** Beam Particles in ImpactX
      *
      * This class stores particles, distributed over MPI ranks.
      */
     class ImpactXParticleContainer
-        : public amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>
+        : public amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>
     {
     public:
         //! amrex iterator for particle boxes
-        using iterator = impactx::ParIter;
+        using iterator = impactx::ParIterSoA;
 
         //! amrex constant iterator for particle boxes (read-only)
-        using const_iterator = impactx::ParConstIter;
+        using const_iterator = impactx::ParConstIterSoA;
 
         //! Construct a new particle container
         ImpactXParticleContainer (amrex::AmrCore* amr_core);

src/particles/ImpactXParticleContainer.cpp

@@ -17,7 +17,6 @@
 #include <AMReX_AmrParGDB.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
-#include <AMReX_ParticleTile.H>
 
 #include <stdexcept>
 
@@ -31,24 +30,24 @@ namespace impactx
         return do_dynamic;
     }
 
-    ParIter::ParIter (ContainerType& pc, int level)
-        : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParIterSoA::ParIterSoA (ContainerType& pc, int level)
+        : amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
                    amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {}
 
-    ParIter::ParIter (ContainerType& pc, int level, amrex::MFItInfo& info)
-        : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParIterSoA::ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info)
+        : amrex::ParIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               info.SetDynamic(do_omp_dynamic())) {}
 
-    ParConstIter::ParConstIter (ContainerType& pc, int level)
-        : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level)
+        : amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {}
 
-    ParConstIter::ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info)
-        : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level,
+    ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info)
+        : amrex::ParConstIterSoA<RealSoA::nattribs, IntSoA::nattribs>(pc, level,
               info.SetDynamic(do_omp_dynamic())) {}
 
     ImpactXParticleContainer::ImpactXParticleContainer (amrex::AmrCore* amr_core)
-        : amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
+        : amrex::ParticleContainerPureSoA<RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
     {
         SetParticleSize();
     }
@@ -105,35 +104,26 @@ namespace impactx
         reserveData();
         resizeData();
 
+        // write Real attributes (SoA) to particle initialized zero
         auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0);
 
         /* Create a temporary tile to obtain data from simulation. This data
          * is then copied to the permanent tile which is stored on the particle
          * (particle_tile).
          */
-        using PinnedTile = amrex::ParticleTile<
-            amrex::Particle<NStructReal, NStructInt>,
-            NArrayReal, NArrayInt,
-            amrex::PinnedArenaAllocator
-        >;
+        using PinnedTile = typename ContainerLike<amrex::PinnedArenaAllocator>::ParticleTileType;
         PinnedTile pinned_tile;
         pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps());
 
         for (int i = 0; i < np; i++)
         {
-            ParticleType p;
-            p.id() = ParticleType::NextID();
-            p.cpu() = amrex::ParallelDescriptor::MyProc();
-            p.pos(RealAoS::x) = x[i];
-            p.pos(RealAoS::y) = y[i];
-            p.pos(RealAoS::t) = t[i];
-            // write position, creating cpu id, and particle id
-            pinned_tile.push_back(p);
+            pinned_tile.push_back_int(IntSoA::id, ParticleType::NextID());
+            pinned_tile.push_back_int(IntSoA::cpu, amrex::ParallelDescriptor::MyProc());
         }
 
-        // write Real attributes (SoA) to particle initialized zero
-        DefineAndReturnParticleTile(0, 0, 0);
-
+        pinned_tile.push_back_real(RealSoA::x, x);
+        pinned_tile.push_back_real(RealSoA::y, y);
+        pinned_tile.push_back_real(RealSoA::t, t);
         pinned_tile.push_back_real(RealSoA::px, px);
         pinned_tile.push_back_real(RealSoA::py, py);
         pinned_tile.push_back_real(RealSoA::pt, pt);