Enable Omega time stepping with diagnostic computations

components/omega/src/ocn/AuxiliaryState.cpp

@@ -20,9 +20,8 @@ static std::string stripDefault(const std::string &Name) {
 // Constructor. Constructs the member auxiliary variables and registers their
 // fields with IOStreams
 AuxiliaryState::AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
-                               Halo *MeshHalo, const VertCoord *VCoord,
-                               VertAdv *VAdv, int NTracers,
-                               TimeInterval TimeStep)
+                               Halo *MeshHalo, VertCoord *VCoord, VertAdv *VAdv,
+                               int NTracers, TimeInterval TimeStep)
     : Mesh(Mesh), MeshHalo(MeshHalo), VCoord(VCoord), VAdv(VAdv),
       Name(stripDefault(Name)), KineticAux(stripDefault(Name), Mesh, VCoord),
       LayerThicknessAux(stripDefault(Name), Mesh, VCoord),
@@ -63,9 +62,50 @@ AuxiliaryState::~AuxiliaryState() {
    FieldGroup::destroy(GroupName);
 }
 
+// Compute auxiliary variables for vertical dynamics
+void AuxiliaryState::computeMomVertAux(const OceanState *State,
+                                       const Array3DReal &TracerArray,
+                                       int ThickTimeLevel,
+                                       int VelTimeLevel) const {
+
+   Pacer::start("AuxState:computeMomVertAux", 2);
+
+   Eos *EosInstance = Eos::getInstance();
+
+   // get layer thickness
+   Array2DReal LayerThickCell = State->getLayerThickness(ThickTimeLevel);
+   // get normal velocity
+   Array2DReal NormalVelEdge = State->getNormalVelocity(VelTimeLevel);
+
+   // get temperature and salinity
+   I4 ConservTempIdx;
+   I4 AbsSalinityIdx;
+   Tracers::getIndex(ConservTempIdx, "Temperature");
+   Tracers::getIndex(AbsSalinityIdx, "Salinity");
+
+   const auto ConservTemp =
+       Kokkos::subview(TracerArray, ConservTempIdx, Kokkos::ALL, Kokkos::ALL);
+   const auto AbsSalinity =
+       Kokkos::subview(TracerArray, AbsSalinityIdx, Kokkos::ALL, Kokkos::ALL);
+
+   // compute pressure
+   const auto &SurfacePressure = VCoord->SurfacePressure;
+   VCoord->computePressure(LayerThickCell, SurfacePressure);
+
+   // compute specific volume
+   const auto &PressureMid = VCoord->PressureMid;
+   EosInstance->computeSpecVol(ConservTemp, AbsSalinity, PressureMid);
+
+   // compute geometric height
+   VCoord->computeZHeight(LayerThickCell, EosInstance->SpecVol);
+
+   Pacer::stop("AuxState:computeMomVertAux", 2);
+}
+
 // Compute the auxiliary variables needed for momentum equation
