Merge branch 'jgfouca/port_gw_prof' into master (PR E3SM-Project#7629)

Port GW gw_prof to Kokkos/C++

[BFB]

Author: jgfouca

components/eamxx/src/physics/gw/gw_functions.hpp

@@ -143,6 +143,19 @@ struct Functions
     Real tndmax; // = huge(1._r8)
   };
 
+  KOKKOS_INLINE_FUNCTION
+  static void midpoint_interp(
+    const MemberType& team,
+    const uview_1d<const Real>& in,
+    const uview_1d<Real>& interp)
+  {
+    EKAT_KERNEL_REQUIRE(in.size() == interp.size() + 1);
+    Kokkos::parallel_for(
+      Kokkos::TeamVectorRange(team, 0, in.extent(0)-1), [&] (const int k) {
+        interp(k) = (in(k)+in(k+1)) / 2;
+    });
+  }
+
   //
   // --------- Init/Finalize Functions ---------
   //
@@ -201,20 +214,27 @@ struct Functions
     const uview_1d<Real>& utgw,
     const uview_1d<Real>& vtgw);
 
+  /*
+   * Compute profiles of background state quantities for the multiple
+   * gravity wave drag parameterization.
+   *
+   * The parameterization is assumed to operate only where water vapor
+   * concentrations are negligible in determining the density.
+   */
   KOKKOS_FUNCTION
   static void gw_prof(
     // Inputs
+    const MemberType& team,
     const Int& pver,
-    const Int& ncol,
-    const Spack& cpair,
-    const uview_1d<const Spack>& t,
-    const uview_1d<const Spack>& pmid,
-    const uview_1d<const Spack>& pint,
+    const Real& cpair,
+    const uview_1d<const Real>& t,
+    const uview_1d<const Real>& pmid,
+    const uview_1d<const Real>& pint,
     // Outputs
-    const uview_1d<Spack>& rhoi,
-    const uview_1d<Spack>& ti,
-    const uview_1d<Spack>& nm,
-    const uview_1d<Spack>& ni);
+    const uview_1d<Real>& rhoi,
+    const uview_1d<Real>& ti,
+    const uview_1d<Real>& nm,
+    const uview_1d<Real>& ni);
 
   KOKKOS_FUNCTION
   static void momentum_energy_conservation(

components/eamxx/src/physics/gw/impl/gw_gw_prof_impl.hpp

@@ -3,6 +3,8 @@
 
 #include "gw_functions.hpp" // for ETI only but harmless for GPU
 
+#include <ekat_subview_utils.hpp>
+
 namespace scream {
 namespace gw {
 
@@ -14,21 +16,59 @@ namespace gw {
 template<typename S, typename D>
 KOKKOS_FUNCTION
 void Functions<S,D>::gw_prof(
-// Inputs
-const Int& pver,
-const Int& ncol,
-const Spack& cpair,
-const uview_1d<const Spack>& t,
-const uview_1d<const Spack>& pmid,
-const uview_1d<const Spack>& pint,
-// Outputs
-const uview_1d<Spack>& rhoi,
-const uview_1d<Spack>& ti,
-const uview_1d<Spack>& nm,
-const uview_1d<Spack>& ni)
+  // Inputs
+  const MemberType& team,
+  const Int& pver,
+  const Real& cpair,
+  const uview_1d<const Real>& t,
+  const uview_1d<const Real>& pmid,
+  const uview_1d<const Real>& pint,
+  // Outputs
+  const uview_1d<Real>& rhoi,
+  const uview_1d<Real>& ti,
+  const uview_1d<Real>& nm,
+  const uview_1d<Real>& ni)
 {
-  // TODO
-  // Note, argument types may need tweaking. Generator is not always able to tell what needs to be packed
+  // Minimum value of Brunt-Vaisalla frequency squared.
+  static constexpr Real n2min = 1.e-8;
+
+  //-----------------------------------------------------------------------
+  // Determine the interface densities and Brunt-Vaisala frequencies.
+  //-----------------------------------------------------------------------
+
+  // The top interface values are calculated assuming an isothermal
+  // atmosphere above the top level.
+  Kokkos::single(Kokkos::PerTeam(team), [&] {
+    ti(0) = t(0);
+    rhoi(0) = pint(0) / (C::Rair*ti(0));
+    ni(0) = sqrt(C::gravit*C::gravit / (cpair*ti(0)));
+  });
+
+  // Interior points use centered differences.
+  midpoint_interp(team, t, ekat::subview(ti, Kokkos::pair<int, int>{1, pver}));
+  team.team_barrier();
+  Kokkos::parallel_for(
+    Kokkos::TeamVectorRange(team, 1, pver), [&] (const int k) {
+    rhoi(k) = pint(k) / (C::Rair*ti(k));
+    const Real dtdp = (t(k)-t(k-1)) / (pmid(k)-pmid(k-1));
+    const Real n2 = C::gravit*C::gravit/ti(k) * (1/cpair - rhoi(k)*dtdp);
+    ni(k) = std::sqrt(ekat::impl::max(n2min, n2));
+  });
+
+  // Bottom interface uses bottom level temperature, density; next interface
+  // B-V frequency.
+  team.team_barrier();
+  Kokkos::single(Kokkos::PerTeam(team), [&] {
+    ti(pver) = t(pver-1);
+    rhoi(pver) = pint(pver) / (C::Rair*ti(pver));
+    ni(pver) = ni(pver-1);
+  });
+
+  //------------------------------------------------------------------------
+  // Determine the midpoint Brunt-Vaisala frequencies.
+  //------------------------------------------------------------------------
+  team.team_barrier();
+  midpoint_interp(team, ni, nm);
 }
 
 } // namespace gw

components/eamxx/src/physics/gw/tests/gw_gw_prof_tests.cpp

@@ -57,24 +57,29 @@ struct UnitWrap::UnitTest<D>::TestGwProf : public UnitWrap::UnitTest<D>::Base {
 
     // Get data from test
     for (auto& d : test_data) {
-      gw_prof(d);
+      if (this->m_baseline_action == GENERATE) {
+        gw_prof_f(d);
+      }
+      else {
+        gw_prof(d);
+      }
     }
 
     // Verify BFB results, all data should be in C layout
     if (SCREAM_BFB_TESTING && this->m_baseline_action == COMPARE) {
       for (Int i = 0; i < num_runs; ++i) {
         GwProfData& d_baseline = baseline_data[i];
         GwProfData& d_test = test_data[i];
+        REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.rhoi));
+        REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.ti));
+        REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.ni));
         for (Int k = 0; k < d_baseline.total(d_baseline.rhoi); ++k) {
-          REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.rhoi));
           REQUIRE(d_baseline.rhoi[k] == d_test.rhoi[k]);
