Fix restart btd (#3734)

* Correct calculation of m_t_lab

* write BTD related restart, read, and init data before and after accordingly

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* write snapshot domain lo and hi in moving window dir

* write lo and hi of snapshot domain for all dimensions

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean up

* more meta data and redefine snapshot geom, buffer mf and write to new dir for plotfile instead of merging

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use virtual functions

* empty else condition

* remove commented line

* fix error in returning k-index. Thanks CI

* doxygen comments

* doxygen comments

* more doxygen

---------

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

13 files changed, 344 insertions, 40 deletions

diff --git a/Source/Diagnostics/BTDiagnostics.H b/Source/Diagnostics/BTDiagnostics.H
index c25a031d7..6a4bd9826 100644
--- a/Source/Diagnostics/BTDiagnostics.H
+++ b/Source/Diagnostics/BTDiagnostics.H
@@ -125,12 +125,20 @@ private:
      */
     bool GetZSliceInDomainFlag (const int i_buffer, const int lev);
 
+    /** whether the k-index corresponding to the z-slice for the ith buffer, i_buffer,
+     *  is within the bounds of the box in index-space
+     * \param[in] i_buffer index of the buffer for which the k-index is checked.
+     * \param[in] lev mesh-refinement level using which the k-index for the z-slice being filled in the ith buffer is computed
+     *            Return true if the k-index of the z-slice is within the bounds of the box for the buffer being filled
+     */
+    bool GetKIndexInSnapshotBoxFlag (const int i_buffer, const int lev);
+
     /** Initialize buffer domain, buffer box and lab-frame parameters such as
      *  m_t_lab, and z-positions for the i^th snapshot, i_buffer, and level, lev.
      * \param[in] i_buffer i^th snapshot or buffer
      * \param[in] lev mesh-refinement level for which the field and/or particle buffer data is initialized.
      */
-    void InitializeBufferData ( int i_buffer , int lev) override;
+    void InitializeBufferData ( int i_buffer , int lev, bool restart=false) override;
     /** Whether to compute back-tranformed values for field-data.
      *  default value is true.
      */
@@ -229,6 +237,9 @@ private:
     /** Vector of counters tracking number of times the buffer of multifab is
      *  flushed out and emptied before being refilled again for each snapshot */
     amrex::Vector<int> m_buffer_flush_counter;
+    /** Vector of k-indices in the hi-end along the moving window direction for the buffer being filled
+     *  for each snapshot.
+     */
     amrex::Vector<int> m_buffer_k_index_hi;
     /** Multi-level cell-centered multifab with all field-data components, namely,
      *  Ex, Ey, Ez, Bx, By, Bz, jx, jy, jz, and rho.
@@ -310,6 +321,19 @@ private:
         return ( m_buffer_counter[i_buffer] == 0) ;
     }
 
+    /** Vector of integers to indicate if is the first flush for a snapshot so that appropriate meta-data and geometry is written (0 if not, 1 if yes). */
+    amrex::Vector<int> m_first_flush_after_restart;
+
+    /** Resets integer for first flush to 0 for the ith snapshot (i_buffer).
+     *  This function is called after flush the first buffer after restarting the simulation
+     */
+    void NullifyFirstFlush(int i_buffer) {m_first_flush_after_restart[i_buffer] = 0;}
+
+    /** Vector of integers to indicate if the snapshot geometry is defined. */
+    amrex::Vector<int> m_snapshot_geometry_defined;
+    /** Vector of integers to indicate if the field buffer multifab, m_mf_output, is defined in DefineFieldBufferMultifab() function */
+    amrex::Vector<int> m_field_buffer_multifab_defined;
+
     /** Reset buffer counter to zero.
      * \param[in] i_buffer snapshot index for which the counter is set to zero.
      */
@@ -388,5 +412,17 @@ private:
     void RedistributeParticleBuffer (const int i_buffer);
     void UpdateVarnamesForRZopenPMD();
 
+    amrex::Real gettlab (int i_buffer) override {return m_t_lab[i_buffer];}
+    void settlab (int i_buffer, amrex::Real tlab) override {m_t_lab[i_buffer] = tlab; }
+    int get_buffer_k_index_hi (int i_buffer) override {return m_buffer_k_index_hi[i_buffer]; }
+    void set_buffer_k_index_hi (int i_buffer, int kindex) override {m_buffer_k_index_hi[i_buffer] = kindex;}
+    amrex::Real get_snapshot_domain_lo (int i_buffer, int idim) override {return m_snapshot_domain_lab[i_buffer].lo(idim); }
+    amrex::Real get_snapshot_domain_hi (int i_buffer, int idim) override {return m_snapshot_domain_lab[i_buffer].hi(idim); }
+    int get_flush_counter (int i_buffer) override {return m_buffer_flush_counter[i_buffer]; }
+    void set_flush_counter ([[maybe_unused]] int i_buffer, int flush_counter) override { m_buffer_flush_counter[i_buffer] = flush_counter; }
+    int get_last_valid_Zslice ( int i_buffer) override { return m_lastValidZSlice[i_buffer]; }
+    void set_last_valid_Zslice ( int i_buffer, int last_valid_Zslice) override { m_lastValidZSlice[i_buffer] = last_valid_Zslice; }
+    int get_snapshot_full_flag ( int i_buffer) override { return m_snapshot_full[i_buffer]; }
+    void set_snapshot_full ( int i_buffer, int snapshot_full) override { m_snapshot_full[i_buffer] = snapshot_full; }
 };
 #endif // WARPX_BTDIAGNOSTICS_H_
