Add vertical advection

components/omega/configs/Default.yml

@@ -23,9 +23,12 @@ Omega:
   State:
     NTimeLevels: 2
   Advection:
+    Coef3rdOrder: 0.25
     FluxThicknessType: Center
     HorzTracerFluxOrder: 2
-    Coef3rdOrder: 0.25
+    FluxTracerType: Center
+    VerticalTracerFluxLimiterEnable: true
+    VerticalTracerFluxOrder: 3
   WindStress:
     InterpType: Isotropic
   VertCoord:
@@ -51,6 +54,9 @@ Omega:
     EddyDiff4: 0.0
     UseCustomTendency: false
     ManufacturedSolutionTendency: false
+    ThicknessVertAdvTendencyEnable: true
+    VelocityVertAdvTendencyEnable: true
+    TracerVertAdvTendencyEnable: true
   Tracers:
     Base: [Temperature, Salinity]
     Debug: [Debug1, Debug2, Debug3]

components/omega/doc/devGuide/VertAdv.md

@@ -0,0 +1,103 @@
+(omega-dev-vert-adv)=
+
+## Vertical Advection
+
+The `VertAdv` class contains methods and arrays for calculating vertical
+velocities and the tendencies of thickness, horizontal velocity, and tracers due
+to vertical advection.
+
+### Initialization
+
+The default `VertAdv` instance is created by calling `VertAdv::init()` method.
+The default `HorzMesh`, `VertCoord`, and `Tracers` objects must be initialized
+first. The default instance is retrieved by:
+```
+VertAdv *DefVertAdv = VertAdv::getDefault();
+```
+
+Additional instances can be created with the `create` method, which calls the
+constructor and places the new instance in a map identified with the label
+`Name`:
+```
+VertAdv *VertAdv::create(const std::string &Name, ///< [in] name for new VertAdv
+                         const HorzMesh *Mesh,    ///< [in] associated HorzMesh
+                         const VertCoord *VCoord, ///< [in] associated VertCoord
+                         Config *Options          ///< [in] configuration options
+)
+```
+This instance is retrieved with the get method:
+```
+VertAdv *NewVertAdv = VertAdv::get("Name");
+```
+The constructor reads configuration options, stores some info from the
+`HorzMesh`, `VertCoord`, and `Tracers` objects, allocates needed arrays,
+and defines `Field` metadata.
+
+### Variables
+
+The following arrays are stored as members of the `VertAdv` class.
+
+| Variable Name | Type | Dimensions |
+| ------------- | ---- | ---------- |
+| VerticalVelocity | Real | NCellsSize, NVertLayersP1 |
+| TotalVerticalVelocity | Real | NCellsSize, NVertLayersP1 |
+| VertFlux | Real | NTracers, NCellsSize, NVertLayersP1 |
+| LowOrderVertFlux | Real | NTracers, NCellsSize, NVertLayersP1 |
+
+The `VerticalVelocity` represents the raw vertical pseudovelocity through the
+top interfaces of cell layers, computed from the divergence of the horizontal
+velocity. The `TotalVerticalVelocity` is the transport velocity and includes
+corrections applied to the `VerticalVelocity`. The `VertFlux` and
+`LowOrderVertFlux` store tracer fluxes at layer interfaces. The specific
+numerical algorithms to compute these fluxes are chosen by the user through
+configuration options.
+
+### Methods
+
+The `VerticalVelocity` and `TotalVerticalVelocity` arrays are computed by
+calling the `computeVerticalVelocity` method:
+```
+VertAdv::computeVerticalVelocity(NormalVelocity, FluxLayerThickEdge);
+```
+This method takes as input the `NormalVelocity` field from the `OceanState`
+object, and the `FluxLayerThickEdge` field from the `AuxiliaryState`. At
+present, `TotalVerticalVelocity` is equivalent to `VerticalVelocity`;
+additional corrections will be added in subsequent updates. The
+`TotalVerticalVelocity` array is used by each of the tendency methods,
+therefore `computeVerticalVelocity` must be called before any tendency
+computations.
+
+There are three methods for computing vertical advection tendencies of
+thickness, horizontal velocity and tracers that are called from the time
+stepper: `computeThicknessVAdvTend`, `computeVelocityVAdvTend`, and
+`computeTracerVAdvTend`. Each method updates the corresponding tendency array
+in place (inout).
+
+Only the thickness tendency array is passed to the `computeThicknessVAdvTend`
+method:
+```
+VertAdv::computeThicknessVAdvTend(ThickTend);
+```
+
+The `computeVelocityVAdvTend` method requires the `NormalVelocity` and
+`FluxLayerThickEdge` fields, in addition to the velocity tendency array:
+```
+VertAdv::computeVelocityVAdvTend(VelTend, NormalVelocity, FluxLayerThickEdge);
+```
+
+The tracer vertical advection tendency depends on the configured settings.
+The `computeTracerVAdvTend` method takes as arguments the full 3D array of
+tracers, a thickness array (selected based on the configured advection
+algorithm) and a `TimeInterval` representing the time step, along with the
+tracer tendency array:
+```
+VertAdv::computeTracerVAdvTend(TracerTend, TracerArray, Thickness, TimeStep);
+```
+For the standard advection algorithm, `LayerThickness` from the `OceanState` is
+passed as the thickness argument; for the FCT algorithm, the `ProvThickness`
+from the `AuxiliaryState` is used instead.
+
+This method first calls `computeVerticalFluxes` to compute the tracer fluxes at
+layer interfaces, using the order of accuracy specified in the configuration
+file. It then dispatches to the selected advection algorithm (standard or FCT)
+to compute the tracer tendencies.

