SoA: InitRandom/RandomPerBox/OnePerCell

Support more init methods on SoA particle containers.

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.

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]);
  }
  }
 

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);
 

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}};
 

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";
  }
  }