Port sphere transport time-varying prescribed velocity fields

components/omega/configs/Default.yml

@@ -22,6 +22,8 @@ Omega:
     IODefaultFormat: pnetcdf
   State:
     NTimeLevels: 2
+    PrescribeThicknessType: None
+    PrescribeVelocityType: None
   Advection:
     Coef3rdOrder: 0.25
     FluxThicknessType: Center

components/omega/src/timeStepping/ForwardBackwardStepper.cpp

@@ -62,6 +62,10 @@ void ForwardBackwardStepper::doStep(
    // u^{n+1} = u^{n} + R_u^{n+1}
    updateVelocityByTend(State, NextLevel, State, CurLevel, TimeStep);
 
+   prescribeThickness(State, NextLevel, State, CurLevel);
+   prescribeVelocity(State, NextLevel, State, CurLevel,
+                     SimTime + TimeStep);
+
    // Update time levels (New -> Old) of prognostic variables with halo
    // exchanges
    const MPI_Comm Comm = MeshHalo->getComm();

components/omega/src/timeStepping/RungeKutta2Stepper.cpp

@@ -56,6 +56,8 @@ void RungeKutta2Stepper::doStep(OceanState *State,   // model state
    updateTracersByTend(NextTracerArray, CurTracerArray, State, NextLevel, State,
                        CurLevel, TimeStep);
 
+   prescribeState(State, NextLevel, State, CurLevel, SimTime + TimeStep);
+
    // Update time levels (New -> Old) of prognostic variables with halo
    // exchanges
    const MPI_Comm Comm = MeshHalo->getComm();

components/omega/src/timeStepping/RungeKutta4Stepper.cpp

@@ -88,6 +88,7 @@ void RungeKutta4Stepper::doStep(OceanState *State,   // model state
                                     CurLevel, CurLevel, StageTime);
          updateStateByTend(State, NextLevel, State, CurLevel,
                            RKB[Stage] * TimeStep);
+         prescribeState(State, NextLevel, State, CurLevel, StageTime);
          accumulateTracersUpdate(NextTracerArray, RKB[Stage] * TimeStep);
       } else {
          // every other stage does:
@@ -96,6 +97,7 @@ void RungeKutta4Stepper::doStep(OceanState *State,   // model state
          // q^{n+1} += RKB[stage] * dt * R^{(s)}
          updateStateByTend(ProvisState, CurLevel, State, CurLevel,
                            RKA[Stage] * TimeStep);
+         prescribeState(State, NextLevel, State, CurLevel, StageTime);
          updateTracersByTend(ProvisTracers, CurTracerArray, ProvisState,
                              CurLevel, State, CurLevel, RKA[Stage] * TimeStep);
 
@@ -114,6 +116,7 @@ void RungeKutta4Stepper::doStep(OceanState *State,   // model state
                                     CurLevel, CurLevel, CurLevel, StageTime);
          updateStateByTend(State, NextLevel, State, NextLevel,
                            RKB[Stage] * TimeStep);
+         prescribeState(State, NextLevel, State, CurLevel, StageTime);
          accumulateTracersUpdate(NextTracerArray, RKB[Stage] * TimeStep);
       }
    }

components/omega/src/timeStepping/TimeStepper.cpp

@@ -10,16 +10,20 @@
 #include "ForwardBackwardStepper.h"
 #include "RungeKutta2Stepper.h"
 #include "RungeKutta4Stepper.h"
+#include "Logging.h"
 
 namespace OMEGA {
-
 //------------------------------------------------------------------------------
 // create the static class members
 // Default model time stepper
 TimeStepper *TimeStepper::DefaultTimeStepper = nullptr;
 // All defined time steppers
 std::map<std::string, std::unique_ptr<TimeStepper>>
     TimeStepper::AllTimeSteppers;
+PrescribeStateType TimeStepper::DefaultPrescribeThicknessMode =
+   PrescribeStateType::None;
+PrescribeStateType TimeStepper::DefaultPrescribeVelocityMode =
+   PrescribeStateType::None;
 
 //------------------------------------------------------------------------------
 // utility functions
@@ -44,6 +48,36 @@ TimeStepperType getTimeStepperFromStr(const std::string &InString) {
    return TimeStepperChoice;
 }
 
+PrescribeStateType getPrescribeThicknessTypeFromStr(const std::string &InString) {
+
+   if (InString == "None") {
+      return PrescribeStateType::None;
+   }
+   if (InString == "Init") {
+      return PrescribeStateType::Init;
+   }
+
+   ABORT_ERROR("PrescribeStateType should be 'None' or 'Init' for thickness but got {}",
+               InString);
+   return PrescribeStateType::Invalid;
+}
+PrescribeStateType getPrescribeVelocityTypeFromStr(const std::string &InString) {
+
+   if (InString == "None") {
+      return PrescribeStateType::None;
+   }else if (InString == "Init") {
+      return PrescribeStateType::Init;
+   }else if (InString == "NonDivergent") {
+      return PrescribeStateType::NonDivergent;
+   }else if (InString == "Divergent") {
+      return PrescribeStateType::Divergent;
+   }
+
+   ABORT_ERROR("PrescribeStateType should be 'None', 'Init', 'NonDivergent' or 'Divergent' for velocity but got {}",
+               InString);
+   return PrescribeStateType::Invalid;
+}
+
 //------------------------------------------------------------------------------
 // Constructors and creation methods.
 
@@ -119,6 +153,11 @@ TimeStepper *TimeStepper::create(
       ABORT_ERROR("Unknown time stepping method");
    }
 
+   NewTimeStepper->PrescribeThicknessMode =
+      DefaultPrescribeThicknessMode;
+   NewTimeStepper->PrescribeVelocityMode =
+      DefaultPrescribeVelocityMode;
+
    // Attach data pointers
    NewTimeStepper->attachData(InTend, InAuxState, InMesh, InVCoord, InMeshHalo);
 
@@ -170,6 +209,11 @@ TimeStepper *TimeStepper::create(
       ABORT_ERROR("Unknown time stepping method");
    }
 
+   NewTimeStepper->PrescribeThicknessMode =
+      DefaultPrescribeThicknessMode;
+   NewTimeStepper->PrescribeVelocityMode =
+      DefaultPrescribeVelocityMode;
+
    // Store instance
    AllTimeSteppers.emplace(InName, NewTimeStepper);
 
@@ -298,6 +342,22 @@ void TimeStepper::init1() {
       StopTime = StopTime2;
    }
 
+   Config StateConfig("State");
+   Error StateErr = OmegaConfig->get(StateConfig);
+   if (StateErr.isSuccess()) {
+      std::string ThicknessMode;
+      if (StateConfig.get("PrescribeThicknessType", ThicknessMode).isSuccess()) {
+         TimeStepper::DefaultPrescribeThicknessMode =
+             getPrescribeThicknessTypeFromStr(ThicknessMode);
+      }
+
+      std::string VelocityMode;
+      if (StateConfig.get("PrescribeVelocityType", VelocityMode).isSuccess()) {
+         TimeStepper::DefaultPrescribeVelocityMode =
+             getPrescribeVelocityTypeFromStr(VelocityMode);
+      }
+   }
+
    // Now that all the inputs are defined, create the default time stepper
    // Use the partial creation function for only the time info. Data
    // pointers will be attached in phase 2 initialization
@@ -460,6 +520,169 @@ void TimeStepper::updateStateByTend(OceanState *State1, int TimeLevel1,
    updateVelocityByTend(State1, TimeLevel1, State2, TimeLevel2, Coeff);
 }
 
+//------------------------------------------------------------------------------
+// Reset state variables to their initial values
+void TimeStepper::prescribeThickness(OceanState *State1, int TimeLevel1,
+                                     OceanState *State2, int TimeLevel2) const {
+
+   if (PrescribeThicknessMode == PrescribeStateType::None) {
+      return;
+   }
+
+   if (PrescribeThicknessMode == PrescribeStateType::Init) {
+      Array2DReal LayerThick1 = State1->getLayerThickness(TimeLevel1);
+      Array2DReal LayerThick2 = State2->getLayerThickness(TimeLevel2);
+
+      OMEGA_SCOPE(MinLayerCell, VCoord->MinLayerCell);
+      OMEGA_SCOPE(MaxLayerCell, VCoord->MaxLayerCell);
+
+      parallelForOuter(
+          "prescribeThickness", {Mesh->NCellsAll},
+          KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {
+             const int KMin   = MinLayerCell(ICell);
+             const int KMax   = MaxLayerCell(ICell);
+             const int KRange = vertRange(KMin, KMax);
+
+             parallelForInner(
+                 Team, KRange, INNER_LAMBDA(int KChunk) {
+                    const int K = KMin + KChunk;
+                    LayerThick1(ICell, K) = LayerThick2(ICell, K);
+                 });
+          });
+      return;
+   }
+}
+
+//------------------------------------------------------------------------------
+void TimeStepper::prescribeVelocity(OceanState *State1, int TimeLevel1,
+                           OceanState *State2, int TimeLevel2,
+                           const TimeInstant &SimTime) const {
+
+   if (PrescribeVelocityMode == PrescribeStateType::None) {
+      return;
+   }
+
+   if (PrescribeVelocityMode == PrescribeStateType::Init) {
+      Array2DReal NormalVel1 = State1->getNormalVelocity(TimeLevel1);
+      Array2DReal NormalVel2 = State2->getNormalVelocity(TimeLevel2);
+
+      OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBot);
+      OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTop);
+
+      parallelForOuter(
+          "prescribeVelocity", {Mesh->NEdgesAll},
+          KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) {
+             const int KMin   = MinLayerEdgeBot(IEdge);
+             const int KMax   = MaxLayerEdgeTop(IEdge);
+             const int KRange = vertRange(KMin, KMax);
+
+             parallelForInner(
+                 Team, KRange, INNER_LAMBDA(int KChunk) {
+                    const int K = KMin + KChunk;
+                    NormalVel2(IEdge, K) = NormalVel1(IEdge, K);
+                 });
+          });
+      return;
+   } else if (PrescribeVelocityMode == PrescribeStateType::NonDivergent) {
+      Array2DReal NormalVel2 = State2->getNormalVelocity(TimeLevel2);
+
+      OMEGA_SCOPE(LatEdge, Mesh->LatEdgeH);
+      OMEGA_SCOPE(LonEdge, Mesh->LonEdgeH);
+      OMEGA_SCOPE(AngleEdge, Mesh->AngleEdgeH);
+      OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBotH);
+      OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTopH);
+
+      const Clock *ModelClock = StepClock.get();
+      R8 ElapsedTimeSec;
+      TimeInterval ElapsedTimeInterval = SimTime - ModelClock->getCurrentTime();
+      ElapsedTimeInterval.get(ElapsedTimeSec, TimeUnits::Seconds);
+
+      const R8 Tau     = 12. * Day2Sec; // 12 days in seconds
+      const R8 TSim    = ElapsedTimeSec;
+
+      parallelForOuter(
+          "prescribeVelocityNonDivergent", {Mesh->NEdgesAll},
+          KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) {
+             const int KMin   = MinLayerEdgeBot(IEdge);
+             const int KMax   = MaxLayerEdgeTop(IEdge);
+             const int KRange = vertRange(KMin, KMax);
+
+             const R8 lon_p =
+                 LonEdge(IEdge) - 2.0 * Pi * TSim / Tau;
+             const R8 u = (1 / Tau) *
+                          (10.0 * Kokkos::pow(sin(lon_p), 2) *
+                           sin(2.0 * LatEdge(IEdge)) *
+                           cos(Pi * TSim / Tau) +
+                           2.0 * Pi * cos(LatEdge(IEdge)));
+             const R8 v = (10.0 / Tau) * sin(2.0 * lon_p) *
+                          cos(LatEdge(IEdge)) * cos(Pi * TSim / Tau);
+             const R8 normalVel = REarth * (
+                 u * cos(AngleEdge(IEdge)) + v * sin(AngleEdge(IEdge)));
+
+             parallelForInner(
+                 Team, KRange, INNER_LAMBDA(int KChunk) {
+                    const int K = KMin + KChunk;
+                    NormalVel2(IEdge, K) = normalVel;
+                 });
+          });
+      return;
+   } else if (PrescribeVelocityMode == PrescribeStateType::Divergent) {
+      Array2DReal NormalVel2 = State2->getNormalVelocity(TimeLevel2);
+
+      OMEGA_SCOPE(LatEdge, Mesh->LatEdgeH);
+      OMEGA_SCOPE(LonEdge, Mesh->LonEdgeH);
+      OMEGA_SCOPE(AngleEdge, Mesh->AngleEdgeH);
+      OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBotH);
+      OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTopH);
+
+      const Clock *ModelClock = StepClock.get();
+      R8 ElapsedTimeSec;
+      TimeInterval ElapsedTimeInterval = SimTime - ModelClock->getCurrentTime();
+      ElapsedTimeInterval.get(ElapsedTimeSec, TimeUnits::Seconds);
+
+      const R8 Tau     = 12. * Day2Sec; // 14 days in seconds
+      const R8 TSim    = ElapsedTimeSec;
+
+      parallelForOuter(
+          "prescribeVelocityDivergent", {Mesh->NEdgesAll},
+          KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) {
+             const int KMin   = MinLayerEdgeBot(IEdge);
+             const int KMax   = MaxLayerEdgeTop(IEdge);
+             const int KRange = vertRange(KMin, KMax);
+
+             const R8 lon_p =
+                 LonEdge(IEdge) - 2.0 * Pi * TSim / Tau;
+             const R8 u = (1.0 / Tau) *
+                          (-5.0 * Kokkos::pow(sin(lon_p / 2), 2) *
+                          sin(2.0 * LatEdge(IEdge)) *
+                          Kokkos::pow(cos(LatEdge(IEdge)), 2) *
+                          cos(Pi * TSim / Tau) +
+                          2.0 * Pi * cos(LatEdge(IEdge)));
+             const R8 v = ((2.5 / Tau) *
+                           sin(lon_p) * 
+                           Kokkos::pow(cos(LatEdge(IEdge)), 3) *
+                           cos(Pi * TSim / Tau));
+             const R8 normalVel = REarth * (
+                 u * cos(AngleEdge(IEdge)) + v * sin(AngleEdge(IEdge)));
+
+             parallelForInner(
+                 Team, KRange, INNER_LAMBDA(int KChunk) {
+                    const int K = KMin + KChunk;
+                    NormalVel2(IEdge, K) = normalVel;
+                 });
+          });
+      return;
+   }
+}
+
+//------------------------------------------------------------------------------
+void TimeStepper::prescribeState(OceanState *State1, int TimeLevel1,
+                                 OceanState *State2, int TimeLevel2,
+                                 const TimeInstant &SimTime) const {
+   prescribeThickness(State1, TimeLevel1, State2, TimeLevel2);
+   prescribeVelocity(State1, TimeLevel1, State2, TimeLevel2, SimTime);
+}
+
 //------------------------------------------------------------------------------
 // Updates tracers
 // NextTracers = (CurTracers * LayerThickness2(TimeLevel2)) +