Initialize multi-physics module (ionization, QED) when restarting simulations from checkpoints (#3194)

* Move the initialization of multi-physics module earlier

* Fix syntax error

* Initialize multi-physics modules also when restarting the simulation

5 files changed, 29 insertions, 16 deletions

diff --git a/Source/Particles/MultiParticleContainer.H b/Source/Particles/MultiParticleContainer.H
index a05cde60b..80a2f362d 100644
--- a/Source/Particles/MultiParticleContainer.H
+++ b/Source/Particles/MultiParticleContainer.H
@@ -93,6 +93,8 @@ public:
 
     void InitData ();
 
+    void InitMultiPhysicsModules ();
+
     ///
     /// This evolves all the particles by one PIC time step, including current deposition, the
     /// field solve, and pushing the particles, for all the species in the MultiParticleContainer.
diff --git a/Source/Particles/MultiParticleContainer.cpp b/Source/Particles/MultiParticleContainer.cpp
index 2e69cd085..67039d8a3 100644
--- a/Source/Particles/MultiParticleContainer.cpp
+++ b/Source/Particles/MultiParticleContainer.cpp
@@ -397,21 +397,42 @@ MultiParticleContainer::AllocData ()
 void
 MultiParticleContainer::InitData ()
 {
+    InitMultiPhysicsModules();
+
     for (auto& pc : allcontainers) {
         pc->InitData();
     }
     pc_tmp->InitData();
+
+}
+
+void
+MultiParticleContainer::PostRestart ()
+{
+    InitMultiPhysicsModules();
+
+    for (auto& pc : allcontainers) {
+        pc->PostRestart();
+    }
+    pc_tmp->PostRestart();
+}
+
+void
+MultiParticleContainer::InitMultiPhysicsModules ()
+{
+    // Init ionization module here instead of in the MultiParticleContainer
+    // constructor because dt is required to compute ionization rate pre-factors
+    for (auto& pc : allcontainers) {
+        pc->InitIonizationModule();
+    }
     // For each species, get the ID of its product species.
     // This is used for ionization and pair creation processes.
     mapSpeciesProduct();
-
     CheckIonizationProductSpecies();
-
 #ifdef WARPX_QED
     CheckQEDProductSpecies();
     InitQED();
 #endif
-
 }
 
 void
@@ -657,15 +678,6 @@ MultiParticleContainer::SetParticleDistributionMap (int lev, DistributionMapping
 }
 
 void
-MultiParticleContainer::PostRestart ()
-{
-    for (auto& pc : allcontainers) {
-        pc->PostRestart();
-    }
-    pc_tmp->PostRestart();
-}
-
-void
 MultiParticleContainer
 ::GetLabFrameData (const std::string& /*snapshot_name*/,
                    const int /*i_lab*/, const int direction,
diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H
index 47ece0d31..aa077d3bc 100644
--- a/Source/Particles/PhysicalParticleContainer.H
+++ b/Source/Particles/PhysicalParticleContainer.H
@@ -64,7 +64,7 @@ public:
 
     virtual void WriteHeader (std::ostream& os) const override;
 
-    void InitIonizationModule ();
+    virtual void InitIonizationModule () override;
 
     /**
      * \brief Evolve is the central function PhysicalParticleContainer that
diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
index 762b5080e..feed808df 100644
--- a/Source/Particles/PhysicalParticleContainer.cpp
+++ b/Source/Particles/PhysicalParticleContainer.cpp
@@ -385,9 +385,6 @@ PhysicalParticleContainer::BackwardCompatibility ()
 
 void PhysicalParticleContainer::InitData ()
 {
-    // Init ionization module here instead of in the PhysicalParticleContainer
-    // constructor because dt is required
-    if (do_field_ionization) {InitIonizationModule();}
     AddParticles(0); // Note - add on level 0
     Redistribute();  // We then redistribute
 }
diff --git a/Source/Particles/WarpXParticleContainer.H b/Source/Particles/WarpXParticleContainer.H
index dc402a1c5..e735927e9 100644
--- a/Source/Particles/WarpXParticleContainer.H
+++ b/Source/Particles/WarpXParticleContainer.H
@@ -130,6 +130,8 @@ public:
 
     virtual void InitData () = 0;
 
+    virtual void InitIonizationModule () {}
+
     /**
      * Evolve is the central WarpXParticleContainer function that advances
      * particles for a time dt (typically one timestep). It is a pure virtual