EAMxx: allow diags in IO to perform some ops at beginning of a timestep

Author: bartgol

components/eamxx/src/control/atmosphere_driver.cpp

@@ -1597,6 +1597,16 @@ void AtmosphereDriver::run (const int dt) {
   //       nano-opt of removing the call for the 1st timestep.
   reset_accumulated_fields();
 
+  // Tell the output managers that we're starting a timestep. This is usually
+  // a no-op, but some diags *may* require to do something. E.g., a diag that
+  // computes tendency of an arbitrary quantity may want to store a copy of
+  // that quantity at the beginning of the timestep. Or they may need to store
+  // the timestamp at the beginning of the timestep, so that we can compute
+  // dt at the end.
+  for (auto it : m_output_managers) {
+    it.init_timestep(m_current_ts);
+  }
+
   // The class AtmosphereProcessGroup will take care of dispatching arguments to
   // the individual processes, which will be called in the correct order.
   m_atm_process_group->run(dt);

components/eamxx/src/share/atm_process/atmosphere_diagnostic.hpp

@@ -51,6 +51,8 @@ class AtmosphereDiagnostic : public AtmosphereProcess
   // Getting the diagnostic output
   Field get_diagnostic () const;
 
+  virtual void init_timestep (const util::TimeStamp& /* start_of_step */) {}
+
   void compute_diagnostic (const double dt = 0);
 protected:
 

components/eamxx/src/share/io/scorpio_output.cpp

@@ -361,6 +361,14 @@ void AtmosphereOutput::init()
   register_views();
 }
 