components/omega/doc/index.md

@@ -51,6 +51,7 @@ userGuide/TridiagonalSolvers
 userGuide/VertCoord
 userGuide/Timing
 userGuide/VerticalMixingCoeff
+userGuide/VertAdv
 ```
 
 ```{toctree}
@@ -94,6 +95,7 @@ devGuide/TridiagonalSolvers
 devGuide/VertCoord
 devGuide/Timing
 devGuide/VerticalMixingCoeff
+devGuide/VertAdv
 ```
 
 ```{toctree}

components/omega/doc/userGuide/HorzMesh.md

@@ -23,7 +23,6 @@ This includes the following variables:
 | XCell, YCell, ZCell | Cartesian coordinates of cell centers | m |
 | XEdge, YEdge, ZEdge | Cartesian coordinates of edge centers | m |
 | XVertex, YVertex, ZVertex | Cartesian coordinates of vertices | m |
-| BottomDepth | Depth of the bottom of the ocean at cell centers | m |
 | FCell, FEdge, FVertex | Coriolis parameter at cell centers/edges/vertices | radians/s |
 | LonCell, LatCell | Longitude/latitude coordinates of cell centers | radians |
 | LonEdge, LatEdge | Longitude/latitude coordinates of edge centers | radians |

components/omega/doc/userGuide/VertAdv.md

@@ -0,0 +1,47 @@
+(omega-user-vert-adv)=
+
+## Vertical Advection
+
+### Overview
+
+In Omega, vertical advection represents the transport of mass, momentum, and
+tracer quantities along the vertical coordinate due to vertical motion within a
+water column. In the context of Omega's
+[governing equations](omega-design-governing-eqns-omega1), vertical advection
+contributes to the tendencies of pseudo-thickness, horizontal velocity, and
+tracer fields. Vertical motion is inferred from the divergence of horizontal
+flow through the continuity equation and is necessary for accurately
+representing vertical transport in three-dimensional ocean simulations.
+
+The `VertAdv` class contains variables and functions for computing the vertical
+velocity and the tendencies of `LayerThickness`, `NormalVelocity`, and
+`Tracers`. The algorithms used are determined by options specified in the
+configuration file.
+
+### Configuration Options
+
+Within the `Tendencies` group of the configuration file, flags are used to
+enable or disable the contribution of vertical advection to the tendencies of
+thickness, velocity, and tracers:
+| Configuration name | Options |
+| ------------------ | ------- |
+| ThicknessVertAdvTendencyEnable | true / false |
+| VelocityVertAdvTendencyEnable | true / false |
+| TracerVertAdvTendencyEnable | true / false |
+
+Within the `Advection` group, options are provided to configure the algorithms
+used to compute the tracer tendencies:
+| Configuration name | Description | Options |
+| ------------------ | ----------- | ------- |
+| VerticalTracerFluxOrder | order of accuracy used for the tracer flux calculation | 2, 3, or 4 |
+| VerticalTracerFluxLimiterEnable | selects the tracer tendency algorithm | true / false |
+| Coef3rdOrder | coefficient used for blending 3rd- and 4th-order fluxes when `VerticalTracerFluxOrder == 3` | 0.0 - 1.0 |
+
+For `VerticalTracerFluxLimiterEnable`, `true` enables flux-corrected transport,
+and `false` selects the standard algorithm. For `Coef3rdOrder`, `1` corresponds
+to purely 3rd-order, `0` purely 4th-order.
+
+The `VertAdv` class provides four fields that can be written to output by
+adding them to the contents of an output stream in the configuration file:
+`VerticalVelocity`, `TotalVerticalVelocity`, `VertFlux`, and `LowOrderVertFlux`.
+These fields collectively make up the `VertAdv` output group.

