Task: Implement Checkpoint, Restart, and Plotfile for pure SOA (AMReX-Codes/amrex#3332)

org: AMReX-Codes
repo: amrex
pr_number: 3332
issue_title: Implement Checkpoint, Restart, and Plotfile for pure SOA
issue_body: 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: number: 2878
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


Reference solution:
fix_patch: diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H
index abffbefc860..3e6a6113fc8 100644
--- a/Src/Particle/AMReX_ParticleContainer.H
+++ b/Src/Particle/AMReX_ParticleContainer.H
@@ -1346,7 +1346,8 @@ protected:
     * \param lev_max
     * \param local_grid
     */
-    bool EnforcePeriodicWhere (ParticleType& prt, ParticleLocData& pld,
+    template <typename P>
+    bool EnforcePeriodicWhere (P& prt, ParticleLocData& pld,
                                int lev_min = 0, int lev_max = -1, int local_grid=-1) const;
 
 public:
@@ -1378,7 +1379,8 @@ private:
     void RedistributeMPI (std::map<int, Vector<char> >& not_ours,
                           int lev_min = 0, int lev_max = 0, int nGrow = 0, int local=0);
 
-    void locateParticle (ParticleType& p, ParticleLocData& pld,
+    template <typename P>
+    void locateParticle (P& p, ParticleLocData& pld,
                          int lev_min, int lev_max, int nGrow, int local_grid=-1) const;
 
     void Initialize ();
diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
index 9a5007f86ec..034bc50cd6f 100644
--- a/Src/Particle/AMReX_ParticleContainerI.H
+++ b/Src/Particle/AMReX_ParticleContainerI.H
@@ -164,9 +164,10 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
 
 template <typename ParticleType, int NArrayReal, int NArrayInt,
           template<class> class Allocator>
+template <typename P>
 bool
 ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
-::EnforcePeriodicWhere (ParticleType&    p,
+::EnforcePeriodicWhere (P&               p,
                         ParticleLocData& pld,
                         int              lev_min,
                         int              lev_max,
@@ -312,8 +313,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::resizeDa
 
 template <typename ParticleType, int NArrayReal, int NArrayInt,
           template<class> class Allocator>
+template <typename P>
 void
-ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::locateParticle (ParticleType& p, ParticleLocData& pld,
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::locateParticle (P& p, ParticleLocData& pld,
                                                                                         int lev_min, int lev_max, int nGrow, int local_grid) const
 {
     bool success;
diff --git a/Src/Particle/AMReX_ParticleIO.H b/Src/Particle/AMReX_ParticleIO.H
index a9849ca8a3c..cfc55dfd550 100644
--- a/Src/Particle/AMReX_ParticleIO.H
+++ b/Src/Particle/AMReX_ParticleIO.H
@@ -919,7 +919,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
     int*   iptr = istuff.dataPtr();
     RTYPE* rptr = rstuff.dataPtr();
 
-    ParticleType p;
+    Particle<NStructReal, NStructInt> ptemp;
     ParticleLocData pld;
 
     Vector<std::map<std::pair<int, int>, Gpu::HostVector<ParticleType> > > host_particles;
@@ -937,41 +937,43 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
     host_int_attribs.resize(finest_level_in_file+1);
 
     for (int i = 0; i < cnt; i++) {
-        if (convert_ids) {
+        // note: for pure SoA particle layouts, we do write the id, cpu and positions as a struct
+        //       for backwards compatibility with readers
+        if (!ParticleType::is_soa_particle && convert_ids) {
             std::int32_t  xi, yi;
             std::uint32_t xu, yu;
             xi = iptr[0];
             yi = iptr[1];
             std::memcpy(&xu, &xi, sizeof(xi));
             std::memcpy(&yu, &yi, sizeof(yi));
-            p.m_idcpu = ((std::uint64_t)xu) << 32 | yu;
+            ptemp.m_idcpu = ((std::uint64_t)xu) << 32 | yu;
         } else {
-            p.id()   = iptr[0];
-            p.cpu()  = iptr[1];
+            ptemp.id()   = iptr[0];
+            ptemp.cpu()  = iptr[1];
         }
         iptr += 2;
 
         for (int j = 0; j < NStructInt; j++)
         {
-            p.idata(j) = *iptr;
+            ptemp.idata(j) = *iptr;
             ++iptr;
         }
 
-        AMREX_ASSERT(p.id() > 0);
+        AMREX_ASSERT(ptemp.id() > 0);
 
-        AMREX_D_TERM(p.pos(0) = ParticleReal(rptr[0]);,
-                     p.pos(1) = ParticleReal(rptr[1]);,
-                     p.pos(2) = ParticleReal(rptr[2]););
+        AMREX_D_TERM(ptemp.pos(0) = ParticleReal(rptr[0]);,
+                     ptemp.pos(1) = ParticleReal(rptr[1]);,
+                     ptemp.pos(2) = ParticleReal(rptr[2]););
 
         rptr += AMREX_SPACEDIM;
 
         for (int j = 0; j < NStructReal; j++)
         {
-            p.rdata(j) = ParticleReal(*rptr);
+            ptemp.rdata(j) = ParticleReal(*rptr);
             ++rptr;
         }
 
-        locateParticle(p, pld, 0, finestLevel(), 0);
+        locateParticle(ptemp, pld, 0, finestLevel(), 0);
 
         std::pair<int, int> ind(grd, pld.m_tile);
 
@@ -979,35 +981,69 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
         host_int_attribs[lev][ind].resize(NumIntComps());
 
         // add the struct
-        host_particles[lev][ind].push_back(p);
+        if constexpr(!ParticleType::is_soa_particle)
+        {
+            host_particles[lev][ind].push_back(ptemp);
 
-        // add the real...
-        for (int icomp = 0; icomp < NumRealComps(); icomp++) {
-            host_real_attribs[lev][ind][icomp].push_back(*rptr);
-            ++rptr;
-        }
+            // add the real...
+            for (int icomp = 0; icomp < NumRealComps(); icomp++) {
+                host_real_attribs[lev][ind][icomp].push_back(*rptr);
+                ++rptr;
+            }
 
-        // ... and int array data
-        for (int icomp = 0; icomp < NumIntComps(); icomp++) {
-            host_int_attribs[lev][ind][icomp].push_back(*iptr);
-            ++iptr;
+            // ... and int array data
+            for (int icomp = 0; icomp < NumIntComps(); icomp++) {
+                host_int_attribs[lev][ind][icomp].push_back(*iptr);
+                ++iptr;
+            }
+        } else {
+            host_particles[lev][ind];
+
+            for (int j = 0; j < AMREX_SPACEDIM; j++) {
+                host_real_attribs[pld.m_lev][ind][j].push_back(ptemp.pos(j));
+            }
+
+            host_int_attribs[pld.m_lev][ind][0].push_back(ptemp.id());
+            host_int_attribs[pld.m_lev][ind][1].push_back(ptemp.cpu());
+
+            // read all other SoA
+            // add the real...
+            for (int icomp = AMREX_SPACEDIM; icomp < NumRealComps(); icomp++) {
+                host_real_attribs[lev][ind][icomp].push_back(*rptr);
+                ++rptr;
+            }
+
+            // ... and int array data
+            for (int icomp = 2; icomp < NumIntComps(); icomp++) {
+                host_int_attribs[lev][ind][icomp].push_back(*iptr);
+                ++iptr;
+            }
         }
     }
 
     for (int host_lev = 0; host_lev < static_cast<int>(host_particles.size()); ++host_lev)
-      {
+    {
         for (auto& kv : host_particles[host_lev]) {
           auto grid = kv.first.first;
           auto tile = kv.first.second;
           const auto& src_tile = kv.second;
 
           auto& dst_tile = DefineAndReturnParticleTile(host_lev, grid, tile);
-          auto old_size = dst_tile.GetArrayOfStructs().size();
-          auto new_size = old_size + src_tile.size();
+          auto old_size = dst_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 < NumRealComps(); ++i) {
               Gpu::copyAsync(Gpu::hostToDevice,
diff --git a/Src/Particle/AMReX_ParticleTile.H b/Src/Particle/AMReX_ParticleTile.H
index 9f1fe712c64..3b5ee15fdb0 100644
--- a/Src/Particle/AMReX_ParticleTile.H
+++ b/Src/Particle/AMReX_ParticleTile.H
@@ -226,6 +226,9 @@ struct ParticleTileData
     {
         AMREX_ASSERT(index < m_size);
         SuperParticleType sp;
+        for (int i = 0; i < AMREX_SPACEDIM; ++i) {sp.pos(i) = m_rdata[i][index];}
+        sp.id()  = m_idata[0][index];
+        sp.cpu() = m_idata[1][index];
         for (int i = 0; i < NAR; ++i)
             sp.rdata(i) = m_rdata[i][index];
         for (int i = 0; i < NAI; ++i)
@@ -595,6 +598,9 @@ struct ConstParticleTileData
     {
         AMREX_ASSERT(index < m_size);
         SuperParticleType sp;
+        for (int i = 0; i < AMREX_SPACEDIM; ++i) {sp.pos(i) = m_rdata[i][index];}
+        sp.id()  = m_idata[0][index];
+        sp.cpu() = m_idata[1][index];
         for (int i = 0; i < NAR; ++i)
             sp.rdata(i) = m_rdata[i][index];
         for (int i = 0; i < NAI; ++i)
diff --git a/Src/Particle/AMReX_WriteBinaryParticleData.H b/Src/Particle/AMReX_WriteBinaryParticleData.H
index 7767972998d..f182eb2b3e8 100644
--- a/Src/Particle/AMReX_WriteBinaryParticleData.H
+++ b/Src/Particle/AMReX_WriteBinaryParticleData.H
@@ -13,7 +13,7 @@ struct KeepValidFilter
     AMREX_GPU_HOST_DEVICE
     int operator() (const SrcData& src, int i) const noexcept
     {
-        return (src.m_aos[i].id() > 0);
+        return (src.id(i) > 0);
     }
 };
 
@@ -197,7 +197,7 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
     for (int i = 0; i < tiles.size(); i++) {
         const auto& ptile = pc.ParticlesAt(lev, grid, tiles[i]);
         const auto& pflags = particle_io_flags[lev].at(std::make_pair(grid, tiles[i]));
-        int np_tile = ptile.GetArrayOfStructs().numParticles();
+        int np_tile = ptile.numParticles();
         typename PC::IntVector offsets(np_tile);
         int num_copies = Scan::ExclusiveSum(np_tile, pflags.begin(), offsets.begin(), Scan::retSum);
 
@@ -283,44 +283,70 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
     const Long rChunkSize = AMREX_SPACEDIM + num_output_real;
     rdata.resize(np*rChunkSize);
 
-    int* iptr = idata.dataPtr();
+        int* iptr = idata.dataPtr();
     ParticleReal* rptr = rdata.dataPtr();
     for (unsigned i = 0; i < tiles.size(); i++) {
         const auto& ptile = pc.ParticlesAt(lev, grid, tiles[i]);
         const auto& pflags = particle_io_flags[lev].at(std::make_pair(grid, tiles[i]));
-        for (int pindex = 0; pindex < ptile.GetArrayOfStructs().numParticles(); ++pindex) {
-            const auto& aos = ptile.GetArrayOfStructs();
-            const auto& p = aos[pindex];
+        for (int pindex = 0; pindex < ptile.numParticles(); ++pindex) {
             if (pflags[pindex]) {
-                packParticleIDs(iptr, p, is_checkpoint);
-                iptr += 2;
-
-                for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = p.pos(j);
-                rptr += AMREX_SPACEDIM;
+                const auto& soa  = ptile.GetStructOfArrays();
 
-                for (int j = 0; j < PC::NStructInt; j++) {
-                    if (write_int_comp[j]) {
-                        *iptr = p.idata(j);
-                        ++iptr;
+                // note: for pure SoA particle layouts, we do write the id, cpu and positions as a struct
+                //       for backwards compatibility with readers
+                if constexpr(!PC::ParticleType::is_soa_particle)
+                {
+                    const auto& aos = ptile.GetArrayOfStructs();
+                    const auto& p = aos[pindex];
+
+                    // Int: id, cpu
+                    packParticleIDs(iptr, p, is_checkpoint);
+                    iptr += 2;
+
+                    // Real: positions
+                    for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = p.pos(j);
+                    rptr += AMREX_SPACEDIM;
+
+                    // extra AoS Int components
+                    for (int j = 0; j < PC::NStructInt; j++) {
+                        if (write_int_comp[j]) {
+                            *iptr = p.idata(j);
+                            ++iptr;
+                        }
+                    }
+                    // extra AoS Real components
+                    for (int j = 0; j < PC::NStructReal; j++) {
+                        if (write_real_comp[j]) {
+                            *rptr = p.rdata(j);
+                            ++rptr;
+                        }
                     }
                 }
+                else {
+                    amrex::ignore_unused(is_checkpoint);
+                    // Int: id, cpu
+                    *iptr = soa.GetIntData(0)[pindex];
+                    iptr += 1;
+                    *iptr = soa.GetIntData(1)[pindex];
+                    iptr += 1;
+
+                    // Real: position
+                    for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = soa.GetRealData(j)[pindex];
+                    rptr += AMREX_SPACEDIM;
+                }
 
-                const auto& soa  = ptile.GetStructOfArrays();
-                for (int j = 0; j < pc.NumIntComps(); j++) {
+                // extra SoA int data
+                const int int_start_offset = PC::ParticleType::is_soa_particle ? 2 : 0; // pure SoA: skip id, cpu
+                for (int j = int_start_offset; j < pc.NumIntComps(); j++) {
                     if (write_int_comp[PC::NStructInt+j]) {
                         *iptr = soa.GetIntData(j)[pindex];
                         ++iptr;
                     }
                 }
 
-                for (int j = 0; j < PC::NStructReal; j++) {
-                    if (write_real_comp[j]) {
-                        *rptr = p.rdata(j);
-                        ++rptr;
-                    }
-                }
-
-                for (int j = 0; j < pc.NumRealComps(); j++) {
+                // extra SoA Real components
+                const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
+                for (int j = real_start_offset; j < pc.NumRealComps(); j++) {
                     if (write_real_comp[PC::NStructReal+j]) {
                         *rptr = (ParticleReal) soa.GetRealData(j)[pindex];
                         ++rptr;
@@ -647,8 +673,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,
         {
             int gid = pti.index();
             const auto& ptile = pc.ParticlesAt(lev, pti);
-            const auto& aos = ptile.GetArrayOfStructs();
-            const auto pstruct = aos().dataPtr();
+            const auto& ptd = ptile.getConstParticleTileData();
             const int np = ptile.numParticles();
 
             ReduceOps<ReduceOpSum> reduce_op;
@@ -658,7 +683,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,
             reduce_op.eval(np, reduce_data,
             [=] AMREX_GPU_DEVICE (int i) -> ReduceTuple
             {
-                return (pstruct[i].id() > 0) ? 1 : 0;
+                return (ptd.id(i) > 0) ? 1 : 0;
             });
 
             int np_valid = amrex::get<0>(reduce_data.value(reduce_op));
@@ -938,30 +963,43 @@ void WriteBinaryParticleDataAsync (PC const& pc,
                 for (unsigned i = 0; i < tile_map[grid].size(); i++) {
                     auto ptile_index = std::make_pair(grid, tile_map[grid][i]);
                     const auto& pbox = (*myptiles)[lev][ptile_index];
-                    for (int pindex = 0;
-                         pindex < pbox.GetArrayOfStructs().numParticles(); ++pindex)
+                    const auto ptd = pbox.getConstParticleTileData();
+                    for (int pindex = 0; pindex < pbox.numParticles(); ++pindex)
                     {
-                        const auto& aos = pbox.GetArrayOfStructs();
-                        const auto& p = aos[pindex];
+                        const auto& soa  = pbox.GetStructOfArrays();
 
+                        const auto& p = make_particle<typename PC::ConstParticleType>{}(ptd, pindex);
                         if (p.id() <= 0) continue;
 
-                        // always write these
-                        particle_detail::packParticleIDs(iptr, p, is_checkpoint);
-                        iptr += 2;
-
-                        // optionally write these
-                        for (int j = 0; j < NStructInt; j++)
+                        // note: for pure SoA particle layouts, we do write the id, cpu and positions as a struct
+                        //       for backwards compatibility with readers
+                        if constexpr(!PC::ParticleType::is_soa_particle)
                         {
-                            if (write_int_comp[j])
+                            // Ints: id, cpu
+                            particle_detail::packParticleIDs(iptr, p, is_checkpoint);
+                            iptr += 2;
+
+                            // extra AoS Int components
+                            for (int j = 0; j < NStructInt; j++)
                             {
-                                *iptr = p.idata(j);
-                                ++iptr;
+                                if (write_int_comp[j])
+                                {
+                                    *iptr = p.idata(j);
+                                    ++iptr;
+                                }
                             }
                         }
+                        else {
+                            // Ints: id, cpu
+                            *iptr = soa.GetIntData(0)[pindex];
+                            iptr += 1;
+                            *iptr = soa.GetIntData(1)[pindex];
+                            iptr += 1;
+                        }
 
-                        const auto& soa  = pbox.GetStructOfArrays();
-                        for (int j = 0; j < nic; j++)
+                        // extra SoA Ints
+                        const int int_start_offset = PC::ParticleType::is_soa_particle ? 2 : 0; // pure SoA: skip id, cpu
+                        for (int j = int_start_offset; j < nic; j++)
                         {
                             if (write_int_comp[NStructInt+j])
                             {
@@ -987,30 +1025,40 @@ void WriteBinaryParticleDataAsync (PC const& pc,
                 for (unsigned i = 0; i < tile_map[grid].size(); i++) {
                     auto ptile_index = std::make_pair(grid, tile_map[grid][i]);
                     const auto& pbox = (*myptiles)[lev][ptile_index];
+                    const auto ptd = pbox.getConstParticleTileData();
                     for (int pindex = 0;
-                         pindex < pbox.GetArrayOfStructs().numParticles(); ++pindex)
+                         pindex < pbox.numParticles(); ++pindex)
                     {
-                        const auto& aos = pbox.GetArrayOfStructs();
-                        const auto& p = aos[pindex];
+                        const auto& soa  = pbox.GetStructOfArrays();
+                        const auto& p = make_particle<typename PC::ConstParticleType>{}(ptd, pindex);
 
                         if (p.id() <= 0) continue;
 
-                        // always write these
-                        for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = p.pos(j);
-                        rptr += AMREX_SPACEDIM;
-
-                        // optionally write these
-                        for (int j = 0; j < NStructReal; j++)
+                        if constexpr(!PC::ParticleType::is_soa_particle)
                         {
-                            if (write_real_comp[j])
+                            // Real: position
+                            for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = p.pos(j);
+                            rptr += AMREX_SPACEDIM;
+
+                            // extra AoS real
+                            for (int j = 0; j < NStructReal; j++)
                             {
-                                *rptr = p.rdata(j);
-                                ++rptr;
+                                if (write_real_comp[j])
+                                {
+                                    *rptr = p.rdata(j);
+                                    ++rptr;
+                                }
                             }
                         }
+                        else {
+                            // Real: position
+                            for (int j = 0; j < AMREX_SPACEDIM; j++) rptr[j] = soa.GetRealData(j)[pindex];
+                            rptr += AMREX_SPACEDIM;
+                        }
 
-                        const auto& soa  = pbox.GetStructOfArrays();
-                        for (int j = 0; j < nrc; j++)
+                        // extra SoA real
+                        const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: positions
+                        for (int j = real_start_offset; j < nrc; j++)
                         {
                             if (write_real_comp[NStructReal+j])
                             {