-          REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.ti));
           REQUIRE(d_baseline.ti[k] == d_test.ti[k]);
-          REQUIRE(d_baseline.total(d_baseline.rhoi) == d_test.total(d_test.ni));
           REQUIRE(d_baseline.ni[k] == d_test.ni[k]);
         }
+        REQUIRE(d_baseline.total(d_baseline.nm) == d_test.total(d_test.nm));
         for (Int k = 0; k < d_baseline.total(d_baseline.nm); ++k) {
-          REQUIRE(d_baseline.total(d_baseline.nm) == d_test.total(d_test.nm));
           REQUIRE(d_baseline.nm[k] == d_test.nm[k]);
         }
 

components/eamxx/src/physics/gw/tests/infra/gw_test_data.cpp

@@ -237,14 +237,84 @@ void gwd_compute_tendencies_from_stress_divergence(GwdComputeTendenciesFromStres
   gw_finalize_cxx(d.init);
 }
 
-void gw_prof(GwProfData& d)
+void gw_prof_f(GwProfData& d)
 {
   gw_init(d.init);
   d.transition<ekat::TransposeDirection::c2f>();
   gw_prof_c(d.ncol, d.cpair, d.t, d.pmid, d.pint, d.rhoi, d.ti, d.nm, d.ni);
   d.transition<ekat::TransposeDirection::f2c>();
 }
 
