From daa00fcb19df7d4bb7263426e41b223e36081714 Mon Sep 17 00:00:00 2001
From: Ben Wibking <ben@wibking.com>
Date: Fri, 17 Jul 2026 21:49:23 -0400
Subject: [PATCH] fix prepost bugs

---
 Src/Extern/HDF5/AMReX_ParticleHDF5.H          | 29 +++++++++----------
 Src/Particle/AMReX_ParticleIO.H               | 14 +--------
 .../CheckpointRestartDualGridHDF5SOA/main.cpp |  5 ++++
 Tests/Particles/CheckpointRestartSOA/inputs   |  3 ++
 Tests/Particles/CheckpointRestartSOA/main.cpp |  2 ++
 5 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/Src/Extern/HDF5/AMReX_ParticleHDF5.H b/Src/Extern/HDF5/AMReX_ParticleHDF5.H
index b70faef6507..4a848d18263 100644
--- a/Src/Extern/HDF5/AMReX_ParticleHDF5.H
+++ b/Src/Extern/HDF5/AMReX_ParticleHDF5.H
@@ -453,19 +453,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     Long  maxnextid  = ParticleType::NextID();
 
     for (int lev = 0; lev < std::ssize(m_particles);  lev++) {
-        const auto& pmap = m_particles[lev];
-        for (const auto& kv : pmap) {
-            const auto& aos = kv.second.GetArrayOfStructs();
-            for (int k = 0; k < aos.numParticles(); ++k) {
-                const ParticleType& p = aos[k];
-                if (p.id().is_valid()) {
-                    //
-                    // Only count (and checkpoint) valid particles.
-                    //
-                    nparticles++;
-                }
-            }
-        }
+        nparticles += NumberOfParticlesAtLevel(lev, true, true);
     }
     ParallelDescriptor::ReduceLongSum(nparticles, IOProcNumber);
 
@@ -729,18 +717,21 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     }
 
     #ifdef AMREX_USE_MPI
-    // Collect the number of mf and total size of mf from each rank
+    // Collect the number of non-empty grids from each rank
     MPI_Allgather(&my_mfi_cnt, 1, ParallelDescriptor::Mpi_typemap<int>::type(), &(all_mfi_cnt[0]), 1,
                   ParallelDescriptor::Mpi_typemap<int>::type(), ParallelDescriptor::Communicator());
+    #else
+    all_mfi_cnt[0] = my_mfi_cnt;
+    #endif
 
     for (int i = 0; i < ParallelDescriptor::NProcs(); i++)
         total_mfi += all_mfi_cnt[i];
 
     // Create the int data
+    #ifdef AMREX_USE_MPI
     MPI_Allgather(&my_mfi_int_total_size, 1, ParallelDescriptor::Mpi_typemap<ULong>::type(),
                   &(all_mfi_int_total_size[0]), 1, ParallelDescriptor::Mpi_typemap<ULong>::type(), ParallelDescriptor::Communicator());
     #else
-    all_mfi_cnt[0] = my_mfi_cnt;
     all_mfi_int_total_size[0] = my_mfi_int_total_size;
     #endif
 
@@ -1295,6 +1286,14 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
             amrex::Abort(msg.c_str());
         }
 
+        htri_t const has_particles = H5Lexists(grp, "nparticles_grid", H5P_DEFAULT);
+        if (has_particles < 0) {
+            amrex::Abort("ParticleContainer::RestartHDF5(): unable to check for nparticles_grid dataset");
+        } else if (has_particles == 0) {
+            H5Gclose(grp);
+            continue;
+        }
+
         dset = H5Dopen(grp, "nparticles_grid", H5P_DEFAULT);
         Vector<int> count(ngrids[lev]);
         ret = H5Dread(dset, H5T_NATIVE_INT, H5S_ALL, H5S_ALL, H5P_DEFAULT, &(count[0]));
diff --git a/Src/Particle/AMReX_ParticleIO.H b/Src/Particle/AMReX_ParticleIO.H
index 6d3493a9234..e47bf9de899 100644
--- a/Src/Particle/AMReX_ParticleIO.H
+++ b/Src/Particle/AMReX_ParticleIO.H
@@ -460,19 +460,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     Long  maxnextid  = ParticleType::NextID();
 
     for (int lev = 0; lev < std::ssize(m_particles);  lev++) {
-        const auto& pmap = m_particles[lev];
-        for (const auto& kv : pmap) {
-            const auto& aos = kv.second.GetArrayOfStructs();
-            for (int k = 0; k < aos.numParticles(); ++k) {
-                const ParticleType& p = aos[k];
-                if (p.id().is_valid()) {
-                    //
-                    // Only count (and checkpoint) valid particles.
-                    //
-                    nparticles++;
-                }
-            }
-        }
+        nparticles += NumberOfParticlesAtLevel(lev, true, true);
     }
     ParallelDescriptor::ReduceLongSum(nparticles, IOProcNumber);
 
diff --git a/Tests/Particles/CheckpointRestartDualGridHDF5SOA/main.cpp b/Tests/Particles/CheckpointRestartDualGridHDF5SOA/main.cpp
index 7deb9429303..967173bc66f 100644
--- a/Tests/Particles/CheckpointRestartDualGridHDF5SOA/main.cpp
+++ b/Tests/Particles/CheckpointRestartDualGridHDF5SOA/main.cpp
@@ -166,12 +166,17 @@ void test_fewer_levels ()
     auto mesh2 = build_mesh(ncells, 2, max_grid_size);
     MyPC pc_write(mesh2.geom, mesh2.dmap, mesh2.ba, mesh2.ref_ratio);
     pc_write.SetVerbose(false);
+    pc_write.SetUsePrePost(true);
     pc_write.InitRandom(nppc * AMREX_D_TERM(ncells, *ncells, *ncells),
                         iseed, pdata, /*serialize=*/false);
 #ifdef AMREX_USE_HDF5
+    pc_write.CheckpointPreHDF5();
     pc_write.CheckpointHDF5(chkdir, "particles", true, real_names, int_names);
+    pc_write.CheckpointPostHDF5();
 #else
+    pc_write.CheckpointPre();
     pc_write.Checkpoint(chkdir, "particles", real_names, int_names);
+    pc_write.CheckpointPost();
 #endif
 
     AsyncOut::Finish();
diff --git a/Tests/Particles/CheckpointRestartSOA/inputs b/Tests/Particles/CheckpointRestartSOA/inputs
index 81855c5947d..0b653c202f7 100644
--- a/Tests/Particles/CheckpointRestartSOA/inputs
+++ b/Tests/Particles/CheckpointRestartSOA/inputs
@@ -24,3 +24,6 @@ nparticlefile = 1
 restart_check = 1
 
 directory=.
+
+# Exercise the pre/post checkpoint path for pure-SoA particles.
+particles.use_prepost = 1
diff --git a/Tests/Particles/CheckpointRestartSOA/main.cpp b/Tests/Particles/CheckpointRestartSOA/main.cpp
index af2ce0f7a0b..b1384ba66f6 100644
--- a/Tests/Particles/CheckpointRestartSOA/main.cpp
+++ b/Tests/Particles/CheckpointRestartSOA/main.cpp
@@ -134,7 +134,9 @@ void test ()
 
             amrex::Print() << "Writing particle file [" << fname << "] ..." << '\n';
 
+            myPC.CheckpointPre();
             myPC.Checkpoint(fname, "particle0", false, particle_realnames, particle_intnames);
+            myPC.CheckpointPost();
 
             amrex::Print() << " done \n";
         }