Update Beam and Plasma to Pure SoA

src/Hipace.cpp

@@ -1043,7 +1043,8 @@ Hipace::Wait (const int step, int it, bool only_ghost)
     {
         const amrex::Long np_total = std::accumulate(np_rcv.begin(), np_rcv.begin()+nbeams, 0);
         if (np_total == 0 && !m_multi_laser.m_use_laser) return;
-        const amrex::Long psize = sizeof(BeamParticleContainer::SuperParticleType);
+        const amrex::Long psize = sizeof(amrex::ParticleReal)*BeamParticleContainer::NAR
+                                  + sizeof(int)*BeamParticleContainer::NAI;
         const amrex::Long buffer_size = psize*np_total;
         auto recv_buffer = (char*)amrex::The_Pinned_Arena()->alloc(buffer_size);
 
@@ -1236,7 +1237,8 @@ Hipace::Notify (const int step, const int it, bool only_ghost)
     {
         const amrex::Long np_total = std::accumulate(np_snd.begin(), np_snd.begin()+nbeams, 0);
         if (np_total == 0 && !m_multi_laser.m_use_laser) return;
-        const amrex::Long psize = sizeof(BeamParticleContainer::SuperParticleType);
+        const amrex::Long psize = sizeof(amrex::ParticleReal)*BeamParticleContainer::NAR
+                                  + sizeof(int)*BeamParticleContainer::NAI;
         const amrex::Long buffer_size = psize*np_total;
         char*& psend_buffer = only_ghost ? m_psend_buffer_ghost : m_psend_buffer;
         psend_buffer = (char*)amrex::The_Pinned_Arena()->alloc(buffer_size);
@@ -1493,18 +1495,19 @@ Hipace::CheckGhostSlice (int it)
 
         // Get pointers to ghost particles
         auto& ptile = m_multi_beam.getBeam(ibeam);
-        auto& aos = ptile.GetArrayOfStructs();
-        const auto& pos_structs = aos.begin() + nreal;
+        auto& soa = ptile.GetStructOfArrays();
+        const amrex::Real * const pos_z = soa.GetRealData(BeamIdx::z).data() + nreal;
+        int * const idp = soa.GetIntData(BeamIdx::id).data() + nreal;
 
         // Invalidate particles out of the ghost slice
         amrex::ParallelFor(
             nghost,
             [=] AMREX_GPU_DEVICE (long idx) {
                 // Get zp of ghost particle
-                const amrex::Real zp = pos_structs[idx].pos(2);
+                const amrex::Real zp = pos_z[idx];
                 // Invalidate ghost particle if not in the ghost slice
                 if ( zp < zmin_leftcell || zp > zmax_leftcell ) {
-                    pos_structs[idx].id() = -1;
+                    idp[idx] = -1;
                 }
             }
             );

src/diagnostics/OpenPMDWriter.H

@@ -96,7 +96,11 @@ private:
 
     /** Named Beam SoA attributes per particle as defined in BeamIdx
      */
-    amrex::Vector<std::string> m_real_names {"weighting", "momentum_x", "momentum_y", "momentum_z"};
+    amrex::Vector<std::string> m_real_names {
+        "position_x", "position_y", "position_z",
+        "weighting",
+        "momentum_x", "momentum_y", "momentum_z"
+    };
 
     /** vector over levels of openPMD-api Series object for output */
     std::unique_ptr< openPMD::Series > m_outputSeries;

src/diagnostics/OpenPMDWriter.cpp

@@ -17,8 +17,8 @@
 
 OpenPMDWriter::OpenPMDWriter ()
 {
-    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_names.size() == BeamIdx::nattribs,
-        "List of real names in openPMD Writer class do not match BeamIdx::nattribs");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_real_names.size() == BeamIdx::real_nattribs,
+        "List of real names in openPMD Writer class do not match BeamIdx::real_nattribs");
     amrex::ParmParse pp("hipace");
     queryWithParser(pp, "openpmd_backend", m_openpmd_backend);
     // pick first available backend if default is chosen
@@ -207,35 +207,13 @@ OpenPMDWriter::WriteBeamParticleData (MultiBeam& beams, openPMD::Iteration itera
             continue;
         }
 
-        // get position and particle ID from aos
-        // note: this implementation iterates the AoS 4x...
-        // if we flush late as we do now, we can also copy out the data in one go
-        const auto& aos = beam.GetArrayOfStructs();  // size =  numParticlesOnTile
-        const auto& pos_structs = aos.begin() + box_offset;
         {
-            // Save positions
-            std::vector< std::string > const positionComponents{"x", "y", "z"};
-
-            for (auto currDim = 0; currDim < AMREX_SPACEDIM; currDim++)
-            {
-                std::shared_ptr< amrex::ParticleReal > curr(
-                    new amrex::ParticleReal[numParticleOnTile],
-                    [](amrex::ParticleReal const *p){ delete[] p; } );
-
-                for (uint64_t i=0; i<numParticleOnTile; i++) {
-                    curr.get()[i] = pos_structs[i].pos(currDim);
-                }
-                std::string const positionComponent = positionComponents[currDim];
-                beam_species["position"][positionComponent].storeChunk(curr, {m_offset[ibeam]},
-                                                                       {numParticleOnTile64});
-            }
-
             // save particle ID
             std::shared_ptr< uint64_t > ids( new uint64_t[numParticleOnTile],
                                              [](uint64_t const *p){ delete[] p; } );
 
             for (uint64_t i=0; i<numParticleOnTile; i++) {
-                ids.get()[i] = pos_structs[i].id();
+                ids.get()[i] = beam.id(i);
             }
             auto const scalar = openPMD::RecordComponent::SCALAR;
             beam_species["id"][scalar].storeChunk(ids, {m_offset[ibeam]}, {numParticleOnTile64});

src/particles/beam/BeamParticleContainer.H

@@ -21,21 +21,34 @@
 struct BeamIdx
 {
     enum {
-        w = 0,      // weight
+        x=0, y, z,  // position
+        w,          // weight
         ux, uy, uz, // momentum
-        nattribs
+        real_nattribs
+    };
+    enum {
+        id=0, cpu,
+        int_nattribs
     };
 };
 
+using BeamTile = amrex::ParticleTile<
+        amrex::SoAParticle<
+            BeamIdx::real_nattribs,
+            BeamIdx::int_nattribs
+        >,
+        BeamIdx::real_nattribs,
+        BeamIdx::int_nattribs
+    >;
+
 struct BeamBins : amrex::DenseBins<
-        amrex::ParticleTile<amrex::Particle<0, 0>, BeamIdx::nattribs, 0>::ParticleType> {
+        BeamTile::ParticleTileDataType> {
 
     template<class...Args>
     void build (Args&&...args) {
         // call build function of the underlying DenseBins object
         // with all of the arguments forwarded
-        amrex::DenseBins<amrex::ParticleTile<
-            amrex::Particle<0, 0>, BeamIdx::nattribs, 0>::ParticleType>::build(args...);
+        amrex::DenseBins<BeamTile::ParticleTileDataType>::build(args...);
 
         // after every build call copy offsets array form GPU to CPU
         const auto offset_size = numBins() + 1;
@@ -58,13 +71,12 @@ struct BeamBins : amrex::DenseBins<
 };
 
 /** \brief Container for particles of 1 beam species. */
-class BeamParticleContainer
-    : public amrex::ParticleTile<amrex::Particle<0, 0>, BeamIdx::nattribs, 0>
+class BeamParticleContainer : public BeamTile
 {
 public:
     /** Constructor */
     explicit BeamParticleContainer (std::string name) :
-        amrex::ParticleTile<amrex::Particle<0, 0>, BeamIdx::nattribs, 0>(),
+        BeamTile(),
         m_name(name)
     {
         ReadParameters();

src/particles/beam/BeamParticleContainer.cpp

@@ -208,13 +208,13 @@ amrex::Long BeamParticleContainer::TotalNumberOfParticles (bool only_valid, bool
         amrex::ReduceData<unsigned long long> reduce_data(reduce_op);
         using ReduceTuple = typename decltype(reduce_data)::Type;
 
-        auto const& ptaos = this->GetArrayOfStructs();
-        ParticleType const* pp = ptaos().data();
+        auto const& ptsoa = this->GetStructOfArrays();
+        const int * const idp = ptsoa.GetIntData(BeamIdx::id).data();
 
-        reduce_op.eval(ptaos.numParticles(), reduce_data,
+        reduce_op.eval(ptsoa.numParticles(), reduce_data,
                        [=] AMREX_GPU_DEVICE (int i) -> ReduceTuple
                        {
-                           return (pp[i].id() > 0) ? 1 : 0;
+                           return (idp[i] > 0) ? 1 : 0;
                        });
         nparticles = static_cast<amrex::Long>(amrex::get<0>(reduce_data.value()));
     }
@@ -239,7 +239,10 @@ void BeamParticleContainer::TagByLevel (const int current_N_level,
     int box_offset = m_box_sorter.boxOffsetsPtr()[m_ibox];
     if (deposit_ghost) box_offset = numParticles()-nghost;
 
-    const auto pos_structs = GetArrayOfStructs().begin() + box_offset;
+    auto& soa = GetStructOfArrays();
+    const amrex::Real * const pos_x = soa.GetRealData(BeamIdx::x).data() + box_offset;
+    const amrex::Real * const pos_y = soa.GetRealData(BeamIdx::y).data() + box_offset;
+    int * const cpup = soa.GetIntData(BeamIdx::cpu).data() + box_offset;
 
     BeamBins::index_type const * const indices = m_slice_bins.permutationPtr();
     BeamBins::index_type const * const offsets = m_slice_bins.offsetsPtrCpu();
@@ -281,22 +284,22 @@ void BeamParticleContainer::TagByLevel (const int current_N_level,
             // Ghost particles are simply contiguous in memory.
             const int ip = deposit_ghost ? cell_start+idx : indices[cell_start+idx];
 
-            const amrex::Real xp = pos_structs[ip].pos(0);
-            const amrex::Real yp = pos_structs[ip].pos(1);
+            const amrex::Real xp = pos_x[ip];
+            const amrex::Real yp = pos_y[ip];
 
             if (current_N_level > 2 &&
                 lo_x_lev2 < xp && xp < hi_x_lev2 &&
                 lo_y_lev2 < yp && yp < hi_y_lev2) {
                 // level 2
-                pos_structs[ip].cpu() = 2;
+                cpup[ip] = 2;
             } else if (current_N_level > 1 &&
                 lo_x_lev1 < xp && xp < hi_x_lev1 &&
                 lo_y_lev1 < yp && yp < hi_y_lev1) {
                 // level 1
-                pos_structs[ip].cpu() = 1;
+                cpup[ip] = 1;
             } else {
                 // level 0
-                pos_structs[ip].cpu() = 0;
+                cpup[ip] = 0;
             }
         }
     );
@@ -324,13 +327,14 @@ BeamParticleContainer::InSituComputeDiags (int islice, int islice_local)
     const amrex::Real clightsq_inv = 1.0_rt/(phys_const.c*phys_const.c);
 
     const int box_offset = m_box_sorter.boxOffsetsPtr()[m_ibox];
-    auto const& ptaos = this->GetArrayOfStructs();
-    const auto& pos_structs = ptaos.begin() + box_offset;
     auto const& soa = this->GetStructOfArrays();
+    const auto pos_x = soa.GetRealData(BeamIdx::x).data() + box_offset;
+    const auto pos_y = soa.GetRealData(BeamIdx::y).data() + box_offset;
     const auto  wp = soa.GetRealData(BeamIdx::w).data() + box_offset;
     const auto uxp = soa.GetRealData(BeamIdx::ux).data() + box_offset;
     const auto uyp = soa.GetRealData(BeamIdx::uy).data() + box_offset;
     const auto uzp = soa.GetRealData(BeamIdx::uz).data() + box_offset;
+    auto idp = soa.GetIntData(BeamIdx::id).data() + box_offset;
 
     BeamBins::index_type const * const indices = m_slice_bins.permutationPtr();
     BeamBins::index_type const * const offsets = m_slice_bins.offsetsPtrCpu();
@@ -356,24 +360,24 @@ BeamParticleContainer::InSituComputeDiags (int islice, int islice_local)
         [=] AMREX_GPU_DEVICE (int i) -> ReduceTuple
         {
             const int ip = indices[cell_start+i];
-            if (pos_structs[ip].id() < 0) {
+            if (idp[ip] < 0) {
                 return{0._rt, 0._rt, 0._rt, 0._rt, 0._rt, 0._rt, 0._rt,
                     0._rt, 0._rt, 0._rt, 0._rt, 0._rt, 0._rt, 0};
             }
             const amrex::Real gamma = std::sqrt(1.0_rt + uxp[ip]*uxp[ip]*clightsq_inv
                                                        + uyp[ip]*uyp[ip]*clightsq_inv
                                                        + uzp[ip]*uzp[ip]*clightsq_inv);
             return {wp[ip],
-                    wp[ip]*pos_structs[ip].pos(0),
-                    wp[ip]*pos_structs[ip].pos(0)*pos_structs[ip].pos(0),
-                    wp[ip]*pos_structs[ip].pos(1),
-                    wp[ip]*pos_structs[ip].pos(1)*pos_structs[ip].pos(1),
+                    wp[ip]*pos_x[ip],
+                    wp[ip]*pos_x[ip]*pos_x[ip],
+                    wp[ip]*pos_y[ip],
+                    wp[ip]*pos_y[ip]*pos_y[ip],
                     wp[ip]*uxp[ip]*clight_inv,
                     wp[ip]*uxp[ip]*uxp[ip]*clightsq_inv,
                     wp[ip]*uyp[ip]*clight_inv,
                     wp[ip]*uyp[ip]*uyp[ip]*clightsq_inv,
-                    wp[ip]*pos_structs[ip].pos(0)*uxp[ip]*clight_inv,
-                    wp[ip]*pos_structs[ip].pos(1)*uyp[ip]*clight_inv,
+                    wp[ip]*pos_x[ip]*uxp[ip]*clight_inv,
+                    wp[ip]*pos_y[ip]*uyp[ip]*clight_inv,
                     wp[ip]*gamma,
                     wp[ip]*gamma*gamma,
                     1};