Merge pull request #219 from mwarusz/omega/wind-forcing-and-bottom-drag

Add wind forcing and bottom drag tendency terms

Author: mark-petersen

components/omega/configs/Default.yml

@@ -16,6 +16,8 @@ Omega:
   Advection:
     FluxThicknessType: Center
     FluxTracerType: Center
+  WindStress:
+    InterpType: Isotropic
   Tendencies:
     ThicknessFluxTendencyEnable: true
     PVTendencyEnable: true
@@ -26,6 +28,10 @@ Omega:
     VelHyperDiffTendencyEnable: true
     ViscDel4: 1.2e11
     DivFactor: 1.0
+    WindForcingTendencyEnable: false
+    Density0: 1026.0
+    BottomDragTendencyEnable: false
+    BottomDragCoeff: 0.0
     TracerHorzAdvTendencyEnable: true
     TracerDiffTendencyEnable: true
     EddyDiff2: 10.0
@@ -91,7 +97,7 @@ Omega:
       Contents:
         - Tracers
         - State
-        - SshCellDefault
+        - SshCell
     Highfreq:
       UsePointerFile: false
       Filename: ocn.hifreq.$Y-$M

components/omega/doc/devGuide/AuxiliaryVariables.md

@@ -101,3 +101,6 @@ The following auxiliary variable groups are currently implemented:
 ||  Del2RelVortVertex ||
 | TracerAuxVars | HTracersEdge | Center or Upwind|
 || Del2TracersCell ||
+| WindForcingAuxVars |  ZonalStressCell ||
+||  MeridStressCell ||
+||  NormalStressEdge ||

components/omega/doc/devGuide/TendencyTerms.md

@@ -38,3 +38,5 @@ implemented:
 - `TracerHorzAdvOnCell`
 - `TracerDiffOnCell`
 - `TracerHyperDiffOnCell`
+- `WindForcingOnEdge`
+- `BottomDragOnEdge`

components/omega/doc/userGuide/AuxiliaryVariables.md

@@ -29,3 +29,6 @@ The following auxiliary variables are currently available:
 | VelDel2RelVortVertex | curl of laplacian of horizontal velocity on cells
 | HTracersEdge | thickness-weighted tracers used for fluxes through edges. May be centered, upwinded or a combination of the two
 | Del2TracersCell | laplacian of tracers on cells
+| ZonalStressCell | zonal component of wind stress on cells
+| MeridStressCell | meridional component of wind stress on cells
+| NormalStressEdge | normal component of wind stress on edge

components/omega/doc/userGuide/TendencyTerms.md

@@ -17,6 +17,8 @@ tendency terms are currently implemented:
 | TracerHorzAdvOnCell | horizontal advection of thickness-weighted tracers
 | TracerDiffOnCell | horizontal diffusion of thickness-weighted tracers
 | TracerHyperDiffOnCell | biharmonic horizontal mixing of thickness-weighted tracers
+| WindForcingOnEdge | forcing by wind stress, defined on edges
+| BottomDragOnEdge | bottom drag, defined on edges
 
 Among the internal data stored by each functor is a `bool` which can enable or
 disable the contribution of that particular term to the tendency. These flags
@@ -45,3 +47,6 @@ the currently available tendency terms:
 | | EddyDiff2 | horizontal diffusion coefficient
 | TracerHyperDiffOnCell | TracerHyperDiffTendencyEnable | enable/disable term
 | | EddyDiff4 | biharmonic horizontal mixing coeffienct for tracers
+| WindForcingOnEdge | WindForcingTendencyEnable | enable/disable term
+| BottomDragOnEdge | BottomDragTendencyEnable | enable/disable term
+| | BottomDragCoeff | bottom drag coefficient

components/omega/src/ocn/AuxiliaryState.cpp

@@ -10,15 +10,21 @@ AuxiliaryState *AuxiliaryState::DefaultAuxState = nullptr;
 std::map<std::string, std::unique_ptr<AuxiliaryState>>
     AuxiliaryState::AllAuxStates;
 
+static std::string stripDefault(const std::string &Name) {
+   return Name != "Default" ? Name : "";
+}
+
 // Constructor. Constructs the member auxiliary variables and registers their
 // fields with IOStreams
 AuxiliaryState::AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
