Finish porting the last of the GW subroutines

components/eamxx/src/physics/gw/CMakeLists.txt

@@ -27,6 +27,7 @@ if (NOT EAMXX_ENABLE_GPU OR Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE OR Kokkos
     eti/gw_vd_lu_decomp.cpp
     eti/gw_vd_lu_solve.cpp
     eti/gw_gw_convect_init.cpp
+    eti/gw_gw_front_init.cpp
   ) # GW ETI SRCS
 endif()
 

components/eamxx/src/physics/gw/eti/gw_gw_front_init.cpp

@@ -0,0 +1,14 @@
+#include "impl/gw_gw_front_init_impl.hpp"
+
+namespace scream {
+namespace gw {
+
+/*
+ * Explicit instantiation for doing gw_front_init on Reals using the
+ * default device.
+ */
+
+template struct Functions<Real,DefaultDevice>;
+
+} // namespace gw
+} // namespace scream

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

@@ -129,6 +129,9 @@ struct Functions
     // Effective horizontal wave number.
     Real kwv; // = huge(1._r8)
 
+    // 1/2 * horizontal wavenumber (for oro)
+    Real oroko2;
+
     // Interface levels for gravity wave sources.
     int ktop; // = huge(1)
     int kbotbg; // = huge(1)
@@ -162,6 +165,24 @@ struct Functions
     view_3d<Real> mfcc;
   };
 
+  struct GwFrontInit {
+    GwFrontInit() : initialized(false) {}
+
+    // Tell us if initialize has been called
+    bool initialized;
+
+    // Frontogenesis function critical threshold.
+    Real frontgfc;
+
+    // Level at which to check the frontogenesis function to determine when
+    // waves will be launched.
+    Int kfront;
+
+    // Average value of gaussian over gravity wave spectrum bins, multiplied by
+    // background source strength (taubgnd).
+    view_1d<Real> fav;
+  };
+
   //
   // --------- Util Functions ---------
   //
@@ -228,15 +249,21 @@ struct Functions
 
   static void gw_convect_init(
     // Inputs
-    const GwCommonInit& init,
     const Real& plev_src_wind,
     const uview_3d<const Real>& mfcc_in);
 
+  static void gw_front_init(
+    // Inputs
+    const Real& taubgnd,
+    const Real& frontgfc_in,
+    const Int& kfront_in);
+
   static void gw_finalize()
   {
     s_common_init.cref  = decltype(s_common_init.cref)();
     s_common_init.alpha = decltype(s_common_init.alpha)();
     s_convect_init.mfcc = decltype(s_convect_init.mfcc)();
+    s_front_init.fav    = decltype(s_front_init.fav)();
   }
 
   //
