From 005552306a8651081bed5644044bb652c48e282d Mon Sep 17 00:00:00 2001
From: Chris Eldred <celdred@sandia.gov>
Date: Wed, 20 Dec 2023 14:07:35 -0700
Subject: [PATCH] adding W clipping and CRM dt adjustment based on vertical cfl

---
 dynamics/spam/Dycore.h                   | 19 +++++++--
 dynamics/spam/src/core/params.h          | 14 +++++++
 dynamics/spam/src/models/extrudedmodel.h | 50 ++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/dynamics/spam/Dycore.h b/dynamics/spam/Dycore.h
index d922c936..973175d2 100644
--- a/dynamics/spam/Dycore.h
+++ b/dynamics/spam/Dycore.h
@@ -241,6 +241,13 @@ class Dycore {
         coupler, prognostic_vars, auxiliary_vars, constant_vars,
         params.couple_wind, params.couple_wind_exact_inverse);
 
+    // adjust crm_per_phys and dtcrm if needed
+
+  if(params.adjust_crm_per_phys_using_vert_cfl)
+  {
+  tendencies.adjust_crm_per_phys_using_vert_cfl(params, prognostic_vars);
+  }
+  //tendencies.a
     // Time stepping loop
     debug_print("start time stepping loop", par.masterproc);
     for (uint nstep = 0; nstep < params.crm_per_phys; nstep++) {
@@ -248,6 +255,15 @@ class Dycore {
 
       time_integrator->step_forward(params.dtcrm);
 
+  if(params.clip_vertical_velocities)
+  {
+tendencies.clip_vertical_velocities(params, prognostic_vars);
+}
+
+    if (params.clip_negative_densities) {
+      tendencies.clip_negative_densities(prognostic_vars);
+    }
+
 #if defined PAMC_AN || defined PAMC_MAN
       if (params.check_anelastic_constraint) {
         real max_div = tendencies.compute_max_anelastic_constraint(
@@ -283,9 +299,6 @@ class Dycore {
     }
     prevstep += params.crm_per_phys;
 
-    if (params.clip_negative_densities) {
-      tendencies.clip_negative_densities(prognostic_vars);
-    }
 
     // convert dynamics state to Coupler state
     tendencies.convert_dynamics_to_coupler_state(
diff --git a/dynamics/spam/src/core/params.h b/dynamics/spam/src/core/params.h
index b2c88fd7..2a076bc6 100644
--- a/dynamics/spam/src/core/params.h
+++ b/dynamics/spam/src/core/params.h
@@ -42,6 +42,11 @@ class Parameters {
   // during conversion to coupler state when coupling winds
   bool couple_wind_exact_inverse = false;
   bool clip_negative_densities = true;
+
+  bool clip_vertical_velocities = false;
+  bool adjust_crm_per_phys_using_vert_cfl = false;
+  real target_cfl = 0.7;
+  real max_w = 50.0;
 };
 
 void readParamsFile(std::string inFile, Parameters &params, Parallel &par,
@@ -82,6 +87,15 @@ void readParamsFile(std::string inFile, Parameters &params, Parallel &par,
   params.clip_negative_densities =
       config["clip_negative_densities"].as<bool>(true);
 
+  params.clip_vertical_velocities =
+      config["clip_vertical_velocities"].as<bool>(false);
+
+  params.adjust_crm_per_phys_using_vert_cfl =
+      config["adjust_crm_per_phys_using_vert_cfl"].as<bool>(false);
+
+  params.target_cfl = config["max_w"].as<real>(0.7);
+  params.max_w = config["max_w"].as<real>(50.0);
+
   // ADD A CHECK HERE THAT TOTAL TIME IS EXACTLY DIVISIBLE BY STAT_FREQ
   if (params.stat_freq >= 0.) {
     params.statSize = params.sim_time / params.stat_freq;
diff --git a/dynamics/spam/src/models/extrudedmodel.h b/dynamics/spam/src/models/extrudedmodel.h
index da8c2433..ea6c0ddf 100644
--- a/dynamics/spam/src/models/extrudedmodel.h
+++ b/dynamics/spam/src/models/extrudedmodel.h
@@ -396,6 +396,56 @@ class ModelTendencies : public ExtrudedTendencies {
 #endif
   }
 
+void clip_vertical_velocities(ModelParameters &params, FieldSet<nprognostic> &x)
+{
+    real max_w = params.max_w;
+
+    const auto &primal_topology = primal_geometry.topology;
+
+    const auto wvar = x.fields_arr[WVAR].data;
+    const int pis = primal_topology.is;
+    const int pjs = primal_topology.js;
+    const int pks = primal_topology.ks;
+
+    parallel_for(
+        "Clip vertical velocity",
+        SimpleBounds<4>(primal_topology.nl, primal_topology.n_cells_y,
+                        primal_topology.n_cells_x, primal_topology.nens),
+        YAKL_LAMBDA(int k, int j, int i, int n) {
+          real wedgelength = primal_geometry.get_area_01entity(k + pks, j + pjs, i + pis, n);
+              wvar(0, k + pks, j + pjs, i + pis, n) =
+                  std::max(-max_w * wedgelength, std::min(max_w * wedgelength, wvar(0, k + pks, j + pjs, i + pis, n)));
+        });
+  }
+
+void adjust_crm_per_phys_using_vert_cfl(ModelParameters &params, FieldSet<nprognostic> &x)
+{
+//cfl = c * dt / dx
+//dt = cfl * dx / c
+
+    const auto &primal_topology = primal_geometry.topology;
+    const auto wvar = x.fields_arr[WVAR].data;
+    const int pis = primal_topology.is;
+    const int pjs = primal_topology.js;
+    const int pks = primal_topology.ks;
+
+    real4d min_dt_4d("min_dt_4d", primal_topology.nl, primal_topology.n_cells_y, primal_topology.n_cells_x, primal_topology.nens);
+
+    real target_cfl = params.target_cfl;
+    parallel_for(
+        "compute adjusted dt",
+        SimpleBounds<4>(primal_topology.nl, primal_topology.n_cells_y,
+                        primal_topology.n_cells_x, primal_topology.nens),
+        YAKL_LAMBDA(int k, int j, int i, int n) {
+          real wedgelength = primal_geometry.get_area_01entity(k + pks, j + pjs, i + pis, n);
+            real wval = std::abs(wvar(0, k + pks, j + pjs, i + pis, n)) / wedgelength;
+            min_dt_4d(k, j, i, n) = target_cfl * wedgelength / wval;
+        });
+  real min_dt = yakl::intrinsics::minval( min_dt_4d );
+  params.crm_per_phys = (int) std::ceil(params.dt_crm_phys / min_dt);
+  params.dtcrm = params.dt_crm_phys / params.crm_per_phys;
+}
+
   void convert_dynamics_to_coupler_state(PamCoupler &coupler,
                                          const FieldSet<nprognostic> &prog_vars,
                                          const FieldSet<nconstant> &const_vars,