Repository: AMReX-Codes/amrex — pull request #3325 (cpp)
Title: Port InitRandom for pure SoA

This is useful for porting and testing other features, like Checkpoint / Restart. I've also added a test.

Follow-up to #2878.

The proposed changes:
- [x] fix a bug or incorrect behavior in AMReX
- [x] add new capabilities to AMReX
- [ ] changes answers in the test suite to more than roundoff level
- [ ] are likely to significantly affect the results of downstream AMReX users
- [ ] include documentation in the code and/or rst files, if appropriate


Resolved issues:
title: Support SoA-Only Particles
body: ## Summary

Further updates of  #2853 where @atmyers and @ax3l can also push to.

## Upgrade Guide for C++ API Changes

These PRs show changes to keep using the *old layout*:
- example in WarpX: https://github.com/ECP-WarpX/WarpX/pull/3848
- example in HiPACE++: https://github.com/Hi-PACE/hipace/pull/921
- example in ImpactX: https://github.com/ECP-WarpX/impactx/pull/342
- example in pyAMReX: https://github.com/AMReX-Codes/pyamrex/pull/119

These PRs show changes to *use the new layout* (pure SoA):
- example in WarpX: https://github.com/ECP-WarpX/WarpX/pull/3850
- example in ImpactX: https://github.com/ECP-WarpX/impactx/pull/348
- example in pyAMReX: https://github.com/AMReX-Codes/pyamrex/pull/124

## Additional background

## Checklist

The proposed changes:
- [ ] fix a bug or incorrect behavior in AMReX
- [x] add new capabilities to AMReX
- [ ] changes answers in the test suite to more than roundoff level
- [ ] are likely to significantly affect the results of downstream AMReX users
- [ ] include documentation in the code and/or rst files, if appropriate


Hints:

Fix patch:
diff --git a/Src/Particle/AMReX_ParticleInit.H b/Src/Particle/AMReX_ParticleInit.H
index 60a20c0e835..28427de32fb 100644
--- a/Src/Particle/AMReX_ParticleInit.H
+++ b/Src/Particle/AMReX_ParticleInit.H
@@ -1046,17 +1046,13 @@ InitRandom (Long                    icount,
 
         for (Long j = 0; j < icount; j++)
         {
-            ParticleType p;
+            Particle<0, 0> ptest;
 
             for (int i = 0; i < AMREX_SPACEDIM; i++) {
-                p.pos(i) = pos[j*AMREX_SPACEDIM + i];
+                ptest.pos(i) = pos[j*AMREX_SPACEDIM + i];
             }
 
-            for (int i = 0; i < NStructReal; i++) {
-                p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
-            }
-
-            if (!Where(p, pld)) {
+            if (!Where(ptest, pld)) {
                 amrex::Abort("ParticleContainer::InitRandom(): invalid particle");
             }
 
@@ -1067,27 +1063,58 @@ InitRandom (Long                    icount,
 
             if (who == MyProc) {
 
-                // We own it. Add it at the appropriate level.
-                p.id()  = ParticleType::NextID();
-                p.cpu() = MyProc;
+                if constexpr(!ParticleType::is_soa_particle)
+                {
+                    ParticleType p;
+                    for (int i = 0; i < AMREX_SPACEDIM; i++) {
+                        p.pos(i) = pos[j*AMREX_SPACEDIM + i];
+                    }
 
-                for (int i = 0; i < NStructInt; i++) {
-                    p.idata(i) = pdata.int_struct_data[i];
-                }
+                    // We own it. Add it at the appropriate level.
+                    p.id()  = ParticleType::NextID();
+                    p.cpu() = MyProc;
 
-                // add the struct
-                host_particles[pld.m_lev][ind].push_back(p);
+                    for (int i = 0; i < NStructInt; i++) {
+                        p.idata(i) = pdata.int_struct_data[i];
+                    }
 
-                // add the real...
-                for (int i = 0; i < NArrayReal; i++) {
-                    host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
-                }
+                    for (int i = 0; i < NStructReal; i++) {
+                        p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
+                    }
 
-                // ... and int array data
-                for (int i = 0; i < NArrayInt; i++) {
-                    host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
+                    // add the struct
+                    host_particles[pld.m_lev][ind].push_back(p);
+
+                    // add the real...
+                    for (int i = 0; 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 = 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]);
+                    }
                 }
-              }
+            }
         }
 
         for (int host_lev = 0; host_lev < static_cast<int>(host_particles.size()); ++host_lev)
@@ -1099,11 +1126,20 @@ 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 + src_tile.size();
+                auto new_size = old_size;
+                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();
+                }
                 dst_tile.resize(new_size);
 