components/omega/src/infra/OmegaKokkos.h

@@ -78,6 +78,8 @@ using ScratchMemSpace = ExecSpace::scratch_memory_space;
 using Kokkos::MemoryUnmanaged;
 using Kokkos::PerTeam;
 using Kokkos::TeamThreadRange;
+using RealScratchArray =
+    Kokkos::View<Real *, ScratchMemSpace, Kokkos::MemoryUnmanaged>;
 
 /// team_size for hierarchical parallelism
 #ifdef OMEGA_TARGET_DEVICE
@@ -321,7 +323,8 @@ KOKKOS_INLINE_FUNCTION void teamBarrier(const TeamMember &Team) {
 // parallelForOuter: with label
 template <int N, class F>
 inline void parallelForOuter(const std::string &Label,
-                             const int (&UpperBounds)[N], F &&Functor) {
+                             const int (&UpperBounds)[N], F &&Functor,
+                             int ScratchValsPerTeam = 0) {
 
    auto LinFunctor = LinearIdxWrapper{std::forward<F>(Functor), UpperBounds};
    int LinBound    = 1;
@@ -330,6 +333,12 @@ inline void parallelForOuter(const std::string &Label,
    }
 
    auto Policy = TeamPolicy(LinBound, OMEGA_TEAMSIZE);
+
+   if (ScratchValsPerTeam > 0) {
+      Policy.set_scratch_size(
+          0, Kokkos::PerTeam(ScratchValsPerTeam * sizeof(Real)));
+   }
+
    Kokkos::parallel_for(
        Label, Policy, KOKKOS_LAMBDA(const TeamMember &Team) {
           const int TeamId = Team.league_rank();
@@ -339,8 +348,10 @@ inline void parallelForOuter(const std::string &Label,
 
 // parallelForOuter: without label
 template <int N, class F>
-inline void parallelForOuter(const int (&UpperBounds)[N], F &&Functor) {
-   parallelForOuter("", UpperBounds, std::forward<F>(Functor));
+inline void parallelForOuter(const int (&UpperBounds)[N], F &&Functor,
+                             int ScratchValsPerTeam = 0) {
+   parallelForOuter("", UpperBounds, std::forward<F>(Functor),
+                    ScratchValsPerTeam);
 }
 
 // parallelForInner

components/omega/src/ocn/AuxiliaryState.cpp

@@ -3,6 +3,7 @@
 #include "Field.h"
 #include "Logging.h"
 #include "Pacer.h"
+#include "TimeStepper.h"
 
 namespace OMEGA {
 
@@ -19,9 +20,10 @@ static std::string stripDefault(const std::string &Name) {
 // fields with IOStreams
 AuxiliaryState::AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
                                Halo *MeshHalo, const VertCoord *VCoord,
-                               int NTracers)
-    : Mesh(Mesh), MeshHalo(MeshHalo), VCoord(VCoord), Name(stripDefault(Name)),
-      KineticAux(stripDefault(Name), Mesh, 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),
       VorticityAux(stripDefault(Name), Mesh, VCoord),
       VelocityDel2Aux(stripDefault(Name), Mesh, VCoord),
@@ -80,6 +82,9 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
    OMEGA_SCOPE(MaxLayerEdgeBot, VCoord->MaxLayerEdgeBot);
    OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTop);
 
+   R8 TimeStepSeconds;
+   TimeStep.get(TimeStepSeconds, TimeUnits::Seconds);
+
    Pacer::start("AuxState:computeMomAux", 1);
 
    Pacer::start("AuxState:vertexAuxState1", 2);
@@ -194,12 +199,20 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
 
           parallelForInner(
               Team, KRange, INNER_LAMBDA(int KChunk) {
-                 LocLayerThicknessAux.computeVarsOnCells(ICell, KChunk,
-                                                         LayerThickCell);
+                 LocLayerThicknessAux.computeVarsOnCells(
+                     ICell, KChunk, LayerThickCell, NormalVelEdge,
+                     TimeStepSeconds);
               });
        });
    Pacer::stop("AuxState:cellAuxState3", 2);
 
+   Pacer::start("AuxState:computeVerticalVelocity", 2);
+
+   const auto &FluxLayerThickEdge = LayerThicknessAux.FluxLayerThickEdge;
+   VAdv->computeVerticalVelocity(NormalVelEdge, FluxLayerThickEdge);
+
+   Pacer::stop("AuxState:computeVerticalVelocity", 2);
+
    Pacer::stop("AuxState:computeMomAux", 1);
 }
 
@@ -212,16 +225,37 @@ void AuxiliaryState::computeAll(const OceanState *State,
 
    const int NTracers = TracerArray.extent_int(0);
 
+   OMEGA_SCOPE(LocLayerThicknessAux, LayerThicknessAux);
    OMEGA_SCOPE(LocTracerAux, TracerAux);
    OMEGA_SCOPE(MinLayerCell, VCoord->MinLayerCell);
    OMEGA_SCOPE(MaxLayerCell, VCoord->MaxLayerCell);
    OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBot);
    OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTop);
 
+   R8 TimeStepSeconds;
+   TimeStep.get(TimeStepSeconds, TimeUnits::Seconds);
+
    Pacer::start("AuxState:computeAll", 1);
 
    computeMomAux(State, ThickTimeLevel, VelTimeLevel);
 
+   Pacer::start("AuxState:cellAuxState3", 2);
+   parallelForOuter(
+       "cellAuxState3", {Mesh->NCellsAll},
+       KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {
+          const int KMin   = MinLayerCell(ICell);
+          const int KMax   = MaxLayerCell(ICell);
+          const int KRange = vertRangeChunked(KMin, KMax);
+
+          parallelForInner(
+              Team, KRange, INNER_LAMBDA(int KChunk) {
+                 LocLayerThicknessAux.computeVarsOnCells(
+                     ICell, KChunk, LayerThickCell, NormalVelEdge,
+                     TimeStepSeconds);
+              });
+       });
+   Pacer::stop("AuxState:cellAuxState3", 2);
+
    const auto &MeanLayerThickEdge = LayerThicknessAux.MeanLayerThickEdge;
 
    Pacer::start("AuxState:cellAuxState4", 2);
@@ -252,33 +286,37 @@ 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,
-                                       const int NTracers) {
+                                       const VertCoord *VCoord, VertAdv *VAdv,
+                                       const int NTracers,
+                                       TimeInterval TimeStep) {
    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, MeshHalo, VCoord, NTracers);
+   auto *NewAuxState = new AuxiliaryState(Name, Mesh, MeshHalo, VCoord, VAdv,
+                                          NTracers, TimeStep);
    AllAuxStates.emplace(Name, NewAuxState);
 
    return NewAuxState;
 }
 