diff --git a/Source/Diagnostics/BTDiagnostics.cpp b/Source/Diagnostics/BTDiagnostics.cpp
index 204b28b07..a5b85c6b6 100644
--- a/Source/Diagnostics/BTDiagnostics.cpp
+++ b/Source/Diagnostics/BTDiagnostics.cpp
@@ -88,10 +88,17 @@ void BTDiagnostics::DerivedInitData ()
     m_snapshot_full.resize( m_num_buffers );
     m_lastValidZSlice.resize( m_num_buffers );
     m_buffer_k_index_hi.resize(m_num_buffers);
+    m_first_flush_after_restart.resize(m_num_buffers);
+    m_snapshot_geometry_defined.resize(m_num_buffers);
+    m_field_buffer_multifab_defined.resize(m_num_buffers);
     for (int i = 0; i < m_num_buffers; ++i) {
         m_geom_snapshot[i].resize(nmax_lev);
         m_snapshot_full[i] = 0;
         m_lastValidZSlice[i] = 0;
+        m_buffer_flush_counter[i] = 0;
+        m_first_flush_after_restart[i] = 1;
+        m_snapshot_geometry_defined[i] = 0;
+        m_field_buffer_multifab_defined[i] = 0;
     }
     for (int lev = 0; lev < nmax_lev; ++lev) {
         // Define cell-centered multifab over the whole domain with
@@ -329,13 +336,21 @@ BTDiagnostics::DoComputeAndPack (int step, bool force_flush)
 }
 
 void
-BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
+BTDiagnostics::InitializeBufferData ( int i_buffer , int lev, bool restart)
 {
     auto & warpx = WarpX::GetInstance();
+
+    // When restarting boosted simulations, the code below needs to take
+    // into account the fact that the position of the box at the beginning
+    // of the simulation, is not the one that we had at t=0 (because of the moving window)
+    amrex::Real boosted_moving_window_v = (warpx.moving_window_v - m_beta_boost*PhysConst::c)
+                                        / (1._rt - m_beta_boost * warpx.moving_window_v/PhysConst::c);
     // Lab-frame time for the i^th snapshot
-    amrex::Real zmax_0 = warpx.Geom(lev).ProbHi(m_moving_window_dir);
-    m_t_lab.at(i_buffer) = m_intervals.GetBTDIteration(i_buffer) * m_dt_snapshots_lab
-        + m_gamma_boost*m_beta_boost*zmax_0/PhysConst::c;
+    if (restart == false) {
+        amrex::Real zmax_0 = warpx.Geom(lev).ProbHi(m_moving_window_dir);
+        m_t_lab.at(i_buffer) = m_intervals.GetBTDIteration(i_buffer) * m_dt_snapshots_lab
+            + m_gamma_boost*m_beta_boost*zmax_0/PhysConst::c;
+    }
 
     // Define buffer domain in boosted frame at level, lev, with user-defined lo and hi
     amrex::RealBox diag_dom;
@@ -392,15 +407,13 @@ BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
 
     // Define buffer_domain in lab-frame for the i^th snapshot.
     // Replace z-dimension with lab-frame co-ordinates.
-    amrex::Real zmin_buffer_lab = diag_dom.lo(m_moving_window_dir)
-                                 / ( (1.0_rt + m_beta_boost) * m_gamma_boost);
-    amrex::Real zmax_buffer_lab = diag_dom.hi(m_moving_window_dir)
-                                 / ( (1.0_rt + m_beta_boost) * m_gamma_boost);
-
+    amrex::Real zmin_buffer_lab = ( diag_dom.lo(m_moving_window_dir) - boosted_moving_window_v * warpx.gett_new(0) )
+                                / ( (1.0_rt + m_beta_boost) * m_gamma_boost);
+    amrex::Real zmax_buffer_lab = ( diag_dom.hi(m_moving_window_dir) - boosted_moving_window_v * warpx.gett_new(0) )
+                                / ( (1.0_rt + m_beta_boost) * m_gamma_boost);
 
     // Initialize buffer counter and z-positions of the  i^th snapshot in
     // boosted-frame and lab-frame
-    m_buffer_flush_counter[i_buffer] = 0;
     m_buffer_counter[i_buffer] = 0;
     m_current_z_lab[i_buffer] = 0._rt;
     m_current_z_boost[i_buffer] = 0._rt;
@@ -450,7 +463,6 @@ BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
     m_snapshot_ncells_lab[i_buffer] = amrex::IntVect(Nz_lab);
 #endif
 
-
     // Box covering the extent of the user-defined diag in the back-transformed frame
     // for the ith snapshot
     // estimating the maximum number of buffer multifabs needed to obtain the
@@ -461,20 +473,23 @@ BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
     // number of cells in z is modified since each buffer multifab always
     // contains a minimum m_buffer_size=256 cells
     int num_z_cells_in_snapshot = m_max_buffer_multifabs[i_buffer] * m_buffer_size;
-    m_snapshot_domain_lab[i_buffer] = diag_dom;
-    m_snapshot_domain_lab[i_buffer].setLo(m_moving_window_dir,
-                                  zmin_buffer_lab + warpx.moving_window_v * m_t_lab[i_buffer]);
-    m_snapshot_domain_lab[i_buffer].setHi(m_moving_window_dir,
-                                  zmax_buffer_lab + warpx.moving_window_v * m_t_lab[i_buffer]);
-    // To prevent round off errors, moving the snapshot domain by half a cell so that all the slices
-    // lie close to the cell-centers in the lab-frame grid instead of on the edge of cell.
-    amrex::Real new_hi = m_snapshot_domain_lab[i_buffer].hi(m_moving_window_dir)
-                       + 0.5_rt * dz_lab(warpx.getdt(lev), ref_ratio[m_moving_window_dir]);
-    m_snapshot_domain_lab[i_buffer].setHi(m_moving_window_dir,new_hi);
-    amrex::Real new_lo = m_snapshot_domain_lab[i_buffer].hi(m_moving_window_dir) -
-                         num_z_cells_in_snapshot *
-                         dz_lab(warpx.getdt(lev), ref_ratio[m_moving_window_dir]);
-    m_snapshot_domain_lab[i_buffer].setLo(m_moving_window_dir, new_lo);
+    if (restart == false) {
+        m_snapshot_domain_lab[i_buffer] = diag_dom;
+        m_snapshot_domain_lab[i_buffer].setLo(m_moving_window_dir,
+                                      zmin_buffer_lab + warpx.moving_window_v * m_t_lab[i_buffer]);
+        m_snapshot_domain_lab[i_buffer].setHi(m_moving_window_dir,
+                                      zmax_buffer_lab + warpx.moving_window_v * m_t_lab[i_buffer]);
+        // To prevent round off errors, moving the snapshot domain by half a cell so that all the slices
+        // lie close to the cell-centers in the lab-frame grid instead of on the edge of cell.
+        amrex::Real new_hi = m_snapshot_domain_lab[i_buffer].hi(m_moving_window_dir)
+                           + 0.5_rt * dz_lab(warpx.getdt(lev), ref_ratio[m_moving_window_dir]);
+        m_snapshot_domain_lab[i_buffer].setHi(m_moving_window_dir,new_hi);
+        amrex::Real new_lo = m_snapshot_domain_lab[i_buffer].hi(m_moving_window_dir) -
+                             num_z_cells_in_snapshot *
+                             dz_lab(warpx.getdt(lev), ref_ratio[m_moving_window_dir]);
+        m_snapshot_domain_lab[i_buffer].setLo(m_moving_window_dir, new_lo);
+    }
+
     // cell-centered index that corresponds to the hi-end of the lab-frame in the z-direction
     // Adding 0.5 dz_lab so that we obtain the cell-centered index consistent to the hi-end
     int snapshot_kindex_hi = static_cast<int>(floor(
@@ -487,8 +502,9 @@ BTDiagnostics::InitializeBufferData ( int i_buffer , int lev)
     m_snapshot_box[i_buffer].setSmall( m_moving_window_dir,
                                        snapshot_kindex_hi - (num_z_cells_in_snapshot-1) );
     // Setting hi k-index for the first buffer
-    m_buffer_k_index_hi[i_buffer] = m_snapshot_box[i_buffer].bigEnd(m_moving_window_dir);
-
+    if (restart == false) {
+        m_buffer_k_index_hi[i_buffer] = m_snapshot_box[i_buffer].bigEnd(m_moving_window_dir);
+    }
 }
 
 void
@@ -809,15 +825,18 @@ BTDiagnostics::PrepareFieldDataForOutput ()
                 // Check if the zslice is in domain
                 bool ZSliceInDomain = GetZSliceInDomainFlag (i_buffer, lev);
                 // Initialize and define field buffer multifab if buffer is empty
-                if (ZSliceInDomain) {
+                bool kindexInSnapshotBox = GetKIndexInSnapshotBoxFlag (i_buffer, lev);
+                if (kindexInSnapshotBox) {
                     if ( buffer_empty(i_buffer) ) {
-                        if ( m_buffer_flush_counter[i_buffer] == 0) {
+                        if ( m_buffer_flush_counter[i_buffer] == 0 || m_first_flush_after_restart[i_buffer] == 1) {
                             // Compute the geometry, snapshot lab-domain extent
                             // and box-indices
                             DefineSnapshotGeometry(i_buffer, lev);
                         }
                         DefineFieldBufferMultiFab(i_buffer, lev);
                     }
+                }
+                if (ZSliceInDomain) {
                     WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
                         m_current_z_lab[i_buffer] >= m_buffer_domain_lab[i_buffer].lo(m_moving_window_dir) and
                         m_current_z_lab[i_buffer] <= m_buffer_domain_lab[i_buffer].hi(m_moving_window_dir),
@@ -877,6 +896,7 @@ BTDiagnostics::SetSnapshotFullStatus (const int i_buffer)
 void
 BTDiagnostics::DefineFieldBufferMultiFab (const int i_buffer, const int lev)
 {
+    if (m_field_buffer_multifab_defined[i_buffer] == 1) return;
     auto & warpx = WarpX::GetInstance();
 
     const int hi_k_lab = m_buffer_k_index_hi[i_buffer];
@@ -931,12 +951,14 @@ BTDiagnostics::DefineFieldBufferMultiFab (const int i_buffer, const int lev)
         m_geom_output[i_buffer][lev] = amrex::refine( m_geom_output[i_buffer][lev-1],
                                                       warpx.RefRatio(lev-1) );
     }
+    m_field_buffer_multifab_defined[i_buffer] = 1;
 }
 
 
 void
 BTDiagnostics::DefineSnapshotGeometry (const int i_buffer, const int lev)
 {
+    if (m_snapshot_geometry_defined[i_buffer] == 1) return;
     auto & warpx = WarpX::GetInstance();
 
     if (lev == 0) {
@@ -954,6 +976,7 @@ BTDiagnostics::DefineSnapshotGeometry (const int i_buffer, const int lev)
         m_geom_snapshot[i_buffer][lev] = amrex::refine( m_geom_snapshot[i_buffer][lev-1],
                                                         warpx.RefRatio(lev-1) );
     }
+    m_snapshot_geometry_defined[i_buffer] = 1;
 }
 
 bool
@@ -975,6 +998,18 @@ BTDiagnostics::GetZSliceInDomainFlag (const int i_buffer, const int lev)
     return true;
 }
 
+
+bool
+BTDiagnostics::GetKIndexInSnapshotBoxFlag (const int i_buffer, const int lev)
+{
+    if (k_index_zlab(i_buffer, lev) >= m_snapshot_box[i_buffer].smallEnd(m_moving_window_dir) and
+        k_index_zlab(i_buffer, lev) <= m_snapshot_box[i_buffer].bigEnd(m_moving_window_dir)) {
+        return true;
+    }
+
+    return false;
+}
+
 void
 BTDiagnostics::Flush (int i_buffer)
 {
@@ -1070,7 +1105,9 @@ BTDiagnostics::Flush (int i_buffer)
 
     // Reset the buffer counter to zero after flushing out data stored in the buffer.
     ResetBufferCounter(i_buffer);
+    m_field_buffer_multifab_defined[i_buffer] = 0;
     IncrementBufferFlushCounter(i_buffer);
+    NullifyFirstFlush(i_buffer);
     // if particles are selected for output then update and reset counters
     if (m_output_species_names.size() > 0) {
         UpdateTotalParticlesFlushed(i_buffer);
@@ -1120,7 +1157,7 @@ void BTDiagnostics::MergeBuffersForPlotfile (int i_snapshot)
         std::string recent_Buffer_Level0_path = recent_Buffer_filepath + "/Level_0";
         std::string recent_Buffer_FabHeaderFilename = recent_Buffer_Level0_path + "/Cell_H";
         // Create directory only when the first buffer is flushed out.
-        if (m_buffer_flush_counter[i_snapshot] == 0 ) {
+        if (m_buffer_flush_counter[i_snapshot] == 0 || m_first_flush_after_restart[i_snapshot] == 1) {
             // Create Level_0 directory to store all Cell_D and Cell_H files
             if (!amrex::UtilCreateDirectory(snapshot_Level0_path, permission_flag_rwxrxrx) )
                 amrex::CreateDirectoryFailed(snapshot_Level0_path);
@@ -1162,7 +1199,7 @@ void BTDiagnostics::MergeBuffersForPlotfile (int i_snapshot)
             // Cell_D_<number> is padded with 5 zeros as that is the default AMReX output
             std::string new_snapshotFabFilename = amrex::Concatenate("Cell_D_", m_buffer_flush_counter[i_snapshot], amrex_fabfile_digits);
 
-            if ( m_buffer_flush_counter[i_snapshot] == 0) {
+            if (m_buffer_flush_counter[i_snapshot] == 0 || m_first_flush_after_restart[i_snapshot] == 1) {
                 std::rename(recent_Header_filename.c_str(), snapshot_Header_filename.c_str());
                 Buffer_FabHeader.SetFabName(0, Buffer_FabHeader.fodPrefix(0),
                                             new_snapshotFabFilename,
@@ -1208,7 +1245,7 @@ void BTDiagnostics::MergeBuffersForPlotfile (int i_snapshot)
                 m_buffer_flush_counter[i_snapshot],
                 amrex_partfile_digits);
 
-            if (m_buffer_flush_counter[i_snapshot] == 0) {
+            if (m_buffer_flush_counter[i_snapshot] == 0 || m_first_flush_after_restart[i_snapshot] == 1) {
                 BufferSpeciesHeader.set_DataIndex(0,0,m_buffer_flush_counter[i_snapshot]);
                 BufferSpeciesHeader.WriteHeader();
 
diff --git a/Source/Diagnostics/BoundaryScrapingDiagnostics.H b/Source/Diagnostics/BoundaryScrapingDiagnostics.H
index 01b0e8305..71ee7bfbe 100644
--- a/Source/Diagnostics/BoundaryScrapingDiagnostics.H
+++ b/Source/Diagnostics/BoundaryScrapingDiagnostics.H
@@ -52,7 +52,7 @@ private:
      *
      * This is not used for BoundaryScrapingDiagnostics: no field to output
      */
-     void InitializeBufferData (int i_buffer, int lev) override;
+     void InitializeBufferData (int i_buffer, int lev, bool restart=false) override;
     /** Initialize functors that point to the fields requested by the user.
      *
      * This is not used for BoundaryScrapingDiagnostics: no field to output
diff --git a/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp b/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp
index 5b5bf4761..792413d98 100644
--- a/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp
+++ b/Source/Diagnostics/BoundaryScrapingDiagnostics.cpp
@@ -74,7 +74,7 @@ BoundaryScrapingDiagnostics::InitializeFieldFunctors (int /*lev*/)
 }
 
 void
-BoundaryScrapingDiagnostics::InitializeBufferData (int /*i_buffer*/, int /*lev*/)
+BoundaryScrapingDiagnostics::InitializeBufferData (int /*i_buffer*/, int /*lev*/, bool /*restart*/)
 {
     // This function is usually used for field output
     // Nothing to do here for boundary scraping output,
diff --git a/Source/Diagnostics/Diagnostics.H b/Source/Diagnostics/Diagnostics.H
index e57447fea..75441d6d9 100644
--- a/Source/Diagnostics/Diagnostics.H
+++ b/Source/Diagnostics/Diagnostics.H
@@ -62,6 +62,8 @@ public:
     virtual void Flush (int i_buffer) = 0;
     /** Initialize pointers to main fields and allocate output multifab m_mf_output. */
     void InitData ();
+    void InitDataBeforeRestart ();
+    void InitDataAfterRestart ();
     /** Initialize functors that store pointers to the fields requested by the user.
      *
      * Derived classes MUST implement this function, and it must allocate m_all_field_functors
@@ -104,6 +106,75 @@ public:
     void FilterComputePackFlush (int step, bool force_flush=false);
     /** Whether the last timestep is always dumped */
     bool DoDumpLastTimestep () const {return  m_dump_last_timestep;}
+    /** Returns the number of snapshots used in BTD. For Full-Diagnostics, the value is 1*/
+    int getnumbuffers() {return m_num_buffers;}
+    /** Time in lab-frame associated with the ith snapshot
+     * \param[in] i_buffer index of the buffer
+     */
+    virtual amrex::Real gettlab ( [[maybe_unused]] int i_buffer) {return 0.;}
+    /** Set time in lab-frame for the ith snapshot
+     * \param[in] i_buffer index of the buffer
+     * \param[in] tlab value used to set the lab-frame time for the ith buffer
+     */
+    virtual void settlab ( [[maybe_unused]] int i_buffer, [[maybe_unused]] amrex::Real tlab) { }
+    /** Returns k-index in big end of the buffer box currently being filled for the ith snapshot in BTD
+     * \param[in] i_buffer index of the buffer
+     */
+    virtual int get_buffer_k_index_hi ( [[maybe_unused]] int i_buffer) {return 0; }
+    /** Set k-index of the buffer in the moving window direction currently being filled for the ith snapshot
+     * \param[in] i_buffer index of the buffer
+     * \param[in] kindex value used to set the k-index for the big end of the buffer box
+     */
+    virtual void set_buffer_k_index_hi ( [[maybe_unused]] int i_buffer, [[maybe_unused]] int kindex) { }
+    /** Returns lo-end of the lab-frame physical domain for the ith snapshot in BTD
+     * \param[in] i_buffer index of the buffer
+     * \param[in] idim dimension along which the lo-end of the lab-frame physical domain is returned
+     */
+    virtual amrex::Real get_snapshot_domain_lo ([[maybe_unused]] int i_buffer, [[maybe_unused]] int idim) {return 0.; }
+    /** Returns hi-end of the lab-frame physical domain for the ith snapshot in BTD
+     * \param[in] i_buffer index of the buffer
+     * \param[in] idim dimension along which the hi-end of the lab-frame physical domain is returned
+     */
+    virtual amrex::Real get_snapshot_domain_hi ([[maybe_unused]] int i_buffer, [[maybe_unused]] int idim) {return 0.; }
+    /** Sets lo-end of the lab-frame physical domain for the ith snapshot in BTD
+     * \param[in] i_buffer index of the buffer
+     * \param[in] idim dimension along which the lo-end of the lab-frame physical domain is set
+     * \param[in] domain_lab_lo value of the lo-end of the lab-frame physical domain for the ith snapshot
+     */
+    void setSnapshotDomainLo(int i_buffer, int idim,  amrex::Real domain_lab_lo) {m_snapshot_domain_lab[i_buffer].setLo(idim, domain_lab_lo); }
+    /** Sets hi-end of the lab-frame physical domain for the ith snapshot in BTD
+     * \param[in] i_buffer index of the buffer
+     * \param[in] idim dimension along which the hi-end of the lab-frame physical domain is set
+     * \param[in] domain_lab_hi value of the hi-end of the lab-frame physical domain for the ith snapshot
+     */
+    void setSnapshotDomainHi(int i_buffer, int idim,  amrex::Real domain_lab_hi) {m_snapshot_domain_lab[i_buffer].setHi(idim, domain_lab_hi); }
+    /** Returns counter for the number of times buffer data for the ith snapshot has been flushed
+     * \param[in] i_buffer index of the buffer
+     */
+    virtual int get_flush_counter ( [[maybe_unused]] int i_buffer) { return 0; }
+    /** Sets counter for the number of times buffer data for the ith snapshot has been flushed to parameter input in the function
+     * \param[in] i_buffer index of the buffer
+     * \param[in] flush_counter value of the counter
+     */
+    virtual void set_flush_counter ( [[maybe_unused]] int i_buffer, [[maybe_unused]] int flush_counter) { }
+    /** Returns 0/1 if the last valid Zslice for the ith snapshot has been filled (1) or not(0)
+     * \param[in] i_buffer index of the buffer
+     */
+    virtual int get_last_valid_Zslice ( [[maybe_unused]] int i_buffer) { return 0; }
+    /** Sets if the last valid Zslice for the ith snapshot has been filled (1) or not(0) with input parameter last_Valid_Zslice
+     * \param[in] i_buffer index of the buffer
+     * \param[in] last_valid_Zslice value to set m_lastValidZSlice in BTD for the ith snapshot (i_buffer)
+     */
+    virtual void set_last_valid_Zslice ( [[maybe_unused]] int i_buffer, [[maybe_unused]] int last_valid_Zslice) { }
+    /** Returns 0/1 if the snapshot is fully filled and BTD for that snapshot is complete(1) or not(0)
+     * \param[in] i_buffer index of the buffer
+     */
+    virtual int get_snapshot_full_flag ( [[maybe_unused]] int i_buffer) { return 0; }
+    /** Sets the value for m_snapshot_full for the ith snapshot in BTD using parameter snapshot_full
+     * \param[in] i_buffer index of the buffer
+     * \param[in] snapshot_full value to set m_snapshot_full in BTD for the ith snapshot (i_buffer)
+     */
+    virtual void set_snapshot_full ( [[maybe_unused]] int i_buffer, [[maybe_unused]] int snapshot_full) { }
 
 protected:
     /** Read Parameters of the base Diagnostics class */
@@ -114,7 +185,7 @@ protected:
      * \param[in] i_buffer index of buffer for which the output MultiFab is defined.
      * \param[in] lev level on which the output MultiFab is defined
      */
-    virtual void InitializeBufferData (int i_buffer, int lev) = 0;
+    virtual void InitializeBufferData (int i_buffer, int lev, bool restart=false) = 0;
     /** Initialize member variables and arrays specific to the diagnostics in the
      *  derived classes.(FullDiagnostics, BTDiagnostics)
      */
@@ -241,6 +312,9 @@ protected:
      *  over the total number of species selected for diagnostics.
      */
     amrex::Vector< amrex::Vector <int> > m_totalParticles_in_buffer;
+    /** Vector of user-defined physical region for diagnostics in lab-frame
+     *  for each back-transformed snapshot */
+    amrex::Vector<amrex::RealBox> m_snapshot_domain_lab;
 
 };
 
diff --git a/Source/Diagnostics/Diagnostics.cpp b/Source/Diagnostics/Diagnostics.cpp
index 25da0f464..f73844f62 100644
--- a/Source/Diagnostics/Diagnostics.cpp
+++ b/Source/Diagnostics/Diagnostics.cpp
@@ -295,6 +295,85 @@ Diagnostics::BaseReadParameters ()
 
 
 void
+Diagnostics::InitDataBeforeRestart ()
+{
+    // initialize member variables and arrays in base class::Diagnostics
+    InitBaseData();
+    // initialize member variables and arrays specific to each derived class
+    // (FullDiagnostics, BTDiagnostics, etc.)
+    DerivedInitData();
+}
+
+void
+Diagnostics::InitDataAfterRestart ()
+{
+    for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer) {
+        // loop over all levels
+        // This includes full diagnostics and BTD as well as cell-center functors for BTD.
+        // Note that the cell-centered data for BTD is computed for all levels and hence
+        // the corresponding functor is also initialized for all the levels
+        for (int lev = 0; lev < nmax_lev; ++lev) {
+            // allocate and initialize m_all_field_functors depending on diag type
+            InitializeFieldFunctors(lev);
+        }
+        // loop over the levels selected for output
+        // This includes all the levels for full diagnostics
+        // and only the coarse level (mother grid) for BTD
+        for (int lev = 0; lev < nlev_output; ++lev) {
+            // Initialize buffer data required for particle and/or fields
+            InitializeBufferData(i_buffer, lev, true);
+        }
+    }
+
+    amrex::ParmParse pp_diag_name(m_diag_name);
+    // default for writing species output is 1
+    int write_species = 1;
+    pp_diag_name.query("write_species", write_species);
+    if (write_species == 1) {
+        // When particle buffers, m_particle_boundary_buffer are included,
+        // they will be initialized here
+        InitializeParticleBuffer();
+        InitializeParticleFunctors();
+    }
+   if (write_species == 0) {
+        WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
+            m_format != "checkpoint",
+            "For checkpoint format, write_species flag must be 1."
+        );
+        // if user-defined value for write_species == 0, then clear species vector
+        for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer ) {
+            m_output_species.at(i_buffer).clear();
+        }
+        m_output_species_names.clear();
+    } else {
+        amrex::Vector <amrex::Real> dummy_val(AMREX_SPACEDIM);
+        if ( utils::parser::queryArrWithParser(
+                pp_diag_name, "diag_lo", dummy_val, 0, AMREX_SPACEDIM) ||
+             utils::parser::queryArrWithParser(
+                pp_diag_name, "diag_hi", dummy_val, 0, AMREX_SPACEDIM) ) {
+            // set geometry filter for particle-diags to true when the diagnostic domain-extent
+            // is specified by the user.
+            // Note that the filter is set for every ith snapshot, and the number of snapshots
+            // for full diagnostics is 1, while for BTD it is user-defined.
+            for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer ) {
+                for (auto& v : m_output_species.at(i_buffer)) {
+                    v.m_do_geom_filter = true;
+                }
+                // Disabling particle-io for reduced domain diagnostics by reducing
+                // the particle-diag vector to zero.
+                // This is a temporary fix until particle_buffer is supported in diagnostics.
+                m_output_species.at(i_buffer).clear();
+            }
+            std::string warnMsg = "For full diagnostics on a reduced domain, particle I/O is not ";
+            warnMsg += "supported, yet! Therefore, particle I/O is disabled for this diagnostics: ";
+            warnMsg += m_diag_name;
+            ablastr::warn_manager::WMRecordWarning("Diagnostics", warnMsg);
+        }
+    }
+}
+
+
+void
 Diagnostics::InitData ()
 {
     // initialize member variables and arrays in base class::Diagnostics
diff --git a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
index 4a8a57bf3..178219107 100644
--- a/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
+++ b/Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp
@@ -10,6 +10,7 @@
 #include "Utils/TextMsg.H"
 #include "Utils/WarpXProfilerWrapper.H"
 #include "WarpX.H"
+#include "Diagnostics/MultiDiagnostics.H"
 
 #include <AMReX.H>
 #include <AMReX_Box.H>
@@ -306,6 +307,27 @@ FlushFormatPlotfile::WriteWarpXHeader(
         HeaderFile << warpx.getdo_moving_window() << "\n";
 
         HeaderFile << warpx.time_of_last_gal_shift << "\n";
+
+        for (int idiag = 0; idiag < warpx.GetMultiDiags().GetTotalDiags(); ++idiag)
+        {
+            if( warpx.GetMultiDiags().diagstypes(idiag) == DiagTypes::BackTransformed )
+            {
+                auto& diag = warpx.GetMultiDiags().GetDiag(idiag);
+                for (int i_buffer=0; i_buffer<diag.getnumbuffers(); ++i_buffer){
+                    HeaderFile << diag.gettlab(i_buffer) << "\n";
+                    HeaderFile << diag.get_buffer_k_index_hi(i_buffer) << "\n";
+                    for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+                        HeaderFile << diag.get_snapshot_domain_lo(i_buffer, idim) << "\n";
+                    }
+                    for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+                        HeaderFile << diag.get_snapshot_domain_hi(i_buffer, idim) << "\n";
+                    }
+                    HeaderFile << diag.get_flush_counter(i_buffer) << "\n";
+                    HeaderFile << diag.get_last_valid_Zslice(i_buffer) << "\n";
+                    HeaderFile << diag.get_snapshot_full_flag(i_buffer) << "\n";
+                }
+            }
+        }
     }
 }
 
diff --git a/Source/Diagnostics/FullDiagnostics.H b/Source/Diagnostics/FullDiagnostics.H
index 8620af75e..a20b44397 100644
--- a/Source/Diagnostics/FullDiagnostics.H
+++ b/Source/Diagnostics/FullDiagnostics.H
@@ -52,7 +52,7 @@ private:
       * \param[in] i_buffer index of a back-transformed snapshot
       * \param[in] lev level on which source multifabs are defined
       */
-    void InitializeBufferData ( int i_buffer, int lev ) override;
+    void InitializeBufferData ( int i_buffer, int lev, bool restart=false ) override;
     /** Initialize functors that store pointers to the RZ fields requested by the user.
      * \param[in] lev level on which the vector of unique_ptrs to field functors is initialized.
      */
diff --git a/Source/Diagnostics/FullDiagnostics.cpp b/Source/Diagnostics/FullDiagnostics.cpp
index 99e28a7f8..332b93986 100644
--- a/Source/Diagnostics/FullDiagnostics.cpp
+++ b/Source/Diagnostics/FullDiagnostics.cpp
@@ -443,7 +443,8 @@ FullDiagnostics::AddRZModesToOutputNames (const std::string& field, int ncomp){
 
 
 void
-FullDiagnostics::InitializeBufferData (int i_buffer, int lev ) {
+FullDiagnostics::InitializeBufferData (int i_buffer, int lev, bool restart ) {
+    amrex::ignore_unused(restart);
     auto & warpx = WarpX::GetInstance();
     amrex::RealBox diag_dom;
     bool use_warpxba = true;
diff --git a/Source/Diagnostics/MultiDiagnostics.H b/Source/Diagnostics/MultiDiagnostics.H
index 06a2b87b4..f9907b8bd 100644
--- a/Source/Diagnostics/MultiDiagnostics.H
+++ b/Source/Diagnostics/MultiDiagnostics.H
@@ -37,6 +37,9 @@ public:
     void InitializeFieldFunctors (int lev);
     /** Start a new iteration, i.e., dump has not been done yet. */
     void NewIteration ();
+    Diagnostics& GetDiag(int idiag) {return *alldiags[idiag]; }
+    int GetTotalDiags() {return ndiags;}
+    DiagTypes diagstypes(int idiag) {return diags_types[idiag];}
 private:
     /** Vector of pointers to all diagnostics */
     amrex::Vector<std::unique_ptr<Diagnostics> > alldiags;
diff --git a/Source/Diagnostics/WarpXIO.cpp b/Source/Diagnostics/WarpXIO.cpp
index 2e55dbc01..01365cba4 100644
--- a/Source/Diagnostics/WarpXIO.cpp
+++ b/Source/Diagnostics/WarpXIO.cpp
@@ -16,6 +16,7 @@
 #include "Utils/TextMsg.H"
 #include "Utils/WarpXProfilerWrapper.H"
 #include "WarpX.H"
+#include "Diagnostics/MultiDiagnostics.H"
 
 #include <ablastr/utils/Communication.H>
 
@@ -223,6 +224,56 @@ WarpX::InitFromCheckpoint ()
 
         is >> time_of_last_gal_shift;
         GotoNextLine(is);
+
+
+        auto & warpx = WarpX::GetInstance();
+        for (int idiag = 0; idiag < warpx.GetMultiDiags().GetTotalDiags(); ++idiag)
+        {
+            if( warpx.GetMultiDiags().diagstypes(idiag) == DiagTypes::BackTransformed )
+            {
+                auto& diag = warpx.GetMultiDiags().GetDiag(idiag);
+                if (diag.getnumbuffers() > 0) {
+                    diag.InitDataBeforeRestart();
+                    for (int i_buffer=0; i_buffer<diag.getnumbuffers(); ++i_buffer){
+                        amrex::Real tlab;
+                        is >> tlab;
+                        diag.settlab(i_buffer, tlab);
+                        int kindex_hi;
+                        is >> kindex_hi;
+                        diag.set_buffer_k_index_hi(i_buffer, kindex_hi);
+
+                        for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+                            amrex::Real snapshot_lo;
+                            is >> snapshot_lo;
+                            diag.setSnapshotDomainLo(i_buffer, idim, snapshot_lo);
+                        }
+                        for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+                            amrex::Real snapshot_hi;
+                            is >> snapshot_hi;
+                            diag.setSnapshotDomainHi(i_buffer, idim, snapshot_hi);
+                        }
+
+                        int flush_counter;
+                        is >> flush_counter;
+                        diag.set_flush_counter(i_buffer, flush_counter);
+
+                        int last_valid_Zslice;
+                        is >> last_valid_Zslice;
+                        diag.set_last_valid_Zslice(i_buffer, last_valid_Zslice);
+
+                        int snapshot_full_flag;
+                        is >> snapshot_full_flag;
+                        diag.set_snapshot_full(i_buffer, snapshot_full_flag);
+
+                    }
+                    diag.InitDataAfterRestart();
+                } else {
+                    diag.InitData();
+                }
+            } else {
+                warpx.GetMultiDiags().GetDiag(idiag).InitData();
+            }
+        }
     }
 
     const int nlevs = finestLevel()+1;
diff --git a/Source/Initialization/WarpXInitData.cpp b/Source/Initialization/WarpXInitData.cpp
index 4b7cb4df6..ede8cc3c8 100644
--- a/Source/Initialization/WarpXInitData.cpp
+++ b/Source/Initialization/WarpXInitData.cpp
@@ -387,12 +387,14 @@ WarpX::InitData ()
         ComputeDt();
         WarpX::PrintDtDxDyDz();
         InitFromScratch();
+        InitDiagnostics();
     }
     else
     {
         InitFromCheckpoint();
         WarpX::PrintDtDxDyDz();
         PostRestart();
+        reduced_diags->InitData();
     }
 
     ComputeMaxStep();
@@ -409,8 +411,6 @@ WarpX::InitData ()
         m_macroscopic_properties->InitData();
     }
 
-    InitDiagnostics();
-
     if (ParallelDescriptor::IOProcessor()) {
         std::cout << "\nGrids Summary:\n";
         printGridSummary(std::cout, 0, finestLevel());
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 1358630b3..10e74671f 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -98,6 +98,7 @@ public:
 
     MultiParticleContainer& GetPartContainer () { return *mypc; }
     MacroscopicProperties& GetMacroscopicProperties () { return *m_macroscopic_properties; }
+    MultiDiagnostics& GetMultiDiags () {return *multi_diags;}
 
     ParticleBoundaryBuffer& GetParticleBoundaryBuffer () { return *m_particle_boundary_buffer; }