[WIP] Additional pkz+rdg_var mods

pyfv3/initialization/init_utils.py

@@ -223,18 +223,19 @@ def initialize_log_pressure_interfaces(pe, ptop):
     return peln
 
 
-def initialize_pkz_dry(delp, pt, delz):
+def initialize_pkz_dry(delp, pt, delz, rdg_var):
+    # TODO: Is WAM-related variable rdg necessary here?
     return np.exp(
         constants.KAPPA
-        * np.log(constants.RDG * delp[:, :, :-1] * pt[:, :, :-1] / delz[:, :, :-1])
+        * np.log(rdg_var * delp[:, :, :-1] * pt[:, :, :-1] / delz[:, :, :-1])
     )
 
 
-def initialize_pkz_moist(delp, pt, qvapor, delz):
+def initialize_pkz_moist(delp, pt, qvapor, delz, rdg_var):
     return np.exp(
         constants.KAPPA
         * np.log(
-            constants.RDG
+            rdg_var[:, :, :-1]  # TODO: Is WAM-related variable rdg necessary here?
             * delp[:, :, :-1]
             * pt[:, :, :-1]
             * (1.0 + constants.ZVIR * qvapor[:, :, :-1])
@@ -282,6 +283,7 @@ def p_var(
     ptop,
     moist_phys,
     make_nh,
+    rdg_var,
 ):
     """
     Computes auxiliary pressure variables for a hydrostatic state.
@@ -297,9 +299,9 @@ def p_var(
     if make_nh:
         delz[:, :, :-1] = initialize_delz(pt, peln)
     if moist_phys:
-        pkz[:, :, :-1] = initialize_pkz_moist(delp, pt, qvapor, delz)
+        pkz[:, :, :-1] = initialize_pkz_moist(delp, pt, qvapor, delz, rdg_var)
     else:
-        pkz[:, :, :-1] = initialize_pkz_dry(delp, pt, delz)
+        pkz[:, :, :-1] = initialize_pkz_dry(delp, pt, delz, rdg_var)
 
 
 def setup_pressure_fields(

pyfv3/initialization/test_cases/initialize_baroclinic.py

@@ -304,6 +304,8 @@ def init_baroclinic_state(
     numpy_state.delz[:] = 1.0e25
     numpy_state.phis[:] = 1.0e25
     numpy_state.ps[:] = SURFACE_PRESSURE
+    numpy_state.grav_var[:] = constants.GRAV
+    numpy_state.rdg_var[:] = -1 / numpy_state.grav_var
     eta = np.zeros(nz)
     eta_v = np.zeros(nz)
     islice, jslice, slice_3d, slice_2d = init_utils.compute_slices(nx, ny)
@@ -365,6 +367,7 @@ def init_baroclinic_state(
         ptop=grid_data.ptop,
         moist_phys=moist_phys,
         make_nh=(not hydrostatic),
+        rdg_var=numpy_state.rdg_var[slice_3d],
     )
     state = DycoreState.init_from_numpy_arrays(
         numpy_state.__dict__,

pyfv3/stencils/dyn_core.py

@@ -10,6 +10,7 @@
 import pyfv3.stencils.pe_halo as pe_halo
 import pyfv3.stencils.ray_fast as ray_fast
 import pyfv3.stencils.temperature_adjust as temperature_adjust
+import pyfv3.stencils.rdg_adjust as rdg_adjust
 import pyfv3.stencils.updatedzc as updatedzc
 import pyfv3.stencils.updatedzd as updatedzd
 from ndsl import (
@@ -101,17 +102,6 @@ def average_gravity(grav_var: FloatField, grav_var_h: FloatField):
         grav_var[0, 0, 0] = 0.5*(grav_var_h[0, 0, 0] + grav_var_h[0, 0, 1])
 
 
-def neg_rdgas_div_gravity(rdg: FloatField, grav_var: FloatField):
-    """
-    # JK TODO: Is there a better name than this?
-    Args:
-        rdg (out): negative radiative gas divided by variable gravity
-        grav_var (in): variable gravity
-    """
-    with computation(FORWARD), interval(...):
-        rdg = - constants.RDGAS / grav_var
-
-
 def gz_from_surface_height_and_thicknesses(
     zs: FloatFieldIJ, delz: FloatField, gz: FloatField
 ):
@@ -682,7 +672,7 @@ def __init__(
             domain=grid_indexing.domain_full(),
         )
         self._neg_rdgas_div_gravity = stencil_factory.from_origin_domain(
-            neg_rdgas_div_gravity,
+            rdg_adjust.neg_rdgas_div_gravity,
             origin=grid_indexing.origin_full(),
             domain=grid_indexing.domain_full(),
         )
@@ -1068,4 +1058,5 @@ def __call__(
                     self._heat_source,
                     state.pt,
                     delt_time_factor,
+                    state.rdg_var
                 )

pyfv3/stencils/fv_dynamics.py

@@ -25,6 +25,7 @@
 from pyfv3.stencils.dyn_core import AcousticDynamics
 from pyfv3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio
 from pyfv3.stencils.remapping import LagrangianToEulerian
+import pyfv3.stencils.rdg_adjust as rdg_adjust
 
 
 def pt_to_potential_density_pt(
@@ -316,6 +317,11 @@ def __init__(
             origin=grid_indexing.origin_full(),
             domain=grid_indexing.domain_full(add=(0,0,1)),
         )
+        self._adjust_rdg = stencil_factory.from_origin_domain(
+            rdg_adjust.neg_rdgas_div_gravity,
+            origin=grid_indexing.origin_full(),
+            domain=grid_indexing.domain_full(),
+        )
         self.acoustic_dynamics = AcousticDynamics(
             comm=comm,
             stencil_factory=stencil_factory,
@@ -544,6 +550,8 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool):
         if self.config.enable_wam:
             self._adjust_gravity(state.grav_var, state.grav_var_h, state.phis, state.delz)
 
+        self._adjust_rdg(state.rdg_var, state.grav_var)
+
         if self._conserve_total_energy > 0:
             raise NotImplementedError(
                 "Dynamical Core (fv_dynamics): compute total energy is not implemented"

pyfv3/stencils/rdg_adjust.py

@@ -0,0 +1,17 @@
+import ndsl.constants as constants
+from ndsl.dsl.gt4py import FORWARD, computation, interval
+from ndsl.dsl.typing import FloatField
+
+
+def neg_rdgas_div_gravity(rdg: FloatField, grav_var: FloatField):
+    """
+    # JK TODO: Is there a better name than this?
+    Adjust rdg to be the negative RDGAS divided by the variable gravity
+    for Whole Atmosphere Modeling
+
+    Args:
+        rdg (out): negative radiative gas divided by variable gravity
+        grav_var (in): variable gravity
+    """
+    with computation(FORWARD), interval(...):
+        rdg = - constants.RDGAS / grav_var

pyfv3/stencils/temperature_adjust.py

@@ -11,6 +11,7 @@ def apply_diffusive_heating(
     heat_source: FloatField,
     pt: FloatField,
     delt_time_factor: Float,
+    rdg_var: FloatField,
 ):
     """
     Adjust air temperature from heating due to vorticity damping.
@@ -25,9 +26,11 @@ def apply_diffusive_heating(
             energy conservation
         pt (inout): Air potential temperature
         delta_time_factor (in): scaled time step
+        rdg_var (in): negative rdgas divided by variable gravity
     """
     with computation(PARALLEL), interval(...):
-        pkz = exp(cappa / (1.0 - cappa) * log(constants.RDG * delp / delz * pt))
+        # JK TODO: Double-check rdg_var here instead of constants.RDG
+        pkz = exp(cappa / (1.0 - cappa) * log(rdg_var * delp / delz * pt))
         dtmp = heat_source / (constants.CV_AIR * delp)
     with computation(PARALLEL):
         with interval(0, 1):

tests/main/test_wam.py

@@ -27,6 +27,7 @@
 from pyfv3.initialization import init_utils
 from pyfv3.initialization.analytic_init import AnalyticCase
 from pyfv3.stencils.dyn_core import AcousticDynamics
+from pyfv3.stencils.rdg_adjust import neg_rdgas_div_gravity
 from pyfv3.stencils.fv_dynamics import adjust_gravity, init_gravity, init_gravity_h
 from pyfv3.stencils.dyn_core import average_gravity, compute_geopotential