Update Particle Container to Pure SoA

Transition particle containers to pure SoA layouts.

Author: ax3l

cmake/dependencies/ABLASTR.cmake

@@ -178,7 +178,7 @@ set(ImpactX_openpmd_src ""
 set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git"
     CACHE STRING
     "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
-set(ImpactX_ablastr_branch "24.01"
+set(ImpactX_ablastr_branch "94ae11900131846e9f8b3704194673f7f02d8959"
     CACHE STRING
     "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_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):

examples/pytorch_surrogate_model/run_ml_surrogate.py

@@ -11,6 +11,14 @@
 from urllib import request
 
 import numpy as np
+
+try:
+    import cupy as cp
+
+    cupy_available = True
+except ImportError:
+    cupy_available = False
+
 from surrogate_model_definitions import surrogate_model
 
 try:
@@ -20,10 +28,11 @@
     sys.exit(0)
 
 from impactx import (
+    Config,
+    CoordSystem,
     ImpactX,
     ImpactXParIter,
     RefPart,
-    TransformationDirection,
     coordinate_transformation,
     distribution,
     elements,
@@ -79,7 +88,18 @@ def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start):
         self.ds = surrogate_length
 
     def surrogate_push(self, pc, step):
-        array = np.array
+        # CPU/GPU logic
+        if Config.have_gpu:
+            if cupy_available:
+                array = cp.array
+            else:
+                print("Warning: GPU found but cupy not available! Try managed...")
+                array = np.array
+            if Config.gpu_backend == "SYCL":
+                print("Warning: SYCL GPU backend not yet implemented for Python")
+
+        else:
+            array = np.array
 
         ref_part = pc.ref_particle()
         ref_z_i = ref_part.z
@@ -100,26 +120,28 @@ def surrogate_push(self, pc, step):
         ref_part_final = torch.tensor([0, 0, ref_z_f, 0, 0, ref_uz_f])
 
         # transform
-        coordinate_transformation(pc, TransformationDirection.to_fixed_t)
+        coordinate_transformation(pc, direction=CoordSystem.t)
 
         for lvl in range(pc.finest_level + 1):
             for pti in ImpactXParIter(pc, level=lvl):
-                aos = pti.aos()
-                aos_arr = array(aos, copy=False)
-
                 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)
-                data_arr = (
-                    torch.tensor(
-                        np.vstack(
-                            [aos_arr["x"], aos_arr["y"], aos_arr["z"], real_arrays[:3]]
-                        )
-                    )
-                    .float()
-                    .T
+                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)
+                data_arr = torch.stack(
+                    (
+                        x,
+                        y,
+                        t,
+                        px,
+                        py,
+                        py,
+                    ),
+                    dim=0,
                 )
 
                 data_arr[:, 0] += ref_part.x
@@ -147,9 +169,9 @@ def surrogate_push(self, pc, step):
                     data_arr_post_model[:, 3 + ii] -= ref_part_final[3 + ii]
                     data_arr_post_model[:, 3 + ii] /= ref_beta_gamma_final
 
-                aos_arr["x"] = data_arr_post_model[:, 0]
-                aos_arr["y"] = data_arr_post_model[:, 1]
-                aos_arr["z"] = data_arr_post_model[:, 2]
+                x[:] = data_arr_post_model[:, 0]
+                y[:] = data_arr_post_model[:, 1]
+                t[:] = data_arr_post_model[:, 2]
                 px[:] = data_arr_post_model[:, 3]
                 py[:] = data_arr_post_model[:, 4]
                 pt[:] = data_arr_post_model[:, 5]
@@ -174,7 +196,7 @@ def surrogate_push(self, pc, step):
         # ref_part.s += pge1.ds
         # ref_part.t += pge1.ds / ref_beta
 
-        coordinate_transformation(pc, TransformationDirection.to_fixed_s)
+        coordinate_transformation(pc, direction=CoordSystem.s)
         ## Done!
 
 

src/particles/CollectLost.cpp

@@ -27,9 +27,9 @@ namespace impactx
         using DstData = ImpactXParticleContainer::ParticleTileType::ParticleTileDataType;
 
         AMREX_GPU_HOST_DEVICE
-        void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept {
-            dst.m_aos[dst_ip] = src.m_aos[src_ip];
-
+        void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept
+        {
+            dst.m_idcpu[dst_ip] = src.m_idcpu[src_ip];
             for (int j = 0; j < SrcData::NAR; ++j)
                 dst.m_rdata[j][dst_ip] = src.m_rdata[j][src_ip];
             for (int j = 0; j < src.m_num_runtime_real; ++j)
@@ -141,8 +141,7 @@ namespace impactx
                                 // move down
                                 int const new_index = ip - n_removed;
 
-                                ptile_src_data.m_aos[new_index] = ptile_src_data.m_aos[ip];
-
+                                ptile_src_data.m_idcpu[new_index] = ptile_src_data.m_idcpu[ip];
                                 for (int j = 0; j < SrcData::NAR; ++j)
                                     ptile_src_data.m_rdata[j][new_index] = ptile_src_data.m_rdata[j][ip];
                                 for (int j = 0; j < ptile_src_data.m_num_runtime_real; ++j)

src/particles/ImpactXParticleContainer.H

@@ -18,6 +18,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>
@@ -35,43 +36,16 @@ namespace impactx
         t   ///< fixed t as the independent variable
     };
 
-    /** 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)
@@ -80,27 +54,28 @@ 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)
+            nattribs ///< the number of attributes above (always last)
         };
     };
 
@@ -109,46 +84,46 @@ namespace impactx
      * 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 (initialization::AmrCoreData* amr_core);
@@ -276,10 +251,6 @@ namespace impactx
         DepositCharge (std::unordered_map<int, amrex::MultiFab> & rho,
                        amrex::Vector<amrex::IntVect> const & ref_ratio);
 
-        /** Get the name of each Real AoS component */
-        std::vector<std::string>
-        RealAoS_names () const;
-
         /** Get the name of each Real SoA component */
         std::vector<std::string>
         RealSoA_names () const;
@@ -311,10 +282,6 @@ namespace impactx
 
     }; // ImpactXParticleContainer
 
-    /** Get the name of each Real AoS component */
-    std::vector<std::string>
-    get_RealAoS_names ();
-
     /** Get the name of each Real SoA component
      *
      * @param num_real_comps number of compile-time + runtime arrays