+void gw_prof(GwProfData& d)
+{
+  gw_init_cxx(d.init);
+
+  // create device views and copy
+  std::vector<view2dr_d> two_d_reals_in(7);
+
+  ekat::host_to_device({d.t, d.pmid, d.nm, d.pint, d.rhoi, d.ti, d.ni},
+                       std::vector<int>(7, d.ncol),
+                       std::vector<int>{    // dim2 sizes
+                         d.init.pver,
+                         d.init.pver,
+                         d.init.pver,
+                         d.init.pver + 1,
+                         d.init.pver + 1,
+                         d.init.pver + 1,
+                         d.init.pver + 1},
+                       two_d_reals_in);
+
+  const auto t    = two_d_reals_in[0];
+  const auto pmid = two_d_reals_in[1];
+  const auto nm   = two_d_reals_in[2];
+  const auto pint = two_d_reals_in[3];
+  const auto rhoi = two_d_reals_in[4];
+  const auto ti   = two_d_reals_in[5];
+  const auto ni   = two_d_reals_in[6];
+
+  auto policy = ekat::TeamPolicyFactory<ExeSpace>::get_default_team_policy(d.ncol, d.init.pver);
+
+  // unpack init because we do not want the lambda to capture it
+  const int pver = d.init.pver;
+  const Real cpair = d.cpair;
+
+  Kokkos::parallel_for(policy, KOKKOS_LAMBDA(const MemberType& team) {
+    const int col = team.league_rank();
+
+    // Get single-column subviews of all inputs, shouldn't need any i-indexing
+    // after this.
+    const auto t_c    = ekat::subview(t, col);
+    const auto pmid_c = ekat::subview(pmid, col);
+    const auto nm_c   = ekat::subview(nm, col);
+    const auto pint_c = ekat::subview(pint, col);
+    const auto rhoi_c = ekat::subview(rhoi, col);
+    const auto ti_c   = ekat::subview(ti, col);
+    const auto ni_c   = ekat::subview(ni, col);
+
+    GWF::gw_prof(
+      team,
+      pver,
+      cpair,
+      t_c,
+      pmid_c,
+      pint_c,
+      rhoi_c,
+      ti_c,
+      nm_c,
+      ni_c
+    );
+  });
+
+  // Get outputs back
+  std::vector<view2dr_d> two_d_reals_out = {rhoi, ti, nm, ni};
+  ekat::device_to_host({d.rhoi, d.ti, d.nm, d.ni},
+                       std::vector<int>(4, d.ncol),
+                       std::vector<int>{d.init.pver+1, d.init.pver+1, d.init.pver, d.init.pver+1},
+                       two_d_reals_out);
+
+  gw_finalize_cxx(d.init);
+}
+
 void momentum_energy_conservation_f(MomentumEnergyConservationData& d)
 {
   gw_init(d.init);

components/eamxx/src/physics/gw/tests/infra/gw_test_data.hpp

@@ -854,6 +854,7 @@ struct GwOroSrcData : public PhysicsTestData {
 // Glue functions to call fortran from from C++ with the Data struct
 void gwd_compute_tendencies_from_stress_divergence(GwdComputeTendenciesFromStressDivergenceData& d);
 void gwd_compute_tendencies_from_stress_divergence_f(GwdComputeTendenciesFromStressDivergenceData& d);
+void gw_prof_f(GwProfData& d);
 void gw_prof(GwProfData& d);
 void momentum_energy_conservation_f(MomentumEnergyConservationData& d);
 void momentum_energy_conservation(MomentumEnergyConservationData& d);