-// Create the default auxiliary state. Assumes that HorzMesh, VertCoord and
-// Halo have been initialized.
+// Create the default auxiliary state. Assumes that HorzMesh, VertCoord,
+// VertAdv, and Halo have been initialized.
 void AuxiliaryState::init() {
-   const HorzMesh *DefMesh    = HorzMesh::getDefault();
-   Halo *DefHalo              = Halo::getDefault();
-   const VertCoord *DefVCoord = VertCoord::getDefault();
+   const HorzMesh *DefMesh           = HorzMesh::getDefault();
+   Halo *DefHalo                     = Halo::getDefault();
+   const VertCoord *DefVCoord        = VertCoord::getDefault();
+   VertAdv *DefVAdv                  = VertAdv::getDefault();
+   const TimeStepper *DefTimeStepper = TimeStepper::getDefault();
 
-   int NTracers = Tracers::getNumTracers();
+   int NTracers          = Tracers::getNumTracers();
+   TimeInterval TimeStep = DefTimeStepper->getTimeStep();
 
-   AuxiliaryState::DefaultAuxState =
-       AuxiliaryState::create("Default", DefMesh, DefHalo, DefVCoord, NTracers);
+   AuxiliaryState::DefaultAuxState = AuxiliaryState::create(
+       "Default", DefMesh, DefHalo, DefVCoord, DefVAdv, NTracers, TimeStep);
 
    Config *OmegaConfig = Config::getOmegaConfig();
    DefaultAuxState->readConfigOptions(OmegaConfig);