Hybrid-PIC: checkpoint the electron temperature under the electron energy equation

Source/Diagnostics/WarpXIO.cpp

@@ -21,6 +21,7 @@
 #include "EmbeddedBoundary/Enabled.H"
 #include "Fields.H"
 #include "FieldIO.H"
+#include "FieldSolver/FiniteDifferenceSolver/HybridPICModel/HybridPICModel.H"
 #include "FieldSolver/ImplicitSolvers/ImplicitSolver.H"
 #include "Particles/MultiParticleContainer.H"
 #include "Particles/WarpXParticleContainer.H"
@@ -30,6 +31,7 @@
 #include <ablastr/fields/MultiFabRegister.H>
 #include <ablastr/profiler/ProfilerWrapper.H>
 #include <ablastr/utils/text/StreamUtils.H>
+#include <ablastr/warn_manager/WarnManager.H>
 
 #ifdef AMREX_USE_SENSEI_INSITU
 #   include <AMReX_AmrMeshInSituBridge.H>
@@ -292,6 +294,11 @@ WarpX::InitFromCheckpoint ()
 
     const int nlevs = finestLevel()+1;
 
+    // Count the levels on which the electron temperature was restored from
+    // the checkpoint, to decide below whether the T_e seeds must be
+    // suppressed (see m_te_restored_from_checkpoint).
+    int n_levels_with_te_restored = 0;
+
     // Initialize the field data
     for (int lev = 0; lev < nlevs; ++lev)
     {
@@ -404,8 +411,38 @@ WarpX::InitFromCheckpoint ()
             }
         }
 
-        m_fields.read_restarts(lev, amrex::MultiFabFileFullPrefix(lev, restart_chkfile, level_prefix, ""));
+        // Read any fields flagged checkpoint_restart in the field register
+        // (mirrors FlushFormatCheckpoint's write_checkpoints call). Flagged
+        // fields absent from an older checkpoint are skipped, not errors.
+        auto const restored_names = m_fields.read_restarts(
+            lev, amrex::MultiFabFileFullPrefix(lev, restart_chkfile, level_prefix, ""));
 
+        std::string const te_name = m_fields.mf_name(
+            amrex::getEnumNameString(FieldType::hybrid_electron_temperature_fp), lev);
+        for (auto const& name : restored_names) {
+            if (name == te_name) { ++n_levels_with_te_restored; }
+        }
+    }
+
+    // With the electron energy equation on, T_e is evolved state that was
+    // just restored: latch the seeds off so neither HybridPICModel::InitData
+    // (uniform elec_temp fill) nor WarpX::HybridPICInitializeRhoJandB
+    // (adiabat seed) -- both of which run after this -- overwrites it.
+    if (m_hybrid_pic_model && m_hybrid_pic_model->m_solve_electron_energy_equation) {
+        if (n_levels_with_te_restored == nlevs) {
+            m_hybrid_pic_model->m_te_restored_from_checkpoint = true;
+            amrex::Print() << Utils::TextMsg::Info(
+                "restart: electron temperature restored from checkpoint "
+                "(adiabat seed suppressed)");
+        } else {
+            ablastr::warn_manager::WMRecordWarning(
+                "HybridPIC",
+                "Restarting with the electron energy equation from a checkpoint "
+                "that does not contain the electron temperature: T_e will be "
+                "re-seeded from the density adiabat, so evolved electron thermal "
+                "structure from before the checkpoint is not preserved.",
+                ablastr::warn_manager::WarnPriority::high);
+        }
     }
 
     InitPML();

Source/FieldSolver/FiniteDifferenceSolver/HybridPICModel/HybridPICModel.H

@@ -462,6 +462,15 @@ public:
      *  because it is a lazily-set cache flag inside a const call chain. */
     mutable bool m_qdsmc_J_plasma_valid = false;
 