-                Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
-                               dst_tile.GetArrayOfStructs().begin() + old_size);
+                if constexpr(!ParticleType::is_soa_particle)
+                {
+                    Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
+                                   dst_tile.GetArrayOfStructs().begin() + old_size);
+                }
 
                 for (int i = 0; i < NArrayReal; ++i) {
                     Gpu::copyAsync(Gpu::hostToDevice,
@@ -1148,7 +1184,7 @@ InitRandom (Long                    icount,
         host_int_attribs.resize(finestLevel()+1);
 
         for (Long icnt = 0; icnt < M; icnt++) {
-            ParticleType p;
+            Particle<0, 0> ptest;
             for (int i = 0; i < AMREX_SPACEDIM; i++) {
                 do {
                     r = amrex::Random();
@@ -1156,42 +1192,71 @@ InitRandom (Long                    icount,
                 }
                 while (static_cast<ParticleReal>(x) < static_cast<ParticleReal>(xlo[i]) || static_cast<ParticleReal>(x) >= static_cast<ParticleReal>(xhi[i]));
 
-                p.pos(i) = static_cast<ParticleReal>(x);
+                ptest.pos(i) = static_cast<ParticleReal>(x);
 
-                AMREX_ASSERT(p.pos(i) < geom.ProbHi(i));
+                AMREX_ASSERT(ptest.pos(i) < geom.ProbHi(i));
             }
 
-            for (int i = 0; i < NStructReal; i++) {
-                p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
-            }
-
-            // 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];
-            }
+            ptest.id()  = ParticleType::NextID();
+            ptest.cpu() = ParallelDescriptor::MyProc();
 
             // locate the particle
-            if (!Where(p, pld))
+            if (!Where(ptest, pld))
             {
                 amrex::Abort("ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt>::InitRandom(): 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
-            host_particles[pld.m_lev][ind].push_back(p);
+            if constexpr(!ParticleType::is_soa_particle)
+            {
+                ParticleType p;
+                for (int i = 0; i < AMREX_SPACEDIM; i++) {
+                    p.pos(i) = ptest.pos(i);;
+                }
 
-            // add the real...
-            for (int i = 0; i < NArrayReal; i++) {
-                host_real_attribs[pld.m_lev][ind][i].push_back(static_cast<ParticleReal>(pdata.real_array_data[i]));
-            }
+                p.id()  = ptest.id();
+                p.cpu() = ptest.cpu();
 
-            // ... and int array data
-            for (int i = 0; i < NArrayInt; i++) {
-                host_int_attribs[pld.m_lev][ind][i].push_back(pdata.int_array_data[i]);
+                for (int i = 0; i < NStructReal; i++) {
+                    p.rdata(i) = static_cast<ParticleReal>(pdata.real_struct_data[i]);
+                }
+
+                for (int i = 0; i < NStructInt; i++) {
+                    p.idata(i) = pdata.int_struct_data[i];
+                }
+
+                // add the struct
+                host_particles[pld.m_lev][ind].push_back(p);
+
+                // add the real...
+                for (int i = 0; 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 = 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(ptest.pos(i));
+                }
+
+                host_int_attribs[pld.m_lev][ind][0].push_back(ptest.id());
+                host_int_attribs[pld.m_lev][ind][1].push_back(ptest.cpu());
+
+                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]);
+                }
             }
         }
 
@@ -1204,11 +1269,20 @@ 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 + src_tile.size();
+                auto new_size = old_size;
+                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();
+                }
                 dst_tile.resize(new_size);
 
-                Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
-                               dst_tile.GetArrayOfStructs().begin() + old_size);
+                if constexpr(!ParticleType::is_soa_particle)
+                {
+                    Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
+                                   dst_tile.GetArrayOfStructs().begin() + old_size);
+                }
 
                 for (int i = 0; i < NArrayReal; ++i) {
                     Gpu::copyAsync(Gpu::hostToDevice,


Test patch:
diff --git a/Tests/Particles/InitRandom/CMakeLists.txt b/Tests/Particles/InitRandom/CMakeLists.txt
new file mode 100644
index 00000000000..c032392c5df
--- /dev/null
+++ b/Tests/Particles/InitRandom/CMakeLists.txt
@@ -0,0 +1,14 @@
+# This tests requires particle support
+if (NOT AMReX_PARTICLES)
+   return()
+endif ()
+
+foreach(D IN LISTS AMReX_SPACEDIM)
+    set(_sources     main.cpp)
+    set(_input_files inputs  )
+
+    setup_test(${D} _sources _input_files)
+
+    unset(_sources)
+    unset(_input_files)
+endforeach()
diff --git a/Tests/Particles/InitRandom/GNUmakefile b/Tests/Particles/InitRandom/GNUmakefile
new file mode 100644
index 00000000000..b3510b88d2e
--- /dev/null
+++ b/Tests/Particles/InitRandom/GNUmakefile
@@ -0,0 +1,31 @@
+AMREX_HOME = ../../../
+
+# DEBUG	= TRUE
+DEBUG	= FALSE
+
+DIM	= 3
+
+COMP    = gnu
+
+PRECISION = DOUBLE
+
+USE_MPI   = TRUE
+MPI_THREAD_MULTIPLE = FALSE
+
+USE_OMP   = FALSE
+
+TINY_PROFILE = TRUE
+
+USE_PARTICLES = TRUE
+
+###################################################
+
+EBASE     = main
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include ./Make.package
+include $(AMREX_HOME)/Src/Base/Make.package
+include $(AMREX_HOME)/Src/Particle/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules
diff --git a/Tests/Particles/InitRandom/Make.package b/Tests/Particles/InitRandom/Make.package
new file mode 100644
index 00000000000..6b4b865e8fc
--- /dev/null
+++ b/Tests/Particles/InitRandom/Make.package
@@ -0,0 +1,1 @@
+CEXE_sources += main.cpp
diff --git a/Tests/Particles/InitRandom/inputs b/Tests/Particles/InitRandom/inputs
new file mode 100644
index 00000000000..59b4212ed5b
--- /dev/null
+++ b/Tests/Particles/InitRandom/inputs
@@ -0,0 +1,13 @@
+# Domain size
+ncells = 64
+
+# Maximum allowable size of each subdomain in the problem domain;
+# this is used to decompose the domain for parallel calculations.
+max_grid_size = 8
+
+# Number of levels
+nlevs = 1
+
+# Number of particles per cell
+nppc = 2
+
diff --git a/Tests/Particles/InitRandom/main.cpp b/Tests/Particles/InitRandom/main.cpp
new file mode 100644
index 00000000000..c3c9f730b53
--- /dev/null
+++ b/Tests/Particles/InitRandom/main.cpp
@@ -0,0 +1,187 @@
+#include <AMReX.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_Particles.H>
+
+#include <cstdio>
+
+using namespace amrex;
+
+void set_grids_nested (Vector<Box>& domains,
+                       Vector<BoxArray>& grids,
+                       Vector<IntVect>& ref_ratio);
+void test ();
+void testSOA ();
+
+int main(int argc, char* argv[])
+{
+    amrex::Initialize(argc,argv);
+    test();
+    testSOA();
+    amrex::Finalize();
+}
+
+void testSOA ()
+{
+    int ncells, max_grid_size, nlevs, nppc;
+
+    ParmParse pp;
+    pp.get("ncells", ncells);
+    pp.get("max_grid_size", max_grid_size);
+    pp.get("nlevs", nlevs);
+    pp.get("nppc", nppc);
+
+    Vector<Box> domains;
+    Vector<BoxArray> ba;
+    Vector<IntVect> ref_ratio;
+
+    set_grids_nested(domains, ba, ref_ratio);
+
+    RealBox real_box;
+    for (int n = 0; n < AMREX_SPACEDIM; n++) {
+        real_box.setLo(n, 0.0);
+        real_box.setHi(n, 1.0);
+    }
+
+    // This sets the boundary conditions to be doubly or triply periodic
+    int is_per[] = {AMREX_D_DECL(1,1,1)};
+
+    // This defines a Geometry object for each level
+    Vector<Geometry> geom(nlevs);
+    geom[0].define(domains[0], &real_box, CoordSys::cartesian, is_per);
+    for (int lev = 1; lev < nlevs; lev++) {
+        geom[lev].define(domains[lev], &real_box, CoordSys::cartesian, is_per);
+    }
+
+    Vector<DistributionMapping> dmap(nlevs);
+
+    for (int lev = 0; lev < nlevs; lev++) {
+        dmap[lev] = DistributionMapping{ba[lev]};
+    }
+
+    // Add some particles
+    constexpr int NReal  = 12;
+    constexpr int NInt   = 4;
+
+    typedef ParticleContainerPureSoA<NReal, NInt> MyPC;
+    MyPC myPC(geom, dmap, ba, ref_ratio);
+    myPC.SetVerbose(false);
+
+    int num_particles = nppc * AMREX_D_TERM(ncells, * ncells, * ncells);
+    bool serialize = false;
+    int iseed = 451;
+    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}};
+
+    myPC.InitRandom(num_particles, iseed, pdata, serialize);
+    amrex::Print() << "Generated " << myPC.TotalNumberOfParticles() << " particles. \n";
+}
+
+void test ()
+{
+    int ncells, max_grid_size, nlevs, nppc;
+
+    ParmParse pp;
+    pp.get("ncells", ncells);
+    pp.get("max_grid_size", max_grid_size);
+    pp.get("nlevs", nlevs);
+    pp.get("nppc", nppc);
+
+    Vector<Box> domains;
+    Vector<BoxArray> ba;
+    Vector<IntVect> ref_ratio;
+
+    set_grids_nested(domains, ba, ref_ratio);
+
+    RealBox real_box;
+    for (int n = 0; n < AMREX_SPACEDIM; n++) {
+        real_box.setLo(n, 0.0);
+        real_box.setHi(n, 1.0);
+    }
+
+    // This sets the boundary conditions to be doubly or triply periodic
+    int is_per[] = {AMREX_D_DECL(1,1,1)};
+
+    // This defines a Geometry object for each level
+    Vector<Geometry> geom(nlevs);
+    geom[0].define(domains[0], &real_box, CoordSys::cartesian, is_per);
+    for (int lev = 1; lev < nlevs; lev++) {
+        geom[lev].define(domains[lev], &real_box, CoordSys::cartesian, is_per);
+    }
+
+    Vector<DistributionMapping> dmap(nlevs);
+
+    for (int lev = 0; lev < nlevs; lev++) {
+        dmap[lev] = DistributionMapping{ba[lev]};
+    }
+
+    // Add some particles
+    constexpr int NStructReal = 4;
+    constexpr int NStructInt  = 1;
+    constexpr int NArrayReal  = 8;
+    constexpr int NArrayInt   = 3;
+
+    typedef ParticleContainer<NStructReal, NStructInt, NArrayReal, NArrayInt> MyPC;
+    MyPC myPC(geom, dmap, ba, ref_ratio);
+    myPC.SetVerbose(false);
+
+    int num_particles = nppc * AMREX_D_TERM(ncells, * ncells, * ncells);
+    bool serialize = false;
+    int iseed = 451;
+    MyPC::ParticleInitData pdata = {{1.0, 2.0, 3.0, 4.0},
+                                    {5},
+                                    {6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0},
+                                    {14, 15, 16}};
+
+    myPC.InitRandom(num_particles, iseed, pdata, serialize);
+    amrex::Print() << "Generated " << myPC.TotalNumberOfParticles() << " particles. \n";
+}
+
+void set_grids_nested (Vector<Box>& domains,
+                       Vector<BoxArray>& grids,
+                       Vector<IntVect>& ref_ratio)
+{
+    int ncells, max_grid_size, nlevs;
+
+    ParmParse pp;
+    pp.get("ncells", ncells);
+    pp.get("max_grid_size", max_grid_size);
+    pp.get("nlevs", nlevs);
+
+    AMREX_ALWAYS_ASSERT(nlevs < 2); // relax this later
+
+    IntVect domain_lo(AMREX_D_DECL(0, 0, 0));
+    IntVect domain_hi(AMREX_D_DECL(ncells-1, ncells-1, ncells-1));
+
+    domains.resize(nlevs);
+    domains[0].setSmall(domain_lo);
+    domains[0].setBig(domain_hi);
+
+    ref_ratio.resize(nlevs-1);
+    for (int lev = 1; lev < nlevs; lev++)
+        ref_ratio[lev-1] = IntVect(AMREX_D_DECL(2, 2, 2));
+
+    grids.resize(nlevs);
+    grids[0].define(domains[0]);
+
+    // Now we make the refined level be the center eighth of the domain
+    if (nlevs > 1) {
+        int n_fine = ncells*ref_ratio[0][0];
+        IntVect refined_lo(AMREX_D_DECL(n_fine/4,n_fine/4,n_fine/4));
+        IntVect refined_hi(AMREX_D_DECL(3*n_fine/4-1,3*n_fine/4-1,3*n_fine/4-1));
+
+        // Build a box for the level 1 domain
+        Box refined_patch(refined_lo, refined_hi);
+        grids[1].define(refined_patch);
+    }
+
+    // break the BoxArrays at both levels into max_grid_size^3 boxes
+    for (int lev = 0; lev < nlevs; lev++) {
+        grids[lev].maxSize(max_grid_size);
+    }
+
+    for (int lev = 1; lev < nlevs; lev++) {
+        domains[lev] = amrex::refine(domains[lev-1], ref_ratio[lev-1]);
+    }
+}