+void AtmosphereOutput::
+init_timestep (const util::TimeStamp& start_of_step)
+{
+  for (auto& it : m_diagnostics) {
+    it.second->init_timestep(start_of_step);
+  }
+}
+
 void AtmosphereOutput::
 run (const std::string& filename,
      const bool output_step, const bool checkpoint_step,

components/eamxx/src/share/io/scorpio_output.hpp

@@ -139,6 +139,8 @@ class AtmosphereOutput
   void reset_dev_views();
   void update_avg_cnt_view(const Field&, view_1d_dev& dev_view);
   void setup_output_file (const std::string& filename, const std::string& fp_precision, const scorpio::FileMode mode);
+
+  void init_timestep (const util::TimeStamp& start_of_step);
   void run (const std::string& filename,
             const bool output_step, const bool checkpoint_step,
             const int nsteps_since_last_output,

components/eamxx/src/share/io/scream_output_manager.cpp

@@ -265,6 +265,9 @@ setup (const ekat::Comm& io_comm, const ekat::ParameterList& params,
   {
     // In order to trigger a t0 write, we need to have next_write_ts matching run_t0
     m_output_control.next_write_ts = m_run_t0;
+    // This is in case some diags need to init the timestep. Most likely, their output
+    // is meaningless at t0, but they may still require the start-of-step timestamp to be valid
+    init_timestep(m_run_t0);
     this->run(m_run_t0);
   }
 
@@ -282,6 +285,22 @@ add_global (const std::string& name, const ekat::any& global) {
 }
 
 /*===============================================================================================*/
+void OutputManager::init_timestep (const util::TimeStamp& start_of_step)
+{
+  // In case output is disabled, no point in doing anything else
+  if (not m_output_control.output_enabled()) {
+    return;
+  }
+
+  if (m_atm_logger) {
+    m_atm_logger->debug("[OutputManager::init_timestep] filename_prefix: " + m_filename_prefix + "\n");
+  }
+
+  for (auto s : m_output_streams) {
+    s->init_timestep(start_of_step);
+  }
+}
+
 void OutputManager::run(const util::TimeStamp& timestamp)
 {
   // In case output is disabled, no point in doing anything else

components/eamxx/src/share/io/scream_output_manager.hpp

@@ -112,6 +112,8 @@ class OutputManager
       m_atm_logger = atm_logger;
   }
   void add_global (const std::string& name, const ekat::any& global);
+
+  void init_timestep (const util::TimeStamp& start_of_step);
   void run (const util::TimeStamp& current_ts);
   void finalize();
 

components/eamxx/src/share/io/tests/io_diags.cpp

@@ -26,15 +26,6 @@
 
 namespace scream {
 
-void multiply (const Field& f, const double v) {
-  auto data = f.get_internal_view_data<Real,Host>();
-  auto nscalars = f.get_header().get_alloc_properties().get_num_scalars();
-  for (int i=0; i<nscalars; ++i) {
-    data[i] *= v;
-  }
-  f.sync_to_dev();
-}
-
 class MyDiag : public AtmosphereDiagnostic
 {
 public:
@@ -69,16 +60,24 @@ class MyDiag : public AtmosphereDiagnostic
     FieldIdentifier fid (name(), lt, units, grid_name);
     m_diagnostic_output = Field(fid);
     m_diagnostic_output.allocate_view();
+    m_one = m_diagnostic_output.clone("one");
+    m_one.deep_copy(1.0);
+  }
+
+  void init_timestep (const util::TimeStamp& start_of_step) {
+    m_t_beg = start_of_step;
   }
 
 protected:
 
   void compute_diagnostic_impl () {
     const auto& f_in  = get_field_in(m_f_in);
 
-    m_diagnostic_output.deep_copy<Host>(f_in);
-    multiply(m_diagnostic_output,2.0);
-    m_diagnostic_output.sync_to_dev();
+    const auto& t = f_in.get_header().get_tracking().get_time_stamp();
+    const double dt = t - m_t_beg;
+
+    m_diagnostic_output.deep_copy(f_in);
+    m_diagnostic_output.update(m_one,dt,2.0);
   }
 
   void initialize_impl (const RunType /* run_type */ ) {
@@ -92,12 +91,19 @@ class MyDiag : public AtmosphereDiagnostic
   int m_num_cols, m_num_levs;
 
   std::string m_f_in;
+
+  util::TimeStamp m_t_beg;
+  Field m_one;
 };
 
 util::TimeStamp get_t0 () {
   return util::TimeStamp({2023,2,17},{0,0,0});
 }
 
+constexpr double get_dt () {
+  return 10;
+};
+
 std::shared_ptr<const GridsManager>
 get_gm (const ekat::Comm& comm)
 {
@@ -167,6 +173,7 @@ void write (const int seed, const ekat::Comm& comm)
 
   // Time advance parameters
   auto t0 = get_t0();
+  auto dt = get_dt();
 
   // Create some fields
   auto fm = get_fm(grid,t0,seed);
@@ -193,7 +200,15 @@ void write (const int seed, const ekat::Comm& comm)
   om.setup(comm,om_pl,fm,gm,t0,t0,false);
 
   // Run output manager
-  om.run (t0);
+  for (auto it : *fm) {
+    auto& f = *it.second;
+    Field one = f.clone("one");
+    one.deep_copy(1.0);
+    f.get_header().get_tracking().update_time_stamp(t0+dt);
+    f.update(one,1.0,1.0);
+  }
+  om.init_timestep(t0);
+  om.run (t0+dt);
 
   // Close file and cleanup
   om.finalize();
@@ -211,6 +226,7 @@ void read (const int seed, const ekat::Comm& comm)
   // Get initial fields
   auto fm0 = get_fm(grid,t0,seed);
   auto fm  = get_fm(grid,t0,seed,true);
+
   std::vector<std::string> fnames;
   std::string f_name;
   for (auto it : *fm) {
@@ -220,6 +236,10 @@ void read (const int seed, const ekat::Comm& comm)
       f_name = fn;
     }
   }
+
+  auto f0 = fm0->get_field(f_name).clone();
+  auto f  = fm->get_field(f_name);
+
   // Sanity check
   REQUIRE (f_name!="");
 
@@ -235,17 +255,25 @@ void read (const int seed, const ekat::Comm& comm)
   reader_pl.set("Field Names",fnames);
   AtmosphereInput reader(reader_pl,fm);
 
-  reader.read_variables();
+  Field one = f0.clone("one");
+  one.deep_copy(1.0);
+  for (int i=0; i<2; ++i) {
+    reader.read_variables(i);
 
-  // Check regular field is correct
-  auto f0 = fm0->get_field(f_name).clone();
-  auto f  = fm->get_field(f_name);
-  REQUIRE (views_are_equal(f,f0));
+    // Check regular field is correct
+    REQUIRE (views_are_equal(f,f0));
 
-  // Check diag field is correct
-  multiply (f0,2.0);
-  auto d = fm->get_field("MyDiag");
-  REQUIRE (views_are_equal(d,f0));
+    // Check diag field is correct
+    const auto t = t0+i*get_dt();
+    const double dt = t-t0;
+    auto d = fm->get_field("MyDiag");
+    auto d0 = f0.clone();
+    d0.update(one,dt,2.0);
+    REQUIRE (views_are_equal(d,d0));
+
+    // Update f0
+    f0.update(one,1.0,1.0);
+  }
 }
 
 TEST_CASE ("io_diags") {