-void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
-                                   int VelTimeLevel) const {
+void AuxiliaryState::computeMomAux(const OceanState *State,
+                                   const Array3DReal &TracerArray,
+                                   int ThickTimeLevel, int VelTimeLevel) const {
    Array2DReal LayerThickCell = State->getLayerThickness(ThickTimeLevel);
    Array2DReal NormalVelEdge  = State->getNormalVelocity(VelTimeLevel);
 
@@ -91,6 +131,8 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
 
    Pacer::start("AuxState:computeMomAux", 1);
 
+   computeMomVertAux(State, TracerArray, ThickTimeLevel, VelTimeLevel);
+
    Pacer::start("AuxState:vertexAuxState1", 2);
    parallelForOuter(
        "vertexAuxState1", {Mesh->NVerticesAll},
@@ -241,7 +283,7 @@ void AuxiliaryState::computeAll(const OceanState *State,
 
    Pacer::start("AuxState:computeAll", 1);
 
-   computeMomAux(State, ThickTimeLevel, VelTimeLevel);
+   computeMomAux(State, TracerArray, ThickTimeLevel, VelTimeLevel);
 
    Pacer::start("AuxState:cellAuxState3", 2);
    parallelForOuter(
@@ -290,7 +332,7 @@ void AuxiliaryState::computeAll(const OceanState *State,
 // Create a non-default auxiliary state
 AuxiliaryState *AuxiliaryState::create(const std::string &Name,
                                        const HorzMesh *Mesh, Halo *MeshHalo,
-                                       const VertCoord *VCoord, VertAdv *VAdv,
+                                       VertCoord *VCoord, VertAdv *VAdv,
                                        const int NTracers,
                                        TimeInterval TimeStep) {
    if (AllAuxStates.find(Name) != AllAuxStates.end()) {
@@ -312,7 +354,7 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
 void AuxiliaryState::init() {
    const HorzMesh *DefMesh           = HorzMesh::getDefault();
    Halo *DefHalo                     = Halo::getDefault();
-   const VertCoord *DefVCoord        = VertCoord::getDefault();
+   VertCoord *DefVCoord              = VertCoord::getDefault();
    VertAdv *DefVAdv                  = VertAdv::getDefault();
    const TimeStepper *DefTimeStepper = TimeStepper::getDefault();
 

components/omega/src/ocn/AuxiliaryState.h

@@ -54,7 +54,7 @@ class AuxiliaryState {
 
    // Create a non-default auxiliary state
    static AuxiliaryState *create(const std::string &Name, const HorzMesh *Mesh,
-                                 Halo *MeshHalo, const VertCoord *VCoord,
+                                 Halo *MeshHalo, VertCoord *VCoord,
                                  VertAdv *VAdv, int NTracers,
                                  TimeInterval TimeStep);
 
@@ -76,9 +76,14 @@ class AuxiliaryState {
    /// Exchange halo
    I4 exchangeHalo();
 
+   // Compute auxiliary variables for vertical dynamics
+   void computeMomVertAux(const OceanState *State,
+                          const Array3DReal &TracerArray, int ThickTimeLevel,
+                          int VelTimeLevel) const;
+
    // Compute all auxiliary variables needed for momentum equation
-   void computeMomAux(const OceanState *State, int ThickTimeLevel,
-                      int VelTimeLevel) const;
+   void computeMomAux(const OceanState *State, const Array3DReal &TracerArray,
+                      int ThickTimeLevel, int VelTimeLevel) const;
 
    /// Compute all auxiliary variables based on an ocean state at a given time
    /// level
@@ -89,15 +94,15 @@ class AuxiliaryState {
 
  private:
    AuxiliaryState(const std::string &Name, const HorzMesh *Mesh, Halo *MeshHalo,
-                  const VertCoord *VCoord, VertAdv *VAdv, int NTracers,
+                  VertCoord *VCoord, VertAdv *VAdv, int NTracers,
                   TimeInterval TimeStep);
 
    AuxiliaryState(const AuxiliaryState &) = delete;
    AuxiliaryState(AuxiliaryState &&)      = delete;
 
    const HorzMesh *Mesh;
    Halo *MeshHalo;
-   const VertCoord *VCoord;
+   VertCoord *VCoord;
    VertAdv *VAdv;
    TimeInterval TimeStep;
 

components/omega/src/ocn/Tendencies.cpp

@@ -664,28 +664,12 @@ void Tendencies::computeVelocityTendenciesOnly(
    // Compute pressure gradient
    if (PGrad->Enabled) {
 
-      // Temporary handling of surface pressure
-      Array1DReal SurfacePressure("SurfacePressure", Mesh->NCellsSize);
-      deepCopy(SurfacePressure, 0.0_Real);
-
       Pacer::start("Tend:pressureGradTerm", 2);
-      Array2DReal LayerThick = State->getLayerThickness(ThickTimeLevel);
-      VCoord->computePressure(LayerThick, SurfacePressure);
-
+      Array2DReal LayerThick        = State->getLayerThickness(ThickTimeLevel);
       const auto &PressureMid       = VCoord->PressureMid;
       const auto &PressureInterface = VCoord->PressureInterface;
-      Array2DReal Temp     = Kokkos::subview(TracerArray, Tracers::IndxTemp,
-                                             Kokkos::ALL, Kokkos::ALL);
-      Array2DReal Salinity = Kokkos::subview(TracerArray, Tracers::IndxSalt,
-                                             Kokkos::ALL, Kokkos::ALL);
-      EqState->computeSpecVol(Temp, Salinity, PressureMid);
-
-      // Temporary: ensure vertical geometric/geopotential fields are updated
-      // for pressure-gradient tendency calculations.
-      const auto &SpecVol = EqState->SpecVol;
-      VCoord->computeZHeight(LayerThick, SpecVol);
-
-      const auto &ZInterface = VCoord->ZInterface;
+      const auto &SpecVol           = EqState->SpecVol;
+      const auto &ZInterface        = VCoord->ZInterface;
       PGrad->computePressureGrad(LocNormalVelocityTend, PressureMid,
                                  PressureInterface, SpecVol, ZInterface,
                                  LayerThick);
@@ -887,7 +871,7 @@ void Tendencies::computeVelocityTendencies(
 ) {
    Pacer::start("Tend:computeVelocityTendencies", 1);
 
-   AuxState->computeMomAux(State, ThickTimeLevel, VelTimeLevel);
+   AuxState->computeMomAux(State, TracerArray, ThickTimeLevel, VelTimeLevel);
    computeVelocityTendenciesOnly(State, AuxState, TracerArray, ThickTimeLevel,
                                  VelTimeLevel, TracerTimeLevel, Time);
 

components/omega/src/ocn/VertCoord.cpp

@@ -166,6 +166,10 @@ VertCoord::VertCoord(const std::string &Name_, //< [in] Name for new VertCoord
    RefLayerThickness =
        Array2DReal("RefLayerThickness", NCellsSize, NVertLayers);
 
+   // TODO: Temporary handling of SurfacePressure
+   SurfacePressure = Array1DReal("SurfacePressure", NCellsSize);
+   deepCopy(SurfacePressure, 0);
+
    // Make host copies for device arrays not being read from file
    PressureInterfaceH    = createHostMirrorCopy(PressureInterface);
    PressureMidH          = createHostMirrorCopy(PressureMid);

components/omega/src/ocn/VertCoord.h

@@ -166,6 +166,11 @@ class VertCoord {
 
    HostArray1DReal BottomDepthH;
 
+   // TODO: Temporary handling of SurfacePressure
+   Array1DReal SurfacePressure;
+
+   HostArray1DReal SurfacePressureH;
+
    // VertCoord instance name and FieldGroup names
    std::string Name;
    std::string InitGroupName;

components/omega/src/timeStepping/ForwardBackwardStepper.cpp

@@ -29,38 +29,44 @@ void ForwardBackwardStepper::doStep(
     TimeInstant &SimTime // current simulation time
 ) const {
 
-   const int CurLevel  = 0;
-   const int NextLevel = 1;
+   const int VelCurLevel    = 0;
+   const int ThickCurLevel  = 0;
+   const int TracerCurLevel = 0;
 
-   Array3DReal CurTracerArray  = Tracers::getAll(CurLevel);
-   Array3DReal NextTracerArray = Tracers::getAll(NextLevel);
+   const int VelNextLevel    = 1;
+   const int ThickNextLevel  = 1;
+   const int TracerNextLevel = 1;
+
+   Array3DReal CurTracerArray  = Tracers::getAll(TracerCurLevel);
+   Array3DReal NextTracerArray = Tracers::getAll(TracerNextLevel);
 
    if (State == nullptr)
       LOG_CRITICAL("Invalid State");
    if (AuxState == nullptr)
       LOG_CRITICAL("Invalid AuxState");
 
-   // R_h^{n} = RHS_h(u^{n}, h^{n}, t^{n})
-   Tend->computeThicknessTendencies(State, AuxState, CurLevel, CurLevel,
-                                    SimTime);
+   // R_u^{n} = RHS_u(u^{n}, h^{n}, t^{n})
+   Tend->computeVelocityTendencies(State, AuxState, CurTracerArray,
+                                   ThickCurLevel, VelCurLevel, TracerCurLevel,
+                                   SimTime + TimeStep);
 
-   // h^{n+1} = h^{n} + R_h^{n}
-   updateThicknessByTend(State, NextLevel, State, CurLevel, TimeStep);
+   // u^{n+1} = u^{n} + R_u^{n}
+   updateVelocityByTend(State, VelNextLevel, State, VelCurLevel, TimeStep);
 
-   // R_phi^{n} = RHS_phi(u^{n}, h^{n}, phi^{n}, t^{n})
-   Tend->computeTracerTendencies(State, AuxState, CurTracerArray, CurLevel,
-                                 CurLevel, SimTime);
+   // R_h^{n} = RHS_h(u^{n+1}, h^{n}, t^{n})
+   Tend->computeThicknessTendencies(State, AuxState, ThickCurLevel,
+                                    VelNextLevel, SimTime);
 
-   // phi^{n+1} = (phi^{n} * h^{n} + R_phi^{n}) / h^{n+1}
-   updateTracersByTend(NextTracerArray, CurTracerArray, State, NextLevel, State,
-                       CurLevel, TimeStep);
+   // h^{n+1} = h^{n} + R_h^{n}
+   updateThicknessByTend(State, ThickNextLevel, State, ThickCurLevel, TimeStep);
 
-   // R_u^{n+1} = RHS_u(u^{n}, h^{n+1}, t^{n+1})
-   Tend->computeVelocityTendencies(State, AuxState, NextTracerArray, NextLevel,
-                                   CurLevel, NextLevel, SimTime + TimeStep);
+   // R_phi^{n} = RHS_phi(u^{n+1}, h^{n+1}, phi^{n}, t^{n})
+   Tend->computeTracerTendencies(State, AuxState, CurTracerArray,
+                                 ThickNextLevel, VelNextLevel, SimTime);
 
-   // u^{n+1} = u^{n} + R_u^{n+1}
-   updateVelocityByTend(State, NextLevel, State, CurLevel, TimeStep);
+   // phi^{n+1} = (phi^{n} * h^{n} + R_phi^{n}) / h^{n+1}
+   updateTracersByTend(NextTracerArray, CurTracerArray, State, ThickNextLevel,
+                       State, ThickCurLevel, TimeStep);
 
    // Update time levels (New -> Old) of prognostic variables with halo
    // exchanges

components/omega/test/ocn/AuxiliaryStateTest.cpp

@@ -3,6 +3,7 @@
 #include "DataTypes.h"
 #include "Decomp.h"
 #include "Dimension.h"
+#include "Eos.h"
 #include "Field.h"
 #include "GlobalConstants.h"
 #include "Halo.h"
@@ -126,6 +127,8 @@ int initAuxStateTest(const std::string &mesh) {
 
    VertAdv::init();
 
+   Eos::init();
+
    return Err;
 }
 
@@ -146,10 +149,10 @@ int testAuxState() {
       return -1;
    }
 
-   const auto *Mesh   = HorzMesh::getDefault();
-   auto *MeshHalo     = Halo::getDefault();
-   const auto *VCoord = VertCoord::getDefault();
-   auto *VAdv         = VertAdv::getDefault();
+   const auto *Mesh = HorzMesh::getDefault();
+   auto *MeshHalo   = Halo::getDefault();
+   auto *VCoord     = VertCoord::getDefault();
+   auto *VAdv       = VertAdv::getDefault();
    TimeInterval TimeStep;
    // test creation of another auxiliary state
    AuxiliaryState::create("AnotherAuxState", Mesh, MeshHalo, VCoord, VAdv, 3,