diff --git a/pyfv3/initialization/init_utils.py b/pyfv3/initialization/init_utils.py index 761d21c8..a7bb9aa9 100644 --- a/pyfv3/initialization/init_utils.py +++ b/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( diff --git a/pyfv3/initialization/test_cases/initialize_baroclinic.py b/pyfv3/initialization/test_cases/initialize_baroclinic.py index ee846a8c..25784570 100644 --- a/pyfv3/initialization/test_cases/initialize_baroclinic.py +++ b/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__, diff --git a/pyfv3/stencils/dyn_core.py b/pyfv3/stencils/dyn_core.py index ac62eea4..fc5ad774 100644 --- a/pyfv3/stencils/dyn_core.py +++ b/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 ) diff --git a/pyfv3/stencils/fv_dynamics.py b/pyfv3/stencils/fv_dynamics.py index 4314c9ce..7a8c67c9 100644 --- a/pyfv3/stencils/fv_dynamics.py +++ b/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" diff --git a/pyfv3/stencils/rdg_adjust.py b/pyfv3/stencils/rdg_adjust.py new file mode 100644 index 00000000..d891fe1f --- /dev/null +++ b/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 diff --git a/pyfv3/stencils/temperature_adjust.py b/pyfv3/stencils/temperature_adjust.py index 2f17bce3..267640b3 100644 --- a/pyfv3/stencils/temperature_adjust.py +++ b/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): diff --git a/tests/main/test_wam.py b/tests/main/test_wam.py index d830802c..e8138c8d 100644 --- a/tests/main/test_wam.py +++ b/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