+    /** True when the evolved electron temperature was restored from a
+     *  checkpoint (set by WarpX::InitFromCheckpoint). With the electron
+     *  energy equation on, T_e is evolved state and is written to the
+     *  checkpoint; the two seeds that would otherwise overwrite it -- the
+     *  uniform elec_temp fill in InitData and the adiabat seed in
+     *  WarpX::HybridPICInitializeRhoJandB -- are skipped when this is set,
+     *  so a restart continues from the restored thermal structure. */
+    bool m_te_restored_from_checkpoint = false;
+
     /** True when the per-species deposited charge densities rho_fp_{spec}
      *  are needed, i.e. when the electron-energy equation is solved (its
      *  Joule and Q_ei sources read the per-species charge densities).

Source/FieldSolver/FiniteDifferenceSolver/HybridPICModel/HybridPICModel.cpp

@@ -185,9 +185,15 @@ void HybridPICModel::AllocateLevelMFs (
     // the energy equation on it is the QDSMC state variable, otherwise it
     // mirrors the closure's implied temperature T_e = P_e / (n_e k_B),
     // filled alongside P_e in CalculateElectronPressure.
+    // With the energy equation on, T_e is flagged into the checkpoint: it is
+    // evolved state that cannot be reconstructed from the restored rho, so
+    // without it a restart would silently discard the evolved electron
+    // thermal structure. With the equation off it is a pure diagnostic
+    // mirror of the closure, refilled every step, and is not checkpointed.
     fields.alloc_init(FieldType::hybrid_electron_temperature_fp,
         lev, amrex::convert(ba, rho_nodal_flag),
-        dm, ncomps, ngRho, 0.0_rt);
+        dm, ncomps, ngRho, 0.0_rt,
+        true, true, m_solve_electron_energy_equation);
 
     // QDSMC electron-energy-equation working fields, only touched (and
     // therefore only allocated) when the energy equation is solved:
@@ -430,10 +436,16 @@ void HybridPICModel::InitData (const ablastr::fields::MultiFabRegister& fields)
     // solve: CalculateElectronPressure overwrites T_e from the closure, both
     // each step on the algebraic path and once from HybridPICInitializeRhoJandB
     // (on the floored density) to seed the energy-equation path.
-    for (int lev = 0; lev <= warpx.finestLevel(); ++lev) {
-        amrex::MultiFab & Te_mf = *warpx.m_fields.get(
-            FieldType::hybrid_electron_temperature_fp, lev);
-        Te_mf.setVal(m_elec_temp / PhysConst::kb);
+    //
+    // Skipped on a restart that restored T_e: this runs AFTER
+    // InitFromCheckpoint in WarpX::InitData, so an unconditional fill would
+    // overwrite the restored evolved temperature with the uniform constant.
+    if (!m_te_restored_from_checkpoint) {
+        for (int lev = 0; lev <= warpx.finestLevel(); ++lev) {
+            amrex::MultiFab & Te_mf = *warpx.m_fields.get(
+                FieldType::hybrid_electron_temperature_fp, lev);
+            Te_mf.setVal(m_elec_temp / PhysConst::kb);
+        }
     }
 
     // QDSMC: lazy-construct the fictitious-particle container and lay one

Source/FieldSolver/WarpXPushFieldsHybridPIC.cpp

@@ -379,10 +379,25 @@ void WarpX::HybridPICInitializeRhoJandB ()
     // right after each deposition (via the closure, or via the QDSMC entropy
     // transport when solve_electron_energy_equation is on).
     // With the energy equation on the closure is evaluated on floored density.
-    // T_e is not checkpointed either, so on restart the seed re-derives it from
-    // the restored rho: evolved T_e structure is not preserved across a restart.
-    m_hybrid_pic_model->CalculateElectronPressure(
-        m_hybrid_pic_model->m_solve_electron_energy_equation);
+    if (m_hybrid_pic_model->m_te_restored_from_checkpoint) {
+        // Restart with the electron energy equation: T_e is evolved state and
+        // was restored by InitFromCheckpoint. Emit Pe from the RESTORED T_e
+        // (with the boundary treatment grad Pe needs) rather than re-running
+        // the adiabat seed, which would overwrite it and discard the evolved
+        // thermal structure. Only reachable with the energy equation on.
+        for (int lev = 0; lev <= finest_level; ++lev) {
+            m_hybrid_pic_model->QDSMCFillElectronPressureFromTe(lev);
+            ApplyElectronPressureBoundary(lev, PatchType::fine);
+            ablastr::utils::communication::FillBoundary(
+                *m_fields.get(FieldType::hybrid_electron_pressure_fp, lev),
+                do_single_precision_comms,
+                Geom(lev).periodicity(),
+                true);
+        }
+    } else {
+        m_hybrid_pic_model->CalculateElectronPressure(
+            m_hybrid_pic_model->m_solve_electron_energy_equation);
+    }
 
     if (restart_chkfile.empty()) {
         // Handle field splitting for Hybrid field push

Source/ablastr/fields/MultiFabRegister.H

@@ -776,11 +776,19 @@ namespace ablastr::fields
         );
 
         /** Read in any (i)MultiFabs that are flagged checkpoint_restart from the checkpoint files
+         *
+         * A flagged field whose file is absent from the checkpoint (e.g. a
+         * checkpoint written before the field was flagged, or by a run with
+         * a different set of flagged fields) is skipped rather than fatal,
+         * so that older checkpoints remain restartable. Callers that need
+         * the field can detect the skip from the returned names and decide
+         * whether it is an error.
          *
          * @param level the MR level of the MF
          * @param dir the pathname to the checkpoint files
+         * @return register names (see mf_name) of the fields actually read
          */
-        void
+        std::vector<std::string>
         read_restarts (
             int level,
             std::string const & dir

Source/ablastr/fields/MultiFabRegister.cpp

@@ -326,12 +326,13 @@ namespace ablastr::fields
         }
     }
 
-    void
+    std::vector<std::string>
     MultiFabRegister::read_restarts (
         int level,
         const std::string & dir
     )
     {
+        std::vector<std::string> names_read;
         for (auto & element : m_mf_register )
         {
             MultiFabOwner & mf_owner = element.second;
@@ -341,9 +342,17 @@ namespace ablastr::fields
                 // only owning MultiFabs are read in
                 amrex::MultiFab & mf = mf_owner.m_mf;
                 const std::string & name = element.first;
+                if (!amrex::VisMF::Exist(dir + name)) {
+                    // The checkpoint predates this field being flagged (or was
+                    // written by a run that did not flag it): keep the runtime
+                    // initialization instead of failing the whole restart.
+                    continue;
+                }
                 amrex::VisMF::Read(mf, dir + name);
+                names_read.push_back(name);
             }
         }
+        return names_read;
     }
 
     bool