@@ -427,47 +454,50 @@ struct Functions
   KOKKOS_FUNCTION
   static void gw_front_project_winds(
     // Inputs
+    const MemberType& team,
     const Int& pver,
-    const Int& ncol,
     const Int& kbot,
-    const uview_1d<const Spack>& u,
-    const uview_1d<const Spack>& v,
+    const uview_1d<const Real>& u,
+    const uview_1d<const Real>& v,
     // Outputs
-    const uview_1d<Spack>& xv,
-    const uview_1d<Spack>& yv,
-    const uview_1d<Spack>& ubm,
-    const uview_1d<Spack>& ubi);
+    Real& xv,
+    Real& yv,
+    const uview_1d<Real>& ubm,
+    const uview_1d<Real>& ubi);
 
   KOKKOS_FUNCTION
   static void gw_front_gw_sources(
     // Inputs
-    const Int& pver,
+    const MemberType& team,
+    const GwFrontInit& finit,
     const Int& pgwv,
-    const Int& ncol,
+    const Int& pver,
     const Int& kbot,
-    const uview_1d<const Spack>& frontgf,
+    const uview_1d<const Real>& frontgf,
     // Outputs
-    const uview_1d<Spack>& tau);
+    const uview_2d<Real>& tau);
 
   KOKKOS_FUNCTION
   static void gw_cm_src(
     // Inputs
+    const MemberType& team,
+    const GwCommonInit& init,
+    const GwFrontInit& finit,
     const Int& pver,
     const Int& pgwv,
-    const Int& ncol,
     const Int& kbot,
-    const uview_1d<const Spack>& u,
-    const uview_1d<const Spack>& v,
-    const uview_1d<const Spack>& frontgf,
+    const uview_1d<const Real>& u,
+    const uview_1d<const Real>& v,
+    const uview_1d<const Real>& frontgf,
     // Outputs
-    const uview_1d<Int>& src_level,
-    const uview_1d<Int>& tend_level,
-    const uview_1d<Spack>& tau,
-    const uview_1d<Spack>& ubm,
-    const uview_1d<Spack>& ubi,
-    const uview_1d<Spack>& xv,
-    const uview_1d<Spack>& yv,
-    const uview_1d<Spack>& c);
+    Int& src_level,
+    Int& tend_level,
+    const uview_2d<Real>& tau,
+    const uview_1d<Real>& ubm,
+    const uview_1d<Real>& ubi,
+    Real& xv,
+    Real& yv,
+    const uview_1d<Real>& c);
 
   KOKKOS_FUNCTION
   static void gw_convect_project_winds(
@@ -616,27 +646,28 @@ struct Functions
   KOKKOS_FUNCTION
   static void gw_oro_src(
     // Inputs
+    const MemberType& team,
+    const GwCommonInit& init,
     const Int& pver,
     const Int& pgwv,
-    const Int& ncol,
-    const uview_1d<const Spack>& u,
-    const uview_1d<const Spack>& v,
-    const uview_1d<const Spack>& t,
-    const uview_1d<const Spack>& sgh,
-    const uview_1d<const Spack>& pmid,
-    const uview_1d<const Spack>& pint,
-    const uview_1d<const Spack>& dpm,
-    const uview_1d<const Spack>& zm,
-    const uview_1d<const Spack>& nm,
+    const uview_1d<const Real>& u,
+    const uview_1d<const Real>& v,
+    const uview_1d<const Real>& t,
+    const Real& sgh,
+    const uview_1d<const Real>& pmid,
+    const uview_1d<const Real>& pint,
+    const uview_1d<const Real>& dpm,
+    const uview_1d<const Real>& zm,
+    const uview_1d<const Real>& nm,
     // Outputs
-    const uview_1d<Int>& src_level,
-    const uview_1d<Int>& tend_level,
-    const uview_1d<Spack>& tau,
-    const uview_1d<Spack>& ubm,
-    const uview_1d<Spack>& ubi,
-    const uview_1d<Spack>& xv,
-    const uview_1d<Spack>& yv,
-    const uview_1d<Spack>& c);
+    Int& src_level,
+    Int& tend_level,
+    const uview_2d<Real>& tau,
+    const uview_1d<Real>& ubm,
+    const uview_1d<Real>& ubi,
+    Real& xv,
+    Real& yv,
+    const uview_1d<Real>& c);
 
   KOKKOS_FUNCTION
   static void vd_lu_decomp(
@@ -678,6 +709,7 @@ struct Functions
   //
   inline static GwCommonInit s_common_init;
   inline static GwConvectInit s_convect_init;
+  inline static GwFrontInit s_front_init;
 
 }; // struct Functions
 
@@ -710,5 +742,6 @@ struct Functions
 # include "impl/gw_vd_lu_decomp_impl.hpp"
 # include "impl/gw_vd_lu_solve_impl.hpp"
 # include "impl/gw_gw_convect_init_impl.hpp"
+# include "impl/gw_gw_front_init_impl.hpp"
 #endif // GPU && !KOKKOS_ENABLE_*_RELOCATABLE_DEVICE_CODE
 #endif // P3_FUNCTIONS_HPP

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

@@ -14,26 +14,45 @@ namespace gw {
 template<typename S, typename D>
 KOKKOS_FUNCTION
 void Functions<S,D>::gw_cm_src(
-// Inputs
-const Int& pver,
-const Int& pgwv,
-const Int& ncol,
-const Int& kbot,
-const uview_1d<const Spack>& u,
-const uview_1d<const Spack>& v,
-const uview_1d<const Spack>& frontgf,
-// Outputs
-const uview_1d<Int>& src_level,
-const uview_1d<Int>& tend_level,
-const uview_1d<Spack>& tau,
-const uview_1d<Spack>& ubm,
-const uview_1d<Spack>& ubi,
-const uview_1d<Spack>& xv,
-const uview_1d<Spack>& yv,
-const uview_1d<Spack>& c)
+  // Inputs
+  const MemberType& team,
+  const GwCommonInit& init,
+  const GwFrontInit& finit,
+  const Int& pver,
+  const Int& pgwv,
+  const Int& kbot,
+  const uview_1d<const Real>& u,
+  const uview_1d<const Real>& v,
+  const uview_1d<const Real>& frontgf,
+  // Outputs
+  Int& src_level,
+  Int& tend_level,
+  const uview_2d<Real>& tau,
+  const uview_1d<Real>& ubm,
+  const uview_1d<Real>& ubi,
+  Real& xv,
+  Real& yv,
+  const uview_1d<Real>& c)
 {
-  // TODO
-  // Note, argument types may need tweaking. Generator is not always able to tell what needs to be packed
+  //------------------------------------------------------------------------
+  // Determine the source layer wind and unit vectors, then project winds.
+  //------------------------------------------------------------------------
+  gw_front_project_winds(team, pver, kbot, u, v, xv, yv, ubm, ubi);
+
+  //-----------------------------------------------------------------------
+  // Gravity wave sources
+  //-----------------------------------------------------------------------
+  gw_front_gw_sources(team, finit, pgwv, pver, kbot, frontgf, tau);
+
+  src_level = kbot;
+  tend_level = kbot;
+
+  // Set phase speeds as reference speeds plus the wind speed at the source
+  // level.
+  Kokkos::parallel_for(
+    Kokkos::TeamVectorRange(team, init.cref.size()), [&] (const int l) {
+      c(l) = init.cref(l) + std::abs(ubi(kbot+1));
+    });
 }
 
 } // namespace gw

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

@@ -28,6 +28,8 @@ void Functions<S,D>::gw_common_init(
   const Real& kwv_in,
   const uview_1d<const Real>& alpha_in)
 {
+  constexpr Real temp1 = 86400;
+  constexpr Real half = 0.5;
   s_common_init.initialized = true;
   s_common_init.pver = pver_in;
   s_common_init.pgwv = pgwv_in;
@@ -42,10 +44,11 @@ void Functions<S,D>::gw_common_init(
   s_common_init.kbotbg = kbotbg_in;
   s_common_init.fcrit2 = fcrit2_in;
   s_common_init.kwv = kwv_in;
+  s_common_init.oroko2 = half * kwv_in;
   s_common_init.alpha = view_1d<Real>("alpha", alpha_in.size());
   Kokkos::deep_copy(s_common_init.alpha, alpha_in);
   s_common_init.effkwv = kwv_in * fcrit2_in;
-  s_common_init.tndmax = orographic_only_in ? 500. / 86400. : 400. / 86400.;
+  s_common_init.tndmax = orographic_only_in ? 500 / temp1 : 400 / temp1;
 }
 
 } // namespace gw

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

@@ -14,12 +14,11 @@ namespace gw {
 template<typename S, typename D>
 void Functions<S,D>::gw_convect_init(
   // Inputs
-  const GwCommonInit& init,
   const Real& plev_src_wind,
   const uview_3d<const Real>& mfcc_in)
 {
   // Just set k_src_wind to pver. We don't have access to pref_edge
-  s_convect_init.k_src_wind = init.pver - 1;
+  s_convect_init.k_src_wind = s_common_init.pver - 1;
 
   // First dimension is maxh.
   s_convect_init.maxh   = mfcc_in.extent(0);

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

@@ -14,17 +14,36 @@ namespace gw {
 template<typename S, typename D>
 KOKKOS_FUNCTION
 void Functions<S,D>::gw_front_gw_sources(
-// Inputs
-const Int& pver,
-const Int& pgwv,
-const Int& ncol,
-const Int& kbot,
-const uview_1d<const Spack>& frontgf,
-// Outputs
-const uview_1d<Spack>& tau)
+  // Inputs
+  const MemberType& team,
+  const GwFrontInit& finit,
+  const Int& pgwv,
+  const Int& pver,
+  const Int& kbot,
+  const uview_1d<const Real>& frontgf,
+  // Outputs
+  const uview_2d<Real>& tau)
 {
-  // TODO
-  // Note, argument types may need tweaking. Generator is not always able to tell what needs to be packed
+  Kokkos::parallel_for(
+    Kokkos::TeamVectorRange(team, tau.size()), [&] (const int k) {
+      tau.data()[k] = 0;
+    });
+
+  team.team_barrier();
+
+  // GW generation depends on frontogenesis at specified level (may be below
+  // actual source level).
+  const bool launch_wave = frontgf(finit.kfront) > finit.frontgfc;
+
+  if (launch_wave) {
+    Kokkos::parallel_for(
+      Kokkos::TeamVectorRange(team, pgwv+1), [&] (const int l) {
+        const Int negl_idx = pgwv - l;
+        const Int posl_idx  = pgwv + l;
+        tau(posl_idx, kbot+1)  = finit.fav(posl_idx);
+        tau(negl_idx, kbot+1)  = finit.fav(posl_idx); // negative for tau only (not fav?)
+      });
+  }
 }
 
 } // namespace gw