-                               int NVertLevels, int NTracers)
-    : Mesh(Mesh), Name(Name), KineticAux(Name, Mesh, NVertLevels),
-      LayerThicknessAux(Name, Mesh, NVertLevels),
-      VorticityAux(Name, Mesh, NVertLevels),
-      VelocityDel2Aux(Name, Mesh, NVertLevels),
-      TracerAux(Name, Mesh, NVertLevels, NTracers) {
+                               Halo *MeshHalo, int NVertLevels, int NTracers)
+    : Mesh(Mesh), MeshHalo(MeshHalo), Name(stripDefault(Name)),
+      KineticAux(stripDefault(Name), Mesh, NVertLevels),
+      LayerThicknessAux(stripDefault(Name), Mesh, NVertLevels),
+      VorticityAux(stripDefault(Name), Mesh, NVertLevels),
+      VelocityDel2Aux(stripDefault(Name), Mesh, NVertLevels),
+      WindForcingAux(stripDefault(Name), Mesh),
+      TracerAux(stripDefault(Name), Mesh, NVertLevels, NTracers) {
 
    GroupName = "AuxiliaryState";
    if (Name != "Default") {
@@ -32,6 +38,7 @@ AuxiliaryState::AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
    LayerThicknessAux.registerFields(GroupName, AuxMeshName);
    VorticityAux.registerFields(GroupName, AuxMeshName);
    VelocityDel2Aux.registerFields(GroupName, AuxMeshName);
+   WindForcingAux.registerFields(GroupName, AuxMeshName);
    TracerAux.registerFields(GroupName, AuxMeshName);
 }
 
@@ -42,6 +49,7 @@ AuxiliaryState::~AuxiliaryState() {
    LayerThicknessAux.unregisterFields();
    VorticityAux.unregisterFields();
    VelocityDel2Aux.unregisterFields();
+   WindForcingAux.unregisterFields();
    TracerAux.unregisterFields();
 
    int Err = FieldGroup::destroy(GroupName);
@@ -64,6 +72,7 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
    OMEGA_SCOPE(LocLayerThicknessAux, LayerThicknessAux);
    OMEGA_SCOPE(LocVorticityAux, VorticityAux);
    OMEGA_SCOPE(LocVelocityDel2Aux, VelocityDel2Aux);
+   OMEGA_SCOPE(LocWindForcingAux, WindForcingAux);
 
    parallelFor(
        "vertexAuxState1", {Mesh->NVerticesAll, NChunks},
@@ -82,7 +91,12 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
    const auto &RelVortVertex   = VorticityAux.RelVortVertex;
 
    parallelFor(
-       "edgeAuxState1", {Mesh->NEdgesAll, NChunks},
+       "edgeAuxState1", {Mesh->NEdgesAll}, KOKKOS_LAMBDA(int IEdge) {
+          LocWindForcingAux.computeVarsOnEdge(IEdge);
+       });
+
+   parallelFor(
+       "edgeAuxState2", {Mesh->NEdgesAll, NChunks},
        KOKKOS_LAMBDA(int IEdge, int KChunk) {
           LocVorticityAux.computeVarsOnEdge(IEdge, KChunk);
           LocLayerThicknessAux.computeVarsOnEdge(IEdge, KChunk, LayerThickCell,
@@ -153,16 +167,17 @@ void AuxiliaryState::computeAll(const OceanState *State,
 
 // Create a non-default auxiliary state
 AuxiliaryState *AuxiliaryState::create(const std::string &Name,
-                                       const HorzMesh *Mesh, int NVertLevels,
-                                       const int NTracers) {
+                                       const HorzMesh *Mesh, Halo *MeshHalo,
+                                       int NVertLevels, const int NTracers) {
    if (AllAuxStates.find(Name) != AllAuxStates.end()) {
       LOG_ERROR("Attempted to create a new AuxiliaryState with name {} but it "
                 "already exists",
                 Name);
       return nullptr;
    }
 
-   auto *NewAuxState = new AuxiliaryState(Name, Mesh, NVertLevels, NTracers);
+   auto *NewAuxState =
+       new AuxiliaryState(Name, Mesh, MeshHalo, NVertLevels, NTracers);
    AllAuxStates.emplace(Name, NewAuxState);
 
    return NewAuxState;
@@ -172,12 +187,13 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
 // initialized.
 void AuxiliaryState::init() {
    const HorzMesh *DefMesh = HorzMesh::getDefault();
+   Halo *DefHalo           = Halo::getDefault();
 
    int NVertLevels = DefMesh->NVertLevels;
    int NTracers    = Tracers::getNumTracers();
 
-   AuxiliaryState::DefaultAuxState =
-       AuxiliaryState::create("Default", DefMesh, NVertLevels, NTracers);
+   AuxiliaryState::DefaultAuxState = AuxiliaryState::create(
+       "Default", DefMesh, DefHalo, NVertLevels, NTracers);
 
    Config *OmegaConfig = Config::getOmegaConfig();
    DefaultAuxState->readConfigOptions(OmegaConfig);
@@ -248,6 +264,37 @@ void AuxiliaryState::readConfigOptions(Config *OmegaConfig) {
    } else {
       ABORT_ERROR("AuxiliaryState: Unknown FluxTracerType requested");
    }
+
+   Config WindStressConfig("WindStress");
+   Err += OmegaConfig->get(WindStressConfig);
+
+   std::string WindStressInterpTypeStr;
+   Err += WindStressConfig.get("InterpType", WindStressInterpTypeStr);
+   CHECK_ERROR_ABORT(
+       Err, "AuxiliaryState: InterpType not found in WindStressConfig");
+
+   if (WindStressInterpTypeStr == "Isotropic") {
+      this->WindForcingAux.InterpChoice = InterpCellToEdgeOption::Isotropic;
+   } else if (WindStressInterpTypeStr == "Anisotropic") {
+      this->WindForcingAux.InterpChoice = InterpCellToEdgeOption::Anisotropic;
+   } else {
+      ABORT_ERROR("AuxiliaryState: Unknown InterpType requested");
+   }
 }
 
+//------------------------------------------------------------------------------
+// Perform auxiliary state halo exchange
+// Note that only non-computed auxiliary variables needs to be exchanged
+I4 AuxiliaryState::exchangeHalo() {
+   I4 Err = 0;
+
+   Err +=
+       MeshHalo->exchangeFullArrayHalo(WindForcingAux.ZonalStressCell, OnCell);
+   Err +=
+       MeshHalo->exchangeFullArrayHalo(WindForcingAux.MeridStressCell, OnCell);
+
+   return Err;
+
+} // end exchangeHalo
+
 } // namespace OMEGA

components/omega/src/ocn/AuxiliaryState.h

@@ -3,6 +3,7 @@
 
 #include "Config.h"
 #include "DataTypes.h"
+#include "Halo.h"
 #include "HorzMesh.h"
 #include "OceanState.h"
 #include "Tracers.h"
@@ -11,6 +12,7 @@
 #include "auxiliaryVars/TracerAuxVars.h"
 #include "auxiliaryVars/VelocityDel2AuxVars.h"
 #include "auxiliaryVars/VorticityAuxVars.h"
+#include "auxiliaryVars/WindForcingAuxVars.h"
 
 #include <memory>
 #include <string>
@@ -36,6 +38,7 @@ class AuxiliaryState {
    TracerAuxVars TracerAux;
    VorticityAuxVars VorticityAux;
    VelocityDel2AuxVars VelocityDel2Aux;
+   WindForcingAuxVars WindForcingAux;
 
    ~AuxiliaryState();
 
@@ -46,7 +49,7 @@ class AuxiliaryState {
 
    // Create a non-default auxiliary state
    static AuxiliaryState *create(const std::string &Name, const HorzMesh *Mesh,
-                                 int NVertLevels, int NTracers);
+                                 Halo *MeshHalo, int NVertLevels, int NTracers);
 
    /// Get the default auxiliary state
    static AuxiliaryState *getDefault();
@@ -63,6 +66,9 @@ class AuxiliaryState {
    /// Read and set config options
    void readConfigOptions(Config *OmegaConfig);
 
+   /// Exchange halo
+   I4 exchangeHalo();
+
    // Compute all auxiliary variables needed for momentum equation
    void computeMomAux(const OceanState *State, int ThickTimeLevel,
                       int VelTimeLevel) const;
@@ -75,13 +81,14 @@ class AuxiliaryState {
                    int TimeLevel) const;
 
  private:
-   AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
+   AuxiliaryState(const std::string &Name, const HorzMesh *Mesh, Halo *MeshHalo,
                   int NVertLevels, int NTracers);
 
    AuxiliaryState(const AuxiliaryState &) = delete;
    AuxiliaryState(AuxiliaryState &&)      = delete;
 
    const HorzMesh *Mesh;
+   Halo *MeshHalo;
    static AuxiliaryState *DefaultAuxState;
    static std::map<std::string, std::unique_ptr<AuxiliaryState>> AllAuxStates;
 };

components/omega/src/ocn/HorzOperators.cpp

@@ -21,4 +21,10 @@ TangentialReconOnEdge::TangentialReconOnEdge(HorzMesh const *Mesh)
     : NEdgesOnEdge(Mesh->NEdgesOnEdge), EdgesOnEdge(Mesh->EdgesOnEdge),
       WeightsOnEdge(Mesh->WeightsOnEdge) {}
 
+InterpCellToEdge::InterpCellToEdge(const HorzMesh *Mesh)
+    : CellsOnEdge(Mesh->CellsOnEdge), VerticesOnEdge(Mesh->VerticesOnEdge),
+      CellsOnVertex(Mesh->CellsOnVertex),
+      KiteAreasOnVertex(Mesh->KiteAreasOnVertex),
+      VertexDegree(Mesh->VertexDegree) {}
+
 } // namespace OMEGA

components/omega/src/ocn/HorzOperators.h

@@ -132,5 +132,57 @@ class TangentialReconOnEdge {
    Array2DReal WeightsOnEdge;
 };
 
+enum class InterpCellToEdgeOption { Anisotropic, Isotropic };
+
+class InterpCellToEdge {
+ public:
+   InterpCellToEdge(const HorzMesh *Mesh);
+
+   KOKKOS_FUNCTION Real operator()(int IEdge, const Array1DReal &ArrayCell,
+                                   InterpCellToEdgeOption Option) const {
+      switch (Option) {
+      case InterpCellToEdgeOption::Anisotropic:
+         return interpolateAnisotropic(IEdge, ArrayCell);
+      case InterpCellToEdgeOption::Isotropic:
+         return interpolateIsotropic(IEdge, ArrayCell);
+      }
+   };
+
+   KOKKOS_FUNCTION Real
+   interpolateAnisotropic(int IEdge, const Array1DReal &ArrayCell) const {
+      const int JCell0 = CellsOnEdge(IEdge, 0);
+      const int JCell1 = CellsOnEdge(IEdge, 1);
+
+      return 0.5_Real * (ArrayCell(JCell0) + ArrayCell(JCell1));
+   };
+
+   KOKKOS_FUNCTION Real
+   interpolateIsotropic(int IEdge, const Array1DReal &ArrayCell) const {
+
+      Real Accum     = 0;
+      Real AreaAccum = 0;
+      for (int J = 0; J < 2; ++J) {
+         const int JVertex = VerticesOnEdge(IEdge, J);
+         for (int L = 0; L < VertexDegree; ++L) {
+            const Real KiteArea = KiteAreasOnVertex(JVertex, L);
+            const int LCell     = CellsOnVertex(JVertex, L);
+
+            Accum += ArrayCell(LCell) * KiteArea;
+            AreaAccum += KiteArea;
+         }
+      }
+
+      const Real InvAreaAccum = 1._Real / AreaAccum;
+      return Accum * InvAreaAccum;
+   };
+
+ private:
+   Array2DI4 CellsOnEdge;
+   Array2DI4 VerticesOnEdge;
+   Array2DI4 CellsOnVertex;
+   Array2DReal KiteAreasOnVertex;
+   I4 VertexDegree;
+};
+
 } // namespace OMEGA
 #endif

components/omega/src/ocn/OceanInit.cpp

@@ -157,13 +157,17 @@ int initOmegaModules(MPI_Comm Comm) {
       }
    }
 
-   // Update Halos and Host arrays with new state and tracer fields
+   // Update Halos and Host arrays with new state, auxiliary state, and tracer
+   // fields
 
    OceanState *DefState = OceanState::getDefault();
    I4 CurTimeLevel      = 0;
    DefState->exchangeHalo(CurTimeLevel);
    DefState->copyToHost(CurTimeLevel);
 
+   AuxiliaryState *DefAuxState = AuxiliaryState::getDefault();
+   DefAuxState->exchangeHalo();
+
    // Now update tracers - assume using same time level index
    Err = Tracers::exchangeHalo(CurTimeLevel);
    if (Err != 0) {