From 6193a36eb2f94ccc27b4070400df4edd0a227e8d Mon Sep 17 00:00:00 2001
From: RevathiJambunathan <revanathan@gmail.com>
Date: Thu, 12 Sep 2024 15:25:48 -0700
Subject: [PATCH] skeleton sprinkled around

---
 Source/Diagnostics/Diagnostics.H        |  2 ++
 Source/Diagnostics/Diagnostics.cpp      | 11 +++++++++++
 Source/Diagnostics/FullDiagnostics.cpp  | 16 +++++++++++++++-
 Source/Diagnostics/MultiDiagnostics.cpp |  2 ++
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/Source/Diagnostics/Diagnostics.H b/Source/Diagnostics/Diagnostics.H
index 20550364fb7..81553433dde 100644
--- a/Source/Diagnostics/Diagnostics.H
+++ b/Source/Diagnostics/Diagnostics.H
@@ -267,6 +267,8 @@ protected:
      */
     amrex::Vector< amrex::Vector< amrex::MultiFab > >  m_mf_output;
 
+    amrex::Vector< amrex::Vector <amrex::MultiFab > > m_sum_mf_output;
+
     /** Geometry that defines the domain attributes corresponding to output multifab.
      *  Specifically, the user-defined physical co-ordinates for the diagnostics
      *  is used to construct the geometry information for each MultiFab at
diff --git a/Source/Diagnostics/Diagnostics.cpp b/Source/Diagnostics/Diagnostics.cpp
index dc28aeda095..e4e4a98f032 100644
--- a/Source/Diagnostics/Diagnostics.cpp
+++ b/Source/Diagnostics/Diagnostics.cpp
@@ -534,6 +534,11 @@ Diagnostics::InitBaseData ()
     for (int i = 0; i < m_num_buffers; ++i) {
         m_mf_output[i].resize( nmax_lev );
     }
+    m_sum_mf_output.resize(m_num_buffers);
+    for (int i = 0; i < m_num_buffers; ++i) {
+        m_sum_mf_output[i].resize( nmax_lev );
+    }
+
 
     // allocate vector of geometry objects corresponding to each output multifab.
     m_geom_output.resize( m_num_buffers );
@@ -558,6 +563,8 @@ Diagnostics::ComputeAndPack ()
 
     auto & warpx = WarpX::GetInstance();
 
+    bool do_sum_now = cur_step >= time_ave 
+
     // compute the necessary fields and store result in m_mf_output.
     for (int i_buffer = 0; i_buffer < m_num_buffers; ++i_buffer) {
         for(int lev=0; lev<nlev_output; lev++){
@@ -568,7 +575,11 @@ Diagnostics::ComputeAndPack ()
                 // a diagnostics and writes in one or more components of the output
                 // multifab m_mf_output[lev].
                 m_all_field_functors[lev][icomp]->operator()(m_mf_output[i_buffer][lev], icomp_dst, i_buffer);
+                if (do_sum_now)
                 // update the index of the next component to fill
+                int scalar_a = 1;
+                // call amrex sax operation to do the following
+                m_sum_mf_output[i_buffer][lev] += scalar_a * m_mf_output[i_buffer][lev]
                 icomp_dst += m_all_field_functors[lev][icomp]->nComp();
             }
             // Check that the proper number of components of mf_avg were updated.
diff --git a/Source/Diagnostics/FullDiagnostics.cpp b/Source/Diagnostics/FullDiagnostics.cpp
index bcf613f49b0..9e207a03bd9 100644
--- a/Source/Diagnostics/FullDiagnostics.cpp
+++ b/Source/Diagnostics/FullDiagnostics.cpp
@@ -142,6 +142,18 @@ FullDiagnostics::Flush ( int i_buffer, bool /* force_flush */ )
         m_output_species.at(i_buffer), nlev_output, m_file_prefix,
         m_file_min_digits, m_plot_raw_fields, m_plot_raw_fields_guards);
 
+    if (time_average) {
+       //call amrex option to 
+       m_sum_mf_output.at(i_buffer).mult(1./time_average_step);
+           m_flush_format->WriteToFile(
+        m_varnames, m_sum_mf_output.at(i_buffer), m_geom_output.at(i_buffer), warpx.getistep(),
+        warpx.gett_new(0),
+        m_output_species.at(i_buffer), nlev_output, m_file_prefix,
+        m_file_min_digits, m_plot_raw_fields, m_plot_raw_fields_guards);
+
+
+    }
+
     FlushRaw();
 }
 
@@ -166,6 +178,7 @@ FullDiagnostics::DoComputeAndPack (int step, bool force_flush)
     // Data must be computed and packed for full diagnostics
     // whenever the data needs to be flushed.
     return (force_flush || m_intervals.contains(step+1));
+    // add logic here to do compute and pack if m_intervals.containes (step+1-time_average_period) or if (cur_step>time_average_startstep)
 }
 
 void
@@ -609,7 +622,8 @@ FullDiagnostics::InitializeBufferData (int i_buffer, int lev, bool restart ) {
     const int ngrow = (m_format == "sensei" || m_format == "ascent") ? 1 : 0;
     int const ncomp = static_cast<int>(m_varnames.size());
     m_mf_output[i_buffer][lev] = amrex::MultiFab(ba, dmap, ncomp, ngrow);
-
+    m_sum_mf_output[i_buffer][lev] = amrex::MultiFab(ba, dmap, ncomp, ngrow);
+    m_sum_mf_output[i_buffer][lev].setVal(0.);
     if (lev == 0) {
         // The extent of the domain covered by the diag multifab, m_mf_output
         //default non-periodic geometry for diags
diff --git a/Source/Diagnostics/MultiDiagnostics.cpp b/Source/Diagnostics/MultiDiagnostics.cpp
index ea14919c713..ecba8831b20 100644
--- a/Source/Diagnostics/MultiDiagnostics.cpp
+++ b/Source/Diagnostics/MultiDiagnostics.cpp
@@ -23,6 +23,8 @@ MultiDiagnostics::MultiDiagnostics ()
     for (int i=0; i<ndiags; i++){
         if ( diags_types[i] == DiagTypes::Full ){
             alldiags[i] = std::make_unique<FullDiagnostics>(i, diags_names[i]);
+        } else if ( diags_types[i] == DiagTypes::TimeAveraged ){
+            alldiags[i] = std::make_unique<FullDiagnostics>(i, diags_names[i]);
         } else if ( diags_types[i] == DiagTypes::BackTransformed ){
             alldiags[i] = std::make_unique<BTDiagnostics>(i, diags_names[i]);
         } else if ( diags_types[i] == DiagTypes::BoundaryScraping ){