SoA: InitRandom/RandomPerBox/OnePerCell

ax3l · ax3l · commit 54d0623298b3 · 2023-05-30T15:29:02.000-07:00
Support more init methods on SoA particle containers.
diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H
@@ -106,11 +106,23 @@ struct ParticleLocData
 * a given component, then the extra values will be set to zero. If more components
 * are specified, it is a compile-time error.
 *
+* Note that the position attributes, first AMREX_SPACEDIM in ParticleReal, and the IDs, first
+* two ints, are skipped here.
+* Note that the counts in the type of the SoAParticle definition are explicit about these default-required
+* attributes.
+*
 * Example usage:
 *
-*     ParticleInitType<0, 2, 4, 1> pdata = {{}, {7, 9}, {1.5, 2.5, 3.5, 4.5}, {11}};
+*     ParticleInitType<Particle<0, 2>, 4, 1> pdata = {{}, {7, 9}, {1.5, 2.5, 3.5, 4.5}, {11}};
+*     ParticleInitTypeSoA<4, 3> pdata = {{1.5, 2.5, 3.5, 4.5}, {7, 9, 11}};  // 3D: SoAParticle<7, 5>
 */
-template<int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
+template <int NArrayReal, int NArrayInt>
+struct ParticleInitTypeSoA
+{
+    std::array<double, NArrayReal > real_array_data;
+    std::array<int,    NArrayInt  > int_array_data;
+};
+template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
 struct ParticleInitType
 {
     std::array<double, NStructReal> real_struct_data;
@@ -177,7 +189,13 @@ public:
 
     using ParticleContainerType = ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>;
     using ParticleTileType = ParticleTile<ParticleType, NArrayReal, NArrayInt, Allocator>;
-    using ParticleInitData = ParticleInitType<NStructReal, NStructInt, NArrayReal, NArrayInt>;
+    using ParticleInitData = typename std::conditional<
+        T_ParticleType::is_soa_particle,
+        // SoA Particle: init w/o positions and id/cpu, but explicitly listed in define
+        ParticleInitTypeSoA<NArrayReal - AMREX_SPACEDIM, NArrayInt - 2>,
+        // legacy particle: init w/o positions and id/cpu, only implicitly listed in define
+        ParticleInitType<T_ParticleType::NReal, T_ParticleType::NInt, NArrayReal, NArrayInt>
+    >::type;
 
     //! A single level worth of particles is indexed (grid id, tile id)
     //! for both SoA and AoS data.
diff --git a/Src/Particle/AMReX_ParticleInit.H b/Src/Particle/AMReX_ParticleInit.H
@@ -1063,8 +1063,27 @@ InitRandom (Long                    icount,
 
             if (who == MyProc) {
 
-                if constexpr(!ParticleType::is_soa_particle)
-                {
+                if constexpr(ParticleType::is_soa_particle) {
+                    for (int i = 0; i < AMREX_SPACEDIM; i++) {
+                        host_real_attribs[pld.m_lev][ind][i].push_back(pos[j*AMREX_SPACEDIM+i]);
+                    }
+
+                    host_int_attribs[pld.m_lev][ind][0].push_back(ParticleType::NextID());
+                    host_int_attribs[pld.m_lev][ind][1].push_back(MyProc);
+
+                    host_particles[pld.m_lev][ind];
+
+                    // add the real...
+                    for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
+                        host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
+                    }
+
+                    // ... and int array data
+                    for (int i = 2; i < NArrayInt; i++) {
+                        host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
+                    }
+                }
+                else {
                     ParticleType p;
                     for (int i = 0; i < AMREX_SPACEDIM; i++) {
                         p.pos(i) = pos[j*AMREX_SPACEDIM + i];
@@ -1094,25 +1113,6 @@ InitRandom (Long                    icount,
                     for (int i = 0; i < NArrayInt; i++) {
                         host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
                     }
-                } else {
-                    for (int i = 0; i < AMREX_SPACEDIM; i++) {
-                        host_real_attribs[pld.m_lev][ind][i].push_back(pos[j*AMREX_SPACEDIM+i]);
-                    }
-
-                    host_int_attribs[pld.m_lev][ind][0].push_back(ParticleType::NextID());
-                    host_int_attribs[pld.m_lev][ind][1].push_back(MyProc);
-
-                    host_particles[pld.m_lev][ind];
-
-                    // add the real...
-                    for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
-                        host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
-                    }
-
-                    // ... and int array data
-                    for (int i = 2; i < NArrayInt; i++) {
-                        host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
-                    }
                 }
             }
         }
@@ -1127,11 +1127,11 @@ InitRandom (Long                    icount,
                 auto& dst_tile = GetParticles(host_lev)[std::make_pair(grid,tile)];
                 auto old_size = dst_tile.GetArrayOfStructs().size();
                 auto new_size = old_size;
-                if constexpr(!ParticleType::is_soa_particle)
+                if constexpr(ParticleType::is_soa_particle)
                 {
-                    new_size += src_tile.size();
-                } else {
                     new_size += host_real_attribs[host_lev][std::make_pair(grid,tile)][0].size();
+                } else {
+                    new_size += src_tile.size();
                 }
                 dst_tile.resize(new_size);
 
@@ -1362,44 +1362,71 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
         for (Long jcnt = 0; jcnt < icount_per_box; jcnt++) {
         for (Long kcnt = 0; kcnt < icount_per_box; kcnt++)
         {
-            p.pos(0) = static_cast<ParticleReal>(grid_box.lo(0) + (dist(mt) + icnt) / icount_per_box * grid_box.length(0));
-            p.pos(1) = static_cast<ParticleReal>(grid_box.lo(1) + (dist(mt) + jcnt) / icount_per_box * grid_box.length(1));
-            p.pos(2) = static_cast<ParticleReal>(grid_box.lo(2) + (dist(mt) + kcnt) / icount_per_box * grid_box.length(2));
+            // the position data
+            for (int d = 0; d < AMREX_SPACEDIM; d++) {
+                p.pos(d) = static_cast<ParticleReal>(grid_box.lo(d) + (dist(mt) + icnt) / icount_per_box * grid_box.length(d));
+            }
 
-            for (int i = 0; i < AMREX_SPACEDIM; i++)
-                AMREX_ASSERT(p.pos(i) < grid_box.hi(i));
+            for (int d = 0; d < AMREX_SPACEDIM; d++)
+                AMREX_ASSERT(p.pos(d) < grid_box.hi(d));
 
-            // the real struct data
-            for (int i = 0; i < NStructReal; i++) {
-                p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
-            }
+            if constexpr(ParticleType::is_soa_particle) {
+                if (!Where(p, pld)) {
+                    amrex::Abort("ParticleContainer::InitRandom(): invalid particle");
+                }
+                AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
+                std::pair<int, int> ind(pld.m_grid, pld.m_tile);
 
-            // the int struct data
-            p.id()  = ParticleType::NextID();
-            p.cpu() = ParallelDescriptor::MyProc();
+                // IDs
+                p.id()  = ParticleType::NextID();
+                p.cpu() = ParallelDescriptor::MyProc();
 
-            for (int i = 0; i < NStructInt; i++) {
-                p.idata(i) = pdata.int_struct_data[i];
-            }
+                // add the real (after position)
+                for (int i = AMREX_SPACEDIM; i < NArrayReal; i++) {
+                    m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
+                }
 
-            // locate the particle
-            if (!Where(p, pld)) {
-                amrex::Abort("ParticleContainer::InitRandomPerBox(): invalid particle");
+                // add the int array data (after id, cpu)
+                for (int i = 2; i < NArrayInt; i++) {
+                    m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
+                }
+
+                // add
+                m_particles[pld.m_lev][ind].push_back(p);
             }
-            AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
-            std::pair<int, int> ind(pld.m_grid, pld.m_tile);
+            else {
+                // the real struct data
+                for (int i = 0; i < NStructReal; i++) {
+                    p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
+                }
 
-            // add the struct
-            m_particles[pld.m_lev][ind].push_back(p);
+                // the int struct data
+                p.id()  = ParticleType::NextID();
+                p.cpu() = ParallelDescriptor::MyProc();
 
-            // add the real...
-            for (int i = 0; i < NArrayReal; i++) {
-                m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
-            }
+                for (int i = 0; i < NStructInt; i++) {
+                    p.idata(i) = pdata.int_struct_data[i];
+                }
 
-            // ... and int array data
-            for (int i = 0; i < NArrayInt; i++) {
-                m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
+                // locate the particle
+                if (!Where(p, pld)) {
+                    amrex::Abort("ParticleContainer::InitRandomPerBox(): invalid particle");
+                }
+                AMREX_ASSERT(pld.m_lev >= 0 && pld.m_lev <= finestLevel());
+                std::pair<int, int> ind(pld.m_grid, pld.m_tile);
+
+                // add the struct
+                m_particles[pld.m_lev][ind].push_back(p);
+
+                // add the real...
+                for (int i = 0; i < NArrayReal; i++) {
+                    m_particles[pld.m_lev][ind].push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
+                }
+
+                // ... and int array data
+                for (int i = 0; i < NArrayInt; i++) {
+                    m_particles[pld.m_lev][ind].push_back_int(i, pdata.int_array_data[i]);
+                }
             }
 
         } } }
@@ -1438,48 +1465,64 @@ InitOnePerCell (Real x_off, Real y_off, Real z_off, const ParticleInitData& pdat
 
     const Real* dx  = geom.CellSize();
 
-    ParticleType p;
-
     // We'll generate the particles in parallel -- but no tiling of the grid here.
     for (MFIter mfi(*m_dummy_mf[0], false); mfi.isValid(); ++mfi) {
         Box grid = ParticleBoxArray(0)[mfi.index()];
         auto ind = std::make_pair(mfi.index(), mfi.LocalTileIndex());
         RealBox grid_box (grid,dx,geom.ProbLo());
+
+        // tile for one particle
         ParticleTile<ParticleType, NArrayReal, NArrayInt, amrex::PinnedArenaAllocator> ptile_tmp;
+        ptile_tmp.resize(1);
+        auto ptd = ptile_tmp.getParticleTileData();
+
         for (IntVect beg = grid.smallEnd(), end=grid.bigEnd(), cell = grid.smallEnd(); cell <= end; grid.next(cell))
         {
-            // the real struct data
-            AMREX_D_TERM(p.pos(0) = static_cast<ParticleReal>(grid_box.lo(0) + (x_off + cell[0]-beg[0])*dx[0]);,
-                         p.pos(1) = static_cast<ParticleReal>(grid_box.lo(1) + (y_off + cell[1]-beg[1])*dx[1]);,
-                         p.pos(2) = static_cast<ParticleReal>(grid_box.lo(2) + (z_off + cell[2]-beg[2])*dx[2]););
+            // particle index
+            constexpr int i = 0;
+
+            // the position data
+            for (int d = 0; d < AMREX_SPACEDIM; d++) {
+                ptile_tmp.pos(i, d) = static_cast<ParticleReal>(grid_box.lo(d) + (x_off + cell[d]-beg[d])*dx[d]);
+            }
 
             for (int d = 0; d < AMREX_SPACEDIM; ++d) {
-                AMREX_ASSERT(p.pos(d) < grid_box.hi(d));
+                AMREX_ASSERT(ptile_tmp.pos(i, d) < grid_box.hi(d));
             }
 
-            for (int i = 0; i < NStructReal; i++) {
-                p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
+            if constexpr(!ParticleType::is_soa_particle) {
+                for (int n = 0; n < NStructReal; n++) {
+                    ptd.rdata(n)[i] = static_cast<ParticleReal>(pdata.real_struct_data[n]);
+                }
             }
 
             // the int struct data
-            p.id()  = ParticleType::NextID();
-            p.cpu() = ParallelDescriptor::MyProc();
-
-            for (int i = 0; i < NStructInt; i++) {
-                p.idata(i) = pdata.int_struct_data[i];
+            if constexpr(ParticleType::is_soa_particle) {
+                ptd.idata(0)[i] = ParticleType::NextID();
+                ptd.idata(1)[i] = ParallelDescriptor::MyProc();
+            }
+            else {
+                auto& p = make_particle<ParticleType>{}(ptd, i);
+                p.id() = ParticleType::NextID();
+                p.cpu() = ParallelDescriptor::MyProc();
             }
 
-            // add the struct
-            ptile_tmp.push_back(p);
+            if constexpr(!ParticleType::is_soa_particle) {
+                for (int n = 0; n < NStructInt; n++) {
+                    ptd.idata(n)[i] = pdata.int_struct_data[n];
+                }
+            }
 
             // add the real...
-            for (int i = 0; i < NArrayReal; i++) {
-                ptile_tmp.push_back_real(i, static_cast<ParticleReal>(pdata.real_array_data[i]));
+            int n_min_real = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // jump over position
+            for (int n = n_min_real; n < NArrayReal; n++) {
+                ptile_tmp.push_back_real(n, static_cast<ParticleReal>(pdata.real_array_data[n]));
             }
 
             // ... and int array data
-            for (int i = 0; i < NArrayInt; i++) {
-                ptile_tmp.push_back_int(i, pdata.int_array_data[i]);
+            int n_min_int = ParticleType::is_soa_particle ? 2 : 0; // jump over cpuid
+            for (int n = n_min_int; n < NArrayInt; n++) {
+                ptile_tmp.push_back_int(n, pdata.int_array_data[n]);
             }
         }
 
diff --git a/Tests/Particles/AssignDensity/main.cpp b/Tests/Particles/AssignDensity/main.cpp
@@ -21,7 +21,7 @@ void test_assign_density(TestParams& parms)
 {
 
   RealBox real_box;
-  for (int n = 0; n < BL_SPACEDIM; n++) {
+  for (int n = 0; n < AMREX_SPACEDIM; n++) {
     real_box.setLo(n, 0.0);
     real_box.setHi(n, 1.0);
   }
@@ -48,10 +48,10 @@ void test_assign_density(TestParams& parms)
   MultiFab density(ba, dmap, 1, 0);
   density.setVal(0.0);
 
-  MultiFab partMF(ba, dmap, 1 + BL_SPACEDIM, 1);
+  MultiFab partMF(ba, dmap, 1 + AMREX_SPACEDIM, 1);
   partMF.setVal(0.0);
 
-  typedef ParticleContainer<1 + BL_SPACEDIM> MyParticleContainer;
+  typedef ParticleContainer<1 + AMREX_SPACEDIM> MyParticleContainer;
   MyParticleContainer myPC(geom, dmap, ba);
   myPC.SetVerbose(false);
 
diff --git a/Tests/Particles/InitRandom/main.cpp b/Tests/Particles/InitRandom/main.cpp
@@ -60,8 +60,8 @@ void testSOA ()
     }
 
     // Add some particles
-    constexpr int NReal  = 12;
-    constexpr int NInt   = 4;
+    constexpr int NReal  = AMREX_SPACEDIM + 12;
+    constexpr int NInt   = 2 + 4;
 
     typedef ParticleContainerPureSoA<NReal, NInt> MyPC;
     MyPC myPC(geom, dmap, ba, ref_ratio);
@@ -70,7 +70,7 @@ void testSOA ()
     int num_particles = nppc * AMREX_D_TERM(ncells, * ncells, * ncells);
     bool serialize = false;
     int iseed = 451;
-    MyPC::ParticleInitData pdata = {{}, {},
+    MyPC::ParticleInitData pdata = {
                       {1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0},
                       {5, 14, 15, 16}};
 
diff --git a/Tests/Particles/ParticleIterator/main.cpp b/Tests/Particles/ParticleIterator/main.cpp
@@ -1,10 +1,11 @@
+#include "AMReX_Vector.H"
+#include "AMReX_FabArray.H"
+#include "AMReX_Particles.H"
+
 #include <iostream>
 #include <map>
 #include <vector>
 
-#include <AMReX_Vector.H>
-#include "AMReX_FabArray.H"
-#include "AMReX_Particles.H"
 
 using namespace amrex;
 
@@ -19,7 +20,7 @@ int main(int argc, char* argv[])
   int coord = 0;
 
   RealBox real_box;
-  for (int n = 0; n < BL_SPACEDIM; n++)
+  for (int n = 0; n < AMREX_SPACEDIM; n++)
     {
       real_box.setLo(n,0.0);
       real_box.setHi(n,1.0);
@@ -49,10 +50,10 @@ int main(int argc, char* argv[])
   for (int lev = 0; lev < nlevs; lev++)
     dmap[lev].define(ba[lev]);
 
-  typedef ParticleContainer<1+BL_SPACEDIM> MyParticleContainer;
+  using MyParticleContainer = ParticleContainer<1+AMREX_SPACEDIM>;
   MyParticleContainer MyPC(geom, dmap, ba, rr);
 
-  MyParticleContainer::ParticleInitData pdata = {{1.0},{},{},{}};
+  MyParticleContainer::ParticleInitData pdata = {{1.0, AMREX_D_DECL(1.0, 2.0, 3.0)},{},{},{}};
   MyPC.InitOnePerCell(0.5, 0.5, 0.5, pdata);
   MyParticleContainer::do_tiling = true;
 
@@ -61,7 +62,7 @@ int main(int argc, char* argv[])
 #ifdef AMREX_USE_OMP
 #pragma omp parallel
 #endif
-  for (ParIter<1+BL_SPACEDIM> mfi(MyPC, 0); mfi.isValid(); ++mfi) {
+  for (ParIter<1+AMREX_SPACEDIM> mfi(MyPC, 0); mfi.isValid(); ++mfi) {
       amrex::AllPrintToFile("particle_iterator_out") << mfi.index() << " " << mfi.tileIndex() << "\n";
   }
   }
