Fix restart for implicit simulations (#5489)

Fix #5482.

Also, with this PR the temporary implicit particle attributes are not
written to the checkpoint file, and also are not communicated when doing
a redistribute of the particles. Note that this requires AMReX PR #4264.

A CI test of this restart is included in PR #5475.

---------

Co-authored-by: David Grote <[email protected]>

Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp

@@ -188,31 +188,49 @@ FlushFormatCheckpoint::CheckpointParticles (
 
         Vector<std::string> real_names;
         Vector<std::string> int_names;
+        Vector<int> write_real_comps;
+        Vector<int> write_int_comps;
 
         // note: positions skipped here, since we reconstruct a plotfile SoA from them
-        real_names.push_back("weight");
-        real_names.push_back("momentum_x");
-        real_names.push_back("momentum_y");
-        real_names.push_back("momentum_z");
-
+        std::vector<std::string> const fixed_names = {"weight",
+                                                      "momentum_x",
+                                                      "momentum_y",
+                                                      "momentum_z"
 #ifdef WARPX_DIM_RZ
-        real_names.push_back("theta");
+                                                      ,"theta"
 #endif
+                                                      };
+
+        for (auto const& name : fixed_names) {
+            real_names.push_back(name);
+            write_real_comps.push_back(1);
+        }
+
+        int const compile_time_comps = static_cast<int>(real_names.size());
 
         // get the names of the real comps
         //   note: skips the mandatory AMREX_SPACEDIM positions for pure SoA
         real_names.resize(pc->NumRealComps() - AMREX_SPACEDIM);
+        write_real_comps.resize(pc->NumRealComps() - AMREX_SPACEDIM);
         auto runtime_rnames = pc->getParticleRuntimeComps();
         for (auto const& x : runtime_rnames) {
-            real_names[x.second + PIdx::nattribs - AMREX_SPACEDIM] = x.first;
+            int const i = x.second + PIdx::nattribs - AMREX_SPACEDIM;
+            real_names[i] = x.first;
+            write_real_comps[i] = pc->h_redistribute_real_comp[i + compile_time_comps];
         }
 
         // and the int comps
         int_names.resize(pc->NumIntComps());
+        write_int_comps.resize(pc->NumIntComps());
         auto runtime_inames = pc->getParticleRuntimeiComps();
-        for (auto const& x : runtime_inames) { int_names[x.second+0] = x.first; }
+        for (auto const& x : runtime_inames) {
+            int const i = x.second + 0;
+            int_names[i] = x.first;
+            write_int_comps[i] = pc->h_redistribute_int_comp[i+AMREX_SPACEDIM];
+        }
 
-        pc->Checkpoint(dir, part_diag.getSpeciesName(), true,
+        pc->Checkpoint(dir, part_diag.getSpeciesName(),
+                       write_real_comps, write_int_comps,
                        real_names, int_names);
     }
 }

Source/Diagnostics/WarpXIO.cpp

@@ -404,4 +404,12 @@ WarpX::InitFromCheckpoint ()
     mypc->AllocData();
     mypc->Restart(restart_chkfile);
 
+    if (m_implicit_solver) {
+
+        m_implicit_solver->Define(this);
+        m_implicit_solver->GetParticleSolverParams( max_particle_its_in_implicit_scheme,
+                                                    particle_tol_in_implicit_scheme );
+        m_implicit_solver->CreateParticleAttributes();
+    }
+
 }

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.H

@@ -57,6 +57,8 @@ public:
         a_particle_tol = m_particle_tolerance;
     }
 
+    void CreateParticleAttributes () const;
+
     /**
      * \brief Advance fields and particles by one time step using the specified implicit algorithm
      */

Source/FieldSolver/ImplicitSolvers/ImplicitSolver.cpp

@@ -1,8 +1,30 @@
 #include "ImplicitSolver.H"
 #include "WarpX.H"
+#include "Particles/MultiParticleContainer.H"
 
 using namespace amrex;
 
+void ImplicitSolver::CreateParticleAttributes () const
+{
+    // Set comm to false to that the attributes are not communicated
+    // nor written to the checkpoint files
+    int const comm = 0;
+
+    // Add space to save the positions and velocities at the start of the time steps
+    for (auto const& pc : m_WarpX->GetPartContainer()) {
+#if (AMREX_SPACEDIM >= 2)
+        pc->NewRealComp("x_n", comm);
+#endif
+#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
+        pc->NewRealComp("y_n", comm);
+#endif
+        pc->NewRealComp("z_n", comm);
+        pc->NewRealComp("ux_n", comm);
+        pc->NewRealComp("uy_n", comm);
+        pc->NewRealComp("uz_n", comm);
+    }
+}
+
 const Geometry& ImplicitSolver::GetGeometry (const int a_lvl) const
 {
     AMREX_ASSERT((a_lvl >= 0) && (a_lvl < m_num_amr_levels));

Source/Initialization/WarpXInitData.cpp

@@ -691,21 +691,7 @@ WarpX::InitFromScratch ()
         m_implicit_solver->Define(this);
         m_implicit_solver->GetParticleSolverParams( max_particle_its_in_implicit_scheme,
                                                     particle_tol_in_implicit_scheme );
-
-        // Add space to save the positions and velocities at the start of the time steps
-        for (auto const& pc : *mypc) {
-#if (AMREX_SPACEDIM >= 2)
-            pc->NewRealComp("x_n");
-#endif
-#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
-            pc->NewRealComp("y_n");
-#endif
-            pc->NewRealComp("z_n");
-            pc->NewRealComp("ux_n");
-            pc->NewRealComp("uy_n");
-            pc->NewRealComp("uz_n");
-        }
-
+        m_implicit_solver->CreateParticleAttributes();
     }
 
     mypc->AllocData();