Merge branch 'bartgol/eamxx/cosp-fixes' (PR #6933)

In COSP, fix a bug related to timestamp handling, and do some upgrades/cleanup.

[non-BFB] only for cosp outputs

Author: bartgol

components/eamxx/src/physics/cosp/cosp_functions.hpp

@@ -31,13 +31,15 @@ namespace scream {
         };
         inline void main(
                 const Int ncol, const Int nsubcol, const Int nlay, const Int ntau, const Int nctp, const Int ncth, const Real emsfc_lw,
-                view_1d<const Real>& sunlit , view_1d<const Real>& skt,
-                view_2d<const Real>& T_mid  , view_2d<const Real>& p_mid  , view_2d<const Real>& p_int,
-                view_2d<const Real>& z_mid  , view_2d<const Real>& qv     , view_2d<const Real>& qc     , view_2d<const Real>& qi,
-                view_2d<const Real>& cldfrac,
-                view_2d<const Real>& reff_qc, view_2d<const Real>& reff_qi,
-                view_2d<const Real>& dtau067, view_2d<const Real>& dtau105,
-                view_1d<Real>& isccp_cldtot , view_3d<Real>& isccp_ctptau, view_3d<Real>& modis_ctptau, view_3d<Real>& misr_cthtau) {
+                const view_1d<const Real>& sunlit , const view_1d<const Real>& skt,
+                const view_2d<const Real>& T_mid  , const view_2d<const Real>& p_mid  ,
+                const view_2d<const Real>& p_int,  const view_2d<const Real>& z_mid,
+                const view_2d<const Real>& qv     , const view_2d<const Real>& qc,
+                const view_2d<const Real>& qi, const view_2d<const Real>& cldfrac,
+                const view_2d<const Real>& reff_qc, const view_2d<const Real>& reff_qi,
+                const view_2d<const Real>& dtau067, const view_2d<const Real>& dtau105,
+                const view_1d<Real>& isccp_cldtot , const view_3d<Real>& isccp_ctptau,
+                const view_3d<Real>& modis_ctptau, const view_3d<Real>& misr_cthtau) {
 
             // Make host copies and permute data as needed
             lview_host_2d

components/eamxx/src/physics/cosp/eamxx_cosp.cpp

@@ -1,6 +1,7 @@
 #include "eamxx_cosp.hpp"
 #include "cosp_functions.hpp"
 #include "share/property_checks/field_within_interval_check.hpp"
+#include "physics/share/physics_constants.hpp"
 
 #include "ekat/ekat_assert.hpp"
 #include "ekat/util/ekat_units.hpp"
@@ -93,6 +94,12 @@ void Cosp::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
   add_field<Computed>("modis_ctptau", scalar4d_ctptau, percent, grid_name, 1);
   add_field<Computed>("misr_cthtau", scalar4d_cthtau, percent, grid_name, 1);
   add_field<Computed>("cosp_sunlit", scalar2d, nondim, grid_name);
+
+  // We can allocate these now
+  m_z_mid = Field(FieldIdentifier("z_mid",scalar3d_mid,m,grid_name));
+  m_z_int = Field(FieldIdentifier("z_int",scalar3d_int,m,grid_name));
+  m_z_mid.allocate_view();
+  m_z_int.allocate_view();
 }
 
 // =========================================================================================
@@ -131,106 +138,127 @@ void Cosp::run_impl (const double dt)
   // Make sure cosp frequency is multiple of rad frequency?
 
   // Compare frequency in steps with current timestep
-  auto ts = timestamp();
-  auto update_cosp = cosp_do(cosp_freq_in_steps, ts.get_num_steps());
-
-  // Get fields from field manager; note that we get host views because this
-  // interface serves primarily as a wrapper to a c++ to f90 bridge for the COSP
-  // all then need to be copied to layoutLeft views to permute the indices for
-  // F90.
-  //
-  // Need to make sure device data is synced to host?
-  get_field_in("qv").sync_to_host();
-  get_field_in("qc").sync_to_host();
-  get_field_in("qi").sync_to_host();
-  get_field_in("sunlit").sync_to_host();
-  get_field_in("surf_radiative_T").sync_to_host();
-  get_field_in("T_mid").sync_to_host();
-  get_field_in("p_mid").sync_to_host();
-  get_field_in("p_int").sync_to_host();
-  get_field_in("cldfrac_rad").sync_to_host();
-  get_field_in("eff_radius_qc").sync_to_host();
-  get_field_in("eff_radius_qi").sync_to_host();
-  get_field_in("dtau067").sync_to_host();
-  get_field_in("dtau105").sync_to_host();
-  get_field_in("phis").sync_to_host();
-  get_field_in("pseudo_density").sync_to_host();
-
-  auto qv      = get_field_in("qv").get_view<const Real**, Host>();
-  auto qc      = get_field_in("qc").get_view<const Real**, Host>();
-  auto qi      = get_field_in("qi").get_view<const Real**, Host>();
-  auto sunlit  = get_field_in("sunlit").get_view<const Real*, Host>();
-  auto skt     = get_field_in("surf_radiative_T").get_view<const Real*, Host>();
-  auto T_mid   = get_field_in("T_mid").get_view<const Real**, Host>();
-  auto p_mid   = get_field_in("p_mid").get_view<const Real**, Host>();
-  auto p_int   = get_field_in("p_int").get_view<const Real**, Host>();
-  auto phis    = get_field_in("phis").get_view<const Real*, Host>();
-  auto pseudo_density = get_field_in("pseudo_density").get_view<const Real**, Host>();
-  auto cldfrac = get_field_in("cldfrac_rad").get_view<const Real**, Host>();
-  auto reff_qc = get_field_in("eff_radius_qc").get_view<const Real**, Host>();
-  auto reff_qi = get_field_in("eff_radius_qi").get_view<const Real**, Host>();
-  auto dtau067 = get_field_in("dtau067").get_view<const Real**, Host>();
-  auto dtau105 = get_field_in("dtau105").get_view<const Real**, Host>();
-  auto isccp_cldtot = get_field_out("isccp_cldtot").get_view<Real*, Host>();
-  auto isccp_ctptau = get_field_out("isccp_ctptau").get_view<Real***, Host>();
-  auto modis_ctptau = get_field_out("modis_ctptau").get_view<Real***, Host>();
-  auto misr_cthtau  = get_field_out("misr_cthtau").get_view<Real***, Host>();
-  auto cosp_sunlit  = get_field_out("cosp_sunlit").get_view<Real*, Host>();  // Copy of sunlit flag with COSP frequency for proper averaging
-
-  // Compute heights
-  const auto z_mid = CospFunc::view_2d<Real>("z_mid", m_num_cols, m_num_levs);
-  const auto z_int = CospFunc::view_2d<Real>("z_int", m_num_cols, m_num_levs+1);
-  const auto dz = z_mid;  // reuse tmp memory for dz
-  const auto ncol = m_num_cols;
-  const auto nlev = m_num_levs;
-  // calculate_z_int contains a team-level parallel_scan, which requires a special policy
-  // TODO: do this on device?
-  const auto scan_policy = ekat::ExeSpaceUtils<KTH::ExeSpace>::get_thread_range_parallel_scan_team_policy(ncol, nlev);
-  Kokkos::parallel_for(scan_policy, KOKKOS_LAMBDA (const KTH::MemberType& team) {
-      const int i = team.league_rank();
-      const auto dz_s    = ekat::subview(dz,    i);
-      const auto p_mid_s = ekat::subview(p_mid, i);
-      const auto T_mid_s = ekat::subview(T_mid, i);
-      const auto qv_s = ekat::subview(qv, i);
-      const auto z_int_s = ekat::subview(z_int, i);
-      const auto z_mid_s = ekat::subview(z_mid, i);
-      const Real z_surf  = phis(i) / 9.81;
-      const auto pseudo_density_s = ekat::subview(pseudo_density, i);
-      PF::calculate_dz(team, pseudo_density_s, p_mid_s, T_mid_s, qv_s, dz_s);
-      team.team_barrier();
-      PF::calculate_z_int(team,nlev,dz_s,z_surf,z_int_s);
-      team.team_barrier();
-      PF::calculate_z_mid(team,nlev,z_int_s,z_mid_s);
-      team.team_barrier();
-  });
+  // NOTE: timestamp() returns the time RIGHT BEFORE the call to run
+  auto end_of_step = timestamp()+dt;
+  auto update_cosp = cosp_do(cosp_freq_in_steps, end_of_step.get_num_steps());
 
   // Call COSP wrapper routines
   if (update_cosp) {
+    // Get fields from field manager; note that we get host views because this
+    // interface serves primarily as a wrapper to a c++ to f90 bridge for the COSP
+    // all then need to be copied to layoutLeft views to permute the indices for
+    // F90.
+
+    // Ensure host data of input fields is up to date
+    get_field_in("qv").sync_to_host();
+    get_field_in("qc").sync_to_host();
+    get_field_in("qi").sync_to_host();
+    get_field_in("sunlit").sync_to_host();
+    get_field_in("surf_radiative_T").sync_to_host();
+    get_field_in("T_mid").sync_to_host();
+    get_field_in("p_int").sync_to_host();
+    get_field_in("cldfrac_rad").sync_to_host();
+    get_field_in("eff_radius_qc").sync_to_host();
+    get_field_in("eff_radius_qi").sync_to_host();
+    get_field_in("dtau067").sync_to_host();
+    get_field_in("dtau105").sync_to_host();
+
+    // Compute z_mid
+    const auto T_mid_d = get_field_in("T_mid").get_view<const Real**>();
+    const auto qv_d  = get_field_in("qv").get_view<const Real**>();
+    const auto p_mid_d = get_field_in("p_mid").get_view<const Real**>();
+    const auto phis_d  = get_field_in("phis").get_view<const Real*>();
+    const auto pseudo_density_d = get_field_in("pseudo_density").get_view<const Real**>();
+    const auto z_mid_d = m_z_mid.get_view<Real**>();
+    const auto z_int_d = m_z_int.get_view<Real**>();
+    const auto ncol = m_num_cols;
+    const auto nlev = m_num_levs;
+
+    using KT       = KokkosTypes<DefaultDevice>;
+    using ExeSpace = typename KT::ExeSpace;
+    using ESU      = ekat::ExeSpaceUtils<ExeSpace>;
+    using PF       = scream::PhysicsFunctions<DefaultDevice>;
+
+    const auto scan_policy = ESU::get_thread_range_parallel_scan_team_policy(ncol, nlev);
+    const auto g = physics::Constants<Real>::gravit;
+    Kokkos::parallel_for(scan_policy, KOKKOS_LAMBDA (const KT::MemberType& team) {
+        const int i = team.league_rank();
+        const auto p_mid_s = ekat::subview(p_mid_d, i);
+        const auto T_mid_s = ekat::subview(T_mid_d, i);
+        const auto qv_s = ekat::subview(qv_d, i);
+        const auto z_int_s = ekat::subview(z_int_d, i);
+        const auto z_mid_s = ekat::subview(z_mid_d, i);
+        const Real z_surf  = phis_d(i) / g;
+        const auto pseudo_density_s = ekat::subview(pseudo_density_d, i);
+
+        // 1. Compute dz (recycle z_mid_s as a temporary)
+        const auto dz_s = z_mid_s; // 
+        PF::calculate_dz(team, pseudo_density_s, p_mid_s, T_mid_s, qv_s, dz_s);
+        team.team_barrier();
+
+        // 2. Compute z_int (vertical scan)
+        PF::calculate_z_int(team,nlev,dz_s,z_surf,z_int_s);
+        team.team_barrier();
+
+        // 3. Compute z_mid (int->mid interpolation)
+        PF::calculate_z_mid(team,nlev,z_int_s,z_mid_s);
+        team.team_barrier();
+    });
+    Kokkos::fence();
+
+    m_z_mid.sync_to_host();
+    const auto z_mid_h = m_z_mid.get_view<const Real**,Host>();
+    const auto T_mid_h   = get_field_in("T_mid").get_view<const Real**, Host>();
+    const auto qv_h      = get_field_in("qv").get_view<const Real**, Host>();
+    const auto p_mid_h   = get_field_in("p_mid").get_view<const Real**,Host>();
+    const auto qc_h      = get_field_in("qc").get_view<const Real**, Host>();
+    const auto qi_h      = get_field_in("qi").get_view<const Real**, Host>();
+    const auto sunlit_h  = get_field_in("sunlit").get_view<const Real*, Host>();
+    const auto skt_h     = get_field_in("surf_radiative_T").get_view<const Real*, Host>();
+    const auto p_int_h   = get_field_in("p_int").get_view<const Real**, Host>();
+    const auto cldfrac_h = get_field_in("cldfrac_rad").get_view<const Real**, Host>();
+    const auto reff_qc_h = get_field_in("eff_radius_qc").get_view<const Real**, Host>();
+    const auto reff_qi_h = get_field_in("eff_radius_qi").get_view<const Real**, Host>();
+    const auto dtau067_h = get_field_in("dtau067").get_view<const Real**, Host>();
+    const auto dtau105_h = get_field_in("dtau105").get_view<const Real**, Host>();
+
+    auto isccp_cldtot_h = get_field_out("isccp_cldtot").get_view<Real*, Host>();
+    auto isccp_ctptau_h = get_field_out("isccp_ctptau").get_view<Real***, Host>();
+    auto modis_ctptau_h = get_field_out("modis_ctptau").get_view<Real***, Host>();
+    auto misr_cthtau_h  = get_field_out("misr_cthtau").get_view<Real***, Host>();
+    auto cosp_sunlit_h  = get_field_out("cosp_sunlit").get_view<Real*, Host>();  // Copy of sunlit flag with COSP frequency for proper averaging
+
     Real emsfc_lw = 0.99;
-    Kokkos::deep_copy(cosp_sunlit, sunlit);
-    CospFunc::view_2d<const Real> z_mid_c = z_mid;  // Need a const version of z_mid for call to CospFunc::main
+    Kokkos::deep_copy(cosp_sunlit_h, sunlit_h);
     CospFunc::main(
-            m_num_cols, m_num_subcols, m_num_levs, m_num_tau, m_num_ctp, m_num_cth,
-            emsfc_lw, sunlit, skt, T_mid, p_mid, p_int, z_mid_c, qv, qc, qi,
-            cldfrac, reff_qc, reff_qi, dtau067, dtau105,
-            isccp_cldtot, isccp_ctptau, modis_ctptau, misr_cthtau
+            m_num_cols, m_num_subcols, m_num_levs, m_num_tau, m_num_ctp, m_num_cth, emsfc_lw,
+            sunlit_h, skt_h, T_mid_h, p_mid_h, p_int_h, z_mid_h, qv_h, qc_h, qi_h,
+            cldfrac_h, reff_qc_h, reff_qi_h, dtau067_h, dtau105_h,
+            isccp_cldtot_h, isccp_ctptau_h, modis_ctptau_h, misr_cthtau_h
     );
     // Remask night values to ZERO since our I/O does not know how to handle masked/missing values
     // in temporal averages; this is all host data, so we can just use host loops like its the 1980s
     for (int i = 0; i < m_num_cols; i++) {
-        if (sunlit(i) == 0) {
-            isccp_cldtot(i) = 0;
-            for (int j = 0; j < m_num_tau; j++) {
-                for (int k = 0; k < m_num_ctp; k++) {
-                    isccp_ctptau(i,j,k) = 0;
-                    modis_ctptau(i,j,k) = 0;
-                }
-                for (int k = 0; k < m_num_cth; k++) {
-                    misr_cthtau (i,j,k) = 0;
-                }
-            }
+      if (sunlit_h(i) == 0) {
+        isccp_cldtot_h(i) = 0;
+        for (int j = 0; j < m_num_tau; j++) {
+          for (int k = 0; k < m_num_ctp; k++) {
+            isccp_ctptau_h(i,j,k) = 0;
+            modis_ctptau_h(i,j,k) = 0;
+          }
+          for (int k = 0; k < m_num_cth; k++) {
+            misr_cthtau_h (i,j,k) = 0;
+          }
         }
+      }
     }
+
+    // Make sure dev data is up to date
+    get_field_out("isccp_cldtot").sync_to_dev();
+    get_field_out("isccp_ctptau").sync_to_dev();
+    get_field_out("modis_ctptau").sync_to_dev();
+    get_field_out("misr_cthtau").sync_to_dev();
+    get_field_out("cosp_sunlit").sync_to_dev();
   } else {
     // If not updating COSP statistics, set these to ZERO; this essentially weights
     // the ISCCP cloud properties by the sunlit mask. What will be output for time-averages
@@ -241,17 +269,12 @@ void Cosp::run_impl (const double dt)
     //     avg(X) = sum(M * X) / sum(M) = (sum(M * X)/N) / (sum(M)/N) = avg(M * X) / avg(M)
     //
     // TODO: mask this when/if the AD ever supports masked averages
-    Kokkos::deep_copy(isccp_cldtot, 0.0);
-    Kokkos::deep_copy(isccp_ctptau, 0.0);
-    Kokkos::deep_copy(modis_ctptau, 0.0);
-    Kokkos::deep_copy(misr_cthtau, 0.0);
-    Kokkos::deep_copy(cosp_sunlit, 0.0);
+    get_field_out("isccp_cldtot").deep_copy(0);
+    get_field_out("isccp_ctptau").deep_copy(0);
+    get_field_out("modis_ctptau").deep_copy(0);
+    get_field_out("misr_cthtau").deep_copy(0);
+    get_field_out("cosp_sunlit").deep_copy(0);
   }
-  get_field_out("isccp_cldtot").sync_to_dev();
-  get_field_out("isccp_ctptau").sync_to_dev();
-  get_field_out("modis_ctptau").sync_to_dev();
-  get_field_out("misr_cthtau").sync_to_dev();
-  get_field_out("cosp_sunlit").sync_to_dev();
 }
 
 // =========================================================================================
@@ -260,6 +283,5 @@ void Cosp::finalize_impl()
   // Finalize COSP wrappers
   CospFunc::finalize();
 }
-// =========================================================================================
 
 } // namespace scream

components/eamxx/src/physics/cosp/eamxx_cosp.hpp

@@ -20,10 +20,6 @@ class Cosp : public AtmosphereProcess
 {
 
 public:
-  using PF  = scream::PhysicsFunctions<HostDevice>;
-  using KT  = KokkosTypes<DefaultDevice>;
-  using KTH = KokkosTypes<HostDevice>;
-
   // Constructors
   Cosp (const ekat::Comm& comm, const ekat::ParameterList& params);
 
@@ -75,6 +71,9 @@ class Cosp : public AtmosphereProcess
 
   std::shared_ptr<const AbstractGrid> m_grid;
 
+  // TODO: use atm buffer instead
+  Field m_z_mid;
+  Field m_z_int;
 }; // class Cosp
 
 } // namespace scream

components/eamxx/tests/single-process/cosp/input.yaml

@@ -10,6 +10,10 @@ time_stepping:
 
 atmosphere_processes:
   atm_procs_list: [cosp]
+  cosp:
+    cosp_frequency: ${NUM_STEPS}
+    cosp_frequency_units: steps
+
 
 grids_manager:
   Type: Mesh Free