diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore index e3b0ea892..8bb4bb9e4 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore @@ -17,3 +17,4 @@ sandbox/ .venv/ site/ test_data/ +_dacecache/ \ No newline at end of file diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/constants.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/constants.py index 7c7dadb2f..3390d769f 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/constants.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/constants.py @@ -1,15 +1,18 @@ -"""File containing all constants used for pyMoist""" +"""File containing constants used in multiple components of pyMoist""" import numpy as np -from ndsl.dsl.typing import Float +from ndsl.dsl.typing import Float, Int _f32 = np.float32 _f64 = np.float64 _i32 = np.int32 +NUMBER_OF_TRACERS = Int(23) + # Define number of tracers in UW +# NOTE depreciate this, change all references of NCNST to NUMBER_OF_TRACERS NCNST = _i32(23) # MAPL_UNDEF is set to 1E15 in the Fortran @@ -62,12 +65,14 @@ MAPL_RDRY = MAPL_RUNIV / MAPL_AIRMW # J/(kg K) MAPL_CPDRY = Float(3.5) * MAPL_RDRY # J/(kg K) MAPL_KAPPA = MAPL_RDRY / MAPL_CPDRY # (2.0/7.0) +MAPL_EPSILON = MAPL_H2OMW / MAPL_AIRMW # -- MAPL_CVDRY = MAPL_CPDRY - MAPL_RDRY # J/(kg K) MAPL_RVAP = MAPL_RUNIV / MAPL_H2OMW # J/(kg K) MAPL_CPVAP = Float(4) * MAPL_RVAP # J/(kg K) MAPL_CVVAP = MAPL_CPVAP - MAPL_RVAP # J/(kg K) MAPL_RGAS = MAPL_RDRY # MAPL_RDRY # J/(kg K) (DEPRECATED) MAPL_CP = MAPL_RGAS / MAPL_KAPPA # J/(kg K) (DEPRECATED) +MAPL_VIREPS = Float(1.0) / MAPL_EPSILON - Float(1.0) # (DEPRECATED) MAPL_P00 = Float(100000.0) # Pa EPSILON = MAPL_H2OMW / MAPL_AIRMW # -- diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/GF_2020.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/GF_2020.py new file mode 100644 index 000000000..27544d8a6 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/GF_2020.py @@ -0,0 +1,130 @@ +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import NDSLRuntime, QuantityFactory, StencilFactory +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.cumulus_parameterization import ( + GF2020CumulusParameterization, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.finalize import GF2020Finalize +from pyMoist.convection.GF_2020.locals import GF2020Locals +from pyMoist.convection.GF_2020.setup import GF2020Setup +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class GF2020(NDSLRuntime): + """Grell-Fritas convection parameterization scheme, 2020 version (GF2020). + + This class has three subcomponents, and each are called at every execution: + Setup - initializes/resets fields, flips K-axis of state variables, and computes derived fields + CumulusParameterization - the core of the scheme, this class contains the all of the science code + related to deep convection. It requires a flipped (level 0 is surface) data structure, + necessitating the flip/unflip routines in Setup and Finalize, respectively. + Finalize - unflips outputs of core and updates the model state + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + saturation_tables: SaturationVaporPressureTable | None, + ): + """Initialize the GF2020 convection parameterization scheme. + + Initializes subclasses, builds stencils, and allocates dataclasses + + Args: + stencil_factory (StencilFactory) + quantity_factory (QuantityFactory) + config (GF2020Config) + cumulus_parameterization_config (GF2020CumulusParameterizationConfig) + saturation_tables (SaturationVaporPressureTable | None) + """ + super().__init__(stencil_factory) + + # make saturation tables visible at runtime + if saturation_tables is None: + saturation_tables = SaturationVaporPressureTable(stencil_factory.backend) + else: + self.saturation_tables = saturation_tables + + # initialize GF2020 locals + self.locals = GF2020Locals.zeros( + quantity_factory, + data_dimensions={ + "plumes": 3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # initialize GF2020 CumulusParameterization state + self.cumulus_parameterization_state = GF2020CumulusParameterizationState.zeros( + quantity_factory, + data_dimensions={ + "plumes": cumulus_parameterization_constants.NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # initialize submodules + self._setup = GF2020Setup( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + saturation_tables=saturation_tables, + ) + + self._cumulus_parameterization_core = GF2020CumulusParameterization( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + saturation_tables=saturation_tables, + ) + + self._finalize = GF2020Finalize( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + saturation_tables=saturation_tables, + ) + + def __call__(self, state: GF2020State, convection_tracers: ConvectionTracers): + """Run the GF2020 convection parameterization scheme. + + Args: + state (GF2020State): State of the overarching model - not all fields from the model are required + convection_tracers (ConvectionTracers): Collection of tracers from the rest of the model which + will be updated within convection. These may come from a variaty of sources, and need to be + collected into the expected ConvectionTracers data type before being passed down. + """ + # flag to stop convection scheme for single column models + # this can be set in setup if surface temperature is very near zero Kelvin + scm_stop = False + + # call the there parts of the scheme + self._setup( + state=state, + locals=self.locals, + cumulus_parameterization_state=self.cumulus_parameterization_state, + convection_tracers=convection_tracers, + scm_stop=scm_stop, + ) + + if not scm_stop: + self._cumulus_parameterization_core( + state=self.cumulus_parameterization_state, + convection_tracers=convection_tracers, + ) + + self._finalize( + state=state, + locals=self.locals, + cumulus_parameterization_state=self.cumulus_parameterization_state, + convection_tracers=convection_tracers, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/config.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/config.py new file mode 100644 index 000000000..926fa473e --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/config.py @@ -0,0 +1,26 @@ +from dataclasses import dataclass + +from ndsl.dsl.typing import Float, Int + + +@dataclass +class GF2020Config: + SINGLE_COLUMN_MODE: bool + DT_MOIST: Float + LHYDROSTATIC: bool + STOCHASTIC_CNV: bool + STOCH_TOP: Float + STOCH_BOT: Float + GF_MIN_AREA: Float + GF_ENV_SETTING: Int + ENTRVERSION: Int + CONVECTION_TRACER: Int + C1: Float + ADV_TRIGGER: Int + AUTOCONV: Int + USE_TRACER_TRANSPORT: Int + SCLM_DEEP: Float + FIX_CONVECTIVE_CLOUD: bool + APPLY_SUBSIDENCE_MICROPHYSICS: Int + NUMBER_OF_TRACERS: Int + USE_MOMENTUM_TRANSPORT: Int diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/air_density.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/air_density.py new file mode 100644 index 000000000..5a5c5ac5c --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/air_density.py @@ -0,0 +1,34 @@ +import pyMoist.constants as constants +from ndsl.dsl.gt4py import PARALLEL, computation, interval +from ndsl.dsl.typing import FloatField, Int +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume + + +def hydrostatic_air_density( + p: FloatField_Plume, + geopotential_height: FloatField, + error_code: IntFieldIJ_Plume, + air_density: FloatField, + plume: Int, +): + """Compute air density, assuming hydrostatic balance. + + Args: + p (FloatField_Plume) + geopotential_height (FloatField) + error_code (IntFieldIJ_Plume) + air_density (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(...): + # prefil with 0 + air_density = 0.0 + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + air_density = ( + 100.0 + * (p[0, 0, 0][plume] - p[0, 0, 1][plume]) + / (geopotential_height[0, 0, 1] - geopotential_height) + / constants.MAPL_GRAV + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/buoyancy.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/buoyancy.py new file mode 100644 index 000000000..ede77754c --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/buoyancy.py @@ -0,0 +1,39 @@ +from ndsl.dsl.gt4py import PARALLEL, K, computation, interval +from ndsl.dsl.typing import FloatField, Int +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import IntFieldIJ_Plume + + +def get_buoyancy( + lcl_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + cloud_moist_static_energy: FloatField, + environment_moist_static_energy: FloatField, + environment_saturation_moist_static_energy: FloatField, + d_buoyancy: FloatField, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """Determine the "d_buoyancy" of a parcel, defined as the difference between the + moist static energy of the parcel and the environment + + Args: + lcl_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + cloud_moist_static_energy (FloatField) + environment_moist_static_energy (FloatField) + environment_saturation_moist_static_energy (FloatField) + d_buoyancy (FloatField) + error_code (IntFieldIJ_Plume) + plume (Int) + """ + + with computation(PARALLEL), interval(...): + d_buoyancy = 0 + + if error_code[0, 0][plume] == 0: + if K <= lcl_level[0, 0][plume]: + d_buoyancy = cloud_moist_static_energy - environment_moist_static_energy + if K > lcl_level[0, 0][plume] and K <= cloud_top_level[0, 0][plume] + 1: + d_buoyancy = cloud_moist_static_energy - environment_saturation_moist_static_energy diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/check_config.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/check_config.py new file mode 100644 index 000000000..17c37efb2 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/check_config.py @@ -0,0 +1,306 @@ +from ndsl import ndsl_log +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig + + +def check_config( + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, +): + """ + Check config and cumulus_parameterization_config for unexpected options. + + Warn about untested options, block execution of unimplemented options. + """ + + if cumulus_parameterization_config.SATURATION_CALCULATION_CHOICE != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->environmental_conditions construced with" + "untested SATURATION_CALCULATION_CHOICE option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.CLOUD_LEVEL_GRID != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->environmental_cloud_levels construced with " + "untested CLOUD_LEVEL_GRID option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.FRAC_MODIS != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->partition_liquid_ice constructed with " + "untested FRAC_MODIS option. Running untested code... proceed with caution" + ) + + if not cumulus_parameterization_config.MELT_GLAC: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->partition_liquid_ice constructed with " + "untested MELT_GLAC option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->find_lcl constructed with " + "untested BOUNDARY_CONDITION_METHOD option. Running untested code... proceed with caution" + ) + + if config.ADV_TRIGGER != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->find_lcl constructed with " + "untested ADV_TRIGGER option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->parcel_moist_static_energy constructed with " + "untested BOUNDARY_CONDITION_METHOD option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.ZERO_DIFF != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->entrainment_rates constructed with " + "untested ZERO_DIFF option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.OVERSHOOT != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->get_convective_cloud_base_level constructed with " + "untested ZERO_DIFF option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.ZERO_DIFF != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->get_convective_cloud_base_level constructed with " + "untested ZERO_DIFF option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.MOIST_TRIGGER != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->get_convective_cloud_base_level constructed with " + "untested MOIST_TRIGGER option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.USE_MEMORY != -1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->get_convective_cloud_base_level constructed with " + "untested USE_MEMORY option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->get_convective_cloud_base_level constructed with " + "untested BOUNDARY_CONDITION_METHOD option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.OVERSHOOT != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->updraft_rates_pdf constructed with " + "untested OVERSHOOT option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->compute_uc_vc constructed with " + "untested BOUNDARY_CONDITION_METHOD option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->updraft_moisture constructed with " + "untested BOUNDARY_CONDITION_METHOD option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.FRAC_MODIS != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->updraft_moisture constructed with " + "untested FRAC_MODIS option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.ZERO_DIFF != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->updraft_moisture constructed with " + "untested ZERO_DIFF option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.ZERO_DIFF != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->updraft_vertical_velocity constructed with " + "untested ZERO_DIFF option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.USE_WETBULB != 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->downdraft_moisture constructed with " + "untested USE_WETBULB option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.SGS_W_TIMESCALE != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->DiurnalCycle initialized with " + "untested SGS_W_TIMESCALE option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.AEROEVAP != 1: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->DowndraftWindshear initialized with " + "untested AEROEVAP option. Running untested code... proceed with caution" + ) + + if cumulus_parameterization_config.USE_TRACER_EVAPORATION == 0: + ndsl_log.warning( + " GF2020-->CumulusParameterization-->AtmosphericComposition-->downdraft_chemistry initialized " + "with untested USE_TRACER_EVAPORATION option. Running untested code... proceed with caution" + ) + + # generate errors after all warnings + # TODO find a way to generate all then fail at once, printing all simultaneously + if cumulus_parameterization_config.OUTPUT_SOUNDING == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with OUTPUT_SOUNDING = 1. This requires" + "an unimplemented output sounding capabilities (in the Sounding class). Please implement, then" + "disable this error manually to proceed." + ) + + if config.CONVECTION_TRACER == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with CONVECTION_TRACER = 1. This requires" + "an unimplemented class ColdPoolParameterization. Please implement, then disable this error" + "manually to proceed." + ) + + if ( + cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES == 1 + and cumulus_parameterization_config.ENABLE_SHALLOW == 1 + ): + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with" + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES == 1 and shallow plume enabled. This combination requires" + "a call to the unimplemented function get_delmix in CumulusParameterization and updraft_moisture." + "Please implement, then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.ZERO_DIFF == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with ZERO_DIFF = 1. This combination" + "requires an unimplemented porion of UpdraftMassFlux. Please implement, then disable this error" + "manually to proceed." + ) + + if cumulus_parameterization_config.ENABLE_SHALLOW == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with shallow plume enabled. This requires" + "an unimplemented porion of UpdraftMassFlux. Please implement, then disable this error" + "manually to proceed." + ) + + if config.AUTOCONV != 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with AUTOCONV != 1. This requires" + "an unimplemented option in updraft_moisture. Please implement, then disable this" + "error manually to proceed." + ) + + if cumulus_parameterization_config.USE_WETBULB == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with USE_WETBULB = 1. This setting requires" + "the unimplemented function get_wetbulb and an unimplemented option in" + "downdraft_moist_static_energy_and_moisture_budget. Please implement, then disable this error" + "manually to proceed." + ) + + if cumulus_parameterization_config.DIURNAL_CYCLE != 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with DIURNAL_CYCLE != 1. This setting" + "requires an unimplemented option for multiple stencils in DiurnalCycle. Please implement, then" + "disable this error manually to proceed." + ) + + if config.ADV_TRIGGER == 3: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with ADV_TRIGGER = 3. This setting requires" + "the unimplemented XieTriggerFunction. Please implement, then disable this error" + "manually to proceed." + ) + + if cumulus_parameterization_config.VERTICAL_DISCRETIZATION_OPTION != 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with VERTICAL_DISCRETIZATION_OPTION != 1." + "This setting requires an unimplemented option in VerticalDiscretization. Please implement," + "then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.USE_FCT != 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with USE_FCT != 0. This setting requires" + "an unimplemented option for multiple stencils in VerticalDiscretization. Please implement," + "then disable this error manually to proceed." + ) + + if config.APPLY_SUBSIDENCE_MICROPHYSICS != 0: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with APPLY_SUBSIDENCE_MICROPHYSICS != 0." + "This setting requires an unimeplemented option in EnvironmentalSubsidence. Please implement," + "then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.DIURNAL_CYCLE not in (1, 6): + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with DIURNAL_CYCLE != 1 or 6. This" + "setting requires one or more unimeplemented options in LargeScaleForcing. Please" + "implement, then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.ENABLE_SHALLOW == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with shallow plume enabled. This requires" + "an unimeplemented functions in LargeScaleForcing. Please implement, then disable" + "this error manually to proceed." + ) + + if cumulus_parameterization_config.ENABLE_SHALLOW == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with shallow plume enabled. This requires" + "an umplemented portion of ensemble_output_and_feedback. Please impelment, then disable this" + "error manually to proceed." + ) + + if not cumulus_parameterization_config.COUPLE_MICROPHYSICS: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with COUPLE_MICROPHYSICS != True. This" + "setting requires an unimeplemented option in cloud_dissipation. Please implement, then" + "disable this error manually to proceed." + ) + + if not cumulus_parameterization_config.LIGHTNING_DIAGNOSTICS: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with LIGHTNING_DIAGNOSTICS != True. This" + "setting requires unimeplemented functions. Please implement, then disable this error manually" + "to proceed." + ) + + if not cumulus_parameterization_config.USE_TRACER_SCAVENGE: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->AtmosphericComposition initialized with" + "USE_TRACER_SCAVENGE != 1. This setting requires unimeplemented options in updraft_chemistry." + "Please implement, then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.USE_FLUX_FORM != 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->AtmosphericComposition initialized with" + "USE_FLUX_FORM != 1. This setting requires one or more unimeplemented options in" + "vertical_transport_part_1. Please implement, then disable this error manually to proceed." + ) + + if cumulus_parameterization_config.ALP1 == 0: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->AtmosphericComposition initialized with" + "ALP1 == 0. This setting requires an unimeplemented option in vertical_transport_part_1." + "Please implement, then disable this error manually to proceed." + ) + + if not cumulus_parameterization_config.WRTGRADS: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization initialized with WRTGRADS = True. This setting requires" + "the unimeplemented GATE sounding capabilities (in the GATESounding class). Please implement," + "then disable this error manually to proceed." + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/config.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/config.py new file mode 100644 index 000000000..5d750d9fa --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/config.py @@ -0,0 +1,103 @@ +from dataclasses import dataclass + +from ndsl.dsl.typing import Bool, Float, Int + + +@dataclass +class GF2020CumulusParameterizationConfig: + # plume dependent (one for each plume) + DOWNDRAFT_MAX_HEIGHT_LAND_SHALLOW: Float + DOWNDRAFT_MAX_HEIGHT_LAND_MID: Float + DOWNDRAFT_MAX_HEIGHT_LAND_DEEP: Float + DOWNDRAFT_MAX_HEIGHT_OCEAN_SHALLOW: Float + DOWNDRAFT_MAX_HEIGHT_OCEAN_MID: Float + DOWNDRAFT_MAX_HEIGHT_OCEAN_DEEP: Float + UPDRAFT_MAX_HEIGHT_LAND_SHALLOW: Float + UPDRAFT_MAX_HEIGHT_LAND_MID: Float + UPDRAFT_MAX_HEIGHT_LAND_DEEP: Float + UPDRAFT_MAX_HEIGHT_OCEAN_SHALLOW: Float + UPDRAFT_MAX_HEIGHT_OCEAN_MID: Float + UPDRAFT_MAX_HEIGHT_OCEAN_DEEP: Float + MINIMUM_EVAP_FRACTION_LAND_SHALLOW: Float + MINIMUM_EVAP_FRACTION_LAND_MID: Float + MINIMUM_EVAP_FRACTION_LAND_DEEP: Float + MINIMUM_EVAP_FRACTION_OCEAN_SHALLOW: Float + MINIMUM_EVAP_FRACTION_OCEAN_MID: Float + MINIMUM_EVAP_FRACTION_OCEAN_DEEP: Float + MAXIMUM_EVAP_FRACTION_LAND_SHALLOW: Float + MAXIMUM_EVAP_FRACTION_LAND_MID: Float + MAXIMUM_EVAP_FRACTION_LAND_DEEP: Float + MAXIMUM_EVAP_FRACTION_OCEAN_SHALLOW: Float + MAXIMUM_EVAP_FRACTION_OCEAN_MID: Float + MAXIMUM_EVAP_FRACTION_OCEAN_DEEP: Float + CLOUD_BASE_MASS_FLUX_FACTOR_SHALLOW: Float + CLOUD_BASE_MASS_FLUX_FACTOR_MID: Float + CLOUD_BASE_MASS_FLUX_FACTOR_DEEP: Float + USE_EXCESS_SHALLOW: Int + USE_EXCESS_MID: Int + USE_EXCESS_DEEP: Int + AVERAGE_LAYER_DEPTH_SHALLOW: Float + AVERAGE_LAYER_DEPTH_MID: Float + AVERAGE_LAYER_DEPTH_DEEP: Float + ENABLE_SHALLOW: Int + ENABLE_MID: Int + ENABLE_DEEP: Int + ENTRAINMENT_RATE_SHALLOW: Float + ENTRAINMENT_RATE_MID: Float + ENTRAINMENT_RATE_DEEP: Float + C0_SHAL: Float + C0_MID: Float + C0_DEEP: Float + TAU_MID: Float + TAU_DEEP: Float + CLOSURE_CHOICE_SHALLOW: Int + CLOSURE_CHOICE_MID: Int + CLOSURE_CHOICE_DEEP: Int + # plume independent + PLUME_ORDER: Float + DTIME: Float + ZERO_DIFF: Int + MOIST_TRIGGER: Int + LAMBDA_DEEP: Float + LAMBDA_SHALLOW_DOWN: Float + PRESSURE_GRADIENT_CONSTANT: Float + CAP_MAXS: Float + OUTPUT_SOUNDING: Int + USE_SCALE_DEP: Int + USE_RANDOM_NUMBER: Float + SATURATION_CALCULATION_CHOICE: Int + ITEST: Int + CLOUD_LEVEL_GRID: Int + MELT_GLAC: Bool + FRAC_MODIS: Int + BOUNDARY_CONDITION_METHOD: Int + OVERSHOOT: Float + USE_MEMORY: Int + DOWNDRAFT: Int + FIRST_GUESS_W: Bool + USE_WETBULB: Int + DIURNAL_CYCLE: Int + USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES: Int + CRITICAL_MIXING_RATIO_OVER_OCEAN: Float + CRITICAL_MIXING_RATIO_OVER_LAND: Float + BETA_SHALLOW: Float + EVAP_FIX: Int + SGS_W_TIMESCALE: Int + AEROEVAP: Int + VERTICAL_DISCRETIZATION_OPTION: Int + ALP1: Float + USE_FCT: Int + MIN_ENTRAINMENT_RATE: Float + USE_SMOOTH_TENDENCIES: Int + COUPLE_MICROPHYSICS: Bool + USE_RAIN_EVAP_BELOW_CLOUD_BASE: Int + USE_CLOUD_DISSIPATION: Float + LIGHTNING_DIAGNOSTICS: Int + WRTGRADS: Bool + USE_TRACER_SCAVENGE: Int + USE_TRACER_EVAPORATION: Int + USE_FLUX_FORM: Int + USE_FCT: Int + MAX_TEMP_VAPOR_TENDENCY: Float + LIQUID_ICE_NUMBER_CONCENTRATION: Int + FEED_3D_MODEL: bool diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/constants.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/constants.py new file mode 100644 index 000000000..98b67f1ca --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/constants.py @@ -0,0 +1,40 @@ +from ndsl.dsl.typing import Float, Int + + +# used to define data_dimension size for quantity factories +NUMBER_OF_PLUMES = Int(3) + +# physical constants used ONLY for GF_2020, else there may be unintended conflicts with global constants +RGAS = Float(287.0) # J K-1 kg-1 +CP = Float(1004.0) # J K-1 kg-1 +RV = Float(461.0) # J K-1 kg-1 +P00 = Float(1.0e5) # hPa +TCRIT = Float(258.0) # K +CPOR = CP / RGAS +XLV = Float(2.5e6) # J kg-1 +AKMIN = Float(1.0) # # +TKMIN = Float(1.0e-5) # m+2 s-2 +CCNCLEAN = Float(250.0) # # cm-3 +T_0 = Float(273.16) # K +T_ICE = Float(235.16) # K +XLF = Float(0.333e6) # latent heat of freezing (J K-1 kg-1) +MAX_QSAT = Float(0.5) # kg/kg +MX_BUOY = CP * Float(5.0) + XLV * Float(2.0e-3) # temp exc=5 K, q deficit=2 g/kg (=> mx_buoy ~ 10 kJ/kg) + +smaller_qv = Float(1.0e-16) # kg/kg + +MAXENS1 = Int(1) # ensemble one on cap_max +MAXENS2 = Int(1) # ensemble two on precip efficiency +MAXENS3 = Int(16) # ensemble three done in cup_forcing_ens16 for G3d + +USE_LCL = False + +AEROEVAP = 1 + +# plume identifiers +SHALLOW = Int(0) +MID = Int(1) +DEEP = Int(2) + +# control if smoothing is performed in updraft_mass_flux and downdraft_mass_flux stencils +DO_SMOOTHING = False diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/convective_tracers.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/convective_tracers.py new file mode 100644 index 000000000..43324b579 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/convective_tracers.py @@ -0,0 +1,995 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, K, computation, exp, interval +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_ConvectionTracers, + FloatField_ConvectionTracers_Plume, + FloatField_Plume, + FloatFieldIJ_ConvectionTracers, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import get_cloud_boundary_conditions +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import tridiag +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.field_types import ( + ConvectionTracerMetaDataTable_Bool, + ConvectionTracerMetaDataTable_Float, + ConvectionTracerMetaDataTable_x4, +) + + +def environment_cloud_levels_chemistry( + error_code: IntFieldIJ_Plume, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_cloud_levels: FloatField_ConvectionTracers, + plume: Int, +): + """Set chemistry tracers at cloud levels. Behavior is determined by CLOUD_LEVEL_OPTION external. + Only options 1 and 2 have been implemented, only option 2 has been tested. + + Args: + error_code (IntFieldIJ_Plume) + chemistry_tracers (FloatField_ConvectionTracers) + chemistry_tracers_cloud_levels (FloatField_ConvectionTracers) + plume (Int) + """ + from __externals__ import CLOUD_LEVEL_OPTION, NUMBER_OF_TRACERS + + with computation(FORWARD), interval(1, -1): + if error_code[0, 0][plume] == 0 and CLOUD_LEVEL_OPTION == 1: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_cloud_levels[0, 0, 0][tracer] = ( + 0.5 * chemistry_tracers[0, 0, -1][tracer] + chemistry_tracers[0, 0, 0][tracer] + ) + tracer += 1 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and CLOUD_LEVEL_OPTION == 1: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_cloud_levels[0, 0, 0][tracer] = chemistry_tracers[0, 0, 0][tracer] + tracer += 1 + + with computation(FORWARD), interval(-1, None): + if error_code[0, 0][plume] == 0 and CLOUD_LEVEL_OPTION == 1: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_cloud_levels[0, 0, 0][tracer] = chemistry_tracers[0, 0, 0][tracer] + tracer += 1 + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0 and CLOUD_LEVEL_OPTION != 1: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_cloud_levels[0, 0, 0][tracer] = chemistry_tracers[0, 0, 0][tracer] + tracer += 1 + + +def updraft_chemistry( + error_code: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + p_forced: FloatField, + p_cloud_levels_forced: FloatField_Plume, + geopotential_height_cloud_levels: FloatField, + vertical_velocity_3d: FloatField, + mass_entrainment_updraft_forced: FloatField_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_sc_updraft: FloatField_ConvectionTracers, + chemistry_tracers_cloud_levels: FloatField_ConvectionTracers, + chemistry_tracers_pw_updraft: FloatField_ConvectionTracers, + chemistry_tracers_total_pw_updraft: FloatFieldIJ_ConvectionTracers, + convection_tracers_vect_hcts: ConvectionTracerMetaDataTable_x4, + convection_tracers_fscav: ConvectionTracerMetaDataTable_Float, + convection_tracers_use_gcc_washout: ConvectionTracerMetaDataTable_Bool, + tracer_cloud_boundary: FloatFieldIJ_ConvectionTracers, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Compute chemistry effects in the updraft. + + Args: + error_code (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + ocean_fraction (FloatFieldIJ) + p_forced (FloatField) + p_cloud_levels_forced (FloatField_Plume) + geopotential_height_cloud_levels (FloatField) + vertical_velocity_3d (FloatField) + mass_entrainment_updraft_forced (FloatField_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + chemistry_tracers (FloatField_ConvectionTracers) + chemistry_tracers_sc_updraft (FloatField_ConvectionTracers) + chemistry_tracers_cloud_levels (FloatField_ConvectionTracers) + chemistry_tracers_pw_updraft (FloatField_ConvectionTracers) + chemistry_tracers_total_pw_updraft (FloatFieldIJ_ConvectionTracers) + convection_tracers_vect_hcts (ConvectionTracerMetaDataTable_x4) + convection_tracers_fscav (ConvectionTracerMetaDataTable_Float) + convection_tracers_use_gcc_washout (ConvectionTracerMetaDataTable_Bool) + tracer_cloud_boundary (FloatFieldIJ_ConvectionTracers) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import BOUNDARY_CONDITION_METHOD, NUMBER_OF_TRACERS, USE_TRACER_SCAVENGE, k_end + + with computation(FORWARD), interval(...): + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + # initialization + chemistry_tracers_sc_updraft[0, 0, 0][tracer] = chemistry_tracers_cloud_levels[0, 0, 0][tracer] + chemistry_tracers_pw_updraft[0, 0, 0][tracer] = 0.0 + chemistry_tracers_total_pw_updraft[0, 0][tracer] = 0.0 + tracer += 1 + + with computation(PARALLEL), interval(...): + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + chemistry_tracers_cloud_levels_3d = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # NOTE this double loop is disgusting. need a better solution for it + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + level = 0 + while level <= k_end: + chemistry_tracers_cloud_levels_3d[0, 0, level] = chemistry_tracers_cloud_levels[ + 0, 0, level + ][tracer] + level += 1 + tracer_cloud_boundary[0, 0][tracer] = get_cloud_boundary_conditions( + field=chemistry_tracers_cloud_levels_3d, + scalar_perturbation=0, + p=p_forced, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + tracer += 1 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + if K <= updraft_origin_level[0, 0][plume]: + chemistry_tracers_sc_updraft[0, 0, 0][tracer] = tracer_cloud_boundary[0, 0][tracer] + tracer += 1 + + with computation(FORWARD), interval(1, None): + if ( + error_code[0, 0][plume] == 0 + and K >= updraft_origin_level[0, 0][plume] + 1 + and K <= cloud_top_level[0, 0][plume] + 1 + ): + # entrainment, detrainment, mass flux + massflux_internal = normalized_massflux_updraft_forced[0, 0, -1][plume] + entrainment_internal = mass_entrainment_updraft_forced[0, 0, -1][plume] + detrainment_internal = 0.5 * mass_detrainment_updraft_forced[0, 0, -1][plume] + denom = massflux_internal - detrainment_internal + entrainment_internal + + # transport + mixing + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + if denom > 0.0: + chemistry_tracers_sc_updraft[0, 0, 0][tracer] = ( + chemistry_tracers_sc_updraft[0, 0, -1][tracer] * massflux_internal + - chemistry_tracers_sc_updraft[0, 0, -1][tracer] * detrainment_internal + + chemistry_tracers[0, 0, -1][tracer] * entrainment_internal + ) / denom + else: + chemistry_tracers_sc_updraft[0, 0, 0][tracer] = chemistry_tracers_sc_updraft[0, 0, -1][ + tracer + ] + tracer += 1 + + # scavenging section - skip if USE_TRACER_SCAVENGE = 0 or on shallow plume + if not (USE_TRACER_SCAVENGE == 0 or plume == cumulus_parameterization_constants.SHALLOW): + dz = geopotential_height_cloud_levels - geopotential_height_cloud_levels[0, 0, -1] + + # in-cloud vert velocity for scavenging formulation 2 + updraft_vertical_velosity_internal = vertical_velocity_3d + + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + is_gcc = convection_tracers_use_gcc_washout.A[tracer] + # aerosol scavenging + if convection_tracers_fscav.A[tracer] > 1.0e-6: + # formulation 1 as in GOCART with RAS conv_par + if USE_TRACER_SCAVENGE == 1: + chemistry_tracers_pw_updraft[0, 0, 0][tracer] = max( + 0.0, + chemistry_tracers_sc_updraft[0, 0, 0][tracer] + * (1.0 - exp(-convection_tracers_fscav.A[tracer] * (dz / 1000.0))), + ) + + # formulation 2 as in GOCART + elif USE_TRACER_SCAVENGE == 2: + option_not_implemented = True + + # formulation 3 - orignal GF conv_par + elif USE_TRACER_SCAVENGE == 3: + option_not_implemented = True + + # (in cloud) total mixing ratio in gas and aqueous phases + chemistry_tracers_sc_updraft[0, 0, 0][tracer] = ( + chemistry_tracers_sc_updraft[0, 0, 0][tracer] + - chemistry_tracers_pw_updraft[0, 0, 0][tracer] + ) + + # tracer gas phase scavenging + elif convection_tracers_vect_hcts.A[tracer, 1] > 1.0e-6: + option_not_implemented = True + if is_gcc == True: + option_not_implemented = True + if USE_TRACER_SCAVENGE == 3: + option_not_implemented = True + else: + if is_gcc == True: + option_not_implemented = True + else: + option_not_implemented = True + + tracer += 1 + + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_total_pw_updraft[0, 0][tracer] = ( + chemistry_tracers_total_pw_updraft[0, 0][tracer] + + chemistry_tracers_pw_updraft[0, 0, 0][tracer] * dp / constants.MAPL_GRAV + ) + tracer += 1 + + +def downdraft_chemistry( + error_code: IntFieldIJ_Plume, + downdraft_origin_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + total_normalized_integrated_evaporate_forced: FloatFieldIJ, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_cloud_levels: FloatField_ConvectionTracers, + chemistry_tracers_sc_downdraft: FloatField_ConvectionTracers, + chemistry_tracers_pw_downdraft: FloatField_ConvectionTracers, + chemistry_tracers_total_pw_downdraft: FloatFieldIJ_ConvectionTracers, + chemistry_tracers_total_pw_updraft: FloatFieldIJ_ConvectionTracers, + plume: Int, +): + """Compute chemistry effects in the downdraft. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_origin_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + total_normalized_integrated_evaporate_forced (FloatFieldIJ) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + chemistry_tracers (FloatField_ConvectionTracers) + chemistry_tracers_cloud_levels (FloatField_ConvectionTracers) + chemistry_tracers_sc_downdraft (FloatField_ConvectionTracers) + chemistry_tracers_pw_downdraft (FloatField_ConvectionTracers) + chemistry_tracers_total_pw_downdraft (FloatFieldIJ_ConvectionTracers) + chemistry_tracers_total_pw_updraft (FloatFieldIJ_ConvectionTracers) + plume (Int) + """ + from __externals__ import NUMBER_OF_TRACERS, USE_TRACER_EVAPORATION + + with computation(FORWARD), interval(...): + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_sc_downdraft[0, 0, 0][tracer] = 0.0 + chemistry_tracers_pw_downdraft[0, 0, 0][tracer] = 0.0 + chemistry_tracers_total_pw_downdraft[0, 0][tracer] = 0.0 + tracer += 1 + + with computation(FORWARD), interval(0, 1): + if plume != cumulus_parameterization_constants.SHALLOW and error_code[0, 0][plume] == 0: + # fration of the total rain that was evaporated + evaporation_fraction: FloatFieldIJ = -total_normalized_integrated_evaporate_forced / ( + 1.0e-16 + total_normalized_integrated_condensate_forced[0, 0][plume] + ) + + # scalar concentration in-cloud - downdraft + + # at downdraft_origin_level + level = downdraft_origin_level[0, 0][plume] + if USE_TRACER_EVAPORATION == 0: + internal_precip: FloatFieldIJ = 0.0 + else: + internal_precip: FloatFieldIJ = ( + evaporate_in_downdraft_forced.at(K=level, ddim=[plume]) + / (1.0e-16 + total_normalized_integrated_evaporate_forced) + * evaporation_fraction + ) + + dp = 100.0 * ( + p_cloud_levels_forced.at(K=level, ddim=[plume]) + - p_cloud_levels_forced.at(K=level + 1, ddim=[plume]) + ) + + # downdrafts will be initiated with a mixture of 50% environmental and in-cloud concentrations + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + chemistry_tracers_sc_downdraft[0, 0, level][tracer] = chemistry_tracers_cloud_levels[ + 0, 0, level + ][tracer] + chemistry_tracers_pw_downdraft[0, 0, level][tracer] = ( + -internal_precip + * chemistry_tracers_total_pw_updraft[0, 0][tracer] + * constants.MAPL_GRAV + / dp + ) + chemistry_tracers_sc_downdraft[0, 0, level][tracer] = ( + chemistry_tracers_sc_downdraft[0, 0, level][tracer] + - chemistry_tracers_pw_downdraft[0, 0, level][tracer] + ) + chemistry_tracers_total_pw_downdraft[0, 0][tracer] = ( + chemistry_tracers_total_pw_downdraft[0, 0][tracer] + + chemistry_tracers_pw_downdraft[0, 0, level][tracer] * dp / constants.MAPL_GRAV + ) + tracer += 1 + + # calculate downdraft mass terms + with computation(BACKWARD), interval(0, -1): + if ( + plume != cumulus_parameterization_constants.SHALLOW + and error_code[0, 0][plume] == 0 + and K <= downdraft_origin_level[0, 0][plume] - 1 + ): + massflux_internal = normalized_massflux_downdraft_forced[0, 0, 1][plume] + entrainment_internal = mass_entrainment_downdraft_forced[0, 0, 0][plume] + detrainment_internal = 0.5 * mass_detrainment_downdraft_forced[0, 0, 0][plume] + denom = massflux_internal - detrainment_internal + entrainment_internal + + # transport + mixing + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + if denom > 0.0: + chemistry_tracers_sc_downdraft[0, 0, 0][tracer] = ( + chemistry_tracers_sc_downdraft[0, 0, 1][tracer] * massflux_internal + - chemistry_tracers_sc_downdraft[0, 0, 1][tracer] * detrainment_internal + + chemistry_tracers[0, 0, 0][tracer] * entrainment_internal + ) / denom + else: + chemistry_tracers_sc_downdraft[0, 0, 0][tracer] = chemistry_tracers_sc_downdraft.at( + K=K + 1, ddim=[tracer] + ) + tracer += 1 + + # evaporation term + if USE_TRACER_EVAPORATION != 0: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + + # fraction of evaporated precip per layer + internal_precip = evaporate_in_downdraft_forced[0, 0, 0][plume] / ( + 1.0e-16 + total_normalized_integrated_evaporate_forced + ) + + # fraction of the total precip that was actually evaporated at layer k + internal_precip = internal_precip * evaporation_fraction + + # sanity check + internal_precip = min(1.0, max(internal_precip, 0.0)) + + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + # amount evaporated by the downdraft from the precipitation + chemistry_tracers_pw_downdraft[0, 0, 0][tracer] = ( + -internal_precip + * chemistry_tracers_total_pw_updraft[0, 0][tracer] + * constants.MAPL_GRAV + / dp + ) + + # final tracer in the downdraft + chemistry_tracers_sc_downdraft[0, 0, 0][tracer] = ( + chemistry_tracers_sc_downdraft[0, 0, 0][tracer] + - chemistry_tracers_pw_downdraft[0, 0, 0][tracer] + ) + + # total evaporated tracer + chemistry_tracers_total_pw_downdraft[0, 0][tracer] = ( + chemistry_tracers_total_pw_downdraft[0, 0][tracer] + + chemistry_tracers_pw_downdraft[0, 0, 0][tracer] * dp / constants.MAPL_GRAV + ) + + tracer += 1 + + +def vertical_transport_part_1( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + environment_massflux: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + epsilon_forced: FloatFieldIJ_Plume, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_output: FloatField_ConvectionTracers_Plume, + chemistry_tracers_cloud_levels: FloatField_ConvectionTracers, + chemistry_tracers_sc_updraft: FloatField_ConvectionTracers, + chemistry_tracers_sc_downdraft: FloatField_ConvectionTracers, + chemistry_tracers_pw_updraft: FloatField_ConvectionTracers, + chemistry_tracers_pw_downdraft: FloatField_ConvectionTracers, + dd_tracers: FloatField_ConvectionTracers, + aa: FloatField, + bb: FloatField, + cc: FloatField, + plume: Int, +): + """Advect tracers up/down the column - part 1/2 (before tridiag). + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + environment_massflux (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + epsilon_forced (FloatFieldIJ_Plume) + chemistry_tracers (FloatField_ConvectionTracers) + chemistry_tracers_output (FloatField_ConvectionTracers_Plume) + chemistry_tracers_cloud_levels (FloatField_ConvectionTracers) + chemistry_tracers_sc_updraft (FloatField_ConvectionTracers) + chemistry_tracers_sc_downdraft (FloatField_ConvectionTracers) + chemistry_tracers_pw_updraft (FloatField_ConvectionTracers) + chemistry_tracers_pw_downdraft (FloatField_ConvectionTracers) + dd_tracers (FloatField_ConvectionTracers) + aa (FloatField) + bb (FloatField) + cc (FloatField) + plume (Int) + """ + from __externals__ import ( + ALP1, + DTIME, + NUMBER_OF_TRACERS, + USE_FLUX_FORM, + USE_TRACER_EVAPORATION, + USE_TRACER_SCAVENGE, + ) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + alp0: FloatFieldIJ = 0.0 + + # flux form + source/sink terms + time explicit + FCT + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and USE_FLUX_FORM == 1 and ALP1 == 0: + option_not_implemented = True + + # flux form + source/sink terms + time explicit/implicit + upstream + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and USE_FLUX_FORM == 1 and ALP1 > 0.0: + alp0 = 1.0 - ALP1 + if K <= cloud_top_level[0, 0][plume] + 1: + fp = 0.5 * (environment_massflux + abs(environment_massflux)) + fm = 0.5 * (environment_massflux - abs(environment_massflux)) + + with computation(FORWARD), interval(...): + if ( + error_code[0, 0][plume] == 0 + and USE_FLUX_FORM == 1 + and ALP1 > 0.0 + and K <= cloud_top_level[0, 0][plume] + ): + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + beta = DTIME * constants.MAPL_GRAV / dp + aa = ALP1 * beta * fm + bb = 1.0 + ALP1 * beta * (fp - fm[0, 0, 1]) + cc = -ALP1 * beta * fp[0, 0, 1] + + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + dd_tracers[0, 0, 0][tracer] = ( + chemistry_tracers[0, 0, 0][tracer] + - ( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * chemistry_tracers_sc_updraft[0, 0, 1][tracer] + - normalized_massflux_updraft_forced[0, 0, 0][plume] + * chemistry_tracers_sc_updraft[0, 0, 0][tracer] + ) + * beta + + ( + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * chemistry_tracers_sc_downdraft[0, 0, 1][tracer] + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + * chemistry_tracers_sc_downdraft[0, 0, 0][tracer] + ) + * beta + * epsilon_forced[0, 0][plume] + ) + ) + tracer += 1 + + # include evaporation + if USE_TRACER_EVAPORATION == 1 and plume != cumulus_parameterization_constants.SHALLOW: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + # evaporated ( chemistry_tracer_pw_downdraft < 0 => E_dn > 0) + chemistry_tracers_output[0, 0, 0][plume, tracer] = ( + chemistry_tracers_output[0, 0, 0][plume, tracer] + - 0.5 + * epsilon_forced[0, 0][plume] + * ( + normalized_massflux_downdraft_forced[0, 0, 0][plume] + * chemistry_tracers_pw_downdraft[0, 0, 0][tracer] + + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * chemistry_tracers_pw_downdraft[0, 0, 1][tracer] + ) + * beta + ) + tracer += 1 + + # include scavenging + if USE_TRACER_SCAVENGE > 0 and plume != cumulus_parameterization_constants.SHALLOW: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + # incorporated in rainfall (<0) + chemistry_tracers_output[0, 0, 0][plume, tracer] = ( + chemistry_tracers_output[0, 0, 0][plume, tracer] + - 0.5 + * ( + normalized_massflux_updraft_forced[0, 0, 0][plume] + * chemistry_tracers_pw_updraft[0, 0, 0][tracer] + + normalized_massflux_updraft_forced[0, 0, 1][plume] + * chemistry_tracers_pw_updraft[0, 0, 1][tracer] + ) + * beta + ) + tracer += 1 + + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + dd_tracers[0, 0, 0][tracer] = ( + dd_tracers[0, 0, 0][tracer] + + chemistry_tracers_output[0, 0, 0][plume, tracer] + + alp0 + * beta + * ( + -fm * chemistry_tracers.at(K=max(0, K - 1), ddim=[tracer]) + + (fm[0, 0, 1] - fp) * chemistry_tracers[0, 0, 0][tracer] + + fp[0, 0, 1] * chemistry_tracers[0, 0, 1][tracer] + ) + ) + tracer += 1 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and (USE_FLUX_FORM == 2 or USE_FLUX_FORM == 3): + option_not_implemented = True + + +def update_after_tridiag( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + dd: FloatField_ConvectionTracers, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_output: FloatField_ConvectionTracers_Plume, + tracer: Int, + plume: Int, +): + from __externals__ import DTIME + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + chemistry_tracers_output[0, 0, 0][plume, tracer] = ( + dd[0, 0, 0][tracer] - chemistry_tracers[0, 0, 0][tracer] + ) / DTIME + + +def vertical_transport_part_2( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + epsilon_forced: FloatFieldIJ_Plume, + chemistry_tracers: FloatField_ConvectionTracers, + chemistry_tracers_output: FloatField_ConvectionTracers_Plume, + chemistry_tracers_pw_updraft: FloatField_ConvectionTracers, + chemistry_tracers_pw_downdraft: FloatField_ConvectionTracers, + dd_tracers: FloatField_ConvectionTracers, + residual: FloatFieldIJ_ConvectionTracers, + plume: Int, +): + """Advect tracers up/down the column - part 2/2 (after tridiag). + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + epsilon_forced (FloatFieldIJ_Plume) + chemistry_tracers (FloatField_ConvectionTracers) + chemistry_tracers_output (FloatField_ConvectionTracers_Plume) + chemistry_tracers_pw_updraft (FloatField_ConvectionTracers) + chemistry_tracers_pw_downdraft (FloatField_ConvectionTracers) + dd_tracers (FloatField_ConvectionTracers) + residual (FloatFieldIJ_ConvectionTracers) + plume (Int) + """ + from __externals__ import NUMBER_OF_TRACERS + + with computation(FORWARD), interval(0, 1): + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + residual[0, 0][tracer] = 0.0 + tracer += 1 + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + evap = ( + -0.5 + * ( + normalized_massflux_downdraft_forced[0, 0, 0][plume] + * chemistry_tracers_pw_downdraft[0, 0, 0][tracer] + + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * chemistry_tracers_pw_downdraft[0, 0, 1][tracer] + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + wetdep = ( + 0.5 + * ( + normalized_massflux_updraft_forced[0, 0, 0][plume] + * chemistry_tracers_pw_updraft[0, 0, 0][tracer] + + normalized_massflux_updraft_forced[0, 0, 1][plume] + * chemistry_tracers_pw_updraft[0, 0, 1][tracer] + ) + * constants.MAPL_GRAV + / dp + ) + + residual[0, 0][tracer] = residual[0, 0][tracer] + (wetdep - evap) * dp / constants.MAPL_GRAV + + tracer += 1 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + tracer = 0 + while tracer < NUMBER_OF_TRACERS: + if residual[0, 0][tracer] < 0: + beta = constants.MAPL_GRAV / ( + p_cloud_levels_forced.at(K=0, ddim=[plume]) + - p_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume] + 1, ddim=[plume]) + ) + chemistry_tracers_output[0, 0, 0][plume, tracer] = ( + chemistry_tracers_output[0, 0, 0][plume, tracer] + residual[0, 0][tracer] * beta + ) + tracer += 1 + + +class ColdPoolParameterization: + def __init__(self, config: GF2020Config): + if config.CONVECTION_TRACER == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->ColdPoolParameterization: The" + "ColdPoolParameterization section has not been implemented. You should have been caught" + "before getting here by the config checker. Beware, something likely failing in the config" + "checker as well - you may be unknowingly calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass + + +class AtmosphericComposition(NDSLRuntime): + """Consider the effects of and on chemistry/tracers by convection.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + quantity_factory.add_data_dimensions({"convection_tracers": config.NUMBER_OF_TRACERS}) + self._aa: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._bb: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._cc: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._dd_tracers: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "convection_tracers"], "n/a") + self._tracer_cloud_boundary: Local = quantity_factory.zeros( + [X_DIM, Y_DIM, "convection_tracers"], "n/a" + ) + self._residual: Local = quantity_factory.zeros([X_DIM, Y_DIM, "convection_tracers"], "n/a") + + # construct stencils and functions + self._environment_cloud_levels_chemistry = stencil_factory.from_dims_halo( + func=environment_cloud_levels_chemistry, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"NUMBER_OF_TRACERS": config.NUMBER_OF_TRACERS, "CLOUD_LEVEL_OPTION": 2}, + ) + + self._updraft_chemistry = stencil_factory.from_dims_halo( + func=updraft_chemistry, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "NUMBER_OF_TRACERS": config.NUMBER_OF_TRACERS, + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + "USE_TRACER_SCAVENGE": cumulus_parameterization_config.USE_TRACER_SCAVENGE, + }, + ) + + self._downdraft_chemistry = stencil_factory.from_dims_halo( + func=downdraft_chemistry, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "NUMBER_OF_TRACERS": config.NUMBER_OF_TRACERS, + "USE_TRACER_EVAPORATION": cumulus_parameterization_config.USE_TRACER_EVAPORATION, + }, + ) + + self._vertical_transport_part_1 = stencil_factory.from_dims_halo( + func=vertical_transport_part_1, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ALP1": cumulus_parameterization_config.ALP1, + "DTIME": cumulus_parameterization_config.DTIME, + "NUMBER_OF_TRACERS": config.NUMBER_OF_TRACERS, + "USE_TRACER_EVAPORATION": cumulus_parameterization_config.USE_TRACER_EVAPORATION, + "USE_TRACER_SCAVENGE": cumulus_parameterization_config.USE_TRACER_SCAVENGE, + "USE_FLUX_FORM": cumulus_parameterization_config.USE_FLUX_FORM, + "USE_FCT": cumulus_parameterization_config.USE_FCT, + }, + ) + + self._tridiag = stencil_factory.from_dims_halo( + func=tridiag, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._update_after_tridiag = stencil_factory.from_dims_halo( + func=update_after_tridiag, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DTIME": cumulus_parameterization_config.DTIME}, + ) + + self._vertical_transport_part_2 = stencil_factory.from_dims_halo( + func=vertical_transport_part_2, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"NUMBER_OF_TRACERS": config.NUMBER_OF_TRACERS}, + ) + + def __call__( + self, + error_code: Quantity, + cloud_top_level: Quantity, + updraft_origin_level: Quantity, + downdraft_origin_level: Quantity, + ocean_fraction: Quantity, + p_forced: Quantity, + p_cloud_levels_forced: Quantity, + geopotential_height_cloud_levels: Quantity, + environment_massflux: Quantity, + normalized_massflux_updraft_forced: Quantity, + normalized_massflux_downdraft_forced: Quantity, + mass_entrainment_updraft_forced: Quantity, + mass_detrainment_updraft_forced: Quantity, + mass_entrainment_downdraft_forced: Quantity, + mass_detrainment_downdraft_forced: Quantity, + vertical_velocity_3d: Quantity, + total_normalized_integrated_condensate_forced: Quantity, + total_normalized_integrated_evaporate_forced: Quantity, + evaporate_in_downdraft_forced: Quantity, + epsilon_forced: Quantity, + chemistry_tracers: Quantity, + chemistry_tracers_output: Quantity, + chemistry_tracers_cloud_levels: Quantity, + chemistry_tracers_sc_updraft: Quantity, + chemistry_tracers_sc_downdraft: Quantity, + chemistry_tracers_pw_updraft: Quantity, + chemistry_tracers_pw_downdraft: Quantity, + chemistry_tracers_total_pw_updraft: Quantity, + chemistry_tracers_total_pw_downdraft: Quantity, + convection_tracers: ConvectionTracers, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + """Apply the effects of convection to the convection/chemistry tracers. + + Args: + error_code (Quantity) + cloud_top_level (Quantity) + updraft_origin_level (Quantity) + downdraft_origin_level (Quantity) + ocean_fraction (Quantity) + p_forced (Quantity) + p_cloud_levels_forced (Quantity) + geopotential_height_cloud_levels (Quantity) + environment_massflux (Quantity) + normalized_massflux_updraft_forced (Quantity) + normalized_massflux_downdraft_forced (Quantity) + mass_entrainment_updraft_forced (Quantity) + mass_detrainment_updraft_forced (Quantity) + mass_entrainment_downdraft_forced (Quantity) + mass_detrainment_downdraft_forced (Quantity) + vertical_velocity_3d (Quantity) + total_normalized_integrated_condensate_forced (Quantity) + total_normalized_integrated_evaporate_forced (Quantity) + evaporate_in_downdraft_forced (Quantity) + epsilon_forced (Quantity) + chemistry_tracers (Quantity) + chemistry_tracers_output (Quantity) + chemistry_tracers_cloud_levels (Quantity) + chemistry_tracers_sc_updraft (Quantity) + chemistry_tracers_sc_downdraft (Quantity) + chemistry_tracers_pw_updraft (Quantity) + chemistry_tracers_pw_downdraft (Quantity) + chemistry_tracers_total_pw_updraft (Quantity) + chemistry_tracers_total_pw_downdraft (Quantity) + convection_tracers (ConvectionTracers): Collection of tracers from the rest of the model which + will be updated within convection. These may come from a variaty of sources, and need to be + collected into the expected ConvectionTracers data type before being passed down. + plume_dependent_constants (GF2020PlumeDependentConstants) + + Raises: + NotImplementedError: _description_ + """ + + # 1) get mass mixing ratios at the cloud levels + self._environment_cloud_levels_chemistry( + error_code=error_code, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_cloud_levels=chemistry_tracers_cloud_levels, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if (convection_tracers.vect_hcts.field[:, 0] > Float(1.0e-6)).any(): + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->AtmosphericComposition called with " + "convection_tracers.vect_hcts > 1.0e-6. This condition requires an unimplemented option in " + "updraft_chemistry. Please implement, then remove this error to continue." + ) + + # 2) determine in-cloud tracer mixing ratios + # a) updraft chemistry + self._updraft_chemistry( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + cloud_top_level=cloud_top_level, + ocean_fraction=ocean_fraction, + p_forced=p_forced, + p_cloud_levels_forced=p_cloud_levels_forced, + geopotential_height_cloud_levels=geopotential_height_cloud_levels, + vertical_velocity_3d=vertical_velocity_3d, + mass_entrainment_updraft_forced=mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=mass_detrainment_updraft_forced, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_sc_updraft=chemistry_tracers_sc_updraft, + chemistry_tracers_cloud_levels=chemistry_tracers_cloud_levels, + chemistry_tracers_pw_updraft=chemistry_tracers_pw_updraft, + chemistry_tracers_total_pw_updraft=chemistry_tracers_total_pw_updraft, + convection_tracers_vect_hcts=convection_tracers.vect_hcts, + convection_tracers_fscav=convection_tracers.fscav, + convection_tracers_use_gcc_washout=convection_tracers.use_gcc_washout, + tracer_cloud_boundary=self._tracer_cloud_boundary, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # b) downdraft chemistry + self._downdraft_chemistry( + error_code=error_code, + downdraft_origin_level=downdraft_origin_level, + p_cloud_levels_forced=p_cloud_levels_forced, + evaporate_in_downdraft_forced=evaporate_in_downdraft_forced, + total_normalized_integrated_evaporate_forced=total_normalized_integrated_evaporate_forced, + total_normalized_integrated_condensate_forced=total_normalized_integrated_condensate_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + mass_entrainment_downdraft_forced=mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=mass_detrainment_downdraft_forced, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_cloud_levels=chemistry_tracers_cloud_levels, + chemistry_tracers_sc_downdraft=chemistry_tracers_sc_downdraft, + chemistry_tracers_pw_downdraft=chemistry_tracers_pw_downdraft, + chemistry_tracers_total_pw_downdraft=chemistry_tracers_total_pw_downdraft, + chemistry_tracers_total_pw_updraft=chemistry_tracers_total_pw_updraft, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # 3) determine the vertical transport including mixing, scavenging and evaporation + self._vertical_transport_part_1( + error_code=error_code, + cloud_top_level=cloud_top_level, + p_cloud_levels_forced=p_cloud_levels_forced, + environment_massflux=environment_massflux, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + epsilon_forced=epsilon_forced, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_output=chemistry_tracers_output, + chemistry_tracers_cloud_levels=chemistry_tracers_cloud_levels, + chemistry_tracers_sc_updraft=chemistry_tracers_sc_updraft, + chemistry_tracers_sc_downdraft=chemistry_tracers_sc_downdraft, + chemistry_tracers_pw_updraft=chemistry_tracers_pw_updraft, + chemistry_tracers_pw_downdraft=chemistry_tracers_pw_downdraft, + dd_tracers=self._dd_tracers, + aa=self._aa, + bb=self._bb, + cc=self._cc, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if ( + self.cumulus_parameterization_config.USE_FLUX_FORM == 1 + and self.cumulus_parameterization_config.ALP1 > 0.0 + ): + for tracer in range(self.config.NUMBER_OF_TRACERS): + # NOTE, this code is an extention of the vertical transport section for option + # which executes when (USE_FLUX_FORM == 1 and ALP1 > 0.0) + self._tridiag( + m=cloud_top_level, + a=self._aa, + b=self._bb, + c=self._cc, + f=self._dd_tracers.data[:, :, :, tracer], + error_code=error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._update_after_tridiag( + error_code=error_code, + cloud_top_level=cloud_top_level, + dd=self._dd_tracers, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_output=chemistry_tracers_output, + tracer=Int(tracer), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._vertical_transport_part_2( + error_code=error_code, + cloud_top_level=cloud_top_level, + p_cloud_levels_forced=p_cloud_levels_forced, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + epsilon_forced=epsilon_forced, + chemistry_tracers=chemistry_tracers, + chemistry_tracers_output=chemistry_tracers_output, + chemistry_tracers_pw_updraft=chemistry_tracers_pw_updraft, + chemistry_tracers_pw_downdraft=chemistry_tracers_pw_downdraft, + dd_tracers=self._dd_tracers, + residual=self._residual, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/cumulus_parameterization.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/cumulus_parameterization.py new file mode 100644 index 000000000..8b5835df8 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/cumulus_parameterization.py @@ -0,0 +1,2162 @@ +from ndsl import NDSLRuntime, QuantityFactory, StencilFactory, ndsl_log +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.air_density import hydrostatic_air_density +from pyMoist.convection.GF_2020.cumulus_parameterization.buoyancy import get_buoyancy +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import MAXENS1, MAXENS2, MAXENS3 +from pyMoist.convection.GF_2020.cumulus_parameterization.convective_tracers import ( + AtmosphericComposition, + ColdPoolParameterization, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.diurnal_cycle import DiurnalCycle +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import ( + DowndraftOriginLevel, + DowndraftWetBlub, + DowndraftWindShear, + downdraft_lateral_massflux, + downdraft_mass_flux, + downdraft_moist_static_energy_and_buoyancy, + downdraft_moisture, + downdraft_temperature, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.entrainment import ( + compute_lateral_massflux, + compute_uc_vc, + downdraft_entrainment_profiles, + entrainment_rates, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import ( + EnvironmentalSubsidence, + environment_cloud_levels, + environment_conditions, + environment_mass_flux, + modify_environment_profiles, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import ( + cloud_top_checks, + find_detrainmet_start_level, + find_highest_moist_static_energy_level, + find_lcl, + find_maximum_updraft_origin_level, + get_cloud_top, + get_convective_cloud_base_level, + set_start_level, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.kinetic_energy_to_heating import ( + kinetic_energy_to_heating, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.large_scale_forcing import LargeScaleForcing +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.mass_conservation import MassConservation +from pyMoist.convection.GF_2020.cumulus_parameterization.moist_static_energy import ( + StaticControl, + first_guess_moist_static_energy, + parcel_moist_static_energy, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.precip import ( + PrecipFactor, + cloud_dissapation, + get_precip_fluxes, + partition_liquid_ice, + rain_evaporation_below_cloud_base, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import ( + LightningFlashDensity, + OutputWorkfunctionsAndPrecipConcentrations, + deep_precipitation_output, + ensemble_output_and_feedback, + output_updraft_temperature, + prepare_output, + total_evaporation_flux, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.profiles import C1DProfile, melting_profile +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.setup import Setup +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import ( + generic_find_level, + updraft_vertical_velocity, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.smoothing import smooth_tendencies +from pyMoist.convection.GF_2020.cumulus_parameterization.sounding import GATESounding, Sounding +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.triggers import ( + XieTriggerFunction, + convection_trigger, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import ( + UpdateWorkfunctionAndPrecipitationEnsemble, + UpdraftCIN, + UpdraftInitialWorkfunctions, + UpdraftMassFlux, + updraft_moist_static_energy_and_momentum_budget, + updraft_moisture, + updraft_temperature, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.vertical_discretization import VerticalDiscretization +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class GF2020CumulusParameterization(NDSLRuntime): + """GF220 cumulus parameterization core. This component of GF2020 contains all of the science code related + to deep conveection, including (but not limited to): + - updraft dynamics + - downdraft dynamics + - entrainment/detrainment + - multi-phase precipitation accumulation + - tracer support + - near-storm environmental effects (optional, not implemented) + - sounding profiles (optional, not implemented) + + Most pieces are called directly, but more complicated parts are buried into their own subclass. + + Not all parts/options are fully implemented, some have been implemetned without full testing. + ALL of these untested/unimplemented paths are locked behind an NotImplementedError, with an + appropriate message pointing to the configuration option that is causing the error. + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + saturation_tables: SaturationVaporPressureTable, + ): + """Initialize the GF2020 cumulus parameterization core. Allocate locals and construct stencils + + Args: + stencil_factory (StencilFactory) + quantity_factory (QuantityFactory) + config (GF2020Config) + cumulus_parameterization_config (GF2020CumulusParameterizationConfig) + saturation_tables (SaturationVaporPressureTable) + """ + super().__init__(stencil_factory) + + # make saturation tables visible at runtime + self.saturation_tables = saturation_tables + + # get config from the rest of the model + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # initialize the plume dependent constants, to be set for each plume within the loop (in Setup) + self.plume_dependent_constants = GF2020PlumeDependentConstants() + + # initialize the local fields needed within this class + self.locals = GF2020CumulusParameterizationLocals.zeros( + quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # initialize all the subclasses + self._setup = Setup( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._environment_conditions = stencil_factory.from_dims_halo( + func=environment_conditions, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SATURATION_CALCULATION_CHOICE": cumulus_parameterization_config.SATURATION_CALCULATION_CHOICE + }, + ) + + self._sounding = Sounding(cumulus_parameterization_config=cumulus_parameterization_config) + + self._environment_cloud_levels = stencil_factory.from_dims_halo( + func=environment_cloud_levels, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"CLOUD_LEVEL_GRID": cumulus_parameterization_config.CLOUD_LEVEL_GRID}, + ) + + self._hydrostatic_air_density = stencil_factory.from_dims_halo( + func=hydrostatic_air_density, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._partition_liquid_ice = stencil_factory.from_dims_halo( + func=partition_liquid_ice, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "MELT_GLAC": cumulus_parameterization_config.MELT_GLAC, + "FRAC_MODIS": cumulus_parameterization_config.FRAC_MODIS, + }, + ) + + self._find_maximum_updraft_origin_level = stencil_factory.from_dims_halo( + func=find_maximum_updraft_origin_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._find_detrainmet_start_level = stencil_factory.from_dims_halo( + func=find_detrainmet_start_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._find_highest_moist_static_energy_level = stencil_factory.from_dims_halo( + func=find_highest_moist_static_energy_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._precip_factor = PrecipFactor() + + self._cold_pool_parameterization = ColdPoolParameterization(config=config) + + self._find_lcl = stencil_factory.from_dims_halo( + func=find_lcl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + "ADV_TRIGGER": config.ADV_TRIGGER, + }, + ) + + self._parcel_moist_static_energy = stencil_factory.from_dims_halo( + func=parcel_moist_static_energy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + self._entrainment_rates = stencil_factory.from_dims_halo( + func=entrainment_rates, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "MIN_ENTRAINMENT_RATE": cumulus_parameterization_config.MIN_ENTRAINMENT_RATE, + }, + ) + + self._set_start_level = stencil_factory.from_dims_halo( + func=set_start_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._get_convective_cloud_base_level = stencil_factory.from_dims_halo( + func=get_convective_cloud_base_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "OVERSHOOT": cumulus_parameterization_config.OVERSHOOT, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "MOIST_TRIGGER": cumulus_parameterization_config.MOIST_TRIGGER, + "USE_MEMORY": cumulus_parameterization_config.USE_MEMORY, + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + self._downdraft_entraiment_profiles = stencil_factory.from_dims_halo( + func=downdraft_entrainment_profiles, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DOWNDRAFT": cumulus_parameterization_config.DOWNDRAFT}, + ) + + self._generic_find_level = stencil_factory.from_dims_halo( + func=generic_find_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._get_cloud_top = stencil_factory.from_dims_halo( + func=get_cloud_top, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"OVERSHOOT": cumulus_parameterization_config.OVERSHOOT}, + ) + + self._cloud_top_checks = stencil_factory.from_dims_halo( + func=cloud_top_checks, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._updraft_mass_flux = UpdraftMassFlux( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._compute_lateral_massflux = stencil_factory.from_dims_halo( + func=compute_lateral_massflux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._compute_uc_vc = stencil_factory.from_dims_halo( + func=compute_uc_vc, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD + }, + ) + + self._first_guess_moist_static_energy = stencil_factory.from_dims_halo( + func=first_guess_moist_static_energy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._get_buoyancy = stencil_factory.from_dims_halo( + func=get_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._c1d_profile = C1DProfile( + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._updraft_moisture_profile = stencil_factory.from_dims_halo( + func=updraft_moisture, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES": cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + "AUTOCONV": config.AUTOCONV, + "CRITICAL_MIXING_RATIO_OVER_OCEAN": cumulus_parameterization_config.CRITICAL_MIXING_RATIO_OVER_OCEAN, + "CRITICAL_MIXING_RATIO_OVER_LAND": cumulus_parameterization_config.CRITICAL_MIXING_RATIO_OVER_LAND, + "FRAC_MODIS": cumulus_parameterization_config.FRAC_MODIS, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + }, + ) + + self._melting_profile = stencil_factory.from_dims_halo( + func=melting_profile, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "MELT_GLAC": cumulus_parameterization_config.MELT_GLAC, + }, + ) + + self._updraft_moist_static_energy_and_momentum_budget = stencil_factory.from_dims_halo( + func=updraft_moist_static_energy_and_momentum_budget, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES": cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + "PRESSURE_GRADIENT_CONSTANT": cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT, + }, + ) + + self._updraft_temperature = stencil_factory.from_dims_halo( + func=updraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._updraft_vertical_velocity = stencil_factory.from_dims_halo( + func=updraft_vertical_velocity, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF}, + ) + + self._downdraft_origin_level = DowndraftOriginLevel( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._downdraft_mass_flux = stencil_factory.from_dims_halo( + func=downdraft_mass_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF}, + ) + + self._downdraft_lateral_mass_flux = stencil_factory.from_dims_halo( + func=downdraft_lateral_massflux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._downdraft_wet_bulb = DowndraftWetBlub( + cumulus_parameterization_config=cumulus_parameterization_config + ) + + self._downdraft_moist_static_energy_and_buoyancy = stencil_factory.from_dims_halo( + func=downdraft_moist_static_energy_and_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_WETBULB": cumulus_parameterization_config.USE_WETBULB, + "PRESSURE_GRADIENT_CONSTANT": cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT, + }, + ) + + self._downdraft_moisture = stencil_factory.from_dims_halo( + func=downdraft_moisture, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_WETBULB": cumulus_parameterization_config.USE_WETBULB, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "EVAP_FIX": cumulus_parameterization_config.EVAP_FIX, + }, + ) + + self._updraft_initial_workfunctions = UpdraftInitialWorkfunctions( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._updraft_cin = UpdraftCIN( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._convection_trigger = stencil_factory.from_dims_halo( + func=convection_trigger, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DICYCLE": cumulus_parameterization_config.DIURNAL_CYCLE}, + ) + + self._downdraft_temperature = stencil_factory.from_dims_halo( + func=downdraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._diurnal_cycle = DiurnalCycle( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._Xie_trigger_function = XieTriggerFunction( + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._downdraft_windshear = DowndraftWindShear( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._environment_mass_flux = stencil_factory.from_dims_halo( + func=environment_mass_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._mass_conservation = MassConservation() + + self._vertical_discretization = VerticalDiscretization( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._environmental_subsidence = EnvironmentalSubsidence( + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._smooth_tendencies = stencil_factory.from_dims_halo( + func=smooth_tendencies, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"USE_SMOOTH_TENDENCIES": cumulus_parameterization_config.USE_SMOOTH_TENDENCIES}, + ) + + self._modify_environment_profiles = stencil_factory.from_dims_halo( + func=modify_environment_profiles, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "COUPLE_MICROPHYSICS": cumulus_parameterization_config.COUPLE_MICROPHYSICS, + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + self._static_control = StaticControl( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._update_workfunction_and_precipitaiton_ensemble = UpdateWorkfunctionAndPrecipitationEnsemble( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._large_scale_forcing = LargeScaleForcing( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._kinetic_energy_to_heating = stencil_factory.from_dims_halo( + func=kinetic_energy_to_heating, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._ensemble_output_and_feedback = stencil_factory.from_dims_halo( + func=ensemble_output_and_feedback, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "DTIME": cumulus_parameterization_config.DTIME, + "MAX_TEMP_VAPOR_TENDENCY": cumulus_parameterization_config.MAX_TEMP_VAPOR_TENDENCY, + "APPLY_SUBSIDENCE_MICROPHYSICS": config.APPLY_SUBSIDENCE_MICROPHYSICS, + "USE_SMOOTH_TENDENCIES": cumulus_parameterization_config.USE_SMOOTH_TENDENCIES, + }, + ) + + self._precipitation_flux = stencil_factory.from_dims_halo( + func=get_precip_fluxes, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._rain_evaporation_below_cloud_base = stencil_factory.from_dims_halo( + func=rain_evaporation_below_cloud_base, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._cloud_dissapation = stencil_factory.from_dims_halo( + func=cloud_dissapation, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_CLOUD_DISSIPATION": cumulus_parameterization_config.USE_CLOUD_DISSIPATION, + "DTIME": cumulus_parameterization_config.DTIME, + "COUPLE_MICROPHYSICS": cumulus_parameterization_config.COUPLE_MICROPHYSICS, + }, + ) + + self._total_evaporation_flux = stencil_factory.from_dims_halo( + func=total_evaporation_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._lightning_flash_density = LightningFlashDensity( + cumulus_parameterization_config=cumulus_parameterization_config + ) + + self._deep_precipitation_output = stencil_factory.from_dims_halo( + func=deep_precipitation_output, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._output_updraft_temperature = stencil_factory.from_dims_halo( + func=output_updraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._prepare_output = stencil_factory.from_dims_halo( + func=prepare_output, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._output_workfunctions_and_precip_concentrations = OutputWorkfunctionsAndPrecipConcentrations( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._atmospheric_composition = AtmosphericComposition( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + self._gate_sounding = GATESounding(cumulus_parameterization_config=cumulus_parameterization_config) + + def __call__( + self, + state: GF2020CumulusParameterizationState, + convection_tracers: ConvectionTracers, + ): + """Run the GF2020 cumulus parameterization core. + + Args: + state (GF2020CumulusParameterizationState): Special state for the cumulus parameterization core. + This state is initialized in the GF2020Setup class. The overarching model state cannot be + passed to the GF2020CumulusParameterizationState because of incompatible K-axis orientation. + convection_tracers (ConvectionTracers): Collection of tracers from the rest of the model which + will be updated within convection. These may come from a variaty of sources, and need to be + collected into the expected ConvectionTracers data type before being passed down. + """ + if self.cumulus_parameterization_config.PLUME_ORDER == 0: + plume_types = ["shallow", "mid", "deep"] + elif self.cumulus_parameterization_config.PLUME_ORDER == 1: + plume_types = ["shallow", "deep", "mid"] + else: + raise NotImplementedError("plume order not impelemented") + + for plume in plume_types: + # setup constants for the current plume, reset necessary fields, prefill necessary fields + # NOTE test GF2020_CumulusParameterization_Setup_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ❌ (cannot serialize output data because the plume is disabled) + self._setup( + error_code=state.output.error_code, + error_code_2=self.locals.error_code_2, + error_code_3=self.locals.error_code_3, + maximum_updraft_origin_level=self.locals.maximum_updraft_origin_level, + kstabm=state.output.kstabm, + t_excess=state.input.t_excess, + t_excess_local=self.locals.t_excess, + vapor_excess=state.input.vapor_excess, + vapor_excess_local=self.locals.vapor_excess, + ocean_fraction=state.input.ocean_fraction, + ocean_fraction_local=self.locals.ocean_fraction, + t_old=state.input_output.t_old, + t_new=self.locals.t_new, + t_new_pbl=self.locals.t_new_pbl, + vapor_old=state.input_output.vapor_old, + vapor_forced=self.locals.vapor_forced, + vapor_forced_pbl=self.locals.vapor_forced_pbl, + downdraft_saturation_vapor_forced=self.locals.downdraft_saturation_vapor_forced, + grid_scale_forcing_t=state.input.grid_scale_forcing_t, + grid_scale_forcing_vapor=state.input.grid_scale_forcing_vapor, + subgrid_scale_forcing_t=state.input.subgrid_scale_forcing_t, + subgrid_scale_forcing_vapor=state.input.subgrid_scale_forcing_vapor, + dmoist_static_energydt=self.locals.dmoist_static_energydt, + cloud_moist_static_energy_downdraft_forced=self.locals.cloud_moist_static_energy_downdraft_forced, + cloud_moist_static_energy_forced_transported=self.locals.cloud_moist_static_energy_forced_transported, + cap_max=self.locals.cap_max, + cap_max_increment=self.locals.cap_max_increment, + geopotential_height=state.input_output.geopotential_height_forced, + geopotential_height_local=self.locals.geopotential_height, + geopotential_height_modified_local=self.locals.geopotential_height_modified, + cloud_workfunction_0=self.locals.cloud_workfunction_0, + cloud_workfunction_1=self.locals.cloud_workfunction_1, + cloud_workfunction_2=self.locals.cloud_workfunction_2, + cloud_workfunction_3=self.locals.cloud_workfunction_3, + cloud_workfunction_0_pbl=self.locals.cloud_workfunction_0_pbl, + cloud_workfunction_1_pbl=self.locals.cloud_workfunction_1_pbl, + cloud_workfunction_1_fa=self.locals.cloud_workfunction_1_fa, + cin_1=self.locals.cin_1, + k_x_modified=self.locals.k_x_modified, + epsilon_forced=state.output.epsilon_forced, + epsilon_local=self.locals.epsilon, + epsilon_min=self.locals.epsilon_min, + epsilon_max=self.locals.epsilon_max, + pbl_time_scale=self.locals.pbl_time_scale, + t_wetbulb=self.locals.t_wetbulb, + vapor_wetbulb=self.locals.vapor_wetbulb, + cape_removal_time_scale=self.locals.cape_removal_time_scale, + f_dicycle_modified=self.locals.f_dicycle_modified, + add_buoyancy=self.locals.add_buoyancy, + scale_dependence_factor=state.output.scale_dependence_factor, + scale_dependence_factor_downdraft=self.locals.scale_dependence_factor_downdraft, + c1d=self.locals.c1d, + evaporation_below_cloud_base=self.locals.evaporation_below_cloud_base, + mass_flux_ensemble=self.locals.mass_flux_ensemble, + precipitation_ensemble=self.locals.precipitation_ensemble, + precip=state.output.precip, + lightning_density=state.output.lightning_density, + seed_convection=state.input.seed_convection, + grid_length=state.input_output.grid_length, + random_number=self.locals.random_number, + lateral_entrainment_rate=state.input.lateral_entrainment_rate, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=self.locals.detrainment_function_updraft, + arbitrary_numerical_parameter=self.locals.arbitrary_numerical_parameter, + plume_dependent_constants=self.plume_dependent_constants, + plume=plume, + ) + + if self.plume_dependent_constants.ENABLE_PLUME == 1: + # environmental conditions, first heights + # calculate moist static energy, heights, environmental saturation mixing ratio + # NOTE test GF2020_CumulusParameterization_EnvironmentConditions_1_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._environment_conditions( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=state.input_output.t_old, + vapor=state.input_output.vapor_old, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=self.locals.environment_moist_static_energy, + saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy, + saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio, + geopotential_height=self.locals.geopotential_height, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + # NOTE test GF2020_CumulusParameterization_EnvironmentConditions_2_{plume}: + # NOTE deep ❌ (two vars, each one point off by two ulp) + # NOTE mid ❌ (two vars, each one point off by two ulp) + # NOTE shallow ✅ + self._environment_conditions( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=self.locals.t_new, + vapor=self.locals.vapor_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=self.locals.environment_moist_static_energy_forced, + saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_forced, + saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio_forced, + geopotential_height=state.input_output.geopotential_height_forced, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # outputs a model sounding for the stand-alone code (part 1) + if self.cumulus_parameterization_config.CLOUD_LEVEL_GRID != 1: + ndsl_log.warning( + " GF2020 cumulus parameterization initialized with unimplemented OUTPUT_SOUNDING option. " + "Output soundings are not currently available. Contact support if this tool is needed." + ) + + # environmental values on cloud levels + # NOTE test GF2020_CumulusParameterization_EnvironmentCloudLevels_1_{plume}: + # NOTE. deep ❌ (worst fail rate 0.01%, worse fail 32 ULP) + # NOTE. mid ❌ (worst fail rate 0.01%, worse fail 32 ULP) + # NOTE. shallow ✅ + self._environment_cloud_levels( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=self.locals.p_cloud_levels, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=self.locals.geopotential_height, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + t=state.input_output.t_old, + t_surface=state.input_output.t_surface, + t_cloud_levels=self.locals.t_cloud_levels, + vapor=state.input_output.vapor_old, + vapor_cloud_levels=self.locals.vapor_cloud_levels, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + environment_saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio, + environment_saturation_mixing_ratio_cloud_levels=self.locals.environment_saturation_mixing_ratio_cloud_levels, + environment_moist_static_energy=self.locals.environment_moist_static_energy, + environment_moist_static_energy_cloud_levels=self.locals.environment_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy, + environment_saturation_moist_static_energy_cloud_levels=self.locals.environment_saturation_moist_static_energy_cloud_levels, + gamma_cloud_levels=self.locals.gamma_cloud_levels, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + # NOTE test GF2020_CumulusParameterization_EnvironmentCloudLevels_2_{plume}: + # NOTE deep ❌ (worst fail rate 0.02%, worse fail 84 ULP) + # NOTE mid ❌ (worst fail rate 0.02%, worse fail 84 ULP) + # NOTE shallow ✅ + p_3d = state.output.p_cloud_levels_forced.field[ + :, :, :, self.plume_dependent_constants.PLUME_INDEX + ] # TODO this should go in a stencil + self._environment_cloud_levels( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=p_3d, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=state.input_output.geopotential_height_forced, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels_forced, + t=self.locals.t_new, + t_surface=state.input_output.t_surface, + t_cloud_levels=self.locals.t_cloud_levels_forced, + vapor=self.locals.vapor_forced, + vapor_cloud_levels=self.locals.vapor_cloud_levels_forced, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + environment_saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio_forced, + environment_saturation_mixing_ratio_cloud_levels=self.locals.environment_saturation_mixing_ratio_cloud_levels_forced, + environment_moist_static_energy=self.locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=self.locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + gamma_cloud_levels=self.locals.gamma_cloud_levels_forced, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + state.output.p_cloud_levels_forced.field[ + :, :, :, self.plume_dependent_constants.PLUME_INDEX + ] = p_3d # TODO this should go in a stencil + + # get air density at full layer (model levels) by hydrostatic balance (kg/m3) + # NOTE test GF2020_CumulusParameterization_HydrostaticAirDensity_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._hydrostatic_air_density( + p=state.output.p_cloud_levels_forced, + geopotential_height=self.locals.geopotential_height_cloud_levels_forced, + error_code=state.output.error_code, + air_density=self.locals.hydrostatic_air_density, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # partition between liq/ice cloud contents + # NOTE test GF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_{plume}: + # NOTE deep ❌ 1 var failing (500 ULP) -> ice_fraction function causing issues + # NOTE mid ✅ + # NOTE shallow ✅ + self._partition_liquid_ice( + t=self.locals.t_new, + p=state.output.p_cloud_levels_forced, + geopotential_height=self.locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + surface_type=state.input.surface_type, + convection_fraction=state.input.convection_fraction, + error_code=state.output.error_code, + melting_layer=self.locals.melting_layer, + part_liquid_ice=self.locals.partition_liquid_ice, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + self._find_maximum_updraft_origin_level( + geopotential_height=self.locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + maximum_updraft_origin_level=self.locals.maximum_updraft_origin_level, + MAX_UPDRAFT_ORIGIN_HEIGHT=self.plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + self._find_detrainmet_start_level( + geopotential_height=self.locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + detrainment_start_level=self.locals.detrainment_start_level, + DETRAINMENT_CRITICAL_DEPTH=self.plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # determine level with highest moist static energy content (k_max_mse) + # NOTE test GF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._find_highest_moist_static_energy_level( + moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels_forced, + error_code=state.output.error_code, + maximum_updraft_origin_level=self.locals.maximum_updraft_origin_level, + updraft_origin_level=state.output.updraft_origin_level, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # get the pickup of ensemble ave prec, following Neelin et al 2009. + # NOTE this runs in fortran but the output is never used - so it is not implemented + self._precip_factor() + + # cold pool parameterization and convective memory + # NOTE CONVECTION TRACER BLOCK + # NOTE not called in experiment used to design this, so it is not implemented + self._cold_pool_parameterization() + + # determine LCL for the air parcels with highest moist static energy + # NOTE test GF2020_CumulusParameterization_GetLCL_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._find_lcl( + p=state.input_output.p_forced, + p_cloud_levels=self.locals.p_cloud_levels, + t_excess=self.locals.t_excess, + t_cloud_levels_forced=self.locals.t_cloud_levels, + t_perturbation=state.output.t_perturbation, + vapor_excess=self.locals.vapor_excess, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels, + omega=state.input_output.omega, + air_density=state.input_output.air_density, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + topography_height_no_negative=state.input_output.topography_height_no_negative, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + grid_length=state.input_output.grid_length, + lcl_level=state.output.lcl_level, + error_code=state.output.error_code, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # determine the moist static energy of air parcels at source level + # NOTE test GF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._parcel_moist_static_energy( + error_code=state.output.error_code, + t_excess=self.locals.t_excess, + vapor_excess=self.locals.vapor_excess, + add_buoyancy=self.locals.add_buoyancy, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + p=state.input_output.p_forced, + environmenet_moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels, + environmenet_moist_static_energy_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + t_perturbation=state.output.t_perturbation, + moist_static_energy_origin_level=self.locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # determine the vertical entrainment/detrainment rates + # NOTE test GF2020_CumulusParameterization_EntrainmentRates_{plume}: + # NOTE deep ❌ (worst fail 2.41% - max 4 ULP) + # NOTE mid ❌ (worst fail 0.88% - max 2 ULP) + # NOTE shallow ✅ + self._entrainment_rates( + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=self.locals.environment_saturation_mixing_ratio_cloud_levels_forced, + lcl_level=state.output.lcl_level, + error_code=state.output.error_code, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=self.locals.detrainment_function_updraft, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # determine level of convective cloud base + # NOTE test GF2020_CumulusParameterization_ConvectiveCloudBaseLevel_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._set_start_level( + lcl_level=state.output.lcl_level, + start_level=self.locals.start_level, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + self._get_convective_cloud_base_level( + error_code=state.output.error_code, + lcl_level=state.output.lcl_level, + cloud_moist_static_energy_forced_transported=self.locals.cloud_moist_static_energy_forced_transported, + cap_max=self.locals.cap_max, + updraft_origin_level=state.output.updraft_origin_level, + start_level=self.locals.start_level, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + maximum_updraft_origin_level=self.locals.maximum_updraft_origin_level, + negative_buoyancy_depth=self.locals.negative_buoyancy_depth, + frh_lfc=self.locals.frh_lfc, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + entrainment_rate=state.output.entrainment_rate, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + t_excess=self.locals.t_excess, + vapor_excess=self.locals.vapor_excess, + add_buoyancy=self.locals.add_buoyancy, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + vapor_forced=self.locals.vapor_forced, + environment_saturation_mixing_ratio_forced=self.locals.environment_saturation_mixing_ratio_forced, + ocean_fraction=self.locals.ocean_fraction, + cap_max_increment=self.locals.cap_max_increment, + t_perturbation=state.output.t_perturbation, + p_forced=state.input_output.p_forced, + cloud_top_level=state.output.cloud_top_level, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # define entrainment/detrainment profiles for downdrafts + # NOTE test GF2020_CumulusParameterization_DowndraftEntrainmentProfiles_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_entraiment_profiles( + lateral_entrainment_rate=state.input.lateral_entrainment_rate, + entrainment_rate_downdraft=self.locals.entrainment_rate_downdraft, + detrainment_function_downdraft=self.locals.detrainment_function_downdraft, + scale_dependence_factor_downdraft=self.locals.scale_dependence_factor_downdraft, + plume_entrainment_rate=self.plume_dependent_constants.ENTRAINMENT_RATE, + ) + + # update unforced & forced moist static energy + # NOTE test GF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._parcel_moist_static_energy( + error_code=state.output.error_code, + t_excess=self.locals.t_excess, + vapor_excess=self.locals.vapor_excess, + add_buoyancy=self.locals.add_buoyancy, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + p=state.input_output.p_forced, + environmenet_moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels, + environmenet_moist_static_energy_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + t_perturbation=state.output.t_perturbation, + moist_static_energy_origin_level=self.locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # increase detrainment in stable layers + # NOTE test GF2020_CumulusParameterization_StableDetrainment_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._generic_find_level( + array=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + start_index=state.output.updraft_lfc_level, + end_index=state.output.kstabm.field[:, :, self.plume_dependent_constants.PLUME_INDEX], + out_index=state.output.kstabi, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # use cloud for plumes + # NOTE test GF2020_CumulusParameterization_CloudTop_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._get_cloud_top( + entrainment_rate=state.output.entrainment_rate, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + cloud_moist_static_energy_forced_transported=self.locals.cloud_moist_static_energy_forced_transported, + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + self._cloud_top_checks( + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + error_code=state.output.error_code, + last_error_code=state.input.last_error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + MINIMUM_DEPTH=self.plume_dependent_constants.MINIMUM_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # determine the normalized mass flux profile for updraft + # NOTE test GF2020_CumulusParameterization_UpdraftMassFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_mass_flux( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + lcl_level=state.output.lcl_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + ocean_fraction=state.input.ocean_fraction, + normalized_massflux_updraft=self.locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_updraft_modified=self.locals.normalized_massflux_updraft_modified, + random_number=self.locals.random_number, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # calculate mass entrainment and detrainment + # NOTE test GF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._compute_lateral_massflux( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + geopotential_height=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=state.output.normalized_massflux_updraft_forced, + detrainment_function_updraft=self.locals.detrainment_function_updraft, + entrainment_rate=state.output.entrainment_rate, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft=self.locals.mass_entrainment_updraft, + mass_detrainment_updraft=self.locals.mass_detrainment_updraft, + updraft_lfc_level=state.output.updraft_lfc_level, + updraft_origin_level=state.output.updraft_origin_level, + pbl_level=state.input_output.pbl_level, + mass_entrainment_u_updraft=self.locals.mass_entrainment_u_updraft, + mass_detrainment_u_updraft=self.locals.mass_detrainment_u_updraft, + LAMBDA_DEEP=self.plume_dependent_constants.LAMBDA_DEEP, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + self._compute_uc_vc( + u_c=self.locals.u_c, + v_c=self.locals.v_c, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + error_code=state.output.error_code, + start_level=self.locals.start_level, + moist_static_energy_origin_level=self.locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + p=state.input_output.p_forced, + updraft_origin_level=state.output.updraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # 1st guess for moist static energy + # NOTE test GF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._first_guess_moist_static_energy( + error_code=state.output.error_code, + start_level=self.locals.start_level, + cloud_top_level=state.output.cloud_top_level, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + normalized_massflux_updraft=self.locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + vapor_excess=self.locals.vapor_excess, + t_excess=self.locals.t_excess, + add_buoyancy=self.locals.add_buoyancy, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # Get buoyancy of updrafts + # NOTE test GF2020_CumulusParameterization_GetBuoyancy_1_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._get_buoyancy( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy_forced, + environment_moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + d_buoyancy=self.locals.d_buoyancy_forced, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # get "c1d" profile + # NOTE test GF2020_CumulusParameterization_C1DProfile_{plume}: + # NOTE deep ⚠️⚠️⚠️ DOES NOT EXECUTE IN CURRENT SIMULATION + # NOTE mid ⚠️⚠️⚠️ DOES NOT EXECUTE IN CURRENT SIMULATION + # NOTE shallow ⚠️⚠️⚠️ DOES NOT EXECUTE IN CURRENT SIMULATION + # NOTE UNFINISHED - MANUALLY DISABLED (see class docstring) + if False: + self._c1d_profile( + state=state, + locals=self.locals, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # calculate moisture properties of updraft + # NOTE test GF2020_CumulusParameterization_UpdraftMoisture_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_moisture_profile( + start_level=self.locals.start_level, + error_code=state.output.error_code, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=self.locals.cloud_total_water_after_entrainment_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + updraft_column_temperature_forced=self.locals.updraft_column_temperature_forced, + ocean_fraction=state.input.ocean_fraction, + convection_fraction=state.input.convection_fraction, + surface_type=state.input.surface_type, + p_forced=state.input_output.p_forced, + cloud_top_level=state.output.cloud_top_level, + d_buoyancy_forced=self.locals.d_buoyancy_forced, + cloud_liquid_before_rain_forced=self.locals.cloud_liquid_before_rain_forced, + t_cloud_levels=self.locals.t_cloud_levels, + vapor_forced=self.locals.vapor_forced, + gamma_cloud_levels_forced=self.locals.gamma_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=self.locals.environment_saturation_mixing_ratio_cloud_levels_forced, + updraft_origin_level=state.output.updraft_origin_level, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + vapor_excess=self.locals.vapor_excess, + ccn=state.input_output.ccn, + mass_entrainment_updraft=self.locals.mass_entrainment_updraft, + mass_detrainment_updraft=self.locals.mass_detrainment_updraft, + psum=self.locals.psum, + psumh=self.locals.psumh, + c1d=self.locals.c1d, + add_buoyancy=self.locals.add_buoyancy, + vertical_velocity_3d=self.locals.vertical_velocity_3d, + C0=self.plume_dependent_constants.C0, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # get melting profile + # NOTE test GF2020_CumulusParameterization_MeltingProfile_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._melting_profile( + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + melting_layer=self.locals.melting_layer, + partition_liquid_ice=self.locals.partition_liquid_ice, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + melting=self.locals.melting, + ) + + # updraft moist static energy + momentum budget + # NOTE test GF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_moist_static_energy_and_momentum_budget( + error_code=state.output.error_code, + start_level=self.locals.start_level, + cloud_top_level=state.output.cloud_top_level, + p_forced=state.input_output.p_forced, + environment_moist_static_energy=self.locals.environment_moist_static_energy, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=self.locals.environment_moist_static_energy_cloud_levels, + environment_moist_static_energy_cloud_levels_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels=self.locals.environment_saturation_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + normalized_massflux_updraft=self.locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + mass_entrainment_updraft=self.locals.mass_entrainment_updraft, + mass_detrainment_updraft=self.locals.mass_detrainment_updraft, + mass_entrainment_u_updraft=self.locals.mass_entrainment_u_updraft, + mass_detrainment_u_updraft=self.locals.mass_detrainment_u_updraft, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + u=state.input_output.u, + v=state.input_output.v, + u_c=self.locals.u_c, + v_c=self.locals.v_c, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + partition_liquid_ice=self.locals.partition_liquid_ice, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + vapor_excess=self.locals.vapor_excess, + t_excess=self.locals.t_excess, + add_buoyancy=self.locals.add_buoyancy, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # Get buoyancy of updrafts + # NOTE test GF2020_CumulusParameterization_GetBuoyancy_2_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._get_buoyancy( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy, + environment_moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_cloud_levels, + d_buoyancy=self.locals.d_buoyancy, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # NOTE test GF2020_CumulusParameterization_GetBuoyancy_3_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._get_buoyancy( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy_forced, + environment_moist_static_energy=self.locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + d_buoyancy=self.locals.d_buoyancy_forced, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + if not self.cumulus_parameterization_config.FIRST_GUESS_W: + # calculate in-cloud/updraft air temperature for vertical velocity + # NOTE test GF2020_CumulusParameterization_UpdraftTemperature_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_temperature( + error_code=state.output.error_code, + updraft_column_temperature_forced=self.locals.updraft_column_temperature_forced, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=self.locals.cloud_total_water_after_entrainment_forced, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # vertical velocity + # NOTE test GF2020_CumulusParameterization_UpdraftVerticalVelocity_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_vertical_velocity( + vertical_velocity_3d=self.locals.vertical_velocity_3d, + vertical_velocity_2d=self.locals.vertical_velocity_2d, + convective_scale_velocity=state.input_output.convective_scale_velocity, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=self.locals.detrainment_function_updraft, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + updraft_column_temperature_forced=self.locals.updraft_column_temperature_forced, + cloud_total_water_after_entrainment_forced=self.locals.cloud_total_water_after_entrainment_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + vapor_forced=self.locals.vapor_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # downdraft origin level + # NOTE test GF2020_CumulusParameterization_DowndraftOriginLevel_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_origin_level( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + updraft_origin_level=state.output.updraft_origin_level, + downdraft_origin_level=state.output.downdraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + detrainment_start_level=self.locals.detrainment_start_level, + melting_layer=self.locals.melting_layer, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # downdraft normalized mass flux + # NOTE test GF2020_CumulusParameterization_DowndraftMassFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_mass_flux( + error_code=state.output.error_code, + detrainment_start_level=self.locals.detrainment_start_level, + downdraft_origin_level=state.output.downdraft_origin_level, + pbl_level=state.input_output.pbl_level, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + lcl_level=state.output.lcl_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + normalized_massflux_downdraft=self.locals.normalized_massflux_downdraft, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + ocean_fraction=state.input.ocean_fraction, + random_number=self.locals.random_number, + DOWNDRAFT_MAX_HEIGHT_LAND=self.plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_LAND, + DOWNDRAFT_MAX_HEIGHT_OCEAN=self.plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_OCEAN, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # lateral mass fluxes associated with downdrafts + # NOTE test GF2020_CumulusParameterization_DowndraftLateralMassFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_lateral_mass_flux( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_downdraft=self.locals.normalized_massflux_downdraft, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + normalized_massflux_downdraft_modified=self.locals.normalized_massflux_downdraft_modified, + detrainment_function_downdraft=self.locals.detrainment_function_downdraft, + entrainment_rate_downdraft=self.locals.entrainment_rate_downdraft, + mass_entrainment_downdraft=self.locals.mass_entrainment_downdraft, + mass_detrainment_downdraft=self.locals.mass_detrainment_downdraft, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + mass_entrainment_u_downdraft=self.locals.mass_entrainment_u_downdraft, + mass_detrainment_u_downdraft=self.locals.mass_detrainment_u_downdraft, + LAMBDA_DOWN=self.plume_dependent_constants.LAMBDA_DOWN, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # wet bulb temperature and moisture at downdraft origin level + # NOTE this section does not run in the test case, and has not been implemented. + # NOTE an error will stop execution during initalization if this would be called + if ( + self.cumulus_parameterization_config.USE_WETBULB + and self.plume_dependent_constants.PLUME_INDEX != 0 + ): + raise NotImplementedError( + "wet bulb functionality is not implemented. You should not be here," + "there are multiple layers of errors that should have caught you first." + "If you are seeing this (at runtime), seek help." + ) + + # downdraft moist static energy + moisture budget + # NOTE test GF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_moist_static_energy_and_buoyancy( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + u=state.input_output.u, + u_cloud_levels=self.locals.u_cloud_levels, + u_c_downdraft=self.locals.u_c_downdraft, + v=state.input_output.v, + v_cloud_levels=self.locals.v_cloud_levels, + v_c_downdraft=self.locals.v_c_downdraft, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy, + cloud_moist_static_energy_downdraft_forced=self.locals.cloud_moist_static_energy_downdraft_forced, + buoyancy_downdraft_forced=self.locals.d_buoyancy_downdraft_forced, + t_wetbulb=self.locals.t_wetbulb, + vapor_wetbulb=self.locals.vapor_wetbulb, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + mass_entrainment_u_downdraft=self.locals.mass_entrainment_u_downdraft, + mass_detrainment_u_downdraft=self.locals.mass_detrainment_u_downdraft, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # calculate moisture properties of downdraft + # NOTE test GF2020_CumulusParameterization_DowndraftMoisture_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_moisture( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + t_wetbulb=self.locals.t_wetbulb, + vapor_forced=self.locals.vapor_forced, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=self.locals.environment_saturation_mixing_ratio_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=self.locals.cloud_total_water_after_entrainment_forced, + cloud_total_water_after_entrainment_downdraft_forced=self.locals.cloud_total_water_after_entrainment_downdraft_forced, + downdraft_saturation_vapor_forced=self.locals.downdraft_saturation_vapor_forced, + vapor_wetbulb=self.locals.vapor_wetbulb, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy_downdraft_forced=self.locals.cloud_moist_static_energy_downdraft_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + gamma_cloud_levels_forced=self.locals.gamma_cloud_levels_forced, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=self.locals.total_normalized_integrated_evaporate_forced, + buoyancy=self.locals.buoyancy, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # calculate workfunctions for updrafts + # NOTE test GF2020_CumulusParameterization_UpdraftInitialWorkfunctions_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_initial_workfunctions( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=self.locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + d_buoyancy=self.locals.d_buoyancy, + d_buoyancy_forced=self.locals.d_buoyancy_forced, + gamma_cloud_levels=self.locals.gamma_cloud_levels, + gamma_cloud_levels_forced=self.locals.gamma_cloud_levels_forced, + t_cloud_levels=self.locals.t_cloud_levels, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + cloud_workfunction_0=self.locals.cloud_workfunction_0, + cloud_workfunction_1=self.locals.cloud_workfunction_1, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # calculate CIN for updrafts + # NOTE test GF2020_CumulusParameterization_UpdraftCIN_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._updraft_cin( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=self.locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + d_buoyancy=self.locals.d_buoyancy, + d_buoyancy_forced=self.locals.d_buoyancy_forced, + gamma_cloud_levels=self.locals.gamma_cloud_levels, + gamma_cloud_levels_forced=self.locals.gamma_cloud_levels_forced, + t_cloud_levels=self.locals.t_cloud_levels, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + cin_0=self.locals.cin_0, + cin_1=self.locals.cin_1, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # trigger function: KE+CIN < 0 --> no convection + # NOTE test GF2020_CumulusParameterization_ConvectionTrigger_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._convection_trigger( + error_code=state.output.error_code, + convective_scale_velosity=state.input_output.convective_scale_velocity, + cin_0=self.locals.cin_0, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # calculate downdraft air temperature for vertical velocitys + # NOTE test GF2020_CumulusParameterization_DowndraftTemperature_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._downdraft_temperature( + error_code=state.output.error_code, + downdraft_column_temperature_forced=self.locals.downdraft_column_temperature_forced, + cloud_moist_static_energy_downdraft_forced=self.locals.cloud_moist_static_energy_downdraft_forced, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_downdraft_forced=self.locals.cloud_total_water_after_entrainment_downdraft_forced, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # diurnal cycle section + # NOTE test GF2020_CumulusParameterization_DiurnalCycle_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._diurnal_cycle( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + grid_length=state.input_output.grid_length, + ocean_fraction=state.input.ocean_fraction, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + t_old=state.input_output.t_old, + t_new=self.locals.t_new, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + vapor_old=state.input_output.vapor_old, + vapor_forced=self.locals.vapor_forced, + u=state.input_output.u, + v=state.input_output.v, + vertical_velocity_2d=self.locals.vertical_velocity_2d, + cape_removal_time_scale=self.locals.cape_removal_time_scale, + cape_removal_time_scale_from_state=state.output.cape_removal_time_scale, + pbl_time_scale=self.locals.pbl_time_scale, + pbl_time_scale_from_state=state.output.pbl_time_scale, + cloud_work_function_1_pbl=self.locals.cloud_workfunction_1_pbl, + cloud_work_function_1_fa=self.locals.cloud_workfunction_1_fa, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # Trigger function based on Xie et al 2019 + # NOTE not implemented, does not run with test config + self._Xie_trigger_function(plume_dependent_constants=self.plume_dependent_constants) + + # determine downdraft strength in terms of windshear + # NOTE test GF2020_CumulusParameterization_DowndraftWindShear_{plume}: + # NOTE deep ✅ + # NOTE mid ❌ one field, one point (0.17%), 4 ULP + # NOTE shallow ✅ + self._downdraft_windshear( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_forced=state.input_output.geopotential_height_forced, + p_forced=state.input_output.p_forced, + u=state.input_output.u, + v=state.input_output.v, + ccn=state.input_output.ccn, + psum=self.locals.psum, + psumh=self.locals.psumh, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=self.locals.total_normalized_integrated_evaporate_forced, + scale_dependence_factor_downdraft=self.locals.scale_dependence_factor_downdraft, + epsilon=self.locals.epsilon, + epsilon_min=self.locals.epsilon_min, + epsilon_max=self.locals.epsilon_max, + epsilon_computed=self.locals.epsilon_computed, + epsilon_forced=state.output.epsilon_forced, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # get the environmental mass flux + # NOTE test GF2020_CumulusParameterization_EnvironmentMassFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._environment_mass_flux( + error_code=state.output.error_code, + epsilon_forced=state.output.epsilon_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_massflux=self.locals.environment_massflux, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # check mass conservation + # NOTE This code runs in the Fortran and only has one output: totmas (total mass). + # totmas has only one use: a conditional log write if total mass is above a 1e-6 + # with a disabled (commented) fatal error call. + # Since the only consequential outcome is disabled, and this port has thus far not + # implemented other log writes, this code not been implemented. + # If totmas is needed in the future, or this fatal call is reimplemented, + # this code will be revisited + self._mass_conservation() + + # change per unit mass that a model cloud would modify the environment + # NOTE test GF2020_CumulusParameterization_VerticalDiscretization_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._vertical_discretization( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels_forced=self.locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_massflux=self.locals.environment_massflux, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + c1d=self.locals.c1d, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + u_c=self.locals.u_c, + v_c=self.locals.v_c, + u_c_downdraft=self.locals.u_c_downdraft, + v_c_downdraft=self.locals.v_c_downdraft, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + cloud_moist_static_energy_downdraft_forced=self.locals.cloud_moist_static_energy_downdraft_forced, + environment_moist_static_energy_cloud_levels_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=self.locals.cloud_total_water_after_entrainment_forced, + cloud_total_water_after_entrainment_downdraft_forced=self.locals.cloud_total_water_after_entrainment_downdraft_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + melting=self.locals.melting, + partition_liquid_ice=self.locals.partition_liquid_ice, + epsilon_forced=state.output.epsilon_forced, + d_buoyancy_downdraft_forced=self.locals.d_buoyancy_downdraft_forced, + del_u_cloud_ensemble=self.locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=self.locals.del_v_cloud_ensemble, + del_moist_static_energy_cloud_ensemble=self.locals.del_moist_static_energy_cloud_ensemble, + del_t_cloud_ensemble=self.locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=self.locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=self.locals.del_cloud_liquid_cloud_ensemble, + del_buoyancy_cloud_ensemble=self.locals.del_buoyancy_cloud_ensemble, + t_tendency_from_environmental_subsidence=self.locals.t_tendency_from_environmental_subsidence, + moist_static_energy_tendency_from_environmental_subsidence=self.locals.moist_static_energy_tendency_from_environmental_subsidence, + vapor_tendency_from_environmental_subsidence=self.locals.vapor_tendency_from_environmental_subsidence, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # apply environmental subsidence on grid-scale ice and + # liq water contents, and cloud fraction (Upwind scheme) + # NOTE not implemented, does not run with test config + self._environmental_subsidence() + + # make the smoothness procedure + # NOTE test GF2020_CumulusParameterization_SmoothTendencies_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._smooth_tendencies( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + del_moist_static_energy_cloud_ensemble=self.locals.del_moist_static_energy_cloud_ensemble, + del_vapor_cloud_ensemble=self.locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=self.locals.del_cloud_liquid_cloud_ensemble, + del_u_cloud_ensemble=self.locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=self.locals.del_v_cloud_ensemble, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # using smoothed tendencies, calculate changed environmental profiles + # NOTE test GF2020_CumulusParameterization_ModifyEnvironmentProfiles_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._modify_environment_profiles( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + updraft_origin_level=state.output.updraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + p_forced=state.input_output.p_forced, + t_new=self.locals.t_new, + t_modified=self.locals.t_modified, + vapor_forced=self.locals.vapor_forced, + vapor_modified=self.locals.vapor_modified, + environment_moist_static_energy_forced=self.locals.environment_moist_static_energy_forced, + environment_moist_static_energy_modified=self.locals.environment_moist_static_energy_modified, + moist_static_energy_origin_level_forced=self.locals.moist_static_energy_origin_level_forced, + moist_static_energy_origin_level_modified=self.locals.moist_static_energy_origin_level_modified, + partition_liquid_ice=self.locals.partition_liquid_ice, + del_moist_static_energy_cloud_ensemble=self.locals.del_moist_static_energy_cloud_ensemble, + del_t_cloud_ensemble=self.locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=self.locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=self.locals.del_cloud_liquid_cloud_ensemble, + del_u_cloud_ensemble=self.locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=self.locals.del_v_cloud_ensemble, + moist_static_energy_tendency_from_environmental_subsidence=self.locals.moist_static_energy_tendency_from_environmental_subsidence, + vapor_tendency_from_environmental_subsidence=self.locals.vapor_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=self.locals.t_tendency_from_environmental_subsidence, + arbitrary_numerical_parameter=self.locals.arbitrary_numerical_parameter, + AVERAGE_LAYER_DEPTH=self.plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # calculate moist static energy, heights, environmental saturation mixing ratio + # NOTE test GF2020_CumulusParameterization_EnvironmentConditions_3_{plume}: + # NOTE deep ✅ + # NOTE mid ❌ one field, one point (0.0%), 2 ULP + # NOTE shallow ✅ + self._environment_conditions( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=self.locals.t_modified, + vapor=self.locals.vapor_modified, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=self.locals.environment_moist_static_energy_modified, + saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_modified, + saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio_modified, + geopotential_height=self.locals.geopotential_height_modified, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # environmental values on cloud levels + # NOTE test GF2020_CumulusParameterization_EnvironmentCloudLevels_3_{plume}: + # NOTE deep ✅ + # NOTE mid ❌ one field, one point (0.0%), 32 ULP + # NOTE shallow ❌ one field, one point (0.0%), 32 ULP + self._environment_cloud_levels( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=state.output.p_cloud_levels_forced.field[ + :, :, :, self.plume_dependent_constants.PLUME_INDEX + ], + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=self.locals.geopotential_height_modified, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels_modified, + t=self.locals.t_modified, + t_surface=state.input_output.t_surface, + t_cloud_levels=self.locals.t_cloud_levels_modified, + vapor=self.locals.vapor_modified, + vapor_cloud_levels=self.locals.vapor_cloud_levels_modified, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=self.locals.u_cloud_levels, + v_cloud_levels=self.locals.v_cloud_levels, + environment_saturation_mixing_ratio=self.locals.environment_saturation_mixing_ratio_modified, + environment_saturation_mixing_ratio_cloud_levels=self.locals.environment_saturation_mixing_ratio_cloud_levels_modified, + environment_moist_static_energy=self.locals.environment_moist_static_energy_modified, + environment_moist_static_energy_cloud_levels=self.locals.environment_moist_static_energy_cloud_levels_modified, + environment_saturation_moist_static_energy=self.locals.environment_saturation_moist_static_energy_modified, + environment_saturation_moist_static_energy_cloud_levels=self.locals.environment_saturation_moist_static_energy_cloud_levels_modified, + gamma_cloud_levels=self.locals.gamma_cloud_levels, + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # static control + # NOTE test GF2020_CumulusParameterization_StaticControl_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._static_control( + error_code=state.output.error_code, + start_level=self.locals.start_level, + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy_modified=self.locals.cloud_moist_static_energy_modified, + moist_static_energy_origin_level_modified=self.locals.moist_static_energy_origin_level_modified, + environment_moist_static_energy_modified=self.locals.environment_moist_static_energy_modified, + environment_moist_static_energy_cloud_levels_modified=self.locals.environment_moist_static_energy_cloud_levels_modified, + environment_saturation_moist_static_energy_cloud_levels_modified=self.locals.environment_saturation_moist_static_energy_cloud_levels_modified, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + normalized_massflux_updraft_modified=self.locals.normalized_massflux_updraft_modified, + partition_liquid_ice=self.locals.partition_liquid_ice, + vapor_excess=self.locals.vapor_excess, + t_excess=self.locals.t_excess, + add_buoyancy=self.locals.add_buoyancy, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + d_buoyancy_modified=self.locals.d_buoyancy_modified, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # workfunctions for updraft + # NOTE test GF2020_CumulusParameterization_UpdraftWorkfunctions_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._update_workfunction_and_precipitaiton_ensemble( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels_modified=self.locals.geopotential_height_cloud_levels_modified, + normalized_massflux_updraft_modified=self.locals.normalized_massflux_updraft_modified, + d_buoyancy_modified=self.locals.d_buoyancy_modified, + gamma_cloud_levels=self.locals.gamma_cloud_levels, + t_cloud_levels_modified=self.locals.t_cloud_levels_modified, + cloud_workfunction_0_modified=self.locals.cloud_workfunction_0_modified, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + epsilon_forced=state.output.epsilon_forced, + precipitation_ensemble=self.locals.precipitation_ensemble, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # large scale forcing + # calculate cloud base mass flux + # NOTE test GF2020_CumulusParameterization_LargeScaleForcing_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._large_scale_forcing( + error_code=state.output.error_code, + error_code_2=self.locals.error_code_2, + error_code_3=self.locals.error_code_3, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + ocean_fraction=self.locals.ocean_fraction, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + vapor_forced=self.locals.vapor_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + effective_condensate_to_fall_forced=self.locals.effective_condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + omega=state.input_output.omega, + convective_scale_velocity=state.input_output.convective_scale_velocity, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + cloud_moist_static_energy=self.locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=self.locals.cloud_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=self.locals.environment_moist_static_energy_cloud_levels, + environment_moist_static_energy_cloud_levels_forced=self.locals.environment_moist_static_energy_cloud_levels_forced, + dmoist_static_energydt=self.locals.dmoist_static_energydt, + cloud_workfunction_0=self.locals.cloud_workfunction_0, + cloud_workfunction_0_modified=self.locals.cloud_workfunction_0_modified, + cloud_workfunction_1=self.locals.cloud_workfunction_1, + cloud_workfunction_1_pbl=self.locals.cloud_workfunction_1_pbl, + arbitrary_numerical_parameter=self.locals.arbitrary_numerical_parameter, + f_dicycle_modified=self.locals.f_dicycle_modified, + cape_removal_time_scale=self.locals.cape_removal_time_scale, + epsilon_forced=state.output.epsilon_forced, + k_x_modified=self.locals.k_x_modified, + mass_flux_ensemble=self.locals.mass_flux_ensemble, + precipitation_ensemble=self.locals.precipitation_ensemble, + xff_mid=self.locals.xff_mid, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # Include kinetic energy dissipation converted to heating + # NOTE test GF2020_CumulusParameterization_KeToHeating_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._kinetic_energy_to_heating( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + u=state.input_output.u, + v=state.input_output.v, + del_u_cloud_ensemble=self.locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=self.locals.del_v_cloud_ensemble, + del_t_cloud_ensemble=self.locals.del_t_cloud_ensemble, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # feedback + # NOTE test GF2020_CumulusParameterization_EnsembleOutputAndFeedback_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._ensemble_output_and_feedback( + error_code=state.output.error_code, + error_code_2=self.locals.error_code_2, + error_code_3=self.locals.error_code_3, + cloud_top_level=state.output.cloud_top_level, + updraft_lfc_level=state.output.updraft_lfc_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + precip=state.output.precip, + effective_condensate_to_fall_forced=self.locals.effective_condensate_to_fall_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + scale_dependence_factor=state.output.scale_dependence_factor, + ocean_fraction=self.locals.ocean_fraction, + f_dicycle_modified=self.locals.f_dicycle_modified, + del_u_cloud_ensemble=self.locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=self.locals.del_v_cloud_ensemble, + del_t_cloud_ensemble=self.locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=self.locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=self.locals.del_cloud_liquid_cloud_ensemble, + del_buoyancy_cloud_ensemble=self.locals.del_buoyancy_cloud_ensemble, + del_convective_ice_cloud_ensemble=self.locals.del_convective_ice_cloud_ensemble, + del_large_scale_ice_cloud_ensemble=self.locals.del_large_scale_ice_cloud_ensemble, + del_convective_liquid_cloud_ensemble=self.locals.del_convective_liquid_cloud_ensemble, + del_large_scale_liquid_cloud_ensemble=self.locals.del_large_scale_liquid_cloud_ensemble, + del_convective_cloud_fraction_cloud_ensemble=self.locals.del_convective_cloud_fraction_cloud_ensemble, + del_large_scale_cloud_fraction_cloud_ensemble=self.locals.del_large_scale_cloud_fraction_cloud_ensemble, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dcloudicedt=state.output.dcloudicedt, + dudt=state.output.dudt, + dvdt=state.output.dvdt, + dbuoyancydt=state.output.dbuoyancydt, + dconvectiveicedt=state.output.dconvectiveicedt, + dlargescaleicedt=state.output.dlargescaleicedt, + dconvectiveliquiddt=state.output.dconvectiveliquiddt, + dlargescaleliquiddt=state.output.dlargescaleliquiddt, + dconvectivecloudfractiondt=state.output.dconvectivecloudfractiondt, + dlargescalecloudfractiondt=state.output.dlargescalecloudfractiondt, + mass_flux_ensemble=self.locals.mass_flux_ensemble, + precipitation_ensemble=self.locals.precipitation_ensemble, + xff_mid=self.locals.xff_mid, + CLOSURE_CHOICE=self.plume_dependent_constants.CLOSURE_CHOICE, + CLOUD_BASE_MASS_FLUX_FACTOR=self.plume_dependent_constants.CLOUD_BASE_MASS_FLUX_FACTOR, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # net precipitation flux (after downdraft evaporation) + # NOTE test GF2020_CumulusParameterization_PrecipitationFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._precipitation_flux( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + epsilon_forced=state.output.epsilon_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + precipitation_flux=self.locals.precipitation_flux, + evaporation_flux=self.locals.evaporation_flux, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # rainfall evap below cloud base + # NOTE test GF2020_CumulusParameterization_RainEvaporationBelowCloudBase_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._rain_evaporation_below_cloud_base( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + ocean_fraction=state.input.ocean_fraction, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + t_cloud_levels=self.locals.t_cloud_levels, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels=self.locals.environment_saturation_mixing_ratio_cloud_levels, + epsilon_forced=state.output.epsilon_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + precip=state.output.precip, + precipitation_flux=self.locals.precipitation_flux, + evaporation_flux=self.locals.evaporation_flux, + evaporation_below_cloud_base=self.locals.evaporation_below_cloud_base, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dbuoyancydt=state.output.dbuoyancydt, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # includes effects of the remained cloud dissipation into the enviroment + # NOTE test GF2020_CumulusParameterization_CloudDissipation_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._cloud_dissapation( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + hydrostatic_air_density=self.locals.hydrostatic_air_density, + geopotential_height_forced=state.input_output.geopotential_height_forced, + t_cloud_levels_forced=self.locals.t_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + vapor_cloud_levels_forced=self.locals.vapor_cloud_levels_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=self.locals.environment_saturation_mixing_ratio_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=self.locals.environment_saturation_moist_static_energy_cloud_levels_forced, + vertical_velocity_3d=self.locals.vertical_velocity_3d, + scale_dependence_factor=state.output.scale_dependence_factor, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dcloudicedt=state.output.dcloudicedt, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # total (deep+mid) evaporation flux for output (units kg/kg/s) + # NOTE test GF2020_CumulusParameterization_TotalEvaporationFlux_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._total_evaporation_flux( + error_code=state.output.error_code, + plume=self.plume_dependent_constants.PLUME_INDEX, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + evaporation_flux=self.locals.evaporation_flux, + evaporation_sublimation_tendency=state.output.evaporation_sublimation_tendency, + ) + + # lightning flashes density (parameterization from Lopez 2016, MWR) + # NOTE this section does not run in the test case, and has not been implemented. + self._lightning_flash_density() + + # output precipitation (only deep plume) + # NOTE test GF2020_CumulusParameterization_DeepPrecipitationOutput_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._deep_precipitation_output( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + precipitation_flux=self.locals.precipitation_flux, + convective_precip_flux=state.output.convective_precip_flux, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # for tracer convective transport / outputs + # NOTE test GF2020_CumulusParameterization_TracerOutput_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._output_updraft_temperature( + error_code=state.output.error_code, + updraft_column_temperature_forced=self.locals.updraft_column_temperature_forced, + t_cloud_levels=self.locals.t_cloud_levels, + t_updraft=state.output.t_updraft, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # convert mass fluxes, etc... + # NOTE test GF2020_CumulusParameterization_PrepareOutput_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._prepare_output( + error_code=state.output.error_code, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=self.locals.total_normalized_integrated_evaporate_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + environment_massflux=self.locals.environment_massflux, + vapor_tendency_from_environmental_subsidence=self.locals.vapor_tendency_from_environmental_subsidence, + moist_static_energy_tendency_from_environmental_subsidence=self.locals.moist_static_energy_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=self.locals.t_tendency_from_environmental_subsidence, + plume=self.plume_dependent_constants.PLUME_INDEX, + ) + + # outputs a model sounding for the stand-alone code (part 2) + # NOTE this section does not run in the test case, and has not been implemented. + self._sounding() + + # NOTE test GF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_{plume}: + # NOTE deep ❌ Fails dnliquiddt and dnicedt. Fortran bug, not carried to python + # NOTE mid ❌ Fails dnliquiddt and dnicedt. Fortran bug, not carried to python + # NOTE shallow ✅ + self._output_workfunctions_and_precip_concentrations( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + convection_fraction=state.input.convection_fraction, + surface_type=state.input.surface_type, + cloud_workfunction_0_output=state.output.cloud_workfunction_0, + cloud_workfunction_1_output=state.output.cloud_workfunction_1, + cloud_workfunction_0=self.locals.cloud_workfunction_0, + cloud_workfunction_1=self.locals.cloud_workfunction_1, + air_density=state.input_output.air_density, + updraft_column_temperature_forced=self.locals.updraft_column_temperature_forced, + dcloudicedt=state.output.dcloudicedt, + dnliquiddt=state.output.dnliquiddt, + dnicedt=state.output.dnicedt, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # section for atmospheric composition + # NOTE test GF2020_CumulusParameterization_AtmosphericComposition_{plume}: + # NOTE deep ✅ + # NOTE mid ✅ + # NOTE shallow ✅ + self._atmospheric_composition( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + updraft_origin_level=state.output.updraft_origin_level, + downdraft_origin_level=state.output.downdraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + p_forced=state.input_output.p_forced, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels=self.locals.geopotential_height_cloud_levels, + environment_massflux=self.locals.environment_massflux, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + vertical_velocity_3d=self.locals.vertical_velocity_3d, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=self.locals.total_normalized_integrated_evaporate_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + epsilon_forced=state.output.epsilon_forced, + chemistry_tracers=state.input_output.chemistry_tracers, + chemistry_tracers_output=state.input_output.chemistry_tracers_output, + chemistry_tracers_cloud_levels=self.locals.chemistry_tracers_cloud_levels, + chemistry_tracers_sc_updraft=self.locals.chemistry_tracers_sc_updraft, + chemistry_tracers_sc_downdraft=self.locals.chemistry_tracers_sc_downdraft, + chemistry_tracers_pw_updraft=self.locals.chemistry_tracers_pw_updraft, + chemistry_tracers_pw_downdraft=self.locals.chemistry_tracers_pw_downdraft, + chemistry_tracers_total_pw_updraft=self.locals.chemistry_tracers_total_pw_updraft, + chemistry_tracers_total_pw_downdraft=self.locals.chemistry_tracers_total_pw_downdraft, + convection_tracers=convection_tracers, + plume_dependent_constants=self.plume_dependent_constants, + ) + + # begin: for GATE soundings + # NOTE this section does not run in the test case, and has not been implemented. + self._gate_sounding() diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/diurnal_cycle.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/diurnal_cycle.py new file mode 100644 index 000000000..d4db997df --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/diurnal_cycle.py @@ -0,0 +1,309 @@ +import pyMoist.constants as constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, computation, interval, sqrt +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import IntFieldIJ_Plume +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.shared_generic_math import sigma + + +def set_time_scales( + error_code: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + grid_length: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + topography_height_no_negative: FloatFieldIJ, + geopotential_height_cloud_levels_forced: FloatField, + u: FloatField, + v: FloatField, + vertical_velocity_2d: FloatFieldIJ, + cape_removal_time_scale: FloatFieldIJ, + cape_removal_time_scale_from_state: FloatFieldIJ, + pbl_time_scale: FloatFieldIJ, + pbl_time_scale_from_state: FloatFieldIJ, + TAU_CAPE_REMOVAL: Float, + plume: Int, +): + """Set the timescale for cape removal and pbl resonance. + + Args: + error_code (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + grid_length (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + topography_height_no_negative (FloatFieldIJ) + geopotential_height_cloud_levels_forced (FloatField) + u (FloatField) + v (FloatField) + vertical_velocity_2d (FloatFieldIJ) + cape_removal_time_scale (FloatFieldIJ) + cape_removal_time_scale_from_state (FloatFieldIJ) + pbl_time_scale (FloatFieldIJ) + pbl_time_scale_from_state (FloatFieldIJ) + TAU_CAPE_REMOVAL (Float) + plume (Int) + """ + from __externals__ import DTIME, SGS_W_TIMESCALE + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if SGS_W_TIMESCALE == 0: + # time-scale cape removal from Bechtold et al. 2008 + dz = geopotential_height_cloud_levels_forced.at( + K=cloud_top_level[0, 0][plume] + ) - geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + else: + # time-scale cape removal from Bechtold et al. 2008 + dz = geopotential_height_cloud_levels_forced.at( + K=cloud_top_level[0, 0][plume] + ) - geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + cape_removal_time_scale = ( + 3600.0 * (sigma(grid_length)) + + 10800.0 * (1.0 - sigma(grid_length)) + + (dz / vertical_velocity_2d) + ) + cape_removal_time_scale = max(DTIME, cape_removal_time_scale) + + if ocean_fraction > 0.99: # over water + umean = 2.0 + sqrt( + 0.5 + * ( + u**2 + + v**2 + + u.at(K=updraft_lfc_level[0, 0][plume]) ** 2 + + v.at(K=updraft_lfc_level[0, 0][plume]) ** 2 + ) + ) + pbl_time_scale = ( + geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + - topography_height_no_negative + ) / umean + else: # over land + pbl_time_scale = ( + geopotential_height_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume]) + - geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + ) / 3.0 # 3.0 m/s is estimated wmean + + with computation(FORWARD), interval(0, 1): + cape_removal_time_scale_from_state = cape_removal_time_scale + pbl_time_scale_from_state = pbl_time_scale + + +def cloud_workfunction_1_pbl( + error_code: IntFieldIJ_Plume, + pbl_level: IntFieldIJ, + geopotential_height_cloud_levels_forced: FloatField, + t_old: FloatField, + t_new: FloatField, + t_cloud_levels_forced: FloatField, + vapor_old: FloatField, + vapor_forced: FloatField, + cloud_work_function_1_pbl: FloatFieldIJ, + cloud_work_function_1_fa: FloatFieldIJ, + plume: Int, +): + """Compute the initial estimates for cloud_workfunction_1_pbl and cloud_workfunction_1_fa. + + Args: + error_code (IntFieldIJ_Plume) + pbl_level (IntFieldIJ) + geopotential_height_cloud_levels_forced (FloatField) + t_old (FloatField) + t_new (FloatField) + t_cloud_levels_forced (FloatField) + vapor_old (FloatField) + vapor_forced (FloatField) + cloud_work_function_1_pbl (FloatFieldIJ) + cloud_work_function_1_fa (FloatFieldIJ) + plume (Int) + """ + from __externals__ import DTIME + + with computation(FORWARD), interval(0, 1): + cloud_work_function_1_pbl = 0.0 + cloud_work_function_1_fa = 0.0 + + top_bound: IntFieldIJ = pbl_level + 1 + + with computation(FORWARD), interval(0, top_bound): + if error_code[0, 0][plume] == 0: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) * constants.MAPL_GRAV + da = dz * (t_new * (1.0 + 0.608 * vapor_forced) - t_old * (1.0 + 0.608 * vapor_old)) / DTIME + + cloud_work_function_1_pbl = cloud_work_function_1_pbl + da # Units : J K / (kg seg) + + +def scale_cloud_workfunction_1_pbl( + error_code: IntFieldIJ_Plume, + pbl_time_scale: FloatFieldIJ, + cloud_work_function_1_pbl: FloatFieldIJ, + T_STAR: Float, + plume: Int, +): + """Scale cloud_workfunction_1_pbl based on the pbl_time_scale. + + Args: + error_code (IntFieldIJ_Plume) + pbl_time_scale (FloatFieldIJ) + cloud_work_function_1_pbl (FloatFieldIJ) + T_STAR (Float) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + cloud_work_function_1_pbl = cloud_work_function_1_pbl / T_STAR * pbl_time_scale + + +class DiurnalCycle(NDSLRuntime): + """Determine the effects of the diurnal cycle on PBL processes.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._set_time_scales = stencil_factory.from_dims_halo( + func=set_time_scales, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SGS_W_TIMESCALE": cumulus_parameterization_config.SGS_W_TIMESCALE, + "DTIME": cumulus_parameterization_config.DTIME, + }, + ) + + self._cloud_workfunction_1_pbl = stencil_factory.from_dims_halo( + func=cloud_workfunction_1_pbl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SGS_W_TIMESCALE": cumulus_parameterization_config.SGS_W_TIMESCALE, + "DTIME": cumulus_parameterization_config.DTIME, + }, + ) + + self._scale_cloud_workfunction_1_pbl = stencil_factory.from_dims_halo( + func=scale_cloud_workfunction_1_pbl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # initialize internal fields + self._tau_ecmwf: Local = quantity_factory.zeros([X_DIM, Y_DIM], "n/a") + self._tau_bl: Local = quantity_factory.zeros([X_DIM, Y_DIM], "n/a") + + def __call__( + self, + error_code: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + pbl_level: Quantity, + grid_length: Quantity, + ocean_fraction: Quantity, + topography_height_no_negative: Quantity, + geopotential_height_cloud_levels_forced: Quantity, + t_old: Quantity, + t_new: Quantity, + t_cloud_levels_forced: Quantity, + vapor_old: Quantity, + vapor_forced: Quantity, + u: Quantity, + v: Quantity, + vertical_velocity_2d: Quantity, + cape_removal_time_scale: Quantity, + cape_removal_time_scale_from_state: Quantity, + pbl_time_scale: Quantity, + pbl_time_scale_from_state: Quantity, + cloud_work_function_1_pbl: Quantity, + cloud_work_function_1_fa: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + # Bechtold et al 2008 time-scale of cape removal + self._set_time_scales( + error_code=error_code, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + grid_length=grid_length, + ocean_fraction=ocean_fraction, + topography_height_no_negative=topography_height_no_negative, + geopotential_height_cloud_levels_forced=geopotential_height_cloud_levels_forced, + u=u, + v=v, + vertical_velocity_2d=vertical_velocity_2d, + cape_removal_time_scale=cape_removal_time_scale, + cape_removal_time_scale_from_state=cape_removal_time_scale_from_state, + pbl_time_scale=pbl_time_scale, + pbl_time_scale_from_state=pbl_time_scale_from_state, + TAU_CAPE_REMOVAL=plume_dependent_constants.TAU_CAPE_REMOVAL, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if ( + self.cumulus_parameterization_config.DIURNAL_CYCLE == 1 + or self.cumulus_parameterization_config.DIURNAL_CYCLE == 6 + ) or ( + self.cumulus_parameterization_config.DIURNAL_CYCLE == 0 + and plume_dependent_constants.PLUME_INDEX == 1 # mid plume + ): + # calculate pcape from BL forcing only + self._cloud_workfunction_1_pbl( + error_code=error_code, + pbl_level=pbl_level, + geopotential_height_cloud_levels_forced=geopotential_height_cloud_levels_forced, + t_old=t_old, + t_new=t_new, + t_cloud_levels_forced=t_cloud_levels_forced, + vapor_old=vapor_old, + vapor_forced=vapor_forced, + cloud_work_function_1_pbl=cloud_work_function_1_pbl, + cloud_work_function_1_fa=cloud_work_function_1_fa, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if self.cumulus_parameterization_config.DIURNAL_CYCLE == 6: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->DiurnalCycle called with an unimplemented" + "path. This should have been caught at initialization, but somehow you made it here." + "Choose another option for DIURNAL_CYCLE or implement to continue." + ) + + self._scale_cloud_workfunction_1_pbl( + error_code=error_code, + pbl_time_scale=pbl_time_scale, + cloud_work_function_1_pbl=cloud_work_function_1_pbl, + T_STAR=plume_dependent_constants.T_STAR, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + elif self.cumulus_parameterization_config.DIURNAL_CYCLE == 4: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->DiurnalCycle called with an unimplemented path." + "This should have been caught at initialization, but somehow you made it here." + "Choose another option or for DIURNAL_CYCLE implement to continue." + ) + + if ( + self.cumulus_parameterization_config.DIURNAL_CYCLE == 5 + or self.cumulus_parameterization_config.DIURNAL_CYCLE == 2 + ): + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->DiurnalCycle called with an unimplemented path." + "This should have been caught at initialization, but somehow you made it here." + "Choose another option or for DIURNAL_CYCLE implement to continue." + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/downdraft.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/downdraft.py new file mode 100644 index 000000000..eb649b9ce --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/downdraft.py @@ -0,0 +1,1229 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, K, computation, interval +from ndsl.dsl.typing import BoolFieldIJ, Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max, column_max_ddim, column_min +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_ensemble_2, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import generic_find_level + + +def get_critical_level( + error_code: IntFieldIJ_Plume, + critical_level: IntFieldIJ, + cloud_top_level: IntFieldIJ_Plume, + geopotential_height_cloud_levels_forced: FloatField, + topography_height_no_negative: FloatFieldIJ, + MAX_DOWNDRAFT_ORIGIN_HEIGHt: Float, + plume: Int, +): + """Determine the maximum height at which a downdraft may start (heavily dependent on, + but not necessarily equal to cloud_top_level). + + Args: + error_code (IntFieldIJ_Plume) + critical_level (IntFieldIJ) + cloud_top_level (IntFieldIJ_Plume) + geopotential_height_cloud_levels_forced (FloatField) + topography_height_no_negative (FloatFieldIJ) + MAX_DOWNDRAFT_ORIGIN_HEIGHt (Float) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + cloud_top_height: FloatFieldIJ = ( + geopotential_height_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume]) + - topography_height_no_negative + ) * 0.6 + cloud_top_height = min( + cloud_top_height + topography_height_no_negative, + MAX_DOWNDRAFT_ORIGIN_HEIGHt + topography_height_no_negative, + ) + + # setup mask to stop the next block + stop_computation: BoolFieldIJ = False + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and stop_computation == False: + if geopotential_height_cloud_levels_forced >= cloud_top_height: + critical_level = K + stop_computation = True + + +def fill_plume_data_dimension( + data: IntFieldIJ, + data_dimension_field: IntFieldIJ_Plume, + plume: Int, +): + with computation(FORWARD), interval(0, 1): + data_dimension_field[0, 0][plume] = data + + +def fill_from_plume_data_dimension( + data: IntFieldIJ, + data_dimension_field: IntFieldIJ_Plume, + plume: Int, +): + with computation(FORWARD), interval(0, 1): + data = data_dimension_field[0, 0][plume] + + +def get_downdraft_origin_level( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + detrainment_start_level: IntFieldIJ, + downdraft_origin_level: IntFieldIJ_Plume, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + geopotential_height_cloud_levels_forced: FloatField, + melting_layer: FloatField, + MINIMUM_DEPTH: Float, + plume: Int, +): + """Determine the leve at which the downdraft begins. + For shallow plume, return 0 (downdraft is disabled). For mid and deep plume, perform full calculation. + + This stencil contains an open-ended vertical solver with a nested vertical loop. + To implement this properly, the aforrementioned computation has been constructed on an interval(0, 1), + and all K read/writes have been done with absolute indexes or relative offsets. The alternative is + to break this into a series of stencils and pass data between them using a much larger number of locals. + + Following a decision made in a prior section of the GF2020 cumulus parameterizaiton core + see (get_convective_cloud_base_level), this has been implemented as one stencil using the inefficient + interval to preserve code readability and facilitate an easy future transition to another implementation, + once the required tool/feature has been implemented. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + detrainment_start_level (IntFieldIJ) + downdraft_origin_level (IntFieldIJ_Plume) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + geopotential_height_cloud_levels_forced (FloatField) + melting_layer (FloatField) + MINIMUM_DEPTH (Float) + plume (Int) + """ + from __externals__ import MELT_GLAC, k_end + + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.SHALLOW: + downdraft_origin_level[0, 0][plume] = 0 + elif plume == cumulus_parameterization_constants.MID: + # setup internal constants + beta: FloatFieldIJ = 0.02 + elif plume == cumulus_parameterization_constants.DEEP: + # setup internal constants + beta: FloatFieldIJ = 0.05 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # predefine field for next block + moist_static_energy_internal: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if plume == cumulus_parameterization_constants.DEEP and MELT_GLAC == True: # noqa + _, max_index = column_max(melting_layer, 0, k_end) + downdraft_origin_level[0, 0][plume] = max(downdraft_origin_level[0, 0][plume], max_index) + + # check whether it would have buoyancy, if there where + # no entrainment/detrainment + downdraft_origin_level_internal = downdraft_origin_level[0, 0][plume] + keep_going = True + while keep_going == True: + keep_going = False + if downdraft_origin_level_internal - 1 < detrainment_start_level: + detrainment_start_level = downdraft_origin_level_internal - 1 + + if downdraft_origin_level_internal >= cloud_top_level[0, 0][plume] - 1: + downdraft_origin_level_internal = cloud_top_level[0, 0][plume] - 2 + + level_initial = downdraft_origin_level_internal + dh = 0.0 + level = level_initial - 1 + stop_k_while_loop = False + while level >= 0 and stop_k_while_loop == False: + moist_static_energy_internal = ( + environment_saturation_moist_static_energy_cloud_levels_forced[ + 0, 0, downdraft_origin_level_internal + ] + ) + dz = ( + geopotential_height_cloud_levels_forced[0, 0, level + 1] + - geopotential_height_cloud_levels_forced[0, 0, level] + ) + dh = dh + dz * ( + moist_static_energy_internal + - environment_saturation_moist_static_energy_cloud_levels_forced[0, 0, level] + ) + if dh > 0: + downdraft_origin_level_internal = downdraft_origin_level_internal - 1 + if downdraft_origin_level_internal > 4: + keep_going = True + else: + error_code[0, 0][plume] = 9 + stop_k_while_loop = True + level -= 1 + + downdraft_origin_level[0, 0][plume] = downdraft_origin_level_internal + if downdraft_origin_level_internal <= 4: + error_code[0, 0][plume] = 4 + + with computation(FORWARD), interval(0, 1): + # must have at least depth_min m between cloud convective base and cloud top. + if error_code[0, 0][plume] == 0: + if downdraft_origin_level[0, 0][plume] - 1 < detrainment_start_level: + detrainment_start_level = downdraft_origin_level[0, 0][plume] - 1 + if ( + -geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + + geopotential_height_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume]) + < MINIMUM_DEPTH + ): + error_code[0, 0][plume] = 6 + + +def downdraft_mass_flux( + error_code: IntFieldIJ_Plume, + detrainment_start_level: IntFieldIJ, + downdraft_origin_level: IntFieldIJ_Plume, + pbl_level: IntFieldIJ, + updraft_origin_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + lcl_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + p_surface: FloatFieldIJ, + normalized_massflux_downdraft: FloatField, + normalized_massflux_downdraft_forced: FloatField_Plume, + ocean_fraction: FloatFieldIJ, + random_number: FloatFieldIJ, + DOWNDRAFT_MAX_HEIGHT_LAND: Float, + DOWNDRAFT_MAX_HEIGHT_OCEAN: Float, + plume: Int, +): + """Handle mass fluxes in the downdraft. For shallow plumes mass flux is forced to zero. + + Args: + error_code (IntFieldIJ_Plume) + detrainment_start_level (IntFieldIJ) + downdraft_origin_level (IntFieldIJ_Plume) + pbl_level (IntFieldIJ) + updraft_origin_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + lcl_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + p_surface (FloatFieldIJ) + normalized_massflux_downdraft (FloatField) + normalized_massflux_downdraft_forced (FloatField_Plume) + ocean_fraction (FloatFieldIJ) + random_number (FloatFieldIJ) + DOWNDRAFT_MAX_HEIGHT_LAND (Float) + DOWNDRAFT_MAX_HEIGHT_OCEAN (Float) + plume (Int) + """ + from __externals__ import ZERO_DIFF, k_end + + with computation(FORWARD), interval(...): + normalized_massflux_downdraft = 0.0 + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + # set internal constants + beta: FloatFieldIJ = 2.5 + + height_down = ( + 1.0 - ocean_fraction + ) * DOWNDRAFT_MAX_HEIGHT_LAND + ocean_fraction * DOWNDRAFT_MAX_HEIGHT_OCEAN + + # non-zero-diff-APR-08-2020 + if ZERO_DIFF == 1: + height_down = 0.5 + # non-zero-diff-APR-08-2020 + + p_max: FloatFieldIJ = ( + height_down * p_cloud_levels_forced.at(K=downdraft_origin_level[0, 0][plume], ddim=[plume]) + + (1.0 - height_down) * p_surface + ) + + with computation(PARALLEL), interval(...): + if plume != 0 and error_code[0, 0][plume] == 0: + p_internal = abs(p_cloud_levels_forced[0, 0, 0][plume] - p_max) + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + _, _min_loc = column_min(p_internal, 0, downdraft_origin_level[0, 0][plume]) + min_loc: IntFieldIJ = _min_loc + + # this alpha constrains the location of the maximun ZU to be at "kb_adj" vertical level + alpha: FloatFieldIJ = 1.0 + (beta - 1.0) * ( + (min_loc + 1) / (downdraft_origin_level[0, 0][plume] + 2) + ) / (1.0 - ((min_loc + 1) / (downdraft_origin_level[0, 0][plume] + 2))) + + with computation(PARALLEL), interval(1, None): + if ( + plume != 0 + and error_code[0, 0][plume] == 0 + and K <= min(downdraft_origin_level[0, 0][plume] + 1, k_end - 1) + ): + ratio = float(K + 1) / (downdraft_origin_level[0, 0][plume] + 2) + normalized_massflux_downdraft_forced[0, 0, 0][plume] = ratio ** (alpha - 1.0) * (1.0 - ratio) ** ( + beta - 1.0 + ) + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + normalized_massflux_downdraft_forced[0, 0, 0][plume] = 0.0 + + # get max value for next block + _max_val, _ = column_max_ddim( + normalized_massflux_downdraft_forced, + plume, + 0, + min(k_end, downdraft_origin_level[0, 0][plume] + 1), + ) + max_val: FloatFieldIJ = _max_val + + with computation(FORWARD), interval(...): + if plume != 0 and error_code[0, 0][plume] == 0: + if max_val <= 0: + normalized_massflux_downdraft_forced[0, 0, 0][plume] = 0.0 + error_code[0, 0][plume] = 51 + else: + if K <= min(k_end, downdraft_origin_level[0, 0][plume] + 1): + normalized_massflux_downdraft_forced[0, 0, 0][plume] = ( + normalized_massflux_downdraft_forced[0, 0, 0][plume] / (1.0e-9 + max_val) + ) + + +def downdraft_lateral_massflux( + error_code: IntFieldIJ_Plume, + downdraft_origin_level: IntFieldIJ_Plume, + geopotential_height_cloud_levels_forced: FloatField, + normalized_massflux_downdraft: FloatField, + normalized_massflux_downdraft_forced: FloatField_Plume, + normalized_massflux_downdraft_modified: FloatField, + detrainment_function_downdraft: FloatField, + entrainment_rate_downdraft: FloatField, + mass_entrainment_downdraft: FloatField, + mass_detrainment_downdraft: FloatField, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + mass_entrainment_u_downdraft: FloatField, + mass_detrainment_u_downdraft: FloatField, + LAMBDA_DOWN: Float, + plume: Int, +): + """Compute the mass entrained/detrained by the downdraft. For shallow plumes, this is zero. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_origin_level (IntFieldIJ_Plume) + geopotential_height_cloud_levels_forced (FloatField) + normalized_massflux_downdraft (FloatField) + normalized_massflux_downdraft_forced (FloatField_Plume) + normalized_massflux_downdraft_modified (FloatField) + detrainment_function_downdraft (FloatField) + entrainment_rate_downdraft (FloatField) + mass_entrainment_downdraft (FloatField) + mass_detrainment_downdraft (FloatField) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + mass_entrainment_u_downdraft (FloatField) + mass_detrainment_u_downdraft (FloatField) + LAMBDA_DOWN (Float) + plume (Int) + """ + from __externals__ import k_end + + with computation(PARALLEL), interval(...): + # set entrainment/detrainment to zero + detrainment_function_downdraft = 0.0 + mass_entrainment_downdraft = 0.0 + mass_detrainment_downdraft = 0.0 + mass_entrainment_downdraft_forced[0, 0, 0][plume] = 0.0 + mass_detrainment_downdraft_forced[0, 0, 0][plume] = 0.0 + mass_entrainment_u_downdraft = 0.0 + mass_detrainment_u_downdraft = 0.0 + + # rest of this stencil does not execute for shallow plumes (plume = 0) + with computation(PARALLEL), interval(...): + if plume != 0 and error_code[0, 0][plume] == 0 and K < downdraft_origin_level[0, 0][plume]: + detrainment_function_downdraft = entrainment_rate_downdraft + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + entrainment_rate_downdraft = 0.0 + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + # get maximum index for next block + _, _max_loc = column_max_ddim(normalized_massflux_downdraft_forced, plume, 0, k_end) + max_loc: IntFieldIJ = _max_loc + + with computation(BACKWARD), interval(0, -1): + if ( + plume != 0 + and error_code[0, 0][plume] == 0 + and K >= max_loc + and K <= downdraft_origin_level[0, 0][plume] + ): + + # from downdraft_origin_level to maximum value of + # normalized_massflux_downdraft, change entrainment + dzo = geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + mass_detrainment_downdraft_forced[0, 0, 0][plume] = ( + detrainment_function_downdraft * dzo * normalized_massflux_downdraft_forced[0, 0, 1][plume] + ) + + mass_entrainment_downdraft_forced[0, 0, 0][plume] = ( + normalized_massflux_downdraft_forced[0, 0, 0][plume] + - normalized_massflux_downdraft_forced[0, 0, 1][plume] + + mass_detrainment_downdraft_forced[0, 0, 0][plume] + ) + mass_entrainment_downdraft_forced[0, 0, 0][plume] = max( + 0.0, mass_entrainment_downdraft_forced[0, 0, 0][plume] + ) + # check dd_massdetro in case of dd_massentro has been changed above + mass_detrainment_downdraft_forced[0, 0, 0][plume] = ( + mass_entrainment_downdraft_forced[0, 0, 0][plume] + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + + normalized_massflux_downdraft_forced[0, 0, 1][plume] + ) + + with computation(BACKWARD), interval(0, -1): + if plume != 0 and error_code[0, 0][plume] == 0 and K < max_loc: + # from maximum value zd to surface -> change detrainment + dzo = geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + mass_entrainment_downdraft_forced[0, 0, 0][plume] = ( + entrainment_rate_downdraft * dzo * normalized_massflux_downdraft_forced[0, 0, 1][plume] + ) + + mass_detrainment_downdraft_forced[0, 0, 0][plume] = ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + + mass_entrainment_downdraft_forced[0, 0, 0][plume] + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + ) + mass_detrainment_downdraft_forced[0, 0, 0][plume] = max( + 0.0, mass_detrainment_downdraft_forced[0, 0, 0][plume] + ) + # check dd_massentro in case of dd_massdetro has been changed above + mass_entrainment_downdraft_forced[0, 0, 0][plume] = ( + mass_detrainment_downdraft_forced[0, 0, 0][plume] + + normalized_massflux_downdraft_forced[0, 0, 0][plume] + - normalized_massflux_downdraft_forced[0, 0, 1][plume] + ) + + with computation(BACKWARD), interval(0, -1): + if plume != 0 and error_code[0, 0][plume] == 0 and K <= downdraft_origin_level[0, 0][plume]: + normalized_massflux_downdraft_modified = normalized_massflux_downdraft_forced[0, 0, 0][plume] + normalized_massflux_downdraft = normalized_massflux_downdraft_forced[0, 0, 0][plume] + mass_entrainment_downdraft = mass_entrainment_downdraft_forced[0, 0, 0][plume] + mass_detrainment_downdraft = mass_detrainment_downdraft_forced[0, 0, 0][plume] + mass_entrainment_u_downdraft = ( + mass_entrainment_downdraft_forced[0, 0, 0][plume] + + LAMBDA_DOWN * mass_detrainment_downdraft_forced[0, 0, 0][plume] + ) + mass_detrainment_u_downdraft = ( + mass_detrainment_downdraft_forced[0, 0, 0][plume] + + LAMBDA_DOWN * mass_detrainment_downdraft_forced[0, 0, 0][plume] + ) + + +def downdraft_moist_static_energy_and_buoyancy( + error_code: IntFieldIJ_Plume, + downdraft_origin_level: IntFieldIJ_Plume, + u: FloatField, + u_cloud_levels: FloatField, + u_c_downdraft: FloatField, + v: FloatField, + v_cloud_levels: FloatField, + v_c_downdraft: FloatField, + environment_moist_static_energy_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + cloud_moist_static_energy: FloatField, + cloud_moist_static_energy_downdraft_forced: FloatField, + buoyancy_downdraft_forced: FloatField, + t_wetbulb: FloatFieldIJ, + vapor_wetbulb: FloatFieldIJ, + geopotential_height_cloud_levels_forced: FloatField, + normalized_massflux_downdraft_forced: FloatField_Plume, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + mass_entrainment_u_downdraft: FloatField, + mass_detrainment_u_downdraft: FloatField, + plume: Int, +): + """Compute moist static energy and buoyancy for the downdraft. + + For shallow plumes the majority of the code is skipped. buoyancy_downdraft_forced is set to zero, + u_c_downdraft and v_c_downdraft are set to u_cloud_levels and v_cloud_levels, respectively, and + cloud_moist_static_energy_downdraft_forced is set to environment_saturation_moist_static_energy. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_origin_level (IntFieldIJ_Plume) + u (FloatField) + u_cloud_levels (FloatField) + u_c_downdraft (FloatField) + v (FloatField) + v_cloud_levels (FloatField) + v_c_downdraft (FloatField) + environment_moist_static_energy_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + cloud_moist_static_energy (FloatField) + cloud_moist_static_energy_downdraft_forced (FloatField) + buoyancy_downdraft_forced (FloatField) + t_wetbulb (FloatFieldIJ) + vapor_wetbulb (FloatFieldIJ) + geopotential_height_cloud_levels_forced (FloatField) + normalized_massflux_downdraft_forced (FloatField_Plume) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + mass_entrainment_u_downdraft (FloatField) + mass_detrainment_u_downdraft (FloatField) + plume (Int) + """ + from __externals__ import PRESSURE_GRADIENT_CONSTANT, USE_WETBULB + + with computation(PARALLEL), interval(...): + cloud_moist_static_energy_downdraft_forced = ( + environment_saturation_moist_static_energy_cloud_levels_forced + ) + u_c_downdraft = u_cloud_levels + v_c_downdraft = v_cloud_levels + buoyancy_downdraft_forced = 0.0 + buoyancy_downdraft = 0.0 + + with computation(FORWARD), interval(...): + buoyancy_downdraft: FloatFieldIJ = 0.0 + if error_code[0, 0][plume] == 0 and plume != 0 and K == downdraft_origin_level[0, 0][plume]: + wetbulb_adjustment: IntFieldIJ = 0 + # for future test) + if USE_WETBULB == 1: + cloud_moist_static_energy_downdraft_forced = 0.5 * ( + cumulus_parameterization_constants.CP * t_wetbulb + + cumulus_parameterization_constants.XLV * vapor_wetbulb + + geopotential_height_cloud_levels_forced * constants.MAPL_GRAV + + cloud_moist_static_energy + ) + wetbulb_adjustment = 1 + + buoyancy_downdraft_forced = ( + cloud_moist_static_energy_downdraft_forced + - environment_saturation_moist_static_energy_cloud_levels_forced + ) + buoyancy_downdraft = buoyancy_downdraft_forced * ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) + + with computation(BACKWARD), interval(...): + if ( + error_code[0, 0][plume] == 0 + and plume != 0 + and K <= downdraft_origin_level[0, 0][plume] - wetbulb_adjustment + ): + denom = ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 * mass_detrainment_downdraft_forced[0, 0, 0][plume] + + mass_entrainment_downdraft_forced[0, 0, 0][plume] + ) + denom_u = ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 * mass_detrainment_u_downdraft + + mass_entrainment_u_downdraft + ) + + # tmp fix for denominator being zero + if denom > 0.0 and denom_u > 0.0: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) + + u_c_downdraft = ( + u_c_downdraft[0, 0, 1] * normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 * mass_detrainment_u_downdraft * u_c_downdraft[0, 0, 1] + + mass_entrainment_u_downdraft * u + - PRESSURE_GRADIENT_CONSTANT + * normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (u[0, 0, 1] - u) + ) / denom_u + v_c_downdraft = ( + v_c_downdraft[0, 0, 1] * normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 * mass_detrainment_u_downdraft * v_c_downdraft[0, 0, 1] + + mass_entrainment_u_downdraft * v + - PRESSURE_GRADIENT_CONSTANT + * normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (v[0, 0, 1] - v) + ) / denom_u + + cloud_moist_static_energy_downdraft_forced = ( + cloud_moist_static_energy_downdraft_forced[0, 0, 1] + * normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 + * mass_detrainment_downdraft_forced[0, 0, 0][plume] + * cloud_moist_static_energy_downdraft_forced[0, 0, 1] + + mass_entrainment_downdraft_forced[0, 0, 0][plume] + * environment_moist_static_energy_forced + ) / denom + + buoyancy_downdraft_forced = ( + cloud_moist_static_energy_downdraft_forced + - environment_saturation_moist_static_energy_cloud_levels_forced + ) + buoyancy_downdraft = buoyancy_downdraft + buoyancy_downdraft_forced * dz + else: + u_c_downdraft = u_c_downdraft[0, 0, 1] + v_c_downdraft = v_c_downdraft[0, 0, 1] + cloud_moist_static_energy_downdraft_forced = cloud_moist_static_energy_downdraft_forced[ + 0, 0, 1 + ] + + with computation(FORWARD), interval(0, 1): + if buoyancy_downdraft > 0: + error_code[0, 0][plume] = 7 + + +def downdraft_moisture( + error_code: IntFieldIJ_Plume, + downdraft_origin_level: IntFieldIJ_Plume, + t_cloud_levels_forced: FloatField, + t_wetbulb: FloatFieldIJ, + vapor_forced: FloatField, + vapor_cloud_levels_forced: FloatField, + environment_saturation_mixing_ratio_cloud_levels_forced: FloatField, + cloud_total_water_after_entrainment_forced: FloatField, + cloud_total_water_after_entrainment_downdraft_forced: FloatField, + downdraft_saturation_vapor_forced: FloatField, + vapor_wetbulb: FloatFieldIJ, + normalized_massflux_downdraft_forced: FloatField_Plume, + environment_moist_static_energy_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + cloud_moist_static_energy_downdraft_forced: FloatField, + evaporate_in_downdraft_forced: FloatField_Plume, + geopotential_height_cloud_levels_forced: FloatField, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + gamma_cloud_levels_forced: FloatField, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + total_normalized_integrated_evaporate_forced: FloatFieldIJ, + buoyancy: FloatFieldIJ, + plume: Int, +): + """Compute the moisture profile for the downdraft. + + For shallow plumes outputs are forced to zero and calculation is terminated. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_origin_level (IntFieldIJ_Plume) + t_cloud_levels_forced (FloatField) + t_wetbulb (FloatFieldIJ) + vapor_forced (FloatField) + vapor_cloud_levels_forced (FloatField) + environment_saturation_mixing_ratio_cloud_levels_forced (FloatField) + cloud_total_water_after_entrainment_forced (FloatField) + cloud_total_water_after_entrainment_downdraft_forced (FloatField) + downdraft_saturation_vapor_forced (FloatField) + vapor_wetbulb (FloatFieldIJ) + normalized_massflux_downdraft_forced (FloatField_Plume) + environment_moist_static_energy_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + cloud_moist_static_energy_downdraft_forced (FloatField) + evaporate_in_downdraft_forced (FloatField_Plume) + geopotential_height_cloud_levels_forced (FloatField) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + gamma_cloud_levels_forced (FloatField) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + total_normalized_integrated_evaporate_forced (FloatFieldIJ) + buoyancy (FloatFieldIJ) + plume (Int) + """ + from __externals__ import EVAP_FIX, USE_WETBULB, ZERO_DIFF + + with computation(FORWARD), interval(0, 1): + internal_loop_constant: IntFieldIJ = 1 + + with computation(FORWARD), interval(0, 1): + # prefill outputs with zero + buoyancy = 0.0 + total_normalized_integrated_evaporate_forced = 0.0 + + with computation(PARALLEL), interval(...): + # prefill outputs with zero + cloud_total_water_after_entrainment_downdraft_forced = 0.0 + downdraft_saturation_vapor_forced = 0.0 + evaporate_in_downdraft_forced[0, 0, 0][plume] = 0.0 + + with computation(FORWARD), interval(0, -1): + # NOTE this K level check should be a dynamic interval for the sake of performance, but the current + # of dynamic intervals does not currently work with single levels e.g. interval(field, field+1) + # will revisit when dynamic intervals are more stable + if error_code[0, 0][plume] == 0 and plume != 0 and K == downdraft_origin_level[0, 0][plume]: + # boundary condition at downdraft_origin_level ('level of free sinking') + dz: FloatFieldIJ = ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) + + cloud_total_water_after_entrainment_downdraft_forced = vapor_cloud_levels_forced + + if USE_WETBULB == 1: + # mixture 50% env air + updraft + cloud_total_water_after_entrainment_downdraft_forced = 0.5 * ( + vapor_wetbulb + cloud_total_water_after_entrainment_forced + ) + + d_moist_static_energy: FloatFieldIJ = ( + cloud_moist_static_energy_downdraft_forced + - environment_saturation_moist_static_energy_cloud_levels_forced + ) + + if d_moist_static_energy < 0: + downdraft_saturation_vapor_forced = ( + environment_saturation_mixing_ratio_cloud_levels_forced + + (1.0 / cumulus_parameterization_constants.XLV) + * (gamma_cloud_levels_forced / (1.0 + gamma_cloud_levels_forced)) + * d_moist_static_energy + ) + else: + downdraft_saturation_vapor_forced = environment_saturation_mixing_ratio_cloud_levels_forced + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0 and plume != 0 and K == downdraft_origin_level[0, 0][plume]: + evaporate_in_downdraft_forced[0, 0, 0][plume] = normalized_massflux_downdraft_forced[0, 0, 0][ + plume + ] * min( + 0.0, cloud_total_water_after_entrainment_downdraft_forced - downdraft_saturation_vapor_forced + ) + cloud_total_water_after_entrainment_downdraft_forced = downdraft_saturation_vapor_forced + total_normalized_integrated_evaporate_forced = ( + total_normalized_integrated_evaporate_forced + evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and plume != 0: + buoyancy = dz * d_moist_static_energy + + with computation(BACKWARD), interval(0, -1): + # NOTE this K level check should be a dynamic interval, but the current version of dynamic + # intervals does not work here - it gets stuck (as if in an infinite loop) + # will revisit when dynamic intervals are more stable + if error_code[0, 0][plume] == 0 and plume != 0 and K < downdraft_origin_level[0, 0][plume]: + dz = geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + + # downward transport + mixing + denom = ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 * mass_detrainment_downdraft_forced[0, 0, 0][plume] + + mass_entrainment_downdraft_forced[0, 0, 0][plume] + ) + if denom == 0.0: + cloud_total_water_after_entrainment_downdraft_forced = ( + cloud_total_water_after_entrainment_downdraft_forced[0, 0, 1] + ) + else: + cloud_total_water_after_entrainment_downdraft_forced = ( + cloud_total_water_after_entrainment_downdraft_forced[0, 0, 1] + * normalized_massflux_downdraft_forced[0, 0, 1][plume] + - 0.5 + * mass_detrainment_downdraft_forced[0, 0, 0][plume] + * cloud_total_water_after_entrainment_downdraft_forced[0, 0, 1] + + mass_entrainment_downdraft_forced[0, 0, 0][plume] * vapor_forced + ) / denom + + # to be negatively buoyant, hcd should be smaller than hes! + # ideally, dh should be negative till dd hits ground, but that is not always the case + d_moist_static_energy = ( + cloud_moist_static_energy_downdraft_forced + - environment_saturation_moist_static_energy_cloud_levels_forced + ) + buoyancy = buoyancy + dz * d_moist_static_energy + downdraft_saturation_vapor_forced = ( + environment_saturation_mixing_ratio_cloud_levels_forced + + (1.0 / cumulus_parameterization_constants.XLV) + * (gamma_cloud_levels_forced / (1.0 + gamma_cloud_levels_forced)) + * d_moist_static_energy + ) + + # rain water evaporation amount at layer k + dq_eva = cloud_total_water_after_entrainment_downdraft_forced - downdraft_saturation_vapor_forced + + if dq_eva > 0.0: + dq_eva = 0.0 + downdraft_saturation_vapor_forced = cloud_total_water_after_entrainment_downdraft_forced + # amount of the evaporated rain water + evaporate_in_downdraft_forced[0, 0, 0][plume] = ( + normalized_massflux_downdraft_forced[0, 0, 0][plume] * dq_eva + ) # kg[water vapor]/kg[air] + + # source term for in-downdraft water vapor mixing ratio + # => equiv to qcd = qcd - dq_eva !( -dq_eva >0 => source term for qcd) + cloud_total_water_after_entrainment_downdraft_forced = downdraft_saturation_vapor_forced + + # total evaporated rain water + total_normalized_integrated_evaporate_forced = ( + total_normalized_integrated_evaporate_forced + evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + + with computation(FORWARD), interval(0, 1): + # check for problems that stop the convective scheme + if error_code[0, 0][plume] == 0 and plume != 0: + if total_normalized_integrated_evaporate_forced >= 0 and internal_loop_constant == 1: + error_code[0, 0][plume] = 70 + + if buoyancy >= 0 and internal_loop_constant == 1: + error_code[0, 0][plume] = 73 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and plume != 0: + # initialize a 2d field to be filled in the next block + abs_check: BoolFieldIJ = False + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and plume != 0: + if ( + ZERO_DIFF == 0 + and EVAP_FIX == 1 + and abs(total_normalized_integrated_evaporate_forced) + > total_normalized_integrated_condensate_forced[0, 0][plume] + and error_code[0, 0][plume] == 0 + ): + # note whether the absolute value check is true so that the conditional can be replicated + # in the next block, even after part of it has been changed + abs_check = True + + fix_evap: FloatFieldIJ = total_normalized_integrated_condensate_forced[0, 0][plume] / ( + 1.0e-16 + abs(total_normalized_integrated_evaporate_forced) + ) + total_normalized_integrated_evaporate_forced = 0.0 + + with computation(BACKWARD), interval(...): + if ( + error_code[0, 0][plume] == 0 + and plume != 0 + and ZERO_DIFF == 0 + and EVAP_FIX == 1 + and abs_check == True + and K <= downdraft_origin_level[0, 0][plume] + ): + evaporate_in_downdraft_forced[0, 0, 0][plume] = ( + evaporate_in_downdraft_forced[0, 0, 0][plume] * fix_evap + ) + total_normalized_integrated_evaporate_forced = ( + total_normalized_integrated_evaporate_forced + evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + dq_eva = evaporate_in_downdraft_forced[0, 0, 0][plume] / ( + 1.0e-16 + normalized_massflux_downdraft_forced[0, 0, 0][plume] + ) + cloud_total_water_after_entrainment_downdraft_forced = downdraft_saturation_vapor_forced + dq_eva + + with computation(FORWARD), interval(0, 1): + if ( + error_code[0, 0][plume] == 0 + and plume != 0 + and ZERO_DIFF == 0 + and EVAP_FIX == 1 + and abs_check == True + ): + if total_normalized_integrated_evaporate_forced >= 0.0: + error_code[0, 0][plume] = 70 + + +def downdraft_temperature( + error_code: IntFieldIJ_Plume, + downdraft_column_temperature_forced: FloatField, + cloud_moist_static_energy_downdraft_forced: FloatField, + geopotential_height_cloud_levels_forced: FloatField, + cloud_total_water_after_entrainment_downdraft_forced: FloatField, + t_cloud_levels_forced: FloatField, + plume: Int, +): + """Compute temperature within the downdraft. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_column_temperature_forced (FloatField) + cloud_moist_static_energy_downdraft_forced (FloatField) + geopotential_height_cloud_levels_forced (FloatField) + cloud_total_water_after_entrainment_downdraft_forced (FloatField) + t_cloud_levels_forced (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + downdraft_column_temperature_forced = (1.0 / cumulus_parameterization_constants.CP) * ( + cloud_moist_static_energy_downdraft_forced + - constants.MAPL_GRAV * geopotential_height_cloud_levels_forced + - cumulus_parameterization_constants.XLV + * cloud_total_water_after_entrainment_downdraft_forced + ) + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] != 0: + downdraft_column_temperature_forced = t_cloud_levels_forced + + +def downdraft_windshear( + error_code: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + geopotential_height_forced: FloatField, + p_forced: FloatField, + u: FloatField, + v: FloatField, + ccn: FloatFieldIJ, + psum: FloatFieldIJ, + psumh: FloatFieldIJ, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + total_normalized_integrated_evaporate_forced: FloatFieldIJ, + epsilon: FloatFieldIJ, + epsilon_min: FloatFieldIJ, + epsilon_max: FloatFieldIJ, + epsilon_computed: FloatFieldIJ_ensemble_2, + plume: Int, +): + """Compute the windshear within the downdraft and quantify its effect on precipitation (epsilon_computed). + + Args: + error_code (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + geopotential_height_forced (FloatField) + p_forced (FloatField) + u (FloatField) + v (FloatField) + ccn (FloatFieldIJ) + psum (FloatFieldIJ) + psumh (FloatFieldIJ) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + total_normalized_integrated_evaporate_forced (FloatFieldIJ) + epsilon (FloatFieldIJ) + epsilon_min (FloatFieldIJ) + epsilon_max (FloatFieldIJ) + epsilon_computed (FloatFieldIJ_ensemble_2) + plume (Int) + """ + from __externals__ import AEROEVAP + + with computation(FORWARD), interval(0, 1): + # initialize internal constants + alpha3: FloatFieldIJ = 1.9 + beta3: FloatFieldIJ = -1.13 + + # zero input fields + epsilon = 0.0 + + # zero every part of the ensemble dimension + count = 0 + while count < cumulus_parameterization_constants.MAXENS2: + epsilon_computed[0, 0][count] = 0.0 + count += 1 + + # initialize internal 2d temporaries + vshear: FloatFieldIJ = 0.0 + sdp: FloatFieldIJ = 0.0 + vws: FloatFieldIJ = 0.0 + + lower_bound: FloatFieldIJ = updraft_lfc_level[0, 0][plume] + upper_bound: FloatFieldIJ = cloud_top_level[0, 0][plume] + 1 + + with computation(FORWARD), interval(lower_bound, upper_bound): + if plume != 0 and error_code[0, 0][plume] == 0: + dp = p_forced - p_forced[0, 0, 1] + vws = ( + vws + + ( + abs((u[0, 0, 1] - u) / (geopotential_height_forced[0, 0, 1] - geopotential_height_forced)) + + abs( + (v[0, 0, 1] - v) / (geopotential_height_forced[0, 0, 1] - geopotential_height_forced) + ) + ) + * dp + ) + sdp = sdp + dp + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + vshear = 1.0e3 * vws / sdp + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + precip_efficiency = 1.591 - 0.639 * vshear + 0.0953 * (vshear**2) - 0.00496 * (vshear**3) + precip_efficiency = min(precip_efficiency, 0.9) + precip_efficiency = max(precip_efficiency, 0.1) + + # cloud base precip efficiency + zkbc = geopotential_height_forced.at(K=updraft_lfc_level[0, 0][plume]) * 3.281e-3 + prezk = 0.02 + if zkbc > 3.0: + prezk = 0.96729352 + zkbc * ( + -0.70034167 + + zkbc * (0.162179896 + zkbc * (-1.2569798e-2 + zkbc * (4.2772e-4 - zkbc * 5.44e-6))) + ) + if zkbc > 25.0: + prezk = 2.4 + + precip_efficiency_b = 1.0 / (1.0 + prezk) + precip_efficiency_b = min(precip_efficiency_b, 0.9) + precip_efficiency_b = max(precip_efficiency_b, 0.1) + + epsilon = 1.0 - 0.5 * (precip_efficiency_b + precip_efficiency) + + if AEROEVAP > 1: + aeroadd = (cumulus_parameterization_constants.CCNCLEAN**beta3) * ( + (psumh) ** (alpha3 - 1) + ) # *1.e6 + prop_c = 0.5 * (precip_efficiency_b + precip_efficiency) / aeroadd + aeroadd = (ccn**beta3) * ((psum) ** (alpha3 - 1)) # *1.e6 + aeroadd = prop_c * aeroadd + precip_efficiency_c = aeroadd + if precip_efficiency_c > 0.9: + precip_efficiency_c = 0.9 + if precip_efficiency_c < 0.1: + precip_efficiency_c = 0.1 + epsilon = 1.0 - precip_efficiency_c + if AEROEVAP == 2: + epsilon = 1.0 - 0.25 * ( + precip_efficiency_b + precip_efficiency + 2.0 * precip_efficiency_c + ) + + # epsilon here is 1-precip_efficiency! + einc = 0.2 * epsilon + count = 0 + while count < cumulus_parameterization_constants.MAXENS2: + epsilon_computed[0, 0][count] = epsilon + (count - 1) * einc + count += 1 + + count = 0 + while count < cumulus_parameterization_constants.MAXENS2: + epsilon_computed[0, 0][count] = ( + -epsilon_computed[0, 0][count] + * total_normalized_integrated_condensate_forced[0, 0][plume] + / total_normalized_integrated_evaporate_forced + ) + if epsilon_computed[0, 0][count] > epsilon_max: + epsilon_computed[0, 0][count] = epsilon_max + if epsilon_computed[0, 0][count] < epsilon_min: + epsilon_computed[0, 0][count] = epsilon_min + count += 1 + + +def update_epsilon_forced( + error_code: IntFieldIJ_Plume, + scale_dependence_factor_downdraft: FloatFieldIJ, + epsilon_computed: FloatFieldIJ_ensemble_2, + epsilon: FloatFieldIJ, + epsilon_forced: FloatFieldIJ_Plume, + plume: Int, +): + """Tune precipitation effeciency (epsilon_computed) based on the scale dependence factor. + + Args: + error_code (IntFieldIJ_Plume) + scale_dependence_factor_downdraft (FloatFieldIJ) + epsilon_computed (FloatFieldIJ_ensemble_2) + epsilon (FloatFieldIJ) + epsilon_forced (FloatFieldIJ_Plume) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + count = 0 + while count < cumulus_parameterization_constants.MAXENS2: + epsilon_forced[0, 0][plume] = ( + scale_dependence_factor_downdraft * epsilon_computed[0, 0][count] + ) + epsilon = epsilon_forced[0, 0][plume] + count += 1 + + +class DowndraftOriginLevel(NDSLRuntime): + """Determine the leve at which the downdraft begins.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make config and cumulus_parameterization_config visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # initialize locals + self._critical_level: Local = self.make_local(quantity_factory, [X_DIM, Y_DIM], Int) + + # construct stencils + self._get_critical_level = stencil_factory.from_dims_halo( + func=get_critical_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._unknown_find_level = stencil_factory.from_dims_halo( + func=generic_find_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._get_downdraft_origin_level = stencil_factory.from_dims_halo( + func=get_downdraft_origin_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"MELT_GLAC": cumulus_parameterization_config.MELT_GLAC}, + ) + + def __call__( + self, + error_code: Quantity, + cloud_top_level: Quantity, + geopotential_height_cloud_levels_forced: Quantity, + topography_height_no_negative: Quantity, + environment_saturation_moist_static_energy_cloud_levels_forced: Quantity, + updraft_origin_level: Quantity, + downdraft_origin_level: Quantity, + updraft_lfc_level: Quantity, + detrainment_start_level: Quantity, + melting_layer: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._get_critical_level( + error_code=error_code, + critical_level=self._critical_level, + cloud_top_level=cloud_top_level, + geopotential_height_cloud_levels_forced=geopotential_height_cloud_levels_forced, + topography_height_no_negative=topography_height_no_negative, + MAX_DOWNDRAFT_ORIGIN_HEIGHt=plume_dependent_constants.MAX_DOWNDRAFT_ORIGIN_HEIGHt, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._unknown_find_level( + array=environment_saturation_moist_static_energy_cloud_levels_forced, + start_index=updraft_origin_level, + end_index=self._critical_level, + out_index=downdraft_origin_level, + error_code=error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._get_downdraft_origin_level( + error_code=error_code, + cloud_top_level=cloud_top_level, + updraft_lfc_level=updraft_lfc_level, + detrainment_start_level=detrainment_start_level, + downdraft_origin_level=downdraft_origin_level, + environment_saturation_moist_static_energy_cloud_levels_forced=environment_saturation_moist_static_energy_cloud_levels_forced, + geopotential_height_cloud_levels_forced=geopotential_height_cloud_levels_forced, + melting_layer=melting_layer, + MINIMUM_DEPTH=plume_dependent_constants.MINIMUM_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class DowndraftWetBlub: + def __init__(self, cumulus_parameterization_config: GF2020CumulusParameterizationConfig): + if cumulus_parameterization_config.USE_WETBULB == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->DowndraftWetBulb: The DowndraftWetBulb" + "section has not been implemented. You should have been caught before getting" + "here by the config checker. Beware, something likely failing in the config" + "checker as well - you may be unknowingly calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass + + +class DowndraftWindShear(NDSLRuntime): + """Calculate wint shear within the downdraft. Quantify the effect of wind shear within the downdraft + on precipitation into a "precip effeciency" field (epsilon_computed). + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._downdraft_windshear = stencil_factory.from_dims_halo( + func=downdraft_windshear, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"AEROEVAP": cumulus_parameterization_config.AEROEVAP}, + ) + + self._update_epsilon_forced = stencil_factory.from_dims_halo( + func=update_epsilon_forced, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + geopotential_height_forced: Quantity, + p_forced: Quantity, + u: Quantity, + v: Quantity, + ccn: Quantity, + psum: Quantity, + psumh: Quantity, + total_normalized_integrated_condensate_forced: Quantity, + total_normalized_integrated_evaporate_forced: Quantity, + scale_dependence_factor_downdraft: Quantity, + epsilon: Quantity, + epsilon_min: Quantity, + epsilon_max: Quantity, + epsilon_computed: Quantity, + epsilon_forced: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._downdraft_windshear( + error_code=error_code, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height_forced=geopotential_height_forced, + p_forced=p_forced, + u=u, + v=v, + ccn=ccn, + psum=psum, + psumh=psumh, + total_normalized_integrated_condensate_forced=total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=total_normalized_integrated_evaporate_forced, + epsilon=epsilon, + epsilon_min=epsilon_min, + epsilon_max=epsilon_max, + epsilon_computed=epsilon_computed, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._update_epsilon_forced( + error_code=error_code, + scale_dependence_factor_downdraft=scale_dependence_factor_downdraft, + epsilon_computed=epsilon_computed, + epsilon=epsilon, + epsilon_forced=epsilon_forced, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/entrainment.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/entrainment.py new file mode 100644 index 000000000..94f94a3d7 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/entrainment.py @@ -0,0 +1,319 @@ +from gt4py.cartesian.gtscript import FORWARD, PARALLEL, K, computation, interval + +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max_ddim +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import get_cloud_boundary_conditions + + +def entrainment_rates( + vapor_cloud_levels_forced: FloatField, + environment_saturation_mixing_ratio_cloud_levels_forced: FloatField, + lcl_level: IntFieldIJ_Plume, + error_code: IntFieldIJ_Plume, + entrainment_rate: FloatField_Plume, + detrainment_function_updraft: FloatField, + plume: Int, +): + """Determine the entrainment/detrainment rates, updating the initial estimate + using data computed inside the core. + + Args: + vapor_cloud_levels_forced (FloatField) + environment_saturation_mixing_ratio_cloud_levels_forced (FloatField) + lcl_level (IntFieldIJ_Plume) + error_code (IntFieldIJ_Plume) + entrainment_rate (FloatField_Plume) + detrainment_function_updraft (FloatField) + plume (Int) + """ + from __externals__ import MIN_ENTRAINMENT_RATE, ZERO_DIFF + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + frh = min( + vapor_cloud_levels_forced + / max( + environment_saturation_mixing_ratio_cloud_levels_forced, + cumulus_parameterization_constants.smaller_qv, + ), + 1.0, + ) + if K >= lcl_level[0, 0][plume]: + entrainment_rate[0, 0, 0][plume] = ( + entrainment_rate[0, 0, 0][plume] + * (1.3 - frh) + * ( + environment_saturation_mixing_ratio_cloud_levels_forced + / environment_saturation_mixing_ratio_cloud_levels_forced.at(K=lcl_level[0, 0][plume]) + ) + ** 3 + ) + else: + entrainment_rate[0, 0, 0][plume] = entrainment_rate[0, 0, 0][plume] * (1.3 - frh) + if ZERO_DIFF == 1: + detrainment_function_updraft = 0.75e-4 * (1.6 - frh) + else: + entrainment_rate[0, 0, 0][plume] = max(entrainment_rate[0, 0, 0][plume], MIN_ENTRAINMENT_RATE) + if plume == cumulus_parameterization_constants.SHALLOW: + detrainment_function_updraft = 0.75 * entrainment_rate[0, 0, 0][plume] + if plume == cumulus_parameterization_constants.MID: + detrainment_function_updraft = 0.5 * entrainment_rate[0, 0, 0][plume] + if plume == cumulus_parameterization_constants.DEEP: + detrainment_function_updraft = 0.1 * entrainment_rate[0, 0, 0][plume] + + +def downdraft_entrainment_profiles( + lateral_entrainment_rate: FloatField, + entrainment_rate_downdraft: FloatField, + detrainment_function_downdraft: FloatField, + scale_dependence_factor_downdraft: FloatFieldIJ, + plume_entrainment_rate: Float, +): + """Get the entrainment and detrainment profiles for the downdraft + + Args: + lateral_entrainment_rate (FloatField) + entrainment_rate_downdraft (FloatField) + detrainment_function_downdraft (FloatField) + scale_dependence_factor_downdraft (FloatFieldIJ) + plume_entrainment_rate (Float) + """ + from __externals__ import DOWNDRAFT + + with computation(PARALLEL), interval(0, -1): + entrainment_rate_downdraft = lateral_entrainment_rate * plume_entrainment_rate * 0.3 + detrainment_function_downdraft = entrainment_rate_downdraft + + with computation(FORWARD), interval(0, 1): + if DOWNDRAFT == 0: + scale_dependence_factor_downdraft = 1.0 + else: + scale_dependence_factor_downdraft = 0.0 + + +def compute_lateral_massflux( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + geopotential_height: FloatField, + normalized_massflux_updraft: FloatField_Plume, + detrainment_function_updraft: FloatField, + entrainment_rate: FloatField_Plume, + p_cloud_levels_forced: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_entrainment_updraft: FloatField, + mass_detrainment_updraft: FloatField, + updraft_lfc_level: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + pbl_level: IntFieldIJ, + mass_entrainment_u_updraft: FloatField, + mass_detrainment_u_updraft: FloatField, + LAMBDA_DEEP: Float, + plume: Int, +): + """Compute massfluxes associated with entrainment and detrainment. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + geopotential_height (FloatField) + normalized_massflux_updraft (FloatField_Plume) + detrainment_function_updraft (FloatField) + entrainment_rate (FloatField_Plume) + p_cloud_levels_forced (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_entrainment_updraft (FloatField) + mass_detrainment_updraft (FloatField) + updraft_lfc_level (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + pbl_level (IntFieldIJ) + mass_entrainment_u_updraft (FloatField) + mass_detrainment_u_updraft (FloatField) + LAMBDA_DEEP (Float) + plume (Int) + """ + from __externals__ import k_end + + with computation(PARALLEL), interval(...): + mass_entrainment_updraft_forced[0, 0, 0][plume] = 0.0 + mass_detrainment_updraft_forced[0, 0, 0][plume] = 0.0 + mass_entrainment_u_updraft = 0.0 + mass_detrainment_u_updraft = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + _, index = column_max_ddim(normalized_massflux_updraft, plume, 0, k_end) + max_index: IntFieldIJ = index + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + # will not allow detrainment below cloud base or in the PBL + if plume == cumulus_parameterization_constants.SHALLOW: + if ( + K + <= max( + updraft_lfc_level[0, 0][plume], + updraft_origin_level[0, 0][plume], + ) + + 1 + ): + detrainment_function_updraft = 0 + else: + if K <= max_index + 1: + detrainment_function_updraft = 0 + + # mass entrainment and detrainment are defined on model levels + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0: + # below location of maximum value normalized_massflux_updraft -> change entrainment + if K <= max_index - 1: + height_massflux_avg = ( + geopotential_height[0, 0, 1] - geopotential_height + ) * normalized_massflux_updraft[0, 0, 0][plume] + mass_detrainment_updraft_forced[0, 0, 0][plume] = ( + detrainment_function_updraft * height_massflux_avg + ) + mass_entrainment_updraft_forced[0, 0, 0][plume] = ( + normalized_massflux_updraft[0, 0, 1][plume] + + -normalized_massflux_updraft[0, 0, 0][plume] + + mass_detrainment_updraft_forced[0, 0, 0][plume] + ) + mass_entrainment_updraft_forced[0, 0, 0][plume] = max( + mass_entrainment_updraft_forced[0, 0, 0][plume], 0.0 + ) + + # check mass_detrainment_updraft_forced in case it has been changed above + mass_detrainment_updraft_forced[0, 0, 0][plume] = ( + -normalized_massflux_updraft[0, 0, 1][plume] + + normalized_massflux_updraft[0, 0, 0][plume] + + mass_entrainment_updraft_forced[0, 0, 0][plume] + ) + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0: + if K >= max_index and K <= cloud_top_level[0, 0][plume]: + # above location of maximum value normalized_massflux_updraft -> change entrainment + height_massflux_avg = ( + geopotential_height[0, 0, 1] - geopotential_height + ) * normalized_massflux_updraft[0, 0, 0][plume] + mass_entrainment_updraft_forced[0, 0, 0][plume] = ( + entrainment_rate[0, 0, 0][plume] * height_massflux_avg + ) + mass_detrainment_updraft_forced[0, 0, 0][plume] = ( + normalized_massflux_updraft[0, 0, 0][plume] + + mass_entrainment_updraft_forced[0, 0, 0][plume] + - normalized_massflux_updraft[0, 0, 1][plume] + ) + mass_detrainment_updraft_forced[0, 0, 0][plume] = max( + mass_detrainment_updraft_forced[0, 0, 0][plume], 0.0 + ) + + # check mass_entrainment_updraft_forced in case it has been changed above + mass_entrainment_updraft_forced[0, 0, 0][plume] = ( + -normalized_massflux_updraft[0, 0, 0][plume] + + mass_detrainment_updraft_forced[0, 0, 0][plume] + + normalized_massflux_updraft[0, 0, 1][plume] + ) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + mass_entrainment_updraft = mass_entrainment_updraft_forced[0, 0, 0][plume] + mass_detrainment_updraft = mass_detrainment_updraft_forced[0, 0, 0][plume] + + with computation(FORWARD), interval(1, None): + if error_code[0, 0][plume] == 0: + # for weaker mixing + mass_entrainment_u_updraft[0, 0, -1] = ( + mass_entrainment_updraft_forced[0, 0, -1][plume] + + LAMBDA_DEEP * mass_detrainment_updraft_forced[0, 0, -1][plume] + ) + mass_detrainment_u_updraft[0, 0, -1] = ( + mass_detrainment_updraft_forced[0, 0, -1][plume] + + LAMBDA_DEEP * mass_detrainment_updraft_forced[0, 0, -1][plume] + ) + + +def compute_uc_vc( + u_c: FloatField, + v_c: FloatField, + cloud_moist_static_energy: FloatField, + cloud_moist_static_energy_forced: FloatField, + error_code: IntFieldIJ_Plume, + start_level: IntFieldIJ, + moist_static_energy_origin_level: FloatFieldIJ, + moist_static_energy_origin_level_forced: FloatFieldIJ, + u_cloud_levels: FloatField, + v_cloud_levels: FloatField, + p: FloatField, + updraft_origin_level: IntFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Compute u and v for the c-grid, update cloud moist static energy + below start_level (nominally LCL level). + + Args: + u_c (FloatField) + v_c (FloatField) + cloud_moist_static_energy (FloatField) + cloud_moist_static_energy_forced (FloatField) + error_code (IntFieldIJ_Plume) + start_level (IntFieldIJ) + moist_static_energy_origin_level (FloatFieldIJ) + moist_static_energy_origin_level_forced (FloatFieldIJ) + u_cloud_levels (FloatField) + v_cloud_levels (FloatField) + p (FloatField) + updraft_origin_level (IntFieldIJ_Plume) + ocean_fraction (FloatFieldIJ) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import BOUNDARY_CONDITION_METHOD, k_end + + with computation(PARALLEL), interval(...): + u_c = 0.0 + v_c = 0.0 + cloud_moist_static_energy = 0.0 + cloud_moist_static_energy_forced = 0.0 + + with computation(PARALLEL), interval(...): + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= start_level: + cloud_moist_static_energy = moist_static_energy_origin_level + cloud_moist_static_energy_forced = moist_static_energy_origin_level_forced + + # get uc and vc as average between layers below k22 + u_c = get_cloud_boundary_conditions( + field=u_cloud_levels, + scalar_perturbation=0, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + + v_c = get_cloud_boundary_conditions( + field=v_cloud_levels, + scalar_perturbation=0, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/environment.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/environment.py new file mode 100644 index 000000000..fc744738a --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/environment.py @@ -0,0 +1,794 @@ +from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, exp, interval + +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.gt4py import function +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import ( + get_cloud_boundary_conditions, + saturation_vapor_pressure, +) + + +@function +def saturation_specific_humidity( + potential_t: Float, + p: Float, +): + """ + Compute saturation specific humidity in kg/kg + + Args: + pt (in): potential temperature in Kelvin + p (in): pressure in mb + """ + # local constants, should not be used/visible outside of this function + # not really a suggestion. don't use them anywhere else, use the global constants + rd = 287.06 + rv = 461.52 + rtt = 273.16 + retv = rv / rd - 1.0 + r2es = 611.21 * rd / rv + r3les = 17.502 + r3ies = 22.587 + r4les = 32.19 + r4ies = -0.7 + rtwat = rtt + rtice = rtt - 23.0 + rticecu = rtt - 23.0 + rtwat_rtice_r = 1.0 / (rtwat - rtice) + rtwat_rticecu_r = 1.0 / (rtwat - rticecu) + + foealfcu = min(1.0, ((max(rticecu, min(rtwat, potential_t)) - rticecu) * rtwat_rticecu_r) ** 2) + foeewmcu = r2es * ( + foealfcu * exp(r3les * (potential_t - rtt) / (potential_t - r4les)) + + (1.0 - foealfcu) * exp(r3ies * (potential_t - rtt) / (potential_t - r4ies)) + ) + + zew = foeewmcu + zqs = zew / (100.0 * p) + if 1.0 - retv * zqs > 0.0: + zcor = 1.0 / (1.0 - retv * zqs) + saturation_specific_humidity = zqs * zcor + else: + saturation_specific_humidity = cumulus_parameterization_constants.MAX_QSAT + + return saturation_specific_humidity + + +def environment_conditions( + p: FloatField, + p_surface: FloatFieldIJ, + t: FloatField, + vapor: FloatField, + topography_height_no_negative: FloatFieldIJ, + moist_static_energy: FloatField, + saturation_moist_static_energy: FloatField, + saturation_mixing_ratio: FloatField, + geopotential_height: FloatField, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """Compute moist static energy, geopotential height and saturation + mixing ratio for a set of environmental fields. + + Args: + p (FloatField) + p_surface (FloatFieldIJ) + t (FloatField) + vapor (FloatField) + topography_height_no_negative (FloatFieldIJ) + moist_static_energy (FloatField) + saturation_moist_static_energy (FloatField) + saturation_mixing_ratio (FloatField) + geopotential_height (FloatField) + error_code (IntFieldIJ_Plume) + plume (Int) + """ + from __externals__ import SATURATION_CALCULATION_CHOICE + + with computation(PARALLEL), interval(...): + itest_internal = -1 + moist_static_energy = 0.0 + saturation_moist_static_energy = 0.0 + saturation_mixing_ratio = 0.0 + with computation(PARALLEL), interval(0, -1): + if SATURATION_CALCULATION_CHOICE == 0: + if error_code[0, 0][plume] == 0: + e = saturation_vapor_pressure(t) + saturation_mixing_ratio = 0.622 * e / max(1.0e-8, (p - e)) + + if saturation_mixing_ratio <= 1.0e-08: + saturation_mixing_ratio = 1.0e-08 + if saturation_mixing_ratio > cumulus_parameterization_constants.MAX_QSAT: + saturation_mixing_ratio = cumulus_parameterization_constants.MAX_QSAT + if saturation_mixing_ratio < vapor: + saturation_mixing_ratio = vapor + + TV = t + 0.608 * vapor * t + else: + if error_code[0, 0][plume] == 0: + saturation_mixing_ratio = saturation_specific_humidity(t, p) + saturation_mixing_ratio = min( + cumulus_parameterization_constants.MAX_QSAT, + max(1.0e-08, saturation_mixing_ratio), + ) + saturation_mixing_ratio = max(saturation_mixing_ratio, vapor) + tv = t + 0.608 * vapor * t + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + moist_static_energy = ( + constants.MAPL_GRAV * geopotential_height + + cumulus_parameterization_constants.CP * t + + cumulus_parameterization_constants.XLV * vapor + ) + + saturation_moist_static_energy = ( + constants.MAPL_GRAV * geopotential_height + + cumulus_parameterization_constants.CP * t + + cumulus_parameterization_constants.XLV * saturation_mixing_ratio + ) + + if moist_static_energy >= saturation_moist_static_energy: + moist_static_energy = saturation_moist_static_energy + + +@function +def get_interp( + p, + t, + vapor, +): + """interpolate temperature and water vapor between a set of pressures. + + Args: + p + t + vapor + + Returns: + t_new (Float) + vapor_forced (Float) + """ + # bunch of interal constants, not to be used outside of this function + rd = 287.06 + rv = 461.52 + rcpd = 1004.71 + rtt = 273.16 + rhoh2o = 1000.0 + rlvtt = 2.5008e6 + rlstt = 2.8345e6 + retv = rv / rd - 1.0 + rlmlt = rlstt - rlvtt + rcpv = 4.0 * rv + r2es = 611.21 * rd / rv + r3les = 17.502 + r3ies = 22.587 + r4les = 32.19 + r4ies = -0.7 + r5les = r3les * (rtt - r4les) + r5ies = r3ies * (rtt - r4ies) + r5alvcp = r5les * rlvtt / rcpd + r5alscp = r5ies * rlstt / rcpd + ralvdcp = rlvtt / rcpd + ralsdcp = rlstt / rcpd + ralfdcp = rlmlt / rcpd + rtwat = rtt + rtber = rtt - 5.0 + rtbercu = rtt - 5.0 + rtice = rtt - 23.0 + rticecu = rtt - 23.0 + rtwat_rtice_r = 1.0 / (rtwat - rtice) + rtwat_rticecu_r = 1.0 / (rtwat - rticecu) + rvtmp2 = rcpv / rcpd - 1.0 + zqmax = 0.5 + + pt = t # k + pq = vapor # kg/kg + psp = p * 100.0 # hpa + + zqp = 1.0 / psp + + iteration = 0 + while iteration <= 1: + ptare = pt + + foealfcu = min(1.0, ((max(rticecu, min(rtwat, ptare)) - rticecu) * rtwat_rticecu_r) ** 2) + foeewmcu = r2es * ( + foealfcu * exp(r3les * (ptare - rtt) / (ptare - r4les)) + + (1.0 - foealfcu) * exp(r3ies * (ptare - rtt) / (ptare - r4ies)) + ) + zqsat = foeewmcu * zqp + + zcor = 1.0 / (1.0 - retv * zqsat) + zqsat = zqsat * zcor + + foedemcu = foealfcu * r5alvcp * (1.0 / (ptare - r4les) ** 2) + (1.0 - foealfcu) * r5alscp * ( + 1.0 / (ptare - r4ies) ** 2 + ) + + zcond1 = (pq - zqsat) / (1.0 + zqsat * zcor * foedemcu) + + foeldcpmcu = foealfcu * ralvdcp + (1.0 - foealfcu) * ralsdcp + pt = pt + foeldcpmcu * zcond1 + pq = pq - zcond1 + + iteration += 1 + + t_new = pt + vapor_forced = pq + + return t_new, vapor_forced + + +def environment_cloud_levels( + p: FloatField, + p_surface: FloatFieldIJ, + p_cloud_levels: FloatField, + topography_height_no_negative: FloatFieldIJ, + geopotential_height: FloatField, + geopotential_height_cloud_levels: FloatField, + t: FloatField, + t_surface: FloatFieldIJ, + t_cloud_levels: FloatField, + vapor: FloatField, + vapor_cloud_levels: FloatField, + u: FloatField, + v: FloatField, + u_cloud_levels: FloatField, + v_cloud_levels: FloatField, + environment_saturation_mixing_ratio: FloatField, + environment_saturation_mixing_ratio_cloud_levels: FloatField, + environment_moist_static_energy: FloatField, + environment_moist_static_energy_cloud_levels: FloatField, + environment_saturation_moist_static_energy: FloatField, + environment_saturation_moist_static_energy_cloud_levels: FloatField, + gamma_cloud_levels: FloatField, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """Compute environmental fields at cloud levels. + + Args: + p (FloatField) + p_surface (FloatFieldIJ) + p_cloud_levels (FloatField) + topography_height_no_negative (FloatFieldIJ) + geopotential_height (FloatField) + geopotential_height_cloud_levels (FloatField) + t (FloatField) + t_surface (FloatFieldIJ) + t_cloud_levels (FloatField) + vapor (FloatField) + vapor_cloud_levels (FloatField) + u (FloatField) + v (FloatField) + u_cloud_levels (FloatField) + v_cloud_levels (FloatField) + environment_saturation_mixing_ratio (FloatField) + environment_saturation_mixing_ratio_cloud_levels (FloatField) + environment_moist_static_energy (FloatField) + environment_moist_static_energy_cloud_levels (FloatField) + environment_saturation_moist_static_energy (FloatField) + environment_saturation_moist_static_energy_cloud_levels (FloatField) + gamma_cloud_levels (FloatField) + error_code (IntFieldIJ_Plume) + plume (Int) + """ + from __externals__ import CLOUD_LEVEL_GRID + + with computation(PARALLEL), interval(...): + # ensure no data from previous call slips through untouched + environment_saturation_mixing_ratio_cloud_levels = 0.0 + vapor_cloud_levels = 0.0 + environment_saturation_moist_static_energy_cloud_levels = 0.0 + environment_moist_static_energy_cloud_levels = 0.0 + geopotential_height_cloud_levels = 0.0 + p_cloud_levels = 0.0 + t_cloud_levels = 0.0 + gamma_cloud_levels = 0.0 + u_cloud_levels = 0.0 + v_cloud_levels = 0.0 + + with computation(PARALLEL), interval(1, -1): + if CLOUD_LEVEL_GRID == 2: + # original formulation + if error_code[0, 0][plume] == 0: + environment_saturation_mixing_ratio_cloud_levels = 0.5 * ( + environment_saturation_mixing_ratio[0, 0, -1] + environment_saturation_mixing_ratio + ) + vapor_cloud_levels = 0.5 * (vapor[0, 0, -1] + vapor) + environment_saturation_moist_static_energy_cloud_levels = 0.5 * ( + environment_saturation_moist_static_energy[0, 0, -1] + + environment_saturation_moist_static_energy + ) + environment_moist_static_energy_cloud_levels = 0.5 * ( + environment_moist_static_energy[0, 0, -1] + environment_moist_static_energy + ) + if ( + environment_moist_static_energy_cloud_levels + > environment_saturation_moist_static_energy_cloud_levels + ): + environment_moist_static_energy_cloud_levels = ( + environment_saturation_moist_static_energy_cloud_levels + ) + geopotential_height_cloud_levels = 0.5 * (geopotential_height[0, 0, -1] + geopotential_height) + p_cloud_levels = 0.5 * (p[0, 0, -1] + p) + t_cloud_levels = 0.5 * (t[0, 0, -1] + t) + gamma_cloud_levels = ( + (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + * ( + cumulus_parameterization_constants.XLV + / (cumulus_parameterization_constants.RV * t_cloud_levels * t_cloud_levels) + ) + * environment_saturation_mixing_ratio_cloud_levels + ) + u_cloud_levels = 0.5 * (u[0, 0, -1] + u) + v_cloud_levels = 0.5 * (v[0, 0, -1] + v) + + with computation(FORWARD), interval(0, 1): + if CLOUD_LEVEL_GRID == 2: + # original formulation + if error_code[0, 0][plume] == 0: + environment_saturation_mixing_ratio_cloud_levels = environment_saturation_mixing_ratio + vapor_cloud_levels = vapor + environment_saturation_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * topography_height_no_negative + + cumulus_parameterization_constants.CP * t + + cumulus_parameterization_constants.XLV * environment_saturation_mixing_ratio + ) + environment_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * topography_height_no_negative + + cumulus_parameterization_constants.CP * t + + cumulus_parameterization_constants.XLV * vapor + ) + geopotential_height_cloud_levels = topography_height_no_negative + p_cloud_levels = p_surface + t_cloud_levels = t + gamma_cloud_levels = ( + cumulus_parameterization_constants.XLV + / cumulus_parameterization_constants.CP + * ( + cumulus_parameterization_constants.XLV + / (cumulus_parameterization_constants.RV * t_cloud_levels * t_cloud_levels) + ) + * environment_saturation_mixing_ratio_cloud_levels + ) + u_cloud_levels = u + v_cloud_levels = v + + with computation(FORWARD), interval(0, 1): + if CLOUD_LEVEL_GRID == 0: + # weighted mean + if error_code[0, 0][plume] == 0: + p_cloud_levels = p_surface + geopotential_height_cloud_levels = topography_height_no_negative + + with computation(FORWARD), interval(0, -2): + if CLOUD_LEVEL_GRID == 0: + # weighted mean + if error_code[0, 0][plume] == 0: + p_cloud_levels[0, 0, 1] = 2.0 * p - p_cloud_levels + geopotential_height_cloud_levels[0, 0, 1] = ( + 2.0 * geopotential_height - geopotential_height_cloud_levels + ) + + with computation(FORWARD), interval(0, -2): + if CLOUD_LEVEL_GRID == 0: + # weighted mean + if error_code[0, 0][plume] == 0: + p1 = abs((p[0, 0, 1] - p_cloud_levels[0, 0, 1]) / (p[0, 0, 1] - p)) + p2 = abs((p_cloud_levels[0, 0, 1] - p) / (p[0, 0, 1] - p)) + t_cloud_levels[0, 0, 1] = p1 * t + p2 * t[0, 0, 1] + u_cloud_levels[0, 0, 1] = p1 * u + p2 * u[0, 0, 1] + v_cloud_levels[0, 0, 1] = p1 * v + p2 * v[0, 0, 1] + vapor_cloud_levels[0, 0, 1] = p1 * vapor + p2 * vapor[0, 0, 1] + environment_moist_static_energy_cloud_levels[0, 0, 1] = ( + p1 * environment_moist_static_energy + p2 * environment_moist_static_energy[0, 0, 1] + ) + environment_saturation_mixing_ratio_cloud_levels[0, 0, 1] = ( + p1 * environment_saturation_mixing_ratio + + p2 * environment_saturation_mixing_ratio[0, 0, 1] + ) + environment_saturation_moist_static_energy_cloud_levels[0, 0, 1] = ( + p1 * environment_saturation_moist_static_energy + + p2 * environment_saturation_moist_static_energy[0, 0, 1] + ) + if ( + environment_moist_static_energy_cloud_levels[0, 0, 1] + > environment_saturation_moist_static_energy_cloud_levels[0, 0, 1] + ): + environment_moist_static_energy_cloud_levels[0, 0, 1] = ( + environment_saturation_moist_static_energy_cloud_levels[0, 0, 1] + ) + + gamma_cloud_levels[0, 0, 1] = ( + (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + * ( + cumulus_parameterization_constants.XLV + / ( + cumulus_parameterization_constants.RV + * t_cloud_levels[0, 0, 1] + * t_cloud_levels[0, 0, 1] + ) + ) + * environment_saturation_mixing_ratio_cloud_levels[0, 0, 1] + ) + + with computation(FORWARD), interval(0, 1): + if CLOUD_LEVEL_GRID == 0: + # weighted mean + # surface level: using level 0 and 1 to determine level 0 + if error_code[0, 0][plume] == 0: + p1 = abs(p - p_cloud_levels) + p2 = abs(p_cloud_levels[0, 0, 1] - p_cloud_levels) + + ct1 = (p1 + p2) / p2 + ct2 = p1 / p2 + + t_cloud_levels = ct1 * t - ct2 * t_cloud_levels[0, 0, 1] + vapor_cloud_levels = ct1 * vapor - ct2 * vapor_cloud_levels[0, 0, 1] + + u_cloud_levels = ct1 * u - ct2 * u_cloud_levels[0, 0, 1] + v_cloud_levels = ct1 * v - ct2 * v_cloud_levels[0, 0, 1] + environment_saturation_mixing_ratio_cloud_levels = ( + ct1 * environment_saturation_mixing_ratio + - ct2 * environment_saturation_mixing_ratio_cloud_levels[0, 0, 1] + ) + + environment_saturation_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * geopotential_height_cloud_levels + + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV + * environment_saturation_mixing_ratio_cloud_levels + ) + environment_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * geopotential_height_cloud_levels + + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV * vapor_cloud_levels + ) + + if ( + environment_moist_static_energy_cloud_levels + > environment_saturation_moist_static_energy_cloud_levels + ): + environment_moist_static_energy_cloud_levels = ( + environment_saturation_moist_static_energy_cloud_levels + ) + + gamma_cloud_levels = ( + cumulus_parameterization_constants.XLV + / cumulus_parameterization_constants.CP + * ( + cumulus_parameterization_constants.XLV + / (cumulus_parameterization_constants.RV * t_cloud_levels * t_cloud_levels) + ) + * environment_saturation_mixing_ratio_cloud_levels + ) + + with computation(BACKWARD), interval(1, -1): + if CLOUD_LEVEL_GRID == 1: + # based on Tiedke (1989) + if error_code[0, 0][plume] == 0: + environment_saturation_mixing_ratio_cloud_levels = environment_saturation_mixing_ratio + vapor_cloud_levels = vapor + p_cloud_levels = 0.5 * (p[0, 0, -1] + p) + geopotential_height_cloud_levels = 0.5 * (geopotential_height[0, 0, -1] + geopotential_height) + t_cloud_levels = ( + max( + cumulus_parameterization_constants.CP * t[0, 0, -1] + + constants.MAPL_GRAV * geopotential_height[0, 0, -1], + cumulus_parameterization_constants.CP * t + constants.MAPL_GRAV * geopotential_height, + ) + - constants.MAPL_GRAV * geopotential_height_cloud_levels + ) / cumulus_parameterization_constants.CP + + if environment_saturation_mixing_ratio < cumulus_parameterization_constants.MAX_QSAT: + t_cloud_levels, environment_saturation_mixing_ratio_cloud_levels = get_interp( + p=p_cloud_levels, + t=t_cloud_levels, + vapor=environment_saturation_mixing_ratio_cloud_levels, + ) + + vapor_cloud_levels = ( + min(vapor, environment_saturation_mixing_ratio) + + environment_saturation_mixing_ratio_cloud_levels + - environment_saturation_mixing_ratio + ) + vapor_cloud_levels = max(vapor_cloud_levels, 0.0) + + with computation(FORWARD), interval(0, 1): + if CLOUD_LEVEL_GRID == 1: + # based on Tiedke (1989) + if error_code[0, 0][plume] == 0: + # surface + environment_saturation_mixing_ratio_cloud_levels = environment_saturation_mixing_ratio + vapor_cloud_levels = vapor + geopotential_height_cloud_levels = topography_height_no_negative + p_cloud_levels = p_surface + + t_cloud_levels = ( + cumulus_parameterization_constants.CP * t + + constants.MAPL_GRAV * geopotential_height + - constants.MAPL_GRAV * geopotential_height_cloud_levels + ) / cumulus_parameterization_constants.CP + + environment_saturation_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * geopotential_height_cloud_levels + + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV + * environment_saturation_mixing_ratio_cloud_levels + ) + environment_moist_static_energy_cloud_levels = ( + constants.MAPL_GRAV * geopotential_height_cloud_levels + + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV * vapor_cloud_levels + ) + + gamma_cloud_levels = ( + cumulus_parameterization_constants.XLV + / cumulus_parameterization_constants.CP + * ( + cumulus_parameterization_constants.XLV + / (cumulus_parameterization_constants.RV * t_cloud_levels * t_cloud_levels) + ) + * environment_saturation_mixing_ratio_cloud_levels + ) + u_cloud_levels = u + v_cloud_levels = v + + with computation(BACKWARD), interval(1, -1): + if CLOUD_LEVEL_GRID == 1: + # based on Tiedke (1989) + if error_code[0, 0][plume] == 0: + p1 = max( + cumulus_parameterization_constants.CP * t_cloud_levels + + constants.MAPL_GRAV * geopotential_height_cloud_levels, + cumulus_parameterization_constants.CP * t_cloud_levels[0, 0, -1] + + constants.MAPL_GRAV * geopotential_height_cloud_levels[0, 0, -1], + ) + t_cloud_levels = ( + p1 - constants.MAPL_GRAV * geopotential_height_cloud_levels + ) / cumulus_parameterization_constants.CP + + environment_saturation_moist_static_energy_cloud_levels = ( + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV + * environment_saturation_mixing_ratio_cloud_levels + + constants.MAPL_GRAV * geopotential_height_cloud_levels + ) + environment_moist_static_energy_cloud_levels = ( + cumulus_parameterization_constants.CP * t_cloud_levels + + cumulus_parameterization_constants.XLV * vapor_cloud_levels + + constants.MAPL_GRAV * geopotential_height_cloud_levels + ) + if ( + environment_moist_static_energy_cloud_levels + > environment_saturation_moist_static_energy_cloud_levels + ): + environment_moist_static_energy_cloud_levels = ( + environment_saturation_moist_static_energy_cloud_levels + ) + + gamma_cloud_levels = ( + (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + * ( + cumulus_parameterization_constants.XLV + / (cumulus_parameterization_constants.RV * t_cloud_levels * t_cloud_levels) + ) + * environment_saturation_mixing_ratio_cloud_levels + ) + u_cloud_levels = u + v_cloud_levels = v + + +def environment_mass_flux( + error_code: IntFieldIJ_Plume, + epsilon_forced: FloatFieldIJ_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + environment_massflux: FloatField, + plume: Int, +): + """Compute environment mass flux. + + Args: + error_code (IntFieldIJ_Plume) + epsilon_forced (FloatFieldIJ_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + environment_massflux (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + environment_massflux = ( + normalized_massflux_updraft_forced[0, 0, 0][plume] + - epsilon_forced[0, 0][plume] * normalized_massflux_downdraft_forced[0, 0, 0][plume] + ) + else: + environment_massflux = 0.0 + + +def modify_environment_profiles( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + p_forced: FloatField, + t_new: FloatField, + t_modified: FloatField, + vapor_forced: FloatField, + vapor_modified: FloatField, + environment_moist_static_energy_forced: FloatField, + environment_moist_static_energy_modified: FloatField, + moist_static_energy_origin_level_forced: FloatFieldIJ, + moist_static_energy_origin_level_modified: FloatFieldIJ, + partition_liquid_ice: FloatField, + del_moist_static_energy_cloud_ensemble: FloatField, + del_t_cloud_ensemble: FloatField, + del_vapor_cloud_ensemble: FloatField, + del_cloud_liquid_cloud_ensemble: FloatField, + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + t_tendency_from_environmental_subsidence: FloatField, + moist_static_energy_tendency_from_environmental_subsidence: FloatField, + vapor_tendency_from_environmental_subsidence: FloatField, + arbitrary_numerical_parameter: FloatFieldIJ, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Update the environment profiles based on the computed tendencies. These updated profiles are only + used within the cumulus parameterization core, and changes to them have no effect on the overarching + model state. The output tendencies are used to update the model state at a later time. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + ocean_fraction (FloatFieldIJ) + p_forced (FloatField) + t_new (FloatField) + t_modified (FloatField) + vapor_forced (FloatField) + vapor_modified (FloatField) + environment_moist_static_energy_forced (FloatField) + environment_moist_static_energy_modified (FloatField) + moist_static_energy_origin_level_forced (FloatFieldIJ) + moist_static_energy_origin_level_modified (FloatFieldIJ) + partition_liquid_ice (FloatField) + del_moist_static_energy_cloud_ensemble (FloatField) + del_t_cloud_ensemble (FloatField) + del_vapor_cloud_ensemble (FloatField) + del_cloud_liquid_cloud_ensemble (FloatField) + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + t_tendency_from_environmental_subsidence (FloatField) + moist_static_energy_tendency_from_environmental_subsidence (FloatField) + vapor_tendency_from_environmental_subsidence (FloatField) + arbitrary_numerical_parameter (FloatFieldIJ) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import BOUNDARY_CONDITION_METHOD, COUPLE_MICROPHYSICS, k_end + + with computation(PARALLEL), interval(...): + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + + with computation(PARALLEL), interval(0, -1): + del_t_cloud_ensemble = 0.0 + + if error_code[0, 0][plume] == 0: + environment_moist_static_energy_modified = ( + del_moist_static_energy_cloud_ensemble * arbitrary_numerical_parameter + + environment_moist_static_energy_forced + ) + vapor_modified = ( + del_vapor_cloud_ensemble + del_cloud_liquid_cloud_ensemble + ) * arbitrary_numerical_parameter + vapor_forced + if vapor_modified <= 0.0: + vapor_modified = 1.0e-08 + + # do not feed del_t_cloud_ensemble with del_cloud_liquid_cloud_ensemble if the + # detrainment of liquid water will be used as a source for cloud microphysics + if COUPLE_MICROPHYSICS == True: + del_t_cloud_ensemble = (1.0 / cumulus_parameterization_constants.CP) * ( + del_moist_static_energy_cloud_ensemble + - cumulus_parameterization_constants.XLV * del_vapor_cloud_ensemble + ) + else: + # melt glac + del_t_cloud_ensemble = (1.0 / cumulus_parameterization_constants.CP) * ( + del_moist_static_energy_cloud_ensemble + - cumulus_parameterization_constants.XLV + * (del_vapor_cloud_ensemble + del_cloud_liquid_cloud_ensemble) + * ( + 1.0 + + (cumulus_parameterization_constants.XLF / cumulus_parameterization_constants.XLV) + * (1.0 - partition_liquid_ice) + ) + ) + + # adding dellaqc to dellaq: + del_vapor_cloud_ensemble = del_vapor_cloud_ensemble + del_cloud_liquid_cloud_ensemble + del_cloud_liquid_cloud_ensemble = 0.0 + + # melt glac + t_modified = ( + (1.0 / cumulus_parameterization_constants.CP) * del_moist_static_energy_cloud_ensemble + - (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + * ( + del_vapor_cloud_ensemble + + del_cloud_liquid_cloud_ensemble + * ( + 1.0 + + (cumulus_parameterization_constants.XLF / cumulus_parameterization_constants.XLV) + * (1.0 - partition_liquid_ice) + ) + ) + ) * arbitrary_numerical_parameter + t_new + + # temp tendency due to the environmental subsidence + t_tendency_from_environmental_subsidence = (1.0 / cumulus_parameterization_constants.CP) * ( + moist_static_energy_tendency_from_environmental_subsidence + - cumulus_parameterization_constants.XLV * vapor_tendency_from_environmental_subsidence + ) + + with computation(FORWARD), interval(-2, -1): + if error_code[0, 0][plume] == 0: + environment_moist_static_energy_modified = environment_moist_static_energy_forced + vapor_modified = vapor_forced + t_modified = t_new + if vapor_modified <= 0.0: + vapor_modified = 1.0e-08 + + with computation(FORWARD), interval(0, 1): + # new way for defining moist static energy at updraft origin level + if error_code[0, 0][plume] == 0: + # note that moist_static_energy_origin_level_forced already contains the contribuition from + # t_excess and vapor_excess + mse_boundary_condition = get_cloud_boundary_conditions( + field=del_moist_static_energy_cloud_ensemble, + scalar_perturbation=0, + p=p_forced, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + moist_static_energy_origin_level_modified = ( + mse_boundary_condition * arbitrary_numerical_parameter + + moist_static_energy_origin_level_forced + ) + + +class EnvironmentalSubsidence: + def __init__( + self, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + def __call__(self, *args, **kwds): + if self.config.APPLY_SUBSIDENCE_MICROPHYSICS != 0: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->EnvironmentalSubsidence this code" + "has not been impemented. You should have been caught before getting here by the config" + "checker. Beware, something likely failing in the config checker as well - you may be" + "unknowingly calling other untested/unimplemented sections." + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/field_types.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/field_types.py new file mode 100644 index 000000000..073cee731 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/field_types.py @@ -0,0 +1,28 @@ +from ndsl.dsl.gt4py import IJ, IJK, Field +from ndsl.dsl.typing import Float, Int +from pyMoist.constants import NUMBER_OF_TRACERS +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) + + +# NOTE must cast to int because numpy types are not acceptable for data dimensions +# plume field types +IntFieldIJ_Plume = Field[IJ, (Int, int(NUMBER_OF_PLUMES))] +FloatFieldIJ_Plume = Field[IJ, (Float, int(NUMBER_OF_PLUMES))] +FloatField_Plume = Field[IJK, (Float, int(NUMBER_OF_PLUMES))] + +# internal ensemble field types +FloatFieldIJ_ensemble_1 = Field[IJ, (Float, int(MAXENS1))] +FloatFieldIJ_ensemble_2 = Field[IJ, (Float, int(MAXENS2))] +FloatFieldIJ_ensemble_3 = Field[IJ, (Float, int(MAXENS3))] +FloatFieldIJ_Ensemble = Field[IJ, (Float, (int(MAXENS1 * MAXENS2 * MAXENS3)))] + +# NOTE THESE SHOULD BE MIGRATED TO PYMOIST GLOBAL FIELD TYPES +# convection tracer field types +FloatFieldIJ_ConvectionTracers = Field[IJ, (Float, int(NUMBER_OF_TRACERS))] +FloatField_ConvectionTracers = Field[IJK, (Float, int(NUMBER_OF_TRACERS))] +FloatField_ConvectionTracers_Plume = Field[IJK, (Float, (int(NUMBER_OF_PLUMES), int(NUMBER_OF_TRACERS)))] diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/get_levels.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/get_levels.py new file mode 100644 index 000000000..ca680976e --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/get_levels.py @@ -0,0 +1,974 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, exp, function, interval +from ndsl.dsl.typing import BoolFieldIJ, Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.logging import ndsl_log +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import ( + compute_dewpoint, + get_cloud_boundary_conditions, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +def find_maximum_updraft_origin_level( + geopotential_height: FloatField, + topography_height_no_negative: FloatFieldIJ, + error_code: IntFieldIJ_Plume, + maximum_updraft_origin_level: IntFieldIJ, + MAX_UPDRAFT_ORIGIN_HEIGHT: Float, + plume: Int, +): + """Determine the highest level from which an updraft may originate. + + Args: + geopotential_height (FloatField) + topography_height_no_negative (FloatFieldIJ) + error_code (IntFieldIJ_Plume) + maximum_updraft_origin_level (IntFieldIJ) + MAX_UPDRAFT_ORIGIN_HEIGHT (Float) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + stop_computation: BoolFieldIJ = False + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0 and stop_computation == 0: + if geopotential_height > MAX_UPDRAFT_ORIGIN_HEIGHT + topography_height_no_negative: + maximum_updraft_origin_level = K + stop_computation: BoolFieldIJ = True + + +def find_detrainmet_start_level( + geopotential_height: FloatField, + topography_height_no_negative: FloatFieldIJ, + error_code: IntFieldIJ_Plume, + detrainment_start_level: IntFieldIJ, + DETRAINMENT_CRITICAL_DEPTH: Float, + plume: Int, +): + """Determine the lowest level at which detrainment may occur. + + Args: + geopotential_height (FloatField) + topography_height_no_negative (FloatFieldIJ) + error_code (IntFieldIJ_Plume) + detrainment_start_level (IntFieldIJ) + DETRAINMENT_CRITICAL_DEPTH (Float) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + stop_computation: BoolFieldIJ = False + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0 and stop_computation == 0: + if geopotential_height > DETRAINMENT_CRITICAL_DEPTH + topography_height_no_negative: + detrainment_start_level = K + stop_computation: BoolFieldIJ = True + + +def find_highest_moist_static_energy_level( + moist_static_energy: FloatField, + error_code: IntFieldIJ_Plume, + maximum_updraft_origin_level: IntFieldIJ, + updraft_origin_level: IntFieldIJ_Plume, + plume: Int, +): + """Detemine the level with the highest moist static energy. This level will + be used as the starting point for any subsequent updrafts + + Args: + moist_static_energy (FloatField) + error_code (IntFieldIJ_Plume) + maximum_updraft_origin_level (IntFieldIJ) + updraft_origin_level (IntFieldIJ_Plume) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + # prefil output + updraft_origin_level[0, 0][plume] = 0 + + if plume == cumulus_parameterization_constants.SHALLOW: + # start at surface for shallow plume + start_level: IntFieldIJ = 0 + else: + # start above surface + start_level: IntFieldIJ = 1 + + if error_code[0, 0][plume] == 0: + level = start_level + while level <= maximum_updraft_origin_level + 1: + if moist_static_energy.at(K=level) > moist_static_energy.at( + K=max(start_level, updraft_origin_level[0, 0][plume]) + ): + updraft_origin_level[0, 0][plume] = level + level += 1 + + updraft_origin_level[0, 0][plume] = updraft_origin_level[0, 0][plume] + start_level - 1 + updraft_origin_level[0, 0][plume] = max(updraft_origin_level[0, 0][plume], start_level) + + if plume == cumulus_parameterization_constants.SHALLOW: + # for shallow plumes + updraft_origin_level[0, 0][plume] = min(2, updraft_origin_level[0, 0][plume]) + if updraft_origin_level[0, 0][plume] > maximum_updraft_origin_level: + error_code[0, 0][plume] = 2 + else: + # other plumes + if updraft_origin_level[0, 0][plume] > maximum_updraft_origin_level: + updraft_origin_level[0, 0][plume] = start_level + + +@function +def get_lcl(p_source, t_source, vapor_source): + # internal constants + cpg = 102.45 + rgas = 287.0 + cp = 1004.0 + p00 = 1.0e5 + g = 9.80 + rocp = rgas / cp + p00i = 1.0 / p00 + cpor = cp / rgas + cpi = 1.0 / cp + p00k = 26.870941 + p00ki = 1.0 / p00k + + # convert to pascals + p_source = 100 * p_source + + # simpler, cheaper method + dewpoint = compute_dewpoint(p_source, vapor_source) + t_lcl = dewpoint - (0.001296 * dewpoint + 0.1963) * (t_source - dewpoint) + p_lcl = p_source * (t_lcl / t_source) ** cpor + dz_lcl = 127 * (t_source - dewpoint) + if dz_lcl <= 0.0: + dz_lcl = -999.0 + + return t_lcl, p_lcl, dz_lcl + + +def find_lcl( + p: FloatField, + p_cloud_levels: FloatField, + t_excess: FloatFieldIJ, + t_cloud_levels_forced: FloatField, + t_perturbation: FloatField, + vapor_excess: FloatFieldIJ, + vapor_cloud_levels_forced: FloatField, + omega: FloatField, + air_density: FloatField, + geopotential_height_cloud_levels: FloatField, + topography_height_no_negative: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + updraft_origin_level: IntFieldIJ_Plume, + grid_length: FloatFieldIJ, + lcl_level: IntFieldIJ_Plume, + error_code: IntFieldIJ_Plume, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Determine the LCL level. For shallow plumes, if LCL is buoyant, calculate a temperature perturbation. + + Args: + p (FloatField) + p_cloud_levels (FloatField) + t_excess (FloatFieldIJ) + t_cloud_levels_forced (FloatField) + t_perturbation (FloatField) + vapor_excess (FloatFieldIJ) + vapor_cloud_levels_forced (FloatField) + omega (FloatField) + air_density (FloatField) + geopotential_height_cloud_levels (FloatField) + topography_height_no_negative (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + updraft_origin_level (IntFieldIJ_Plume) + grid_length (FloatFieldIJ) + lcl_level (IntFieldIJ_Plume) + error_code (IntFieldIJ_Plume) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import ADV_TRIGGER, BOUNDARY_CONDITION_METHOD, k_end + + with computation(PARALLEL), interval(...): + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # default value + lcl_level[0, 0][plume] = updraft_origin_level[0, 0][plume] + + # get conditions for source parcel + vapor_source = get_cloud_boundary_conditions( + field=vapor_cloud_levels_forced, + scalar_perturbation=vapor_excess, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + t_source = get_cloud_boundary_conditions( + field=t_cloud_levels_forced, + scalar_perturbation=t_excess, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + p_source = get_cloud_boundary_conditions( + field=p_cloud_levels, + scalar_perturbation=0, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + + # initialize 2d temporaries + p_lcl: FloatFieldIJ = 0.0 + t_lcl: FloatFieldIJ = 0.0 + dz_lcl: FloatFieldIJ = 0.0 + z_source: FloatFieldIJ = 0.0 + stop_computation: BoolFieldIJ = False + + p_lcl, t_lcl, dz_lcl = get_lcl(p_source=p_source, t_source=t_source, vapor_source=vapor_source) + + if dz_lcl >= 0.0: + z_source = get_cloud_boundary_conditions( + field=geopotential_height_cloud_levels, + scalar_perturbation=0, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0: + if dz_lcl >= 0.0: + if geopotential_height_cloud_levels >= z_source + dz_lcl and stop_computation == False: + lcl_level[0, 0][plume] = max(K, updraft_origin_level[0, 0][plume]) + stop_computation = True + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + lcl_level[0, 0][plume] = min(lcl_level[0, 0][plume], k_end - 5) + if ( + cumulus_parameterization_constants.USE_LCL == True + and plume == cumulus_parameterization_constants.MID + ): + error_code[0, 0][plume] = 21 + + if ADV_TRIGGER == 1 and plume == cumulus_parameterization_constants.SHALLOW: + wkf: FloatFieldIJ = 0.02 # m/s + + level = lcl_level[0, 0][plume] + dz_lcl = geopotential_height_cloud_levels - topography_height_no_negative + + ckf = wkf + if dz_lcl <= 2.0e3: + ckf = wkf * dz_lcl / 2000 + + wkflcl = ( + -(omega.at(K=max(0, level - 1))) / air_density.at(K=max(0, level - 1)) + + omega / air_density + + omega[0, 0, 1] / air_density[0, 0, 1] / (3.0 * constants.MAPL_GRAV) + ) + + # check to see if cloud is buoyant using fritsch-chappell trigger + # function described in kain and fritsch (1992)...w0avg is an + # aproximate value for the running-mean grid-scale vertical + # velocity, which gives smoother fields of convective initiation + # than the instantaneous value...formula relating temperature + # perturbation to vertical velocity has been used with the most + # success at grid lengths near 25 km. for different grid-lengths, + # adjust vertical velocity to equivalent value for 25 km grid + # length, assuming linear dependence of w on grid length... + if grid_length >= 25.0e3: + wkflcl = wkflcl * grid_length / 25.0e3 - ckf + else: + wkflcl = wkflcl - ckf + + with computation(FORWARD), interval(0, 1): + if ( + error_code[0, 0][plume] == 0 + and ADV_TRIGGER == 1 + and plume == cumulus_parameterization_constants.SHALLOW + ): + # Kain (2004) Eq. 1 + t_perturbation = 4.64 * wkflcl ** (1.0 / 3.0) + + if t_perturbation > 2.0: + t_perturbation = 2.0 + elif t_perturbation < -2.0: + t_perturbation = -2.0 + + +def set_start_level(lcl_level: IntFieldIJ_Plume, start_level: IntFieldIJ, plume: Int): + """Copy LCL level data into another field called "start_level", which will modified later when searching + for an equilibrium level. + + Args: + lcl_level (IntFieldIJ_Plume) + start_level (IntFieldIJ) + plume (Int) + """ + with computation(FORWARD), interval(...): + start_level = lcl_level[0, 0][plume] + + +def get_convective_cloud_base_level( + error_code: IntFieldIJ_Plume, + lcl_level: IntFieldIJ_Plume, + cloud_moist_static_energy_forced_transported: FloatField, + cap_max: FloatFieldIJ, + updraft_origin_level: IntFieldIJ_Plume, + start_level: IntFieldIJ, + moist_static_energy_origin_level_forced: FloatFieldIJ, + updraft_lfc_level: IntFieldIJ_Plume, + maximum_updraft_origin_level: IntFieldIJ, + negative_buoyancy_depth: FloatFieldIJ, + frh_lfc: FloatFieldIJ, + geopotential_height_cloud_levels_forced: FloatField, + entrainment_rate: FloatField_Plume, + environment_moist_static_energy_forced: FloatField, + environment_moist_static_energy_cloud_levels_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + t_excess: FloatFieldIJ, + vapor_excess: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + p_cloud_levels_forced: FloatField_Plume, + vapor_forced: FloatField, + environment_saturation_mixing_ratio_forced: FloatField, + ocean_fraction: FloatFieldIJ, + cap_max_increment: FloatFieldIJ, + t_perturbation: FloatField, + p_forced: FloatField, + cloud_top_level: IntFieldIJ_Plume, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Determine level of free convection (LFC) for the source parcel. + + May also modify cloud_top_level, depending on LFC location. + + This stencil contains an open-ended vertical solver with multiple nested K-intervals. + To implement this properly, the entire stencil has been constructed on an interval(0, 1), + and all K read/writes have been done with absolute indexes or relative offsets. The alternative is + to break this into a series of stencils (at least seven, based on current understanding of the structure + of this code) and pass data between them using a much larger number of locals. + + The current method (single stencil) is horrible for optimization and speed, but was chosen nonetheless + because it allows for more readable code and should allow for a more seamless implementation of a proper + solution once the necessary tool/feature has been developed. + + Args: + error_code (IntFieldIJ_Plume): _description_ + lcl_level (IntFieldIJ_Plume): _description_ + cloud_moist_static_energy_forced_transported (FloatField): _description_ + cap_max (FloatFieldIJ): _description_ + updraft_origin_level (IntFieldIJ_Plume): _description_ + start_level (IntFieldIJ): _description_ + moist_static_energy_origin_level_forced (FloatFieldIJ): _description_ + updraft_lfc_level (IntFieldIJ_Plume): _description_ + maximum_updraft_origin_level (IntFieldIJ): _description_ + negative_buoyancy_depth (FloatFieldIJ): _description_ + frh_lfc (FloatFieldIJ): _description_ + geopotential_height_cloud_levels_forced (FloatField): _description_ + entrainment_rate (FloatField_Plume): _description_ + environment_moist_static_energy_forced (FloatField): _description_ + environment_moist_static_energy_cloud_levels_forced (FloatField): _description_ + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField): _description_ + t_excess (FloatFieldIJ): _description_ + vapor_excess (FloatFieldIJ): _description_ + add_buoyancy (FloatFieldIJ): _description_ + p_cloud_levels_forced (FloatField_Plume): _description_ + vapor_forced (FloatField): _description_ + environment_saturation_mixing_ratio_forced (FloatField): _description_ + ocean_fraction (FloatFieldIJ): _description_ + cap_max_increment (FloatFieldIJ): _description_ + t_perturbation (FloatField): _description_ + p_forced (FloatField): _description_ + cloud_top_level (IntFieldIJ_Plume): _description_ + AVERAGE_LAYER_DEPTH (Float): _description_ + plume (Int): _description_ + """ + from __externals__ import ( + BOUNDARY_CONDITION_METHOD, + MOIST_TRIGGER, + OVERSHOOT, + USE_MEMORY, + ZERO_DIFF, + k_end, + ) + + with computation(PARALLEL), interval(...): + # prefill some fields + cloud_moist_static_energy_forced_transported = 0.0 + dby = 0.0 + + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + + with computation(FORWARD), interval(0, 1): + # internal constants + frh_crit_O = 0.7 + frh_crit_L = 0.7 + + # prefill some fields + start_level_internal: IntFieldIJ = 0 + cap_max_internal = cap_max + + # initialize some 2d temporaries + dzh: FloatFieldIJ = 0.0 + found_level: BoolFieldIJ = False + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and ZERO_DIFF == 1: + start_level_internal = lcl_level[0, 0][plume] + else: + start_level_internal = start_level + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= start_level_internal: + cloud_moist_static_energy_forced_transported = ( + moist_static_energy_origin_level_forced # assumed no entraiment between these layers + ) + + # determine the level of convective cloud base (updraft_lfc_level) + + with computation(FORWARD), interval(0, 1): + skip_last_check = False + updraft_lfc_level[0, 0][plume] = maximum_updraft_origin_level + 3 + negative_buoyancy_depth = 0.0 + frh_lfc = 0.0 + if error_code[0, 0][plume] == 0: + continue_outer_while_loop = True + while error_code[0, 0][plume] == 0 and continue_outer_while_loop == True: + updraft_lfc_level[0, 0][plume] = start_level_internal + level = start_level_internal + 1 + while level <= maximum_updraft_origin_level + 3: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, level] + - geopotential_height_cloud_levels_forced[0, 0, level - 1] + ) + cloud_moist_static_energy_forced_transported[0, 0, level] = ( + (1.0 - 0.5 * entrainment_rate[0, 0, level - 1][plume] * dz) + * cloud_moist_static_energy_forced_transported[0, 0, level - 1] + + entrainment_rate[0, 0, level - 1][plume] + * dz + * environment_moist_static_energy_forced[0, 0, level - 1] + ) / (1.0 + 0.5 * entrainment_rate[0, 0, level - 1][plume] * dz) + if level == start_level_internal + 1: + modification = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + cloud_moist_static_energy_forced_transported[0, 0, level] = ( + cloud_moist_static_energy_forced_transported[0, 0, level] + modification + ) + level += 1 + + continue_inner_while_loop = True + while ( + cloud_moist_static_energy_forced_transported.at(K=updraft_lfc_level[0, 0][plume]) + < environment_saturation_moist_static_energy_cloud_levels_forced.at( + K=updraft_lfc_level[0, 0][plume] + ) + ) and continue_inner_while_loop == True: + updraft_lfc_level[0, 0][plume] = updraft_lfc_level[0, 0][plume] + 1 + if updraft_lfc_level[0, 0][plume] > maximum_updraft_origin_level + 2: + error_code[0, 0][plume] = 3 + continue_inner_while_loop = False + + if error_code[0, 0][plume] != 0: + continue_outer_while_loop = False + skip_last_check = True + + if continue_outer_while_loop == True: + # cloud base pressure and max moist static energy pressure + # i.e., the depth (in mb) of the layer of negative buoyancy + negative_buoyancy_depth = -( + p_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume], ddim=[plume]) + - p_cloud_levels_forced.at(K=start_level_internal, ddim=[plume]) + ) + + if MOIST_TRIGGER == 1: + frh_lfc = 0.0 + dzh = 0 + level = updraft_origin_level[0, 0][plume] + while level <= updraft_lfc_level[0, 0][plume]: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, level] + - geopotential_height_cloud_levels_forced[0, 0, max(level - 1, 0)] + ) + frh_lfc = frh_lfc + dz * ( + vapor_forced[0, 0, level] + / environment_moist_static_energy_forced[0, 0, level] + ) + dzh = dzh + dz + + frh_lfc = frh_lfc / (dzh + 1.0e-16) + frh_crit = frh_crit_O * ocean_fraction + frh_crit_L * (1.0 - ocean_fraction) + + fx = ( + (2.0 / 0.78) * exp(-((frh_lfc - frh_crit) ** 2)) * (frh_lfc - frh_crit) + ) # exponential + fx = max(-1.0, min(1.0, fx)) + + del_cap_max = fx * cap_max_increment + + cap_max_internal = min(max(cap_max + del_cap_max, 10.0), 150.0) + + # test if the air parcel has enough energy to reach the positive buoyant region + if cap_max_internal > negative_buoyancy_depth: + continue_outer_while_loop = False + skip_last_check = True + + if continue_outer_while_loop == True: + # if am here -> updraft_lfc_level not found for air parcels from updraft_origin_level + updraft_origin_level[0, 0][plume] = updraft_origin_level[0, 0][plume] + 1 + + if USE_MEMORY == 20: + cap_max_internal = cap_max_internal + cap_max_increment + + # get new moist_static_energy_origin_level_forced + modification = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + moist_static_energy_origin_level_forced = get_cloud_boundary_conditions( + field=environment_moist_static_energy_cloud_levels_forced, + scalar_perturbation=modification, + p=p_forced, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=True, + perturbation_field=t_perturbation, + ) + + start_level_internal = start_level_internal + 1 + cloud_moist_static_energy_forced_transported[0, 0, start_level_internal] = ( + moist_static_energy_origin_level_forced + ) + + if skip_last_check == True: + # last check for updraft_lfc_level + if updraft_lfc_level[0, 0][plume] == 0: + error_code[0, 0][plume] = 33 + + # determine the level of neutral buoyancy - ktop + + with computation(FORWARD), interval(0, 1): + cloud_top_level[0, 0][plume] = k_end - 2 + + if error_code[0, 0][plume] == 0: + start_level_internal = updraft_lfc_level[0, 0][plume] + + with computation(FORWARD), interval(1, -2): + if error_code[0, 0][plume] == 0 and K > start_level_internal: + dz = geopotential_height_cloud_levels_forced - geopotential_height_cloud_levels_forced[0, 0, -1] + denom = 1.0 + 0.5 * entrainment_rate[0, 0, -1][plume] * dz + if denom == 0.0: + cloud_moist_static_energy_forced_transported = cloud_moist_static_energy_forced_transported[ + 0, 0, -1 + ] + else: + cloud_moist_static_energy_forced_transported = ( + (1.0 - 0.5 * entrainment_rate[0, 0, -1][plume] * dz) + * cloud_moist_static_energy_forced_transported[0, 0, -1] + + entrainment_rate[0, 0, -1][plume] + * dz + * environment_moist_static_energy_forced[0, 0, -1] + ) / denom + + with computation(FORWARD), interval(0, -2): + if error_code[0, 0][plume] == 0 and K > start_level_internal and found_level == False: + if ( + cloud_moist_static_energy_forced_transported + < environment_saturation_moist_static_energy_cloud_levels_forced + ): + cloud_top_level[0, 0][plume] = K - 1 + found_level = True + + with computation(FORWARD), interval(0, 1): + if ( + error_code[0, 0][plume] == 0 + and cloud_top_level[0, 0][plume] <= updraft_lfc_level[0, 0][plume] + 1 + ): + error_code[0, 0][plume] = 41 + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and OVERSHOOT >= 1.0e-6: + z_overshoot: FloatFieldIJ = (1.0 + OVERSHOOT) * geopotential_height_cloud_levels_forced.at( + K=cloud_top_level[0, 0][plume] + ) + + with computation(FORWARD), interval(0, 1): + # reset mask + found_level = False + + with computation(FORWARD), interval(0, -3): + if ( + error_code[0, 0][plume] == 0 + and OVERSHOOT >= 1.0e-6 + and K >= cloud_top_level[0, 0][plume] + and found_level == False + ): + if z_overshoot < geopotential_height_cloud_levels_forced: + cloud_top_level[0, 0][plume] = min(K - 1, k_end - 2) + found_level = True + + +def get_cloud_top( + entrainment_rate: FloatField_Plume, + environment_moist_static_energy_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + moist_static_energy_origin_level_forced: FloatFieldIJ, + updraft_lfc_level: IntFieldIJ_Plume, + geopotential_height_cloud_levels_forced: FloatField, + cloud_moist_static_energy_forced_transported: FloatField, + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + plume: Int, +): + """Determine the initial estimate for the cloud top level. This can include an overshooting top above the + equilibrium level, if the configuration option OVERSHOOT == True. + + Args: + entrainment_rate (FloatField_Plume) + environment_moist_static_energy_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + moist_static_energy_origin_level_forced (FloatFieldIJ) + updraft_lfc_level (IntFieldIJ_Plume) + geopotential_height_cloud_levels_forced (FloatField) + cloud_moist_static_energy_forced_transported (FloatField) + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + plume (Int) + """ + from __externals__ import OVERSHOOT, k_end + + with computation(PARALLEL), interval(...): + # default value + cloud_moist_static_energy_forced_transported = 0.0 + + with computation(FORWARD), interval(0, 1): + if plume != 0: + # default value + cloud_top_level[0, 0][plume] = k_end - 3 + + if plume != 0 and error_code[0, 0][plume] == 0: + start_level: IntFieldIJ = updraft_lfc_level[0, 0][plume] + + with computation(PARALLEL), interval(...): + if plume != 0 and error_code[0, 0][plume] == 0: + if K <= start_level: + cloud_moist_static_energy_forced_transported = moist_static_energy_origin_level_forced + + with computation(FORWARD), interval(1, None): + if plume != 0 and error_code[0, 0][plume] == 0: + if K > start_level and K < k_end - 2: + dz = ( + geopotential_height_cloud_levels_forced + - geopotential_height_cloud_levels_forced[0, 0, -1] + ) + + cloud_moist_static_energy_forced_transported = ( + (1.0 - 0.5 * entrainment_rate[0, 0, -1][plume] * dz) + * cloud_moist_static_energy_forced_transported[0, 0, -1] + + entrainment_rate[0, 0, -1][plume] + * dz + * environment_moist_static_energy_forced[0, 0, -1] + ) / (1.0 + 0.5 * entrainment_rate[0, 0, -1][plume] * dz) + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + # set up mask for next computation + stop_computation: BoolFieldIJ = False + + with computation(FORWARD), interval(...): + if plume != 0 and error_code[0, 0][plume] == 0: + if K > start_level and K < k_end - 2 and stop_computation == False: + # find the height where the parcel is no longer saturated + if ( + cloud_moist_static_energy_forced_transported + < environment_saturation_moist_static_energy_cloud_levels_forced + ): + cloud_top_level[0, 0][plume] = K - 1 + stop_computation = True + + with computation(FORWARD), interval(0, 1): + if plume != 0 and error_code[0, 0][plume] == 0: + if cloud_top_level[0, 0][plume] <= updraft_lfc_level[0, 0][plume] + 1: + error_code[0, 0][plume] = 41 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and plume != 0 and OVERSHOOT > 0: + z_overshoot = (1.0 + OVERSHOOT) * geopotential_height_cloud_levels_forced.at( + K=cloud_top_level[0, 0][plume] + ) + + with computation(FORWARD), interval(0, 1): + # set up mask for next computation + stop_computation: BoolFieldIJ = False + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and plume != 0 and OVERSHOOT > 0: + if K >= cloud_top_level[0, 0][plume] and K < k_end - 1 and stop_computation == False: + if z_overshoot < geopotential_height_cloud_levels_forced: + cloud_top_level[0, 0][plume] = min(K - 1, k_end - 3) + stop_computation = True + + +def cloud_top_checks( + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + geopotential_height_cloud_levels: FloatField, + error_code: IntFieldIJ_Plume, + last_error_code: IntFieldIJ, + updraft_lfc_level: IntFieldIJ_Plume, + MINIMUM_DEPTH: Float, + plume: Int, +): + """Perform a series of checks on the initial estimate of cloud top level. + + Args: + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + geopotential_height_cloud_levels (FloatField) + error_code (IntFieldIJ_Plume) + last_error_code (IntFieldIJ) + updraft_lfc_level (IntFieldIJ_Plume) + MINIMUM_DEPTH (Float) + plume (Int) + """ + # check if cloud_top_level is too low for deep convection + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.DEEP and error_code[0, 0][plume] == 0: + if p_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume], ddim=[plume]) > 400: + error_code[0, 0][plume] = 22 + + # check if cloud_top_level is too high for mid convection + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.MID and error_code[0, 0][plume] == 0: + if p_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume], ddim=[plume]) < 400: + error_code[0, 0][plume] = 22 + + # check if cloud_top_level is too high for shallow convection + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.SHALLOW and error_code[0, 0][plume] == 0: + if p_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume], ddim=[plume]) < 400: + error_code[0, 0][plume] = 23 + + # avoid double-counting plumes + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and last_error_code == 0: + if p_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume], ddim=[plume]) < 700: + error_code[0, 0][plume] = 27 + + # last checks for cloud_top_level + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if ( + geopotential_height_cloud_levels.at(K=cloud_top_level[0, 0][plume]) + - geopotential_height_cloud_levels.at(K=updraft_lfc_level[0, 0][plume]) + < MINIMUM_DEPTH + ): + error_code[0, 0][plume] = 5 + if cloud_top_level[0, 0][plume] <= updraft_lfc_level[0, 0][plume]: + error_code[0, 0][plume] = 5 + + +class MaximumUpdraftOriginLevel: + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._find_maximum_updraft_origin_level = stencil_factory.from_dims_halo( + func=find_maximum_updraft_origin_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + state: GF2020CumulusParameterizationState, + locals: GF2020CumulusParameterizationLocals, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._find_maximum_updraft_origin_level( + geopotential_height=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + MAX_UPDRAFT_ORIGIN_HEIGHT=plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class DowndraftDetrainmentLevel: + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._find_detrainmet_start_level = stencil_factory.from_dims_halo( + func=find_detrainmet_start_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + state: GF2020CumulusParameterizationState, + locals: GF2020CumulusParameterizationLocals, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._find_detrainmet_start_level( + geopotential_height=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + detrainment_start_level=locals.detrainment_start_level, + DETRAINMENT_CRITICAL_DEPTH=plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class HighestMoistStaticEnergyLevel: + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._find_highest_moist_static_energy_level = stencil_factory.from_dims_halo( + func=find_highest_moist_static_energy_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + state: GF2020CumulusParameterizationState, + locals: GF2020CumulusParameterizationLocals, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._find_highest_moist_static_energy_level( + moist_static_energy=locals.environment_moist_static_energy_cloud_levels_forced, + error_code=state.output.error_code, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + updraft_origin_level=state.output.updraft_origin_level, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class CloudTop: + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._updraft_rates_pdf = stencil_factory.from_dims_halo( + func=get_cloud_top, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"OVERSHOOT": cumulus_parameterization_config.OVERSHOOT}, + ) + + self._cloud_top_checks = stencil_factory.from_dims_halo( + func=cloud_top_checks, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + state: GF2020CumulusParameterizationState, + locals: GF2020CumulusParameterizationLocals, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + + if self.cumulus_parameterization_config.OVERSHOOT != 0: + ndsl_log.warning( + " GF2020 cumulus parameterization called CloudTop with " + "untested OVERSHOOT option. Running untested code... proceed with caution" + ) + + self._updraft_rates_pdf( + entrainment_rate=state.output.entrainment_rate, + moist_static_energy=locals.environment_moist_static_energy_forced, + saturation_moist_static_energy=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + moist_static_energy_origin_level=locals.moist_static_energy_origin_level_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + geopotential_height=locals.geopotential_height_cloud_levels_forced, + cloud_moist_static_energy=locals.cloud_moist_static_energy_forced_transported, + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._cloud_top_checks( + cloud_top_level=state.output.cloud_top_level, + p=state.output.p_cloud_levels_forced, + geopotential_height=locals.geopotential_height_cloud_levels, + error_code=state.output.error_code, + last_error_code=state.input.last_error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + MINIMUM_DEPTH=plume_dependent_constants.MINIMUM_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating.py new file mode 100644 index 000000000..72f7caef7 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating.py @@ -0,0 +1,71 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval, sqrt +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Int +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume + + +def kinetic_energy_to_heating( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + u: FloatField, + v: FloatField, + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + del_t_cloud_ensemble: FloatField, + plume: Int, +): + """Modify the output temperature tendency (which will modify the overarching model state) + based on kinetic energy in the updraft. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + u (FloatField) + v (FloatField) + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + del_t_cloud_ensemble (FloatField) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + dts: FloatFieldIJ = 0.0 + fpi: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + dp = (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) * 100.0 + + dts = dts - ( + ((del_u_cloud_ensemble * u) + (del_v_cloud_ensemble * v)) * dp / constants.MAPL_GRAV + ) + + fpi = fpi + ( + sqrt( + (del_u_cloud_ensemble * del_u_cloud_ensemble) + + (del_v_cloud_ensemble * del_v_cloud_ensemble) + ) + * dp + ) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + if fpi > 0.0: + if K <= cloud_top_level[0, 0][plume]: + fp = ( + sqrt( + ( + del_u_cloud_ensemble * del_u_cloud_ensemble + + del_v_cloud_ensemble * del_v_cloud_ensemble + ) + ) + / fpi + ) + + del_t_cloud_ensemble = del_t_cloud_ensemble + ( + fp * dts * constants.MAPL_GRAV / cumulus_parameterization_constants.CP + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/large_scale_forcing.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/large_scale_forcing.py new file mode 100644 index 000000000..36d3c7d25 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/large_scale_forcing.py @@ -0,0 +1,598 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, K, computation, interval +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Int, IntFieldIJ +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Ensemble, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) + + +def ddim_copy( + field_in: IntFieldIJ_Plume, + field_out: IntFieldIJ, + plume: Int, +): + """Copy a field with a data dimension. The index to be read form the data_dim field must be specified. + + Args: + field_in (IntFieldIJ_Plume): _description_ + field_out (IntFieldIJ): _description_ + plume (Int): _description_ + """ + with computation(FORWARD), interval(0, 1): + field_out = field_in[0, 0][plume] + + +def ensemble_forcing( + error_code: IntFieldIJ_Plume, + error_code_2: IntFieldIJ, + error_code_3: IntFieldIJ, + updraft_lfc_level: IntFieldIJ_Plume, + vapor_forced: FloatField, + ocean_fraction: FloatFieldIJ, + omega: FloatField, + cape_removal_time_scale: FloatFieldIJ, + arbitrary_numerical_parameter: FloatFieldIJ, + f_dicycle_modified: FloatFieldIJ, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_0_modified: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + cloud_workfunction_1_pbl: FloatFieldIJ, + mass_flux_ensemble: FloatFieldIJ_Ensemble, + internal_mass_flux_ensemble: FloatFieldIJ_Ensemble, + precipitation_ensemble: FloatFieldIJ_Ensemble, + CLOSURE_CHOICE: Int, + plume: Int, +): + """Fill the mass flux ensemble based on the cloud workfunctions computed at the start of convection + (cloud_workfunction_1) and the end (cloud_workfunction_0). These values may then be updated based on + member number and level (below, within, or above cloud level). + + Args: + error_code (IntFieldIJ_Plume) + error_code_2 (IntFieldIJ) + error_code_3 (IntFieldIJ) + updraft_lfc_level (IntFieldIJ_Plume) + vapor_forced (FloatField) + ocean_fraction (FloatFieldIJ) + omega (FloatField) + cape_removal_time_scale (FloatFieldIJ) + arbitrary_numerical_parameter (FloatFieldIJ) + f_dicycle_modified (FloatFieldIJ) + cloud_workfunction_0 (FloatFieldIJ) + cloud_workfunction_0_modified (FloatFieldIJ) + cloud_workfunction_1 (FloatFieldIJ) + cloud_workfunction_1_pbl (FloatFieldIJ) + mass_flux_ensemble (FloatFieldIJ_Ensemble) + internal_mass_flux_ensemble (FloatFieldIJ_Ensemble) + precipitation_ensemble (FloatFieldIJ_Ensemble) + CLOSURE_CHOICE (Int) + plume (Int) + """ + from __externals__ import DIURNAL_CYCLE, DTIME, ENSEMBLE_MEMBERS, ZERO_DIFF + + with computation(FORWARD), interval(0, 1): + ensemble_adjustment: FloatFieldIJ = 1.0 + member = 0 + while member < ENSEMBLE_MEMBERS: + mass_flux_ensemble[0, 0][member] = 0.0 + member += 1 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + workfunction_diff_1: FloatFieldIJ = (cloud_workfunction_1 - cloud_workfunction_0) / DTIME + + internal_mass_flux_ensemble[0, 0][0] = max( + 0.0, (cloud_workfunction_1 - cloud_workfunction_0) / DTIME + ) + + internal_mass_flux_ensemble[0, 0][1] = internal_mass_flux_ensemble[0, 0][0] + internal_mass_flux_ensemble[0, 0][2] = internal_mass_flux_ensemble[0, 0][0] + internal_mass_flux_ensemble[0, 0][15] = internal_mass_flux_ensemble[0, 0][0] + + # more like Brown (1979), or Frank-Cohen (199?) + # omega is in Pa/s + omega_modified: FloatFieldIJ = 0.0 + + level_counter: IntFieldIJ = 0 + + internal_mass_flux_ensemble[0, 0][3] = 0.0 + + with computation(FORWARD), interval(...): + if ( + error_code[0, 0][plume] == 0 + and K >= max(0, updraft_lfc_level[0, 0][plume] - 1) + and K <= updraft_lfc_level[0, 0][plume] + 1 + ): + beta = 1.0 + omega_modified = omega_modified - omega / constants.MAPL_GRAV / beta + level_counter = level_counter + 1 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if level_counter > 0: + internal_mass_flux_ensemble[0, 0][3] = omega_modified / level_counter # kg[air]/m^3 * m/s + internal_mass_flux_ensemble[0, 0][3] = max(0.0, internal_mass_flux_ensemble[0, 0][3]) + internal_mass_flux_ensemble[0, 0][4] = internal_mass_flux_ensemble[0, 0][3] + internal_mass_flux_ensemble[0, 0][5] = internal_mass_flux_ensemble[0, 0][3] + internal_mass_flux_ensemble[0, 0][13] = internal_mass_flux_ensemble[0, 0][3] + + # more like Krishnamurti et al.; + # assuming that omega.at(K=cloud_top_level)*vapor_forced.at(K=cloud_top_level)<< + # omega.at(K=updraft_lfc_level)*vapor_forced.at(K=updraft_lfc_level) + moisture_convergence: FloatFieldIJ = ( + -omega.at(K=updraft_lfc_level[0, 0][plume]) + * vapor_forced.at(K=updraft_lfc_level[0, 0][plume]) + / constants.MAPL_GRAV + ) + + moisture_convergence = max(0.0, moisture_convergence) + internal_mass_flux_ensemble[0, 0][6] = moisture_convergence + internal_mass_flux_ensemble[0, 0][7] = internal_mass_flux_ensemble[0, 0][6] + internal_mass_flux_ensemble[0, 0][8] = internal_mass_flux_ensemble[0, 0][6] + internal_mass_flux_ensemble[0, 0][14] = internal_mass_flux_ensemble[0, 0][6] + + # more like Betchold et al (2014). Note that AA1 already includes the forcings tendencies + internal_mass_flux_ensemble[0, 0][9] = cloud_workfunction_1 / cape_removal_time_scale + + internal_mass_flux_ensemble[0, 0][10] = internal_mass_flux_ensemble[0, 0][9] + internal_mass_flux_ensemble[0, 0][11] = internal_mass_flux_ensemble[0, 0][9] + internal_mass_flux_ensemble[0, 0][12] = internal_mass_flux_ensemble[0, 0][9] + + if CLOSURE_CHOICE == 0 and workfunction_diff_1 < 0.0: + internal_mass_flux_ensemble[0, 0][0] = 0.0 + internal_mass_flux_ensemble[0, 0][1] = 0.0 + internal_mass_flux_ensemble[0, 0][2] = 0.0 + internal_mass_flux_ensemble[0, 0][15] = 0.0 + internal_mass_flux_ensemble[0, 0][9] = 0.0 + internal_mass_flux_ensemble[0, 0][10] = 0.0 + internal_mass_flux_ensemble[0, 0][11] = 0.0 + internal_mass_flux_ensemble[0, 0][12] = 0.0 + + workfunction_diff_2: FloatFieldIJ = ( + cloud_workfunction_0_modified - (cloud_workfunction_1) + ) / arbitrary_numerical_parameter + + if workfunction_diff_2 <= 0.0 and workfunction_diff_2 > -0.1 * arbitrary_numerical_parameter: + workfunction_diff_2 = -0.1 * arbitrary_numerical_parameter + + if workfunction_diff_2 > 0.0 and workfunction_diff_2 < 1.0e-2: + workfunction_diff_2 = 1.0e-2 + + # over water, enforce small cap for some of the closures + if ocean_fraction < 0.1 and (error_code_2 > 0.0 or error_code_3 > 0): + member = 0 + while member < 16: + internal_mass_flux_ensemble[0, 0][member] = ( + ensemble_adjustment * internal_mass_flux_ensemble[0, 0][member] + ) + member += 1 + + # special treatment for stability closures + if workfunction_diff_2 < 0.0: + if internal_mass_flux_ensemble[0, 0][0] > 0.0: + mass_flux_ensemble[0, 0][0] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][0] / workfunction_diff_2 + ) + + if internal_mass_flux_ensemble[0, 0][1] > 0.0: + mass_flux_ensemble[0, 0][1] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][1] / workfunction_diff_2 + ) + + if internal_mass_flux_ensemble[0, 0][2] > 0.0: + mass_flux_ensemble[0, 0][2] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][2] / workfunction_diff_2 + ) + + if internal_mass_flux_ensemble[0, 0][15] > 0.0: + mass_flux_ensemble[0, 0][15] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][15] / workfunction_diff_2 + ) + else: + internal_mass_flux_ensemble[0, 0][0] = 0.0 + internal_mass_flux_ensemble[0, 0][1] = 0.0 + internal_mass_flux_ensemble[0, 0][2] = 0.0 + internal_mass_flux_ensemble[0, 0][15] = 0.0 + + mass_flux_ensemble[0, 0][3] = max(0.0, internal_mass_flux_ensemble[0, 0][3]) + mass_flux_ensemble[0, 0][4] = max(0.0, internal_mass_flux_ensemble[0, 0][4]) + mass_flux_ensemble[0, 0][5] = max(0.0, internal_mass_flux_ensemble[0, 0][5]) + mass_flux_ensemble[0, 0][13] = max(0.0, internal_mass_flux_ensemble[0, 0][13]) + + adjustment = max(1.0e-3, precipitation_ensemble[0, 0][6]) + mass_flux_ensemble[0, 0][6] = max(0.0, internal_mass_flux_ensemble[0, 0][6] / adjustment) + + adjustment = max(1.0e-3, precipitation_ensemble[0, 0][7]) + mass_flux_ensemble[0, 0][7] = max(0.0, internal_mass_flux_ensemble[0, 0][7] / adjustment) + + adjustment = max(1.0e-3, precipitation_ensemble[0, 0][8]) + mass_flux_ensemble[0, 0][8] = max(0.0, internal_mass_flux_ensemble[0, 0][8] / adjustment) + + adjustment = max(1.0e-3, precipitation_ensemble[0, 0][14]) + mass_flux_ensemble[0, 0][14] = max(0.0, internal_mass_flux_ensemble[0, 0][14] / adjustment) + + if workfunction_diff_2 < 0.0: + mass_flux_ensemble[0, 0][9] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][9] / workfunction_diff_2 + ) + mass_flux_ensemble[0, 0][10] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][10] / workfunction_diff_2 + ) + mass_flux_ensemble[0, 0][11] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][11] / workfunction_diff_2 + ) + mass_flux_ensemble[0, 0][12] = max( + 0.0, -internal_mass_flux_ensemble[0, 0][12] / workfunction_diff_2 + ) + else: + mass_flux_ensemble[0, 0][9] = 0.0 + mass_flux_ensemble[0, 0][10] = 0.0 + mass_flux_ensemble[0, 0][11] = 0.0 + mass_flux_ensemble[0, 0][12] = 0.0 + + if CLOSURE_CHOICE >= 1: + member = 0 + while member < 16: + mass_flux_ensemble[0, 0][member] = mass_flux_ensemble[0, 0][CLOSURE_CHOICE] + member += 1 + + # over the land, only applies closure 9 + if ZERO_DIFF == 0 and CLOSURE_CHOICE == 0: + member = 0 + while member < 16: + mass_flux_ensemble[0, 0][member] = (1.0 - ocean_fraction) * mass_flux_ensemble[0, 0][ + 9 + ] + ocean_fraction * mass_flux_ensemble[0, 0][member] + member += 1 + + with computation(FORWARD), interval(0, 1): + if DIURNAL_CYCLE == 1 or DIURNAL_CYCLE == 6: + f_dicycle_modified = 0.0 + if error_code[0, 0][plume] == 0: + workfunction_diff_3 = ( + cloud_workfunction_1 - cloud_workfunction_1_pbl + ) / cape_removal_time_scale + if workfunction_diff_2 < 0: + f_dicycle_modified = max(0.0, -workfunction_diff_3 / workfunction_diff_2) + f_dicycle_modified = mass_flux_ensemble[0, 0][9] - f_dicycle_modified + + +def ensemble_forcing_mid_plume( + error_code: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + pbl_level: IntFieldIJ, + p_cloud_levels_forced: FloatField_Plume, + cloud_moist_static_energy_forced: FloatField, + environment_moist_static_energy_cloud_levels_forced: FloatField, + dmoist_static_energydt: FloatField, + convective_scale_velocity: FloatFieldIJ, + cape_removal_time_scale: FloatFieldIJ, + arbitrary_numerical_parameter: FloatFieldIJ, + f_dicycle_modified: FloatFieldIJ, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_0_modified: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + cloud_workfunction_1_pbl: FloatFieldIJ, + xff_mid: FloatFieldIJ_Ensemble, + CLOSURE_CHOICE: Int, + plume: Int, +): + """Compute a special mass flux ensemble for the mid plume. Similar to the generic ensemble computed + in ensemble_forcing, with extra considerations of the PBL and diurnal cycle. + + Args: + error_code (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + pbl_level (IntFieldIJ) + p_cloud_levels_forced (FloatField_Plume) + cloud_moist_static_energy_forced (FloatField) + environment_moist_static_energy_cloud_levels_forced (FloatField) + dmoist_static_energydt (FloatField) + convective_scale_velocity (FloatFieldIJ) + cape_removal_time_scale (FloatFieldIJ) + arbitrary_numerical_parameter (FloatFieldIJ) + f_dicycle_modified (FloatFieldIJ) + cloud_workfunction_0 (FloatFieldIJ) + cloud_workfunction_0_modified (FloatFieldIJ) + cloud_workfunction_1 (FloatFieldIJ) + cloud_workfunction_1_pbl (FloatFieldIJ) + xff_mid (FloatFieldIJ_Ensemble) + CLOSURE_CHOICE (Int) + plume (Int) + """ + from __externals__ import DIURNAL_CYCLE + + with computation(FORWARD), interval(0, 1): + # initialization + member = 0 + while member < 16: # 16 should come from config file + xff_mid[0, 0][member] = 0.0 + member += 1 + f_dicycle_modified = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + workfunction_diff_1: FloatFieldIJ = ( + cloud_workfunction_0_modified - cloud_workfunction_1 + ) / arbitrary_numerical_parameter + + if workfunction_diff_1 <= 0.0 and workfunction_diff_1 > -0.1 * arbitrary_numerical_parameter: + workfunction_diff_1 = -0.1 * arbitrary_numerical_parameter + + if workfunction_diff_1 > 0.0 and workfunction_diff_1 < 0.01: + workfunction_diff_1 = 0.01 + + # diurnal cycle mass flux + if DIURNAL_CYCLE == 1 or DIURNAL_CYCLE == 6 or DIURNAL_CYCLE == 0: + internal_workfunction_modified = cloud_workfunction_1_pbl / cape_removal_time_scale + f_dicycle_modified = max(0.0, -internal_workfunction_modified / workfunction_diff_1) + + # closures 3 and 4 for mid + if workfunction_diff_1 < 0.0: + xff_mid[0, 0][2] = max( + 0.0, -(cloud_workfunction_1 / cape_removal_time_scale) / workfunction_diff_1 + ) + xff_mid[0, 0][3] = f_dicycle_modified + + with computation(FORWARD), interval(0, 1): + # boundary layer quasi-equilibrium (Raymond 1995) + if error_code[0, 0][plume] == 0 and updraft_origin_level[0, 0][plume] < pbl_level + 1: + blqe: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(0, -1): + if ( + error_code[0, 0][plume] == 0 + and updraft_origin_level[0, 0][plume] < pbl_level + 1 + and K < updraft_lfc_level[0, 0][plume] + ): + blqe = ( + blqe + + 100.0 + * dmoist_static_energydt + * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + / constants.MAPL_GRAV + ) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if updraft_origin_level[0, 0][plume] < pbl_level + 1: + trash = max( + ( + cloud_moist_static_energy_forced.at(K=updraft_lfc_level) + - environment_moist_static_energy_cloud_levels_forced.at(K=updraft_lfc_level) + ), + 1.0e1, + ) + xff_mid[0, 0][1] = max(0.0, blqe / trash) + + # W* closure (Grant,2001) + xff_mid[0, 0][0] = 0.03 * convective_scale_velocity + + with computation(FORWARD), interval(0, 1): + # set f_dicycle_modified=0 in case the CLOSURE_CHOICE is 4 and DIURNAL_CYCLE closure will be applied + if DIURNAL_CYCLE > 0 and CLOSURE_CHOICE == 4: + f_dicycle_modified = 0.0 + + with computation(FORWARD), interval(...): + if CLOSURE_CHOICE == 5: + if DIURNAL_CYCLE > 0: + if error_code[0, 0][plume] == 0: + xff_mid[0, 0][4] = 0.05 * xff_mid[0, 0][3] + f_dicycle_modified = 0.0 + + elif DIURNAL_CYCLE == 0: + if error_code[0, 0][plume] == 0: + xff_mid[0, 0][4] = 0.05 * xff_mid[0, 0][3] + f_dicycle_modified = 0.0 + + +def effective_precipitation( + error_code: IntFieldIJ_Plume, + epsilon_forced: FloatFieldIJ_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + effective_condensate_to_fall_forced: FloatField, + plume: Int, +): + """Compute the effective precipitation based on evaporation in the downdraft and precipitation effeciency + (which is dependent on wind shear). + + Args: + error_code (IntFieldIJ_Plume) + epsilon_forced (FloatFieldIJ_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + effective_condensate_to_fall_forced (FloatField) + plume (Int) + """ + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0: + effective_condensate_to_fall_forced = ( + condensate_to_fall_forced[0, 0, 0][plume] + + epsilon_forced[0, 0][plume] * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + + +class LargeScaleForcing(NDSLRuntime): + """Fill the mass flux ensemble based on cloud workfunction data, compute an additional mass flux ensemble + for the mid plume, and compute effective precipitation. + + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # initialize locals + quantity_factory.add_data_dimensions( + { + "ensemble_members": cumulus_parameterization_constants.MAXENS1 + * cumulus_parameterization_constants.MAXENS2 + * cumulus_parameterization_constants.MAXENS3 + } + ) + self._internal_mass_flux_ensemble: Local = quantity_factory.zeros( + [X_DIM, Y_DIM, "ensemble_members"], "n/a" + ) + + # construct stencils + self._copy = stencil_factory.from_dims_halo( + func=ddim_copy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self.ensemble_forcing = stencil_factory.from_dims_halo( + func=ensemble_forcing, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "DIURNAL_CYCLE": cumulus_parameterization_config.DIURNAL_CYCLE, + "DTIME": cumulus_parameterization_config.DTIME, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "ENSEMBLE_MEMBERS": cumulus_parameterization_constants.MAXENS1 + * cumulus_parameterization_constants.MAXENS2 + * cumulus_parameterization_constants.MAXENS3, + }, + ) + + self._ensemble_forcing_mid_plume = stencil_factory.from_dims_halo( + func=ensemble_forcing_mid_plume, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "DIURNAL_CYCLE": cumulus_parameterization_config.DIURNAL_CYCLE, + "DTIME": cumulus_parameterization_config.DTIME, + # "ENSEMBLE_MEMBERS": cumulus_parameterization_config.ENSEMBLE_MEMBERS, + }, + ) + + self._effective_precipitation = stencil_factory.from_dims_halo( + func=effective_precipitation, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + error_code_2: Quantity, + error_code_3: Quantity, + updraft_origin_level: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + pbl_level: Quantity, + ocean_fraction: Quantity, + p_cloud_levels_forced: Quantity, + vapor_forced: Quantity, + condensate_to_fall_forced: Quantity, + effective_condensate_to_fall_forced: Quantity, + evaporate_in_downdraft_forced: Quantity, + omega: Quantity, + convective_scale_velocity: Quantity, + normalized_massflux_updraft_forced: Quantity, + normalized_massflux_downdraft_forced: Quantity, + cloud_moist_static_energy: Quantity, + cloud_moist_static_energy_forced: Quantity, + environment_moist_static_energy_cloud_levels: Quantity, + environment_moist_static_energy_cloud_levels_forced: Quantity, + dmoist_static_energydt: Quantity, + cloud_workfunction_0: Quantity, + cloud_workfunction_0_modified: Quantity, + cloud_workfunction_1: Quantity, + cloud_workfunction_1_pbl: Quantity, + arbitrary_numerical_parameter: Quantity, + f_dicycle_modified: Quantity, + cape_removal_time_scale: Quantity, + epsilon_forced: Quantity, + k_x_modified: Quantity, + mass_flux_ensemble: Quantity, + precipitation_ensemble: Quantity, + xff_mid: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + # copy error codes + self._copy(field_in=error_code, field_out=error_code_2, plume=plume_dependent_constants.PLUME_INDEX) + self._copy(field_in=error_code, field_out=error_code_3, plume=plume_dependent_constants.PLUME_INDEX) + + if plume_dependent_constants.PLUME_INDEX == cumulus_parameterization_constants.DEEP: + + self.ensemble_forcing( + error_code=error_code, + error_code_2=error_code_2, + error_code_3=error_code_3, + updraft_lfc_level=updraft_lfc_level, + vapor_forced=vapor_forced, + ocean_fraction=ocean_fraction, + omega=omega, + cape_removal_time_scale=cape_removal_time_scale, + arbitrary_numerical_parameter=arbitrary_numerical_parameter, + f_dicycle_modified=f_dicycle_modified, + cloud_workfunction_0=cloud_workfunction_0, + cloud_workfunction_0_modified=cloud_workfunction_0_modified, + cloud_workfunction_1=cloud_workfunction_1, + cloud_workfunction_1_pbl=cloud_workfunction_1_pbl, + mass_flux_ensemble=mass_flux_ensemble, + internal_mass_flux_ensemble=self._internal_mass_flux_ensemble, + precipitation_ensemble=precipitation_ensemble, + CLOSURE_CHOICE=plume_dependent_constants.CLOSURE_CHOICE, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if plume_dependent_constants.PLUME_INDEX == cumulus_parameterization_constants.MID: + + self._ensemble_forcing_mid_plume( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + pbl_level=pbl_level, + p_cloud_levels_forced=p_cloud_levels_forced, + cloud_moist_static_energy_forced=cloud_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels_forced=environment_moist_static_energy_cloud_levels_forced, + dmoist_static_energydt=dmoist_static_energydt, + convective_scale_velocity=convective_scale_velocity, + cape_removal_time_scale=cape_removal_time_scale, + arbitrary_numerical_parameter=arbitrary_numerical_parameter, + f_dicycle_modified=f_dicycle_modified, + cloud_workfunction_0=cloud_workfunction_0, + cloud_workfunction_0_modified=cloud_workfunction_0_modified, + cloud_workfunction_1=cloud_workfunction_1, + cloud_workfunction_1_pbl=cloud_workfunction_1_pbl, + xff_mid=xff_mid, + CLOSURE_CHOICE=plume_dependent_constants.CLOSURE_CHOICE, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if plume_dependent_constants.PLUME_INDEX == cumulus_parameterization_constants.SHALLOW: + raise NotImplementedError( + "Shallow plume options not implemented in LargeScaleForcing, but you" + "should have been caught before getting here. Something is wrong with the config checker.," + "Beware, more untested paths may be executing without warning." + ) + + self._effective_precipitation( + error_code=error_code, + epsilon_forced=epsilon_forced, + condensate_to_fall_forced=condensate_to_fall_forced, + evaporate_in_downdraft_forced=evaporate_in_downdraft_forced, + effective_condensate_to_fall_forced=effective_condensate_to_fall_forced, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/locals.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/locals.py new file mode 100644 index 000000000..6dfc64ea4 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/locals.py @@ -0,0 +1,1474 @@ +import dataclasses + +from ndsl import Quantity, State +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import Float, Int + + +@dataclasses.dataclass +class GF2020CumulusParameterizationLocals(State): + t_new: Quantity = dataclasses.field( + metadata={ + "name": "t_new", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_modified: Quantity = dataclasses.field( + metadata={ + "name": "t_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "t_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "t_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "t_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_excess: Quantity = dataclasses.field( + metadata={ + "name": "t_excess", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_new_pbl: Quantity = dataclasses.field( + metadata={ + "name": "t_new_pbl", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_forced: Quantity = dataclasses.field( + metadata={ + "name": "vapor_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_modified: Quantity = dataclasses.field( + metadata={ + "name": "vapor_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "vapor_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "vapor_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "vapor_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_excess: Quantity = dataclasses.field( + metadata={ + "name": "vapor_excess", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_forced_pbl: Quantity = dataclasses.field( + metadata={ + "name": "vapor_forced_pbl", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dmoist_static_energydt: Quantity = dataclasses.field( + metadata={ + "name": "dmoist_static_energydt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + maximum_updraft_origin_level: Quantity = dataclasses.field( + metadata={ + "name": "maximum_updraft_origin_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + ocean_fraction: Quantity = dataclasses.field( + metadata={ + "name": "ocean_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cap_max: Quantity = dataclasses.field( + metadata={ + "name": "cap_max", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + error_code_2: Quantity = dataclasses.field( + metadata={ + "name": "error_code_2", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + error_code_3: Quantity = dataclasses.field( + metadata={ + "name": "error_code_3", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + cap_max_increment: Quantity = dataclasses.field( + metadata={ + "name": "cap_max_increment", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_modified: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_0: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_0", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_0_modified: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_0_modified", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_0_pbl: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_0_pbl", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_1", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_fa: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_1_fa", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_pbl: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_1_pbl", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_2: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_2", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_3: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_3", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cin_0: Quantity = dataclasses.field( + metadata={ + "name": "cin_0", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cin_1: Quantity = dataclasses.field( + metadata={ + "name": "cin_1", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + k_x_modified: Quantity = dataclasses.field( + metadata={ + "name": "k_x_modified", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + epsilon: Quantity = dataclasses.field( + metadata={ + "name": "epsilon", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cape_removal_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "cape_removal_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "pbl_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_wetbulb: Quantity = dataclasses.field( + metadata={ + "name": "t_wetbulb", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_wetbulb: Quantity = dataclasses.field( + metadata={ + "name": "vapor_wetbulb", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + f_dicycle_modified: Quantity = dataclasses.field( + metadata={ + "name": "f_dicycle_modified", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + add_buoyancy: Quantity = dataclasses.field( + metadata={ + "name": "add_buoyancy", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + downdraft_saturation_vapor_forced: Quantity = dataclasses.field( + metadata={ + "name": "downdraft_saturation_vapor_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + c1d: Quantity = dataclasses.field( + metadata={ + "name": "c1d", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_normalized_integrated_evaporate_forced: Quantity = dataclasses.field( + metadata={ + "name": "total_normalized_integrated_evaporate_forced", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + effective_condensate_to_fall_forced: Quantity = dataclasses.field( + metadata={ + "name": "effective_condensate_to_fall_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation_below_cloud_base: Quantity = dataclasses.field( + metadata={ + "name": "evaporation_below_cloud_base", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_massflux: Quantity = dataclasses.field( + metadata={ + "name": "environment_massflux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_ensemble", + "dims": [X_DIM, Y_DIM, "ensemble_members"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + precipitation_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "precipitation_ensemble", + "dims": [X_DIM, Y_DIM, "ensemble_members"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + xff_mid: Quantity = dataclasses.field( + metadata={ + "name": "xff_mid", + "dims": [X_DIM, Y_DIM, "ensemble_members"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + scale_dependence_factor_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "scale_dependence_factor_downdraft", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + random_number: Quantity = dataclasses.field( + metadata={ + "name": "random_number", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + detrainment_function_updraft: Quantity = dataclasses.field( + metadata={ + "name": "detrainment_function_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + detrainment_function_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "detrainment_function_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + entrainment_rate_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "entrainment_rate_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + epsilon_min: Quantity = dataclasses.field( + metadata={ + "name": "epsilon_min", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + epsilon_max: Quantity = dataclasses.field( + metadata={ + "name": "epsilon_max", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + arbitrary_numerical_parameter: Quantity = dataclasses.field( + metadata={ + "name": "arbitrary_numerical_parameter", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_moist_static_energy_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_moist_static_energy_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_moist_static_energy_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_moist_static_energy_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + environment_saturation_mixing_ratio_cloud_levels_modified: Quantity = dataclasses.field( + metadata={ + "name": "environment_saturation_mixing_ratio_cloud_levels_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "p_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "u_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u_c: Quantity = dataclasses.field( + metadata={ + "name": "u_c", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u_c_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "u_c_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "v_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v_c: Quantity = dataclasses.field( + metadata={ + "name": "v_c", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v_c_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "v_c_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + gamma_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "gamma_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + gamma_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "gamma_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + hydrostatic_air_density: Quantity = dataclasses.field( + metadata={ + "name": "hydrostatic_air_density", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + partition_liquid_ice: Quantity = dataclasses.field( + metadata={ + "name": "partition_liquid_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + melting_layer: Quantity = dataclasses.field( + metadata={ + "name": "melting_layer", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + detrainment_start_level: Quantity = dataclasses.field( + metadata={ + "name": "detrainment_start_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + moist_static_energy_origin_level: Quantity = dataclasses.field( + metadata={ + "name": "moist_static_energy_origin_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + moist_static_energy_origin_level_forced: Quantity = dataclasses.field( + metadata={ + "name": "moist_static_energy_origin_level_forced", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + moist_static_energy_origin_level_modified: Quantity = dataclasses.field( + metadata={ + "name": "moist_static_energy_origin_level_modified", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + start_level: Quantity = dataclasses.field( + metadata={ + "name": "start_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + cloud_moist_static_energy: Quantity = dataclasses.field( + metadata={ + "name": "cloud_moist_static_energy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_moist_static_energy_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_moist_static_energy_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_moist_static_energy_modified: Quantity = dataclasses.field( + metadata={ + "name": "cloud_moist_static_energy_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_moist_static_energy_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_moist_static_energy_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_moist_static_energy_forced_transported: Quantity = dataclasses.field( + metadata={ + "name": "cloud_moist_static_energy_forced_transported", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + negative_buoyancy_depth: Quantity = dataclasses.field( + metadata={ + "name": "negative_buoyancy_depth", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + frh_lfc: Quantity = dataclasses.field( + metadata={ + "name": "frh_lfc", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_updraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_updraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_u_updraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_u_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_u_updraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_u_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_updraft: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_updraft_modified: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_updraft_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + buoyancy: Quantity = dataclasses.field( + metadata={ + "name": "buoyancy", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + d_buoyancy: Quantity = dataclasses.field( + metadata={ + "name": "d_buoyancy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + d_buoyancy_forced: Quantity = dataclasses.field( + metadata={ + "name": "d_buoyancy_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + d_buoyancy_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "d_buoyancy_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + d_buoyancy_modified: Quantity = dataclasses.field( + metadata={ + "name": "d_buoyancy_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + melting: Quantity = dataclasses.field( + metadata={ + "name": "melting", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + updraft_column_temperature_forced: Quantity = dataclasses.field( + metadata={ + "name": "updraft_column_temperature_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + downdraft_column_temperature_forced: Quantity = dataclasses.field( + metadata={ + "name": "downdraft_column_temperature_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_total_water_after_entrainment_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_total_water_after_entrainment_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_total_water_after_entrainment_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_total_water_after_entrainment_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_liquid_before_rain_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_liquid_before_rain_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vertical_velocity_3d: Quantity = dataclasses.field( + metadata={ + "name": "vertical_velocity_3d", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vertical_velocity_2d: Quantity = dataclasses.field( + metadata={ + "name": "vertical_velocity_2d", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_downdraft_modified: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_downdraft_modified", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_u_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_u_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_u_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_u_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + epsilon_computed: Quantity = dataclasses.field( + metadata={ + "name": "epsilon_computed", + "dims": [X_DIM, Y_DIM, "ensemble_2"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_u_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_u_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_v_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_v_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_moist_static_energy_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_moist_static_energy_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_t_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_t_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_vapor_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_vapor_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_cloud_liquid_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_cloud_liquidu_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_buoyancy_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_buoyancy_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_convective_ice_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_convective_ice_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_large_scale_ice_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_large_scale_ice_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_convective_liquid_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_convective_liquid_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_large_scale_liquid_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_large_scale_liquid_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_convective_cloud_fraction_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_convective_cloud_fraction_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + del_large_scale_cloud_fraction_cloud_ensemble: Quantity = dataclasses.field( + metadata={ + "name": "del_large_scale_cloud_fraction_cloud_ensemble", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_tendency_from_environmental_subsidence: Quantity = dataclasses.field( + metadata={ + "name": "t_tendency_from_environmental_subsidence", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + moist_static_energy_tendency_from_environmental_subsidence: Quantity = dataclasses.field( + metadata={ + "name": "moist_static_energy_tendency_from_environmental_subsidence", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_tendency_from_environmental_subsidence: Quantity = dataclasses.field( + metadata={ + "name": "vapor_tendency_from_environmental_subsidence", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + precipitation_flux: Quantity = dataclasses.field( + metadata={ + "name": "prec_flux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation_flux: Quantity = dataclasses.field( + metadata={ + "name": "evap_flux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_cloud_levels: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_cloud_levels", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_sc_updraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_sc_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_sc_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_sc_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_pw_updraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_pw_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_pw_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_pw_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_total_pw_updraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_pw_updraft", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_total_pw_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_pw_downdraft", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ####################################################################### + # NOTE AtmosphericComposition variables **CAN BE MOVED AND/OR RENAMED** + ####################################################################### + ddtr: Quantity = dataclasses.field( + metadata={ + "name": "ddtr", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + trash_: Quantity = dataclasses.field( + metadata={ + "name": "trash_", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + trash2_: Quantity = dataclasses.field( + metadata={ + "name": "trash2_", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evap_: Quantity = dataclasses.field( + metadata={ + "name": "evap_", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + wetdep_: Quantity = dataclasses.field( + metadata={ + "name": "wetdep_", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + residu_: Quantity = dataclasses.field( + metadata={ + "name": "residu_", + "dims": [X_DIM, Y_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + + # NOTE these can potentially be removed with a better analysis of what they do/where they go + + # NOTE these can potentially be removed with a better analysis of what they do/where they go + psum: Quantity = dataclasses.field( + metadata={ + "name": "NEED BETTER NAME. WHAT IS THIS FIELD", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + psumh: Quantity = dataclasses.field( + metadata={ + "name": "NEED BETTER NAME. WHAT IS THIS FIELD", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/mass_conservation.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/mass_conservation.py new file mode 100644 index 000000000..7851d9157 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/mass_conservation.py @@ -0,0 +1,6 @@ +class MassConservation: + def __init__(self): + pass + + def __call__(self, *args, **kwds): + pass diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/moist_static_energy.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/moist_static_energy.py new file mode 100644 index 000000000..05cddb189 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/moist_static_energy.py @@ -0,0 +1,327 @@ +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.buoyancy import get_buoyancy +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import get_cloud_boundary_conditions + + +def parcel_moist_static_energy( + error_code: IntFieldIJ_Plume, + t_excess: FloatFieldIJ, + vapor_excess: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + updraft_origin_level: IntFieldIJ_Plume, + p: FloatField, + environmenet_moist_static_energy: FloatField, + environmenet_moist_static_energy_forced: FloatField, + t_perturbation: FloatField, + moist_static_energy_origin_level: FloatFieldIJ, + moist_static_energy_origin_level_forced: FloatFieldIJ, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Determine the moist static energy of the forced and unforced parcel. + + Args: + error_code (IntFieldIJ_Plume) + t_excess (FloatFieldIJ) + vapor_excess (FloatFieldIJ) + add_buoyancy (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + updraft_origin_level (IntFieldIJ_Plume) + p (FloatField) + environmenet_moist_static_energy (FloatField) + environmenet_moist_static_energy_forced (FloatField) + t_perturbation (FloatField) + moist_static_energy_origin_level (FloatFieldIJ) + moist_static_energy_origin_level_forced (FloatFieldIJ) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import BOUNDARY_CONDITION_METHOD, k_end + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + modification = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + + moist_static_energy_origin_level = get_cloud_boundary_conditions( + field=environmenet_moist_static_energy, + scalar_perturbation=modification, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=t_perturbation, + ) + moist_static_energy_origin_level_forced = get_cloud_boundary_conditions( + field=environmenet_moist_static_energy_forced, + scalar_perturbation=modification, + p=p, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=t_perturbation, + ) + + +def first_guess_moist_static_energy( + error_code: IntFieldIJ_Plume, + start_level: IntFieldIJ, + cloud_top_level: IntFieldIJ_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + normalized_massflux_updraft: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + environment_moist_static_energy_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + cloud_moist_static_energy_forced: FloatField, + vapor_excess: FloatFieldIJ, + t_excess: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + plume: Int, +): + """Compute moist static energy within the cloud (between LCL/start_level + and equilibrium level/cloud_top_level) + + Args: + error_code (IntFieldIJ_Plume) + start_level (IntFieldIJ) + cloud_top_level (IntFieldIJ_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + normalized_massflux_updraft (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + environment_moist_static_energy_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + cloud_moist_static_energy_forced (FloatField) + vapor_excess (FloatFieldIJ) + t_excess (FloatFieldIJ) + add_buoyancy (FloatFieldIJ) + plume (Int) + """ + with computation(FORWARD), interval(1, None): + if error_code[0, 0][plume] == 0: + if K >= start_level + 1 and K <= cloud_top_level[0, 0][plume] + 1: # mass cons option + denom: FloatFieldIJ = ( + normalized_massflux_updraft[0, 0, -1] + - 0.5 * mass_detrainment_updraft_forced[0, 0, -1][plume] + + mass_entrainment_updraft_forced[0, 0, -1][plume] + ) + if denom > 0.0: + cloud_moist_static_energy_forced = ( + cloud_moist_static_energy_forced[0, 0, -1] + * normalized_massflux_updraft_forced[0, 0, -1][plume] + - 0.5 + * mass_detrainment_updraft_forced[0, 0, -1][plume] + * cloud_moist_static_energy_forced[0, 0, -1] + + mass_entrainment_updraft_forced[0, 0, -1][plume] + * environment_moist_static_energy_forced[0, 0, -1] + ) / denom + if K == start_level + 1: + perturbation: FloatFieldIJ = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + cloud_moist_static_energy_forced = ( + cloud_moist_static_energy_forced + + perturbation * mass_entrainment_updraft_forced[0, 0, -1][plume] / denom + ) + else: + cloud_moist_static_energy_forced = cloud_moist_static_energy_forced[0, 0, -1] + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + if K >= cloud_top_level[0, 0][plume] + 2: + cloud_moist_static_energy_forced = ( + environment_saturation_moist_static_energy_cloud_levels_forced + ) + + +def moist_static_energy_inside_cloud( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + start_level: IntFieldIJ, + moist_static_energy_origin_level: FloatFieldIJ, + cloud_moist_static_energy: FloatField, + environment_moist_static_energy_modified: FloatField, + environment_saturation_moist_static_energy_cloud_levels_modified: FloatField, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + normalized_massflux_updraft_modified: FloatField, + partition_liquid_ice: FloatField, + vapor_excess: FloatFieldIJ, + t_excess: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + cloud_liquid_after_rain_forced: FloatField_Plume, + plume: Int, +): + """Compute moist static energy within the cloud. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + start_level (IntFieldIJ) + moist_static_energy_origin_level (FloatFieldIJ) + cloud_moist_static_energy (FloatField) + environment_moist_static_energy_modified (FloatField) + environment_saturation_moist_static_energy_cloud_levels_modified (FloatField) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + normalized_massflux_updraft_modified (FloatField) + partition_liquid_ice (FloatField) + vapor_excess (FloatFieldIJ) + t_excess (FloatFieldIJ) + add_buoyancy (FloatFieldIJ) + cloud_liquid_after_rain_forced (FloatField_Plume) + plume (Int) + """ + with computation(PARALLEL), interval(...): + cloud_moist_static_energy = 0.0 + if error_code[0, 0][plume] == 0: + if K <= start_level: + cloud_moist_static_energy = moist_static_energy_origin_level + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K >= start_level + 1 and K <= cloud_top_level[0, 0][plume] + 1: + denom: FloatFieldIJ = ( + normalized_massflux_updraft_modified.at(K=K - 1) + - 0.5 * mass_detrainment_updraft_forced.at(K=K - 1, ddim=[plume]) + + mass_entrainment_updraft_forced.at(K=K - 1, ddim=[plume]) + ) + if denom == 0.0: + cloud_moist_static_energy = cloud_moist_static_energy.at(K=K - 1) + else: + cloud_moist_static_energy = ( + cloud_moist_static_energy.at(K=K - 1) + * normalized_massflux_updraft_modified.at(K=K - 1) + - 0.5 + * mass_detrainment_updraft_forced.at(K=K - 1, ddim=[plume]) + * cloud_moist_static_energy.at(K=K - 1) + + mass_entrainment_updraft_forced.at(K=K - 1, ddim=[plume]) + * environment_moist_static_energy_modified.at(K=K - 1) + ) / denom + if K == start_level + 1: + x_add: FloatFieldIJ = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + cloud_moist_static_energy = ( + cloud_moist_static_energy + + x_add * mass_entrainment_updraft_forced.at(K=K - 1, ddim=[plume]) / denom + ) + + # include glaciation effects on cloud_moist_static_energy + cloud_moist_static_energy = ( + cloud_moist_static_energy + + cumulus_parameterization_constants.XLF + * (1.0 - partition_liquid_ice) + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + if K >= cloud_top_level[0, 0][plume] + 2: + cloud_moist_static_energy = environment_saturation_moist_static_energy_cloud_levels_modified + normalized_massflux_updraft_modified = 0.0 + + +class StaticControl(NDSLRuntime): + """Update cloud moist static energy and compute buoyancy remaining after convection.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._moist_static_energy_inside_cloud = stencil_factory.from_dims_halo( + func=moist_static_energy_inside_cloud, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._get_buoyancy = stencil_factory.from_dims_halo( + func=get_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + start_level: Quantity, + lcl_level: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + cloud_moist_static_energy_modified: Quantity, + moist_static_energy_origin_level_modified: Quantity, + environment_moist_static_energy_modified: Quantity, + environment_moist_static_energy_cloud_levels_modified: Quantity, + environment_saturation_moist_static_energy_cloud_levels_modified: Quantity, + mass_detrainment_updraft_forced: Quantity, + mass_entrainment_updraft_forced: Quantity, + normalized_massflux_updraft_modified: Quantity, + partition_liquid_ice: Quantity, + vapor_excess: Quantity, + t_excess: Quantity, + add_buoyancy: Quantity, + cloud_liquid_after_rain_forced: Quantity, + d_buoyancy_modified: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._moist_static_energy_inside_cloud( + error_code=error_code, + cloud_top_level=cloud_top_level, + start_level=start_level, + moist_static_energy_origin_level=moist_static_energy_origin_level_modified, + cloud_moist_static_energy=cloud_moist_static_energy_modified, + environment_moist_static_energy_modified=environment_moist_static_energy_modified, + environment_saturation_moist_static_energy_cloud_levels_modified=environment_saturation_moist_static_energy_cloud_levels_modified, + mass_detrainment_updraft_forced=mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=mass_entrainment_updraft_forced, + normalized_massflux_updraft_modified=normalized_massflux_updraft_modified, + partition_liquid_ice=partition_liquid_ice, + vapor_excess=vapor_excess, + t_excess=t_excess, + add_buoyancy=add_buoyancy, + cloud_liquid_after_rain_forced=cloud_liquid_after_rain_forced, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._get_buoyancy( + lcl_level=lcl_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + cloud_moist_static_energy=cloud_moist_static_energy_modified, + environment_moist_static_energy=environment_moist_static_energy_cloud_levels_modified, + environment_saturation_moist_static_energy=environment_saturation_moist_static_energy_cloud_levels_modified, + d_buoyancy=d_buoyancy_modified, + error_code=error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/plume_dependent_constants.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/plume_dependent_constants.py new file mode 100644 index 000000000..e8e25c8a3 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/plume_dependent_constants.py @@ -0,0 +1,32 @@ +import dataclasses + +from ndsl.dsl.typing import Float, Int + + +@dataclasses.dataclass +class GF2020PlumeDependentConstants: + PLUME_INDEX: Int = Int(0) + DOWNDRAFT_MAX_HEIGHT_LAND: Float = Float(0.0) + DOWNDRAFT_MAX_HEIGHT_OCEAN: Float = Float(0.0) + UPDRAFT_MAX_HEIGHT_LAND: Float = Float(0.0) + UPDRAFT_MAX_HEIGHT_OCEAN: Float = Float(0.0) + MINIMUM_EVAP_FRACTION_LAND: Float = Float(0.0) + MINIMUM_EVAP_FRACTION_OCEAN: Float = Float(0.0) + MAXIMUM_EVAP_FRACTION_LAND: Float = Float(0.0) + MAXIMUM_EVAP_FRACTION_OCEAN: Float = Float(0.0) + CLOUD_BASE_MASS_FLUX_FACTOR: Float = Float(0.0) + USE_EXCESS: Int = Int(0) + ENTRAINMENT_RATE: Float = Float(0.0) + ENABLE_PLUME: Int = Int(0) + AVERAGE_LAYER_DEPTH = Float(0.0) + CAP_MAX_INC: Float = Float(0.0) + LAMBDA_DEEP: Float = Float(0.0) + LAMBDA_DOWN: Float = Float(0.0) + MINIMUM_DEPTH: Float = Float(0.0) + MAX_UPDRAFT_ORIGIN_HEIGHT: Float = Float(0.0) + MAX_DOWNDRAFT_ORIGIN_HEIGHt: Float = Float(0.0) + DETRAINMENT_CRITICAL_DEPTH: Float = Float(0.0) + C0: Float = Float(0.0) + T_STAR: Float = Float(0.0) + TAU_CAPE_REMOVAL: Float = Float(0.0) + CLOSURE_CHOICE: Int = Int(0.0) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/precip.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/precip.py new file mode 100644 index 000000000..97c97082f --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/precip.py @@ -0,0 +1,439 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, K, computation, function, interval, sqrt +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Int, IntFieldIJ +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.shared_incloud_processes import ice_fraction + + +@function +def liquid_fraction( + t, + convection_fraction, + surface_type, + FRAC_MODIS, +): + """ + Get the fraction of liquid condensates + + Args: + t (in): temperature + convection_fraction (in) + surface_type (in) + FRAC_MODIS (in): use fraction liq/ice content derived from MODIS/CALIPO sensors + """ + if FRAC_MODIS == 1: + liquid_fraction = 1.0 - ice_fraction(t, convection_fraction, surface_type) + else: + liquid_fraction = min( + 1.0, + ( + max(0.0, (t - cumulus_parameterization_constants.T_ICE)) + / (cumulus_parameterization_constants.T_0 - cumulus_parameterization_constants.T_ICE) + ) + ** 2, + ) + + return liquid_fraction + + +def partition_liquid_ice( + t: FloatField, + p: FloatField_Plume, + geopotential_height: FloatField, + topography_height_no_negative: FloatFieldIJ, + surface_type: FloatFieldIJ, + convection_fraction: FloatFieldIJ, + error_code: IntFieldIJ_Plume, + melting_layer: FloatField, + part_liquid_ice: FloatField, + plume: Int, +): + """Partition total condensate into liquid and ice phases. + + Args: + t (FloatField) + p (FloatField_Plume) + geopotential_height (FloatField) + topography_height_no_negative (FloatFieldIJ) + surface_type (FloatFieldIJ) + convection_fraction (FloatFieldIJ) + error_code (IntFieldIJ_Plume) + melting_layer (FloatField) + part_liquid_ice (FloatField) + plume (Int) + """ + from __externals__ import FRAC_MODIS, MELT_GLAC, k_end + + with computation(PARALLEL), interval(...): + # constants, set internally because they may differ from global constants + # and need to only exist inside this stencil + t1 = 276.16 + z_meltlayer1 = 4000.0 + z_meltlayer2 = 6000.0 + del_t = 3.0 + + # prefill some fields + part_liquid_ice = 1.0 + melting_layer = 0.0 + + with computation(PARALLEL), interval(0, -1): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + if error_code[0, 0][plume] == 0: + # get function of T for partition of total condensate into liq and ice phases + part_liquid_ice = liquid_fraction(t, convection_fraction, surface_type, FRAC_MODIS) + + with computation(PARALLEL), interval(0, -1): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + if error_code[0, 0][plume] == 0: + # define the melting layer (the layer will be between T_0+1 < TEMP < T_1 + if t <= (cumulus_parameterization_constants.T_0 - del_t): + melting_layer = 0.0 + + elif t < (cumulus_parameterization_constants.T_0 + del_t) and t > ( + cumulus_parameterization_constants.T_0 - del_t + ): + melting_layer = ( + (t - (cumulus_parameterization_constants.T_0 - del_t)) / (2.0 * del_t) + ) ** 2 + + else: + melting_layer = 1.0 + + melting_layer = melting_layer * (1.0 - melting_layer) + + with computation(FORWARD), interval(0, 1): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + # normalize vertical integral of melting_layer to 1 + norm: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(0, -2): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + if error_code[0, 0][plume] == 0: + # normalize vertical integral of melting_layer to 1 + dp = 100.0 * (p[0, 0, 0][plume] - p[0, 0, 1][plume]) + norm = norm + melting_layer * dp / constants.MAPL_GRAV + + with computation(PARALLEL), interval(...): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + if error_code[0, 0][plume] == 0: + # normalize vertical integral of melting_layer to 1 + melting_layer = ( + melting_layer + / (norm + 1.0e-6) + * ( + 100 + * (p.at(K=0, ddim=[plume]) - p.at(K=k_end - 1, ddim=[plume])) + / constants.MAPL_GRAV + ) + ) + + +class PrecipFactor: + """ + Get the pickup of ensemble average precipitation, following Neelin et al 2009. + + Fortran --> Python translation note: this runs in the fortran, but it does not modify inputs + and the output is never used. For this reason, it is not implemented in the Python version. + """ + + def __init__(self): + pass + + def __call__(self, *args, **kwds): + pass + + +def get_precip_fluxes( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + epsilon_forced: FloatFieldIJ_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + precipitation_flux: FloatField, + evaporation_flux: FloatField, + plume: Int, +): + """Compute the evaporation and precipitation flux throughout the entire column. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + epsilon_forced (FloatFieldIJ_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + precipitation_flux (FloatField) + evaporation_flux (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(...): + precipitation_flux = 0.0 + evaporation_flux = 0.0 + + with computation(BACKWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + precipitation_flux = precipitation_flux[0, 0, 1] + cloud_base_mass_flux_modified[0, 0][ + plume + ] * ( + condensate_to_fall_forced[0, 0, 0][plume] + + epsilon_forced[0, 0][plume] * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + precipitation_flux = max(0.0, precipitation_flux) + + evaporation_flux = ( + evaporation_flux[0, 0, 1] + - cloud_base_mass_flux_modified[0, 0][plume] + * epsilon_forced[0, 0][plume] + * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + evaporation_flux = max(0.0, evaporation_flux) + + +def rain_evaporation_below_cloud_base( + error_code: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + p_cloud_levels_forced: FloatField_Plume, + p_surface: FloatFieldIJ, + t_cloud_levels: FloatField, + vapor_cloud_levels_forced: FloatField, + environment_saturation_mixing_ratio_cloud_levels: FloatField, + epsilon_forced: FloatFieldIJ_Plume, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + precip: FloatFieldIJ_Plume, + precipitation_flux: FloatField, + evaporation_flux: FloatField, + evaporation_below_cloud_base: FloatField, + dtdt: FloatField_Plume, + dvapordt: FloatField_Plume, + dbuoyancydt: FloatField_Plume, + plume: Int, +): + """Allow rain the evaporate below the cloud base. + + Args: + error_code (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + ocean_fraction (FloatFieldIJ) + p_cloud_levels_forced (FloatField_Plume) + p_surface (FloatFieldIJ) + t_cloud_levels (FloatField) + vapor_cloud_levels_forced (FloatField) + environment_saturation_mixing_ratio_cloud_levels (FloatField) + epsilon_forced (FloatFieldIJ_Plume) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + precip (FloatFieldIJ_Plume) + precipitation_flux (FloatField) + evaporation_flux (FloatField) + evaporation_below_cloud_base (FloatField) + dtdt (FloatField_Plume) + dvapordt (FloatField_Plume) + dbuoyancydt (FloatField_Plume) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + # setup internal constants + alpha1: FloatFieldIJ = 5.44e-4 + alpha2: FloatFieldIJ = 5.09e-3 + alpha3: FloatFieldIJ = 0.5777 + c_conv: FloatFieldIJ = 0.05 + + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.SHALLOW: + critical_rh_ocean: FloatFieldIJ = 1.0 + critical_rh_land: FloatFieldIJ = 1.0 + eff_c_conv: FloatFieldIJ = min(0.2, max(cloud_base_mass_flux_modified[0, 0][plume], c_conv)) + else: + critical_rh_ocean: FloatFieldIJ = 0.95 + critical_rh_land: FloatFieldIJ = 0.85 + eff_c_conv: FloatFieldIJ = c_conv + + total_evaporation_below_cloud_base: FloatFieldIJ = 0.0 + + with computation(PARALLEL), interval(...): + precipitation_flux = 0.0 + evaporation_flux = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + critical_rh: FloatFieldIJ = critical_rh_ocean * ocean_fraction + critical_rh_land * ( + 1.0 - ocean_fraction + ) + + with computation(BACKWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + dp: FloatFieldIJ = 100.0 * ( + p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume] + ) + + if K <= updraft_lfc_level[0, 0][plume]: + vapor_deficit: FloatFieldIJ = max( + 0.0, + ( + critical_rh * environment_saturation_mixing_ratio_cloud_levels + - vapor_cloud_levels_forced + ), + ) + + evaporation_below_cloud_base = ( + eff_c_conv + * alpha1 + * vapor_deficit + * ( + sqrt(p_cloud_levels_forced[0, 0, 0][plume] / p_surface) + / alpha2 + * precipitation_flux[0, 0, 1] + / eff_c_conv + ) + ** alpha3 + ) + + evaporation_below_cloud_base = evaporation_below_cloud_base * dp / constants.MAPL_GRAV + + else: + + evaporation_below_cloud_base = 0.0 + + precipitation_flux = ( + precipitation_flux[0, 0, 1] + - evaporation_below_cloud_base + + cloud_base_mass_flux_modified[0, 0][plume] + * ( + condensate_to_fall_forced[0, 0, 0][plume] + + epsilon_forced[0, 0][plume] * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + ) + precipitation_flux = max(0.0, precipitation_flux) + + evaporation_flux = ( + evaporation_flux[0, 0, 1] + + evaporation_below_cloud_base + - cloud_base_mass_flux_modified[0, 0][plume] + * epsilon_forced[0, 0][plume] + * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + evaporation_flux = max(0.0, evaporation_flux) + + total_evaporation_below_cloud_base = ( + total_evaporation_below_cloud_base + evaporation_below_cloud_base + ) + + del_vapor = evaporation_below_cloud_base * constants.MAPL_GRAV / dp + del_t = ( + -evaporation_below_cloud_base + * constants.MAPL_GRAV + / dp + * (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + ) + + dvapordt[0, 0, 0][plume] = dvapordt[0, 0, 0][plume] + del_vapor + dtdt[0, 0, 0][plume] = dtdt[0, 0, 0][plume] + del_t + dbuoyancydt[0, 0, 0][plume] = ( + dbuoyancydt[0, 0, 0][plume] + + cumulus_parameterization_constants.CP * del_t + + cumulus_parameterization_constants.XLV * del_vapor + ) + + precip[0, 0][plume] = precip[0, 0][plume] - evaporation_below_cloud_base + + +def cloud_dissapation( + error_code: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + hydrostatic_air_density: FloatField, + geopotential_height_forced: FloatField, + t_cloud_levels_forced: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + vapor_cloud_levels_forced: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + environment_saturation_mixing_ratio_cloud_levels_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + vertical_velocity_3d: FloatField, + scale_dependence_factor: FloatFieldIJ_Plume, + dtdt: FloatField_Plume, + dvapordt: FloatField_Plume, + dcloudicedt: FloatField_Plume, + plume: Int, +): + """After excess water has precipitated, dissipate clouds which are no longer saturated. + + Args: + error_code (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + hydrostatic_air_density (FloatField) + geopotential_height_forced (FloatField) + t_cloud_levels_forced (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + vapor_cloud_levels_forced (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + environment_saturation_mixing_ratio_cloud_levels_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + vertical_velocity_3d (FloatField) + scale_dependence_factor (FloatFieldIJ_Plume) + dtdt (FloatField_Plume) + dvapordt (FloatField_Plume) + dcloudicedt (FloatField_Plume) + plume (Int) + """ + from __externals__ import COUPLE_MICROPHYSICS, DTIME, USE_CLOUD_DISSIPATION + + with computation(FORWARD), interval(0, 1): + # setup internal constants + cloud_lifetime: FloatFieldIJ = 1800.0 + version_x: IntFieldIJ = 2 + + with computation(BACKWARD), interval(...): + if ( + error_code[0, 0][plume] == 0 + and USE_CLOUD_DISSIPATION >= 0.0 + and K <= cloud_top_level[0, 0][plume] + and K >= updraft_lfc_level[0, 0][plume] + ): + # cloud liq/ice remained in the convection plume + precip_dissipation = max( + 0.0, cloud_liquid_after_rain_forced[0, 0, 0][plume] - dcloudicedt[0, 0, 0][plume] * DTIME + ) + + # get relative humidity + f_rh = 0.0 + + # estimation of the fractional area + fractional_area = ( + (cloud_base_mass_flux_modified[0, 0][plume] / scale_dependence_factor[0, 0][plume]) + * normalized_massflux_updraft_forced[0, 0, 0][plume] + / (hydrostatic_air_density * vertical_velocity_3d) + ) + + # source of enviroment moistening/cooling due to the 'remained' cloud dissipation into it. + out_precip_dissipation = (precip_dissipation * (1.0 - f_rh)) / cloud_lifetime + + # NOTE other option (if this is true) is not implemented + if not (version_x == 1 or COUPLE_MICROPHYSICS == False): + dcloudicedt[0, 0, 0][plume] = ( + dcloudicedt[0, 0, 0][plume] + + out_precip_dissipation * fractional_area * USE_CLOUD_DISSIPATION + ) + + cloud_liquid_after_rain_forced[0, 0, 0][plume] = max( + 0.0, + cloud_liquid_after_rain_forced[0, 0, 0][plume] + - out_precip_dissipation * USE_CLOUD_DISSIPATION * fractional_area * DTIME, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/prepare_output.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/prepare_output.py new file mode 100644 index 000000000..71b635d15 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/prepare_output.py @@ -0,0 +1,661 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Ensemble, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import liquid_fraction +from pyMoist.shared_incloud_processes import ( + G_RATIO, + RADIATIVE_EFFECTIVE_RADIUS, + G_RATIO_Table_Type, + RADIATIVE_EFFECTIVE_RADIUS_Table_Type, + make_droplet_number, + make_ice_number, +) + + +def ensemble_output_and_feedback( + error_code: IntFieldIJ_Plume, + error_code_2: IntFieldIJ, + error_code_3: IntFieldIJ, + cloud_top_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + precip: FloatFieldIJ_Plume, + effective_condensate_to_fall_forced: FloatField, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + scale_dependence_factor: FloatFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + f_dicycle_modified: FloatFieldIJ, + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + del_t_cloud_ensemble: FloatField, + del_vapor_cloud_ensemble: FloatField, + del_cloud_liquid_cloud_ensemble: FloatField, + del_buoyancy_cloud_ensemble: FloatField, + del_convective_ice_cloud_ensemble: FloatField, + del_large_scale_ice_cloud_ensemble: FloatField, + del_convective_liquid_cloud_ensemble: FloatField, + del_large_scale_liquid_cloud_ensemble: FloatField, + del_convective_cloud_fraction_cloud_ensemble: FloatField, + del_large_scale_cloud_fraction_cloud_ensemble: FloatField, + dtdt: FloatField_Plume, + dvapordt: FloatField_Plume, + dcloudicedt: FloatField_Plume, + dudt: FloatField_Plume, + dvdt: FloatField_Plume, + dbuoyancydt: FloatField_Plume, + dconvectiveicedt: FloatField_Plume, + dlargescaleicedt: FloatField_Plume, + dconvectiveliquiddt: FloatField_Plume, + dlargescaleliquiddt: FloatField_Plume, + dconvectivecloudfractiondt: FloatField_Plume, + dlargescalecloudfractiondt: FloatField_Plume, + mass_flux_ensemble: FloatFieldIJ_Ensemble, + precipitation_ensemble: FloatFieldIJ_Ensemble, + xff_mid: FloatFieldIJ_Ensemble, + CLOSURE_CHOICE: Int, + CLOUD_BASE_MASS_FLUX_FACTOR: Float, + plume: Int, +): + """Compute output tendencies for microphysical properties (ice/liquid/vapor concentrations), wind (u/v), + buoyancy, and temperature. + + Behavior of this stencil is heavily dependent on the value of the CLOSURE_CHOICE external. + + Args: + error_code (IntFieldIJ_Plume) + error_code_2 (IntFieldIJ) + error_code_3 (IntFieldIJ) + cloud_top_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + precip (FloatFieldIJ_Plume) + effective_condensate_to_fall_forced (FloatField) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + scale_dependence_factor (FloatFieldIJ_Plume) + ocean_fraction (FloatFieldIJ) + f_dicycle_modified (FloatFieldIJ) + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + del_t_cloud_ensemble (FloatField) + del_vapor_cloud_ensemble (FloatField) + del_cloud_liquid_cloud_ensemble (FloatField) + del_buoyancy_cloud_ensemble (FloatField) + del_convective_ice_cloud_ensemble (FloatField) + del_large_scale_ice_cloud_ensemble (FloatField) + del_convective_liquid_cloud_ensemble (FloatField) + del_large_scale_liquid_cloud_ensemble (FloatField) + del_convective_cloud_fraction_cloud_ensemble (FloatField) + del_large_scale_cloud_fraction_cloud_ensemble (FloatField) + dtdt (FloatField_Plume) + dvapordt (FloatField_Plume) + dcloudicedt (FloatField_Plume) + dudt (FloatField_Plume) + dvdt (FloatField_Plume) + dbuoyancydt (FloatField_Plume) + dconvectiveicedt (FloatField_Plume) + dlargescaleicedt (FloatField_Plume) + dconvectiveliquiddt (FloatField_Plume) + dlargescaleliquiddt (FloatField_Plume) + dconvectivecloudfractiondt (FloatField_Plume) + dlargescalecloudfractiondt (FloatField_Plume) + mass_flux_ensemble (FloatFieldIJ_Ensemble) + precipitation_ensemble (FloatFieldIJ_Ensemble) + xff_mid (FloatFieldIJ_Ensemble) + CLOSURE_CHOICE (Int) + CLOUD_BASE_MASS_FLUX_FACTOR (Float) + plume (Int) + """ + from __externals__ import ( + APPLY_SUBSIDENCE_MICROPHYSICS, + DTIME, + MAX_TEMP_VAPOR_TENDENCY, + USE_SMOOTH_TENDENCIES, + k_end, + ) + + # ensure outputs are all zero + with computation(PARALLEL), interval(...): + dtdt[0, 0, 0][plume] = 0.0 + dvapordt[0, 0, 0][plume] = 0.0 + dcloudicedt[0, 0, 0][plume] = 0.0 + dudt[0, 0, 0][plume] = 0.0 + dvdt[0, 0, 0][plume] = 0.0 + dbuoyancydt[0, 0, 0][plume] = 0.0 + + with computation(FORWARD), interval(0, 1): + precip[0, 0][plume] = 0.0 + cloud_base_mass_flux_modified[0, 0][plume] = 0.0 + + if error_code[0, 0][plume] == 0: + member = 0 + while member < cumulus_parameterization_constants.MAXENS3: + if precipitation_ensemble[0, 0][member] <= 0.0: + mass_flux_ensemble[0, 0][member] = 0.0 + member += 1 + + # calculate ensemble average mass fluxes + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0.0: + average_mass_flux = 0.0 + if plume == cumulus_parameterization_constants.DEEP: + count = 0 + member = 0 + while member < cumulus_parameterization_constants.MAXENS3: + count = count + 1 + average_mass_flux = average_mass_flux + mass_flux_ensemble[0, 0][member] + member += 1 + # ensemble average mass flux + average_mass_flux = average_mass_flux / count + + elif plume == cumulus_parameterization_constants.MID: + if CLOSURE_CHOICE <= 5: + if CLOSURE_CHOICE == 0: + average_mass_flux = 0.3333 * (xff_mid[0, 0][0] + xff_mid[0, 0][1] + xff_mid[0, 0][2]) + else: + average_mass_flux = xff_mid[0, 0][CLOSURE_CHOICE - 1] + + elif plume == cumulus_parameterization_constants.SHALLOW: + option_not_implemented = True + + # set the updradt mass flux and do not allow negative values and apply the diurnal cycle closure + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # mass flux of updradt at cloud base + cloud_base_mass_flux_modified[0, 0][plume] = average_mass_flux + + # apply the adjust factor for tunning + cloud_base_mass_flux_modified[0, 0][plume] = ( + CLOUD_BASE_MASS_FLUX_FACTOR * cloud_base_mass_flux_modified[0, 0][plume] + ) + + # diurnal cycle closure + cloud_base_mass_flux_modified[0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] - f_dicycle_modified + ) + + if cloud_base_mass_flux_modified[0, 0][plume] <= 0.0: + error_code[0, 0][plume] = 13 + cloud_base_mass_flux_modified[0, 0][plume] = 0.0 + + # apply the scale-dependence Arakawa's approach + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # scale dependence + cloud_base_mass_flux_modified[0, 0][plume] = ( + scale_dependence_factor[0, 0][plume] * cloud_base_mass_flux_modified[0, 0][plume] + ) + + if cloud_base_mass_flux_modified[0, 0][plume] == 0.0: + error_code[0, 0][plume] = 14 + if cloud_base_mass_flux_modified[0, 0][plume] > 100.0: + error_code[0, 0][plume] = 15 + + # sanity check for mass flux + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + max_mass_flux = ( + 100.0 + * ( + p_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume], ddim=[plume]) + - p_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume] + 1, ddim=[plume]) + ) + / (constants.MAPL_GRAV * DTIME) + ) + cloud_base_mass_flux_modified[0, 0][plume] = min( + cloud_base_mass_flux_modified[0, 0][plume], max_mass_flux + ) + + with computation(PARALLEL), interval(...): + # prepare field from which max value is pulled + abs_del_t_cloud_ensemble = abs(del_t_cloud_ensemble) + abs_del_vapor_cloud_ensemble = abs(del_vapor_cloud_ensemble) + + # check dtdt and and dvapordt for high values + # criteria: if abs (dtdt or dvapordt) > MAX_TEMP_VAPOR_TENDENCY (default 100 K/day) => fix mass flux + with computation(FORWARD), interval(0, 1): + if MAX_TEMP_VAPOR_TENDENCY > 0.0 and error_code[0, 0][plume] == 0: + max_val_t, _ = column_max(abs_del_t_cloud_ensemble, 0, cloud_top_level[0, 0][plume]) + max_val_vapor, _ = column_max(abs_del_vapor_cloud_ensemble, 0, cloud_top_level[0, 0][plume]) + fixouts = ( + cloud_base_mass_flux_modified[0, 0][plume] + * 86400.0 + * max( + max_val_t, + (cumulus_parameterization_constants.XLV / cumulus_parameterization_constants.CP) + * max_val_vapor, + ) + ) + + if fixouts > MAX_TEMP_VAPOR_TENDENCY: # K/day + fixouts = MAX_TEMP_VAPOR_TENDENCY / fixouts + cloud_base_mass_flux_modified[0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * fixouts + ) + member = 0 + while member < cumulus_parameterization_constants.MAXENS3: + mass_flux_ensemble[0, 0][member] = mass_flux_ensemble[0, 0][member] * fixouts + member += 1 + + # now do feedback + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + precip[0, 0][plume] = ( + precip[0, 0][plume] + + effective_condensate_to_fall_forced * cloud_base_mass_flux_modified[0, 0][plume] + ) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + dtdt[0, 0, 0][plume] = del_t_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + dvapordt[0, 0, 0][plume] = ( + del_vapor_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dcloudicedt[0, 0, 0][plume] = ( + del_cloud_liquid_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dudt[0, 0, 0][plume] = del_u_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + dvdt[0, 0, 0][plume] = del_v_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + dbuoyancydt[0, 0, 0][plume] = ( + del_buoyancy_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + + if APPLY_SUBSIDENCE_MICROPHYSICS == 1: + if K <= cloud_top_level[0, 0][plume]: + dconvectiveicedt[0, 0, 0][plume] = ( + del_convective_ice_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dlargescaleicedt[0, 0, 0][plume] = ( + del_large_scale_ice_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dconvectiveliquiddt[0, 0, 0][plume] = ( + del_convective_liquid_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dlargescaleliquiddt[0, 0, 0][plume] = ( + del_large_scale_liquid_cloud_ensemble * cloud_base_mass_flux_modified[0, 0][plume] + ) + dconvectivecloudfractiondt[0, 0, 0][plume] = ( + del_convective_cloud_fraction_cloud_ensemble + * cloud_base_mass_flux_modified[0, 0][plume] + ) + dlargescalecloudfractiondt[0, 0, 0][plume] = ( + del_large_scale_cloud_fraction_cloud_ensemble + * cloud_base_mass_flux_modified[0, 0][plume] + ) + + if K >= cloud_top_level[0, 0][plume] and K < k_end: + dconvectiveicedt[0, 0, 0][plume] = 0.0 + dlargescaleicedt[0, 0, 0][plume] = 0.0 + dconvectiveliquiddt[0, 0, 0][plume] = 0.0 + dlargescaleliquiddt[0, 0, 0][plume] = 0.0 + dconvectivecloudfractiondt[0, 0, 0][plume] = 0.0 + dlargescalecloudfractiondt[0, 0, 0][plume] = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + member = 0 + while ( + member + < cumulus_parameterization_constants.MAXENS1 + * cumulus_parameterization_constants.MAXENS2 + * cumulus_parameterization_constants.MAXENS3 + ): + mass_flux_ensemble[0, 0][member] = ( + scale_dependence_factor[0, 0][plume] * mass_flux_ensemble[0, 0][member] + ) + member += 1 + + # smooth the tendencies (future work: include outbuoy, outmpc* and tracers) + with computation(PARALLEL), interval(...): + if USE_SMOOTH_TENDENCIES < 0: + option_not_implemented = 0.0 + + +def total_evaporation_flux( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + evaporation_flux: FloatField, + evaporation_sublimation_tendency: FloatField, + plume: Int, +): + """Compute evaporation of rain below cloud_top_level. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + evaporation_flux (FloatField) + evaporation_sublimation_tendency (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + if K <= cloud_top_level[0, 0][plume]: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + evaporation_sublimation_tendency = ( + evaporation_sublimation_tendency + evaporation_flux * constants.MAPL_GRAV / dp + ) + + +def deep_precipitation_output( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + precipitation_flux: FloatField, + convective_precip_flux: FloatField, + plume: Int, +): + """Compute precipitation specifically from deep convection. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + precipitation_flux (FloatField) + convective_precip_flux (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(...): + if ( + error_code[0, 0][plume] == 0 + and plume == cumulus_parameterization_constants.DEEP + and K <= cloud_top_level[0, 0][plume] + 1 + ): + convective_precip_flux = precipitation_flux + + +def output_updraft_temperature( + error_code: IntFieldIJ_Plume, + updraft_column_temperature_forced: FloatField, + t_cloud_levels: FloatField, + t_updraft: FloatField_Plume, + plume: Int, +): + """Copy the internal updraft temperature to an output array so that it can be + fed back to the rest of the model. + + Args: + error_code (IntFieldIJ_Plume) + updraft_column_temperature_forced (FloatField) + t_cloud_levels (FloatField) + t_updraft (FloatField_Plume) + plume (Int) + """ + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + t_updraft[0, 0, 0][plume] = updraft_column_temperature_forced + with computation(PARALLEL), interval(-1, None): + if error_code[0, 0][plume] == 0: + t_updraft[0, 0, 0][plume] = t_cloud_levels + + +def prepare_output( + error_code: IntFieldIJ_Plume, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + total_normalized_integrated_evaporate_forced: FloatFieldIJ, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + environment_massflux: FloatField, + vapor_tendency_from_environmental_subsidence: FloatField, + moist_static_energy_tendency_from_environmental_subsidence: FloatField, + t_tendency_from_environmental_subsidence: FloatField, + plume: Int, +): + """Scale output mass fluxes based on the cloud base mass flux before output. + + Args: + error_code (IntFieldIJ_Plume) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + total_normalized_integrated_evaporate_forced (FloatFieldIJ) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + environment_massflux (FloatField) + vapor_tendency_from_environmental_subsidence (FloatField) + moist_static_energy_tendency_from_environmental_subsidence (FloatField) + t_tendency_from_environmental_subsidence (FloatField) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + total_normalized_integrated_condensate_forced[0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] + * total_normalized_integrated_condensate_forced[0, 0][plume] + ) + total_normalized_integrated_evaporate_forced = ( + cloud_base_mass_flux_modified[0, 0][plume] * total_normalized_integrated_evaporate_forced + ) + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + normalized_massflux_updraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] + * normalized_massflux_updraft_forced[0, 0, 0][plume] + ) + normalized_massflux_downdraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] + * normalized_massflux_downdraft_forced[0, 0, 0][plume] + ) + condensate_to_fall_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * condensate_to_fall_forced[0, 0, 0][plume] + ) + evaporate_in_downdraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + mass_entrainment_updraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * mass_entrainment_updraft_forced[0, 0, 0][plume] + ) + mass_detrainment_updraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * mass_detrainment_updraft_forced[0, 0, 0][plume] + ) + mass_entrainment_downdraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * mass_entrainment_downdraft_forced[0, 0, 0][plume] + ) + mass_detrainment_downdraft_forced[0, 0, 0][plume] = ( + cloud_base_mass_flux_modified[0, 0][plume] * mass_detrainment_downdraft_forced[0, 0, 0][plume] + ) + environment_massflux = cloud_base_mass_flux_modified[0, 0][plume] * environment_massflux + + with computation(PARALLEL), interval(...): + vapor_tendency_from_environmental_subsidence = ( + cloud_base_mass_flux_modified[0, 0][plume] * vapor_tendency_from_environmental_subsidence + ) + moist_static_energy_tendency_from_environmental_subsidence = ( + cloud_base_mass_flux_modified[0, 0][plume] + * moist_static_energy_tendency_from_environmental_subsidence + ) + t_tendency_from_environmental_subsidence = ( + cloud_base_mass_flux_modified[0, 0][plume] * t_tendency_from_environmental_subsidence + ) + + +def output_workfunctions_and_precip_concentrations( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + convection_fraction: FloatFieldIJ, + surface_type: FloatFieldIJ, + cloud_workfunction_0_output: FloatFieldIJ, + cloud_workfunction_1_output: FloatFieldIJ, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + air_density: FloatField, + updraft_column_temperature_forced: FloatField, + dcloudicedt: FloatField_Plume, + dnliquiddt: FloatField_Plume, + dnicedt: FloatField_Plume, + G_RATIO: G_RATIO_Table_Type, + RADIATIVE_EFFECTIVE_RADIUS: RADIATIVE_EFFECTIVE_RADIUS_Table_Type, + plume: Int, +): + """Push the internal copy of workfunctions 0 and 1 to the state fields for output, and get precipitate + concentrations using the functions make_droplet_number and make_ice_number + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + convection_fraction (FloatFieldIJ) + surface_type (FloatFieldIJ) + cloud_workfunction_0_output (FloatFieldIJ) + cloud_workfunction_1_output (FloatFieldIJ) + cloud_workfunction_0 (FloatFieldIJ) + cloud_workfunction_1 (FloatFieldIJ) + air_density (FloatField) + updraft_column_temperature_forced (FloatField) + dcloudicedt (FloatField_Plume) + dnliquiddt (FloatField_Plume) + dnicedt (FloatField_Plume) + G_RATIO (G_RATIO_Table_Type) + RADIATIVE_EFFECTIVE_RADIUS (RADIATIVE_EFFECTIVE_RADIUS_Table_Type) + plume (Int) + """ + from __externals__ import DTIME, FRAC_MODIS + + with computation(PARALLEL), interval(...): + n_water_friendly_aerosols = 99.0e7 # in the future set this as NCPL + + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.DEEP: + cloud_workfunction_0_output = 0.0 + cloud_workfunction_1_output = 0.0 + + if error_code[0, 0][plume] == 0: + cloud_workfunction_0_output = cloud_workfunction_0 + cloud_workfunction_1_output = cloud_workfunction_1 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume] + 1: + fraction = liquid_fraction( + updraft_column_temperature_forced, convection_fraction, surface_type, FRAC_MODIS + ) + cloud_liquid = DTIME * dcloudicedt[0, 0, 0][plume] * air_density * fraction + cloud_ice = DTIME * dcloudicedt[0, 0, 0][plume] * air_density * (1.0 - fraction) + + dnicedt[0, 0, 0][plume] = max( + 0.0, + make_ice_number(cloud_ice, updraft_column_temperature_forced, RADIATIVE_EFFECTIVE_RADIUS) + / air_density, + ) + dnliquiddt[0, 0, 0][plume] = max( + 0.0, make_droplet_number(cloud_liquid, n_water_friendly_aerosols, G_RATIO) / air_density + ) + + # convert to tendencies + dnicedt[0, 0, 0][plume] = dnicedt[0, 0, 0][plume] * (1 / DTIME) # unit [1/s] + dnliquiddt[0, 0, 0][plume] = dnliquiddt[0, 0, 0][plume] * (1 / DTIME) # unit [1/s] + + +class OutputWorkfunctionsAndPrecipConcentrations(NDSLRuntime): + """Push the internal copy of workfunctions 0 and 1 to the state fields for output, and get precipitate + concentrations using the functions make_droplet_number and make_ice_number + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # add dimension to quantityfactory and create classes for constants + quantity_factory.add_data_dimensions({"G_RATIO_Table": len(G_RATIO)}) + quantity_factory.add_data_dimensions( + {"RADIATIVE_EFFECTIVE_RADIUS_Table": len(RADIATIVE_EFFECTIVE_RADIUS)} + ) + + self._G_RATIO: Local = quantity_factory.zeros(["G_RATIO_Table"], "n/a") + self._RADIATIVE_EFFECTIVE_RADIUS: Local = quantity_factory.zeros( + ["RADIATIVE_EFFECTIVE_RADIUS_Table"], "n/a" + ) + + self._G_RATIO.field[:] = G_RATIO + self._RADIATIVE_EFFECTIVE_RADIUS.field[:] = RADIATIVE_EFFECTIVE_RADIUS + + # construct stencils + self._output_workfunctions_and_precip_concentrations = stencil_factory.from_dims_halo( + func=output_workfunctions_and_precip_concentrations, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "FRAC_MODIS": cumulus_parameterization_config.FRAC_MODIS, + "DTIME": cumulus_parameterization_config.DTIME, + }, + ) + + def __call__( + self, + error_code: Quantity, + cloud_top_level: Quantity, + convection_fraction: Quantity, + surface_type: Quantity, + cloud_workfunction_0_output: Quantity, + cloud_workfunction_1_output: Quantity, + cloud_workfunction_0: Quantity, + cloud_workfunction_1: Quantity, + air_density: Quantity, + updraft_column_temperature_forced: Quantity, + dcloudicedt: Quantity, + dnliquiddt: Quantity, + dnicedt: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._output_workfunctions_and_precip_concentrations( + error_code=error_code, + cloud_top_level=cloud_top_level, + convection_fraction=convection_fraction, + surface_type=surface_type, + cloud_workfunction_0_output=cloud_workfunction_0_output, + cloud_workfunction_1_output=cloud_workfunction_1_output, + cloud_workfunction_0=cloud_workfunction_0, + cloud_workfunction_1=cloud_workfunction_1, + air_density=air_density, + updraft_column_temperature_forced=updraft_column_temperature_forced, + dcloudicedt=dcloudicedt, + dnliquiddt=dnliquiddt, + dnicedt=dnicedt, + G_RATIO=self._G_RATIO, + RADIATIVE_EFFECTIVE_RADIUS=self._RADIATIVE_EFFECTIVE_RADIUS, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class LightningFlashDensity: + def __init__(self, cumulus_parameterization_config: GF2020CumulusParameterizationConfig): + if cumulus_parameterization_config.LIGHTNING_DIAGNOSTICS: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->LightningFlashDensity: lightning output has not" + "been implemented. You should have been caught before getting here by the config checker." + "Beware, something likely failing in the config checker as well - you may be unknowingly" + "calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/profiles.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/profiles.py new file mode 100644 index 000000000..513e6b99f --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/profiles.py @@ -0,0 +1,102 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, interval +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Int +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume + + +def melting_profile( + error_code: IntFieldIJ_Plume, + plume: Int, + melting_layer: FloatField, + partition_liquid_ice: FloatField, + p_cloud_levels_forced: FloatField_Plume, + condensate_to_fall_forced: FloatField_Plume, + melting: FloatField, +): + """Generate the melting profile, given the excess water in the updraft. + + Args: + error_code (IntFieldIJ_Plume) + plume (Int) + melting_layer (FloatField) + partition_liquid_ice (FloatField) + p_cloud_levels_forced (FloatField_Plume) + condensate_to_fall_forced (FloatField_Plume) + melting (FloatField) + """ + from __externals__ import MELT_GLAC, k_end + + with computation(FORWARD), interval(...): + if MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP: + solid_phase_precipitable_water = 0.0 + effective_precipitable_water = 0.0 + melting = 0.0 + + total_solid_phase_precipitable_water: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(0, -2): + if ( + MELT_GLAC == True + and plume == cumulus_parameterization_constants.DEEP + and error_code[0, 0][plume] == 0 + ): + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + + effective_precipitable_water = 0.5 * ( + condensate_to_fall_forced[0, 0, 0][plume] + condensate_to_fall_forced[0, 0, 1][plume] + ) + + solid_phase_precipitable_water = (1.0 - partition_liquid_ice) * effective_precipitable_water + + total_solid_phase_precipitable_water = ( + total_solid_phase_precipitable_water + + solid_phase_precipitable_water * dp / constants.MAPL_GRAV + ) + + with computation(PARALLEL), interval(0, -1): + if ( + MELT_GLAC == True + and plume == cumulus_parameterization_constants.DEEP + and error_code[0, 0][plume] == 0 + ): + melting = melting_layer * ( + total_solid_phase_precipitable_water + / ( + 100 + * ( + p_cloud_levels_forced.at(K=0, ddim=[plume]) + - p_cloud_levels_forced.at(K=k_end - 1, ddim=[plume]) + ) + / constants.MAPL_GRAV + ) + ) + + with computation(PARALLEL), interval(...): + if not (MELT_GLAC == True and plume == cumulus_parameterization_constants.DEEP): + melting = 0.0 + + +class C1DProfile: + """ + C1D Profile generator. This code is manually diables with a + "do not enable" note in fortran, so it is not going to be completed until specifically requested. + """ + + def __init__( + self, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + if cumulus_parameterization_config.FIRST_GUESS_W or config.AUTOCONV == 4: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->C1DProfile: C1DProfile option has not been" + "implemented. You should have been caught before getting here by the config checker." + "Beware, something likely failing in the config checker as well - you may be unknowingly" + "calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/set_constants.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/set_constants.py new file mode 100644 index 000000000..70f210020 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/set_constants.py @@ -0,0 +1,255 @@ +from ndsl.dsl.typing import Float, Int +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) + + +def set_constants( + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + plume_dependent_constants: GF2020PlumeDependentConstants, + plume: str, +): + if plume == "shallow": + # set a number of plume dependent constants + plume_dependent_constants.PLUME_INDEX = Int(0) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_LAND_SHALLOW + ) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_OCEAN_SHALLOW + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_LAND_SHALLOW + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_OCEAN_SHALLOW + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_LAND_SHALLOW + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_OCEAN_SHALLOW + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_LAND_SHALLOW + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_OCEAN_SHALLOW + ) + plume_dependent_constants.CLOUD_BASE_MASS_FLUX_FACTOR = ( + cumulus_parameterization_config.CLOUD_BASE_MASS_FLUX_FACTOR_SHALLOW + ) + plume_dependent_constants.USE_EXCESS = cumulus_parameterization_config.USE_EXCESS_SHALLOW + plume_dependent_constants.ENTRAINMENT_RATE = cumulus_parameterization_config.ENTRAINMENT_RATE_SHALLOW + plume_dependent_constants.ENABLE_PLUME = cumulus_parameterization_config.ENABLE_SHALLOW + plume_dependent_constants.AVERAGE_LAYER_DEPTH = ( + cumulus_parameterization_config.AVERAGE_LAYER_DEPTH_SHALLOW + ) + + # maximum depth (mb) of capping inversion (larger cap = no convection) + if ( + cumulus_parameterization_config.ZERO_DIFF == 1 + or cumulus_parameterization_config.MOIST_TRIGGER == 0 + ): + plume_dependent_constants.CAP_MAX_INC = Float(25.0) + else: + plume_dependent_constants.CAP_MAX_INC = Float(10.0) + + # lambda_U parameter for momentum transport + if cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT != 0.0: + plume_dependent_constants.LAMBDA_DEEP = Float(0.0) + plume_dependent_constants.LAMBDA_DOWN = Float(0.0) + else: + plume_dependent_constants.LAMBDA_DEEP = cumulus_parameterization_config.LAMBDA_DEEP + plume_dependent_constants.LAMBDA_DOWN = cumulus_parameterization_config.LAMBDA_SHALLOW_DOWN + + # minimum depth (m) clouds must have + plume_dependent_constants.MINIMUM_DEPTH = Float(500.0) + + # max height(m) above ground where updraft air can originate + plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT = Float(2000.0) + + # height(m) above which no downdrafts are allowed to originate + plume_dependent_constants.MAX_DOWNDRAFT_ORIGIN_HEIGHt = Float(3000.0) + + # depth(m) over which downdraft detrains all its mass + plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH = ( + Float(0.5) * plume_dependent_constants.MINIMUM_DEPTH + ) + + plume_dependent_constants.C0 = cumulus_parameterization_config.C0_SHAL + + # temperature for diurnal cycle section + plume_dependent_constants.T_STAR = Float(40.0) + + # timescale of cape removal + plume_dependent_constants.TAU_CAPE_REMOVAL = ( + -999 + ) # not used - ideally should not exist for this plume + + # closure choice + plume_dependent_constants.CLOSURE_CHOICE = cumulus_parameterization_config.CLOSURE_CHOICE_SHALLOW + + elif plume == "mid": + # set a number of plume dependent constants + plume_dependent_constants.PLUME_INDEX = Int(1) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_LAND_MID + ) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_OCEAN_MID + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_LAND_MID + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_OCEAN_MID + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_LAND_MID + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_OCEAN_MID + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_LAND_MID + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_OCEAN_MID + ) + plume_dependent_constants.CLOUD_BASE_MASS_FLUX_FACTOR = ( + cumulus_parameterization_config.CLOUD_BASE_MASS_FLUX_FACTOR_MID + ) + plume_dependent_constants.USE_EXCESS = cumulus_parameterization_config.USE_EXCESS_MID + plume_dependent_constants.ENTRAINMENT_RATE = cumulus_parameterization_config.ENTRAINMENT_RATE_MID + plume_dependent_constants.ENABLE_PLUME = cumulus_parameterization_config.ENABLE_MID + plume_dependent_constants.AVERAGE_LAYER_DEPTH = ( + cumulus_parameterization_config.AVERAGE_LAYER_DEPTH_MID + ) + + # maximum depth (mb) of capping inversion (larger cap = no convection) + if ( + cumulus_parameterization_config.ZERO_DIFF == 1 + or cumulus_parameterization_config.MOIST_TRIGGER == 0 + ): + plume_dependent_constants.CAP_MAX_INC = Float(10.0) + else: + plume_dependent_constants.CAP_MAX_INC = Float(90.0) + + # lambda_U parameter for momentum transport + if cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT != 0.0: + plume_dependent_constants.LAMBDA_DEEP = Float(0.0) + plume_dependent_constants.LAMBDA_DOWN = Float(0.0) + else: + plume_dependent_constants.LAMBDA_DEEP = cumulus_parameterization_config.LAMBDA_SHALLOW_DOWN + plume_dependent_constants.LAMBDA_DOWN = cumulus_parameterization_config.LAMBDA_SHALLOW_DOWN + + # minimum depth (m) clouds must have + plume_dependent_constants.MINIMUM_DEPTH = Float(1000.0) + + # max height(m) above ground where updraft air can originate + plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT = Float(3000.0) + + # height(m) above which no downdrafts are allowed to originate + plume_dependent_constants.MAX_DOWNDRAFT_ORIGIN_HEIGHt = Float(3000.0) + + # depth(m) over which downdraft detrains all its mass + plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH = ( + Float(0.5) * plume_dependent_constants.MINIMUM_DEPTH + ) + + plume_dependent_constants.C0 = cumulus_parameterization_config.C0_MID + + # temperature for diurnal cycle section + plume_dependent_constants.T_STAR = Float(40.0) + + # timescale of cape removal + plume_dependent_constants.TAU_CAPE_REMOVAL = cumulus_parameterization_config.TAU_MID + + # closure choice + plume_dependent_constants.CLOSURE_CHOICE = cumulus_parameterization_config.CLOSURE_CHOICE_MID + + elif plume == "deep": + # set a number of plume dependent constants + plume_dependent_constants.PLUME_INDEX = Int(2) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_LAND_DEEP + ) + plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.DOWNDRAFT_MAX_HEIGHT_OCEAN_DEEP + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_LAND = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_LAND_DEEP + ) + plume_dependent_constants.UPDRAFT_MAX_HEIGHT_OCEAN = ( + cumulus_parameterization_config.UPDRAFT_MAX_HEIGHT_OCEAN_DEEP + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_LAND_DEEP + ) + plume_dependent_constants.MINIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MINIMUM_EVAP_FRACTION_OCEAN_DEEP + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_LAND = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_LAND_DEEP + ) + plume_dependent_constants.MAXIMUM_EVAP_FRACTION_OCEAN = ( + cumulus_parameterization_config.MAXIMUM_EVAP_FRACTION_OCEAN_DEEP + ) + plume_dependent_constants.CLOUD_BASE_MASS_FLUX_FACTOR = ( + cumulus_parameterization_config.CLOUD_BASE_MASS_FLUX_FACTOR_DEEP + ) + plume_dependent_constants.USE_EXCESS = cumulus_parameterization_config.USE_EXCESS_DEEP + plume_dependent_constants.ENTRAINMENT_RATE = cumulus_parameterization_config.ENTRAINMENT_RATE_DEEP + plume_dependent_constants.ENABLE_PLUME = cumulus_parameterization_config.ENABLE_DEEP + plume_dependent_constants.AVERAGE_LAYER_DEPTH = ( + cumulus_parameterization_config.AVERAGE_LAYER_DEPTH_DEEP + ) + + # maximum depth (mb) of capping inversion (larger cap = no convection) + if ( + cumulus_parameterization_config.ZERO_DIFF == 1 + or cumulus_parameterization_config.MOIST_TRIGGER == 0 + ): + plume_dependent_constants.CAP_MAX_INC = Float(20.0) + else: + plume_dependent_constants.CAP_MAX_INC = Float(90.0) + + # lambda_U parameter for momentum transport + if cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT != 0.0: + plume_dependent_constants.LAMBDA_DEEP = Float(0.0) + plume_dependent_constants.LAMBDA_DOWN = Float(0.0) + else: + plume_dependent_constants.LAMBDA_DEEP = cumulus_parameterization_config.LAMBDA_DEEP + plume_dependent_constants.LAMBDA_DOWN = cumulus_parameterization_config.LAMBDA_SHALLOW_DOWN + + # minimum depth (m) clouds must have + plume_dependent_constants.MINIMUM_DEPTH = Float(1000.0) + + # max height(m) above ground where updraft air can originate + plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT = Float(4000.0) + + # height(m) above which no downdrafts are allowed to originate + plume_dependent_constants.MAX_DOWNDRAFT_ORIGIN_HEIGHt = Float(3000.0) + + # depth(m) over which downdraft detrains all its mass + plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH = ( + Float(0.5) * plume_dependent_constants.MINIMUM_DEPTH + ) + + plume_dependent_constants.C0 = cumulus_parameterization_config.C0_DEEP + + # temperature for diurnal cycle section + plume_dependent_constants.T_STAR = Float(5.0) # temp scale in original paper = 1 K + + # timescale of cape removal + plume_dependent_constants.TAU_CAPE_REMOVAL = cumulus_parameterization_config.TAU_DEEP + + # closure choice + plume_dependent_constants.CLOSURE_CHOICE = cumulus_parameterization_config.CLOSURE_CHOICE_DEEP + + else: + raise NotImplementedError("Unknown plume specified, corresponding constants are unavailable.") + + return plume_dependent_constants diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/setup.py new file mode 100644 index 000000000..da15d2a17 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/setup/setup.py @@ -0,0 +1,616 @@ +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, interval +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import MAXENS1, MAXENS2, MAXENS3 +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Ensemble, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.shared_generic_math import sigma + + +def set_plume_dependent_fields( + t_excess: FloatFieldIJ, + t_excess_local: FloatFieldIJ, + vapor_excess: FloatFieldIJ, + vapor_excess_local: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + use_excess: Int, + t_old: FloatField, + vapor_old: FloatField, + grid_scale_forcing_t: FloatField, + grid_scale_forcing_vapor: FloatField, + subgrid_scale_forcing_t: FloatField, + subgrid_scale_forcing_vapor: FloatField, + t_new: FloatField, + vapor_forced: FloatField, + t_new_pbl: FloatField, + vapor_forced_pbl: FloatField, + dmoist_static_energydt: FloatField, +): + """Fill or modify existing values of temperature/vapor excess before cumulus parameterization + + Args: + t_excess (FloatFieldIJ) + t_excess_local (FloatFieldIJ) + vapor_excess (FloatFieldIJ) + vapor_excess_local (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + use_excess (Int): trigger to control behavior + 0: fill with zero + 1: no change (retain existing values) + 2: enforce min/max everywhere + other: enforce min/max only over ocean + t_old (FloatField) + vapor_old (FloatField) + grid_scale_forcing_t (FloatField) + grid_scale_forcing_vapor (FloatField) + subgrid_scale_forcing_t (FloatField) + subgrid_scale_forcing_vapor (FloatField) + t_new (FloatField) + vapor_forced (FloatField) + t_new_pbl (FloatField) + vapor_forced_pbl (FloatField) + dmoist_static_energydt (FloatField) + """ + from __externals__ import DT_MOIST + + with computation(FORWARD), interval(0, 1): + if use_excess == 0: + t_excess_local = 0.0 + vapor_excess_local = 0.0 + elif use_excess == 2: + t_excess_local = min(0.5, max(0.2, t_excess)) # Kelvin + vapor_excess_local = min(5.0e-4, max(1.0e-4, vapor_excess)) # kg kg^-1 + else: + if ocean_fraction > 0.98: # ocean + t_excess_local = min(0.5, max(0.2, t_excess)) # Kelvin + vapor_excess_local = min(5.0e-4, max(1.0e-4, vapor_excess)) # kg kg^-1 + + with computation(PARALLEL), interval(0, -1): + t_new = t_old + (subgrid_scale_forcing_t + grid_scale_forcing_t) * DT_MOIST + vapor_forced = vapor_old + (subgrid_scale_forcing_vapor + grid_scale_forcing_vapor) * DT_MOIST + vapor_forced = max(cumulus_parameterization_constants.smaller_qv, vapor_forced) + + # temp/water vapor modified only by bl processes + t_new_pbl = t_old + (subgrid_scale_forcing_t) * DT_MOIST + vapor_forced_pbl = vapor_old + (subgrid_scale_forcing_vapor) * DT_MOIST + + # moist static energy + dmoist_static_energydt = cumulus_parameterization_constants.CP * ( + subgrid_scale_forcing_t + grid_scale_forcing_t + ) + cumulus_parameterization_constants.XLV * (subgrid_scale_forcing_vapor + grid_scale_forcing_vapor) + + +def prefil_internal_fields( + plume: Int, + maximum_updraft_origin_level: IntFieldIJ, + kstabm: IntFieldIJ_Plume, + ocean_fraction: FloatFieldIJ, + ocean_fraction_local: FloatFieldIJ, + cap_max: FloatFieldIJ, + error_code_2: IntFieldIJ, + error_code_3: IntFieldIJ, + CAP_MAX_INC: Float, + cap_max_increment: FloatFieldIJ, + geopotential_height: FloatField, + geopotential_height_local: FloatField, + geopotential_height_modified_local: FloatField, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + cloud_workfunction_2: FloatFieldIJ, + cloud_workfunction_3: FloatFieldIJ, + cloud_workfunction_0_pbl: FloatFieldIJ, + cloud_workfunction_1_pbl: FloatFieldIJ, + cloud_workfunction_1_fa: FloatFieldIJ, + cin_1: FloatFieldIJ, + scale_dependence_factor: FloatFieldIJ_Plume, + scale_dependence_factor_downdraft: FloatFieldIJ, + epsilon_forced: FloatFieldIJ_Plume, + epsilon_local: FloatFieldIJ, + cloud_moist_static_energy_downdraft_forced: FloatField, + cloud_moist_static_energy_forced_transported: FloatField, + k_x_modified: FloatFieldIJ, + pbl_time_scale: FloatFieldIJ, + t_wetbulb: FloatFieldIJ, + vapor_wetbulb: FloatFieldIJ, + cape_removal_time_scale: FloatFieldIJ, + f_dicycle_modified: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + downdraft_saturation_vapor_forced: FloatField, + c1d: FloatField, + evaporation_below_cloud_base: FloatField, + mass_flux_ensemble: FloatFieldIJ_Ensemble, + precipitation_ensemble: FloatFieldIJ_Ensemble, + precip: FloatFieldIJ_Plume, + lightning_density: FloatFieldIJ, +): + """Fill internal fields (from the cumulus parameterization state) to ensure no data is left over from + the previous call. Most fields are filled with zero, but not all. + + Args: + plume (Int): _description_ + maximum_updraft_origin_level (IntFieldIJ): _description_ + kstabm (IntFieldIJ_Plume): _description_ + ocean_fraction (FloatFieldIJ): _description_ + ocean_fraction_local (FloatFieldIJ): _description_ + cap_max (FloatFieldIJ): _description_ + error_code_2 (IntFieldIJ): _description_ + error_code_3 (IntFieldIJ): _description_ + CAP_MAX_INC (Float): _description_ + cap_max_increment (FloatFieldIJ): _description_ + geopotential_height (FloatField): _description_ + geopotential_height_local (FloatField): _description_ + geopotential_height_modified_local (FloatField): _description_ + cloud_workfunction_0 (FloatFieldIJ): _description_ + cloud_workfunction_1 (FloatFieldIJ): _description_ + cloud_workfunction_2 (FloatFieldIJ): _description_ + cloud_workfunction_3 (FloatFieldIJ): _description_ + cloud_workfunction_0_pbl (FloatFieldIJ): _description_ + cloud_workfunction_1_pbl (FloatFieldIJ): _description_ + cloud_workfunction_1_fa (FloatFieldIJ): _description_ + cin_1 (FloatFieldIJ): _description_ + scale_dependence_factor (FloatFieldIJ_Plume): _description_ + scale_dependence_factor_downdraft (FloatFieldIJ): _description_ + epsilon_forced (FloatFieldIJ_Plume): _description_ + epsilon_local (FloatFieldIJ): _description_ + cloud_moist_static_energy_downdraft_forced (FloatField): _description_ + cloud_moist_static_energy_forced_transported (FloatField): _description_ + k_x_modified (FloatFieldIJ): _description_ + pbl_time_scale (FloatFieldIJ): _description_ + t_wetbulb (FloatFieldIJ): _description_ + vapor_wetbulb (FloatFieldIJ): _description_ + cape_removal_time_scale (FloatFieldIJ): _description_ + f_dicycle_modified (FloatFieldIJ): _description_ + add_buoyancy (FloatFieldIJ): _description_ + downdraft_saturation_vapor_forced (FloatField): _description_ + c1d (FloatField): _description_ + evaporation_below_cloud_base (FloatField): _description_ + mass_flux_ensemble (FloatFieldIJ_Ensemble): _description_ + precipitation_ensemble (FloatFieldIJ_Ensemble): _description_ + precip (FloatFieldIJ_Plume): _description_ + lightning_density (FloatFieldIJ): _description_ + """ + from __externals__ import CAP_MAXS, ENSEMBLE_MEMBERS, k_end + + # reset to zero manually. cannot use locals.fill(0) because not all fields are reset to zero b/t plumes + # internal fields + with computation(FORWARD), interval(0, 1): + maximum_updraft_origin_level = 0 + kstabm[0, 0][plume] = k_end - 2 + ocean_fraction_local = ocean_fraction + cap_max = CAP_MAXS + error_code_2 = 0 + error_code_3 = 0 + cap_max_increment = CAP_MAX_INC + cloud_workfunction_0 = 0.0 + cloud_workfunction_1 = 0.0 + cloud_workfunction_2 = 0.0 + cloud_workfunction_3 = 0.0 + cloud_workfunction_0_pbl = 0.0 + cloud_workfunction_1_pbl = 0.0 + cloud_workfunction_1_fa = 0.0 + cin_1 = 0.0 + k_x_modified = 0.0 + epsilon_local = 0.0 + pbl_time_scale = 0.0 + t_wetbulb = 0.0 + vapor_wetbulb = 0.0 + cape_removal_time_scale = 0.0 + f_dicycle_modified = 0.0 + add_buoyancy = 0.0 + scale_dependence_factor_downdraft = 0.0 + + # ensure locals are zero, no lingering data from previous call + with computation(PARALLEL), interval(...): + geopotential_height_local = geopotential_height + geopotential_height_modified_local = geopotential_height + cloud_moist_static_energy_downdraft_forced = 0.0 + downdraft_saturation_vapor_forced = 0.0 + cloud_moist_static_energy_forced_transported = 0.0 + c1d = 0.0 + evaporation_below_cloud_base = 0.0 + + # ensure locals are zero, no lingering data from previous call + # reset all ensemble members + with computation(FORWARD), interval(0, 1): + member = 0 + while member < ENSEMBLE_MEMBERS: + mass_flux_ensemble[0, 0][member] = 0.0 + precipitation_ensemble[0, 0][member] = 0.0 + member = member + 1 + + # reset state fields to zero + with computation(FORWARD), interval(0, 1): + epsilon_forced[0, 0][plume] = 0.0 + precip[0, 0][plume] = 0.0 + scale_dependence_factor[0, 0][plume] = 0.0 + lightning_density = 0.0 + + +def compute_scale_dependence_factor( + plume: Int, + scale_dependence_factor: FloatFieldIJ_Plume, + seed_convection: FloatFieldIJ, + error_code: IntFieldIJ_Plume, + grid_length: FloatFieldIJ, +): + """Compute scale dependence factor for use later. + + Args: + plume (Int) + scale_dependence_factor (FloatFieldIJ_Plume) + seed_convection (FloatFieldIJ) + error_code (IntFieldIJ_Plume) + grid_length (FloatFieldIJ) + """ + from __externals__ import USE_SCALE_DEP + + # prepare mask to stop loop in next computation + with computation(FORWARD), interval(0, 1): + error_at_point = False + + with computation(FORWARD), interval(0, 1): + if USE_SCALE_DEP == 0: + scale_dependence_factor[0, 0][plume] = 1.0 + elif USE_SCALE_DEP == 1: + if plume == cumulus_parameterization_constants.SHALLOW: + scale_dependence_factor[0, 0][plume] = 1.0 + else: + if seed_convection < 0.0: + error_code[0, 0][plume] = 1 + error_at_point = True + if error_at_point == False: + scale_dependence_factor[0, 0][plume] = sigma(grid_length) + if seed_convection != 1.0: + scale_dependence_factor[0, 0][plume] = scale_dependence_factor[0, 0][plume] ** ( + seed_convection * max(1.0, scale_dependence_factor[0, 0][plume]) + ) + scale_dependence_factor[0, 0][plume] = max( + 0.1, min(scale_dependence_factor[0, 0][plume], 1.0) + ) + if scale_dependence_factor[0, 0][plume] <= 0.1: + error_code[0, 0][plume] = 1 + error_at_point = True + + +def get_random_number( + plume: Int, + random_number: FloatFieldIJ, +): + """Generate a random number for each column. + + THIS IS UNFINISHED, RIGHT NOW IT USES SERIALIZED FORTRAN DATA. + THIS WILL NOT WORK WITH A PROPER INTEGRATION, NEED A BETTER SOLUTION. + + Args: + plume (Int) + random_number (FloatFieldIJ) + """ + from __externals__ import USE_RANDOM_NUMBER + + with computation(FORWARD), interval(0, 1): + if plume == cumulus_parameterization_constants.DEEP and USE_RANDOM_NUMBER > 1.0e-6: + # need to figure out how to get system clock data + random_number = random_number # keep input data from fortran for now + else: + random_number = 0.0 + + +def initial_entrainment_detrainment( + plume: Int, + lateral_entrainment_rate: FloatField, + current_plume_rate: Float, + entrainment_rate: FloatField_Plume, + detrainment_function_updraft: FloatField, +): + """Get initial entrainment/detrainment estimates based on data from the overarching model. + + Args: + plume (Int) + lateral_entrainment_rate (FloatField) + current_plume_rate (Float) + entrainment_rate (FloatField_Plume) + detrainment_function_updraft (FloatField) + """ + with computation(PARALLEL), interval(0, -1): + entrainment_rate[0, 0, 0][plume] = lateral_entrainment_rate * current_plume_rate + detrainment_function_updraft = lateral_entrainment_rate * current_plume_rate + + +def epsilon_min_max( + ocean_fraction: FloatFieldIJ, + epsilon_min: FloatFieldIJ, + epsilon_max: FloatFieldIJ, + MINIMUM_EVAP_FRACTION_OCEAN: Float, + MAXIMUM_EVAP_FRACTION_OCEAN: Float, + MINIMUM_EVAP_FRACTION_LAND: Float, + MAXIMUM_EVAP_FRACTION_LAND: Float, +): + """Set min/max epsilon for the current plume. + + Args: + ocean_fraction (FloatFieldIJ) + epsilon_min (FloatFieldIJ) + epsilon_max (FloatFieldIJ) + MINIMUM_EVAP_FRACTION_OCEAN (Float) + MAXIMUM_EVAP_FRACTION_OCEAN (Float) + MINIMUM_EVAP_FRACTION_LAND (Float) + MAXIMUM_EVAP_FRACTION_LAND (Float) + """ + with computation(FORWARD), interval(0, 1): + if ocean_fraction > 0.99: # water + epsilon_min = MINIMUM_EVAP_FRACTION_OCEAN + epsilon_max = MAXIMUM_EVAP_FRACTION_OCEAN + else: # land + epsilon_min = MINIMUM_EVAP_FRACTION_LAND + epsilon_max = MAXIMUM_EVAP_FRACTION_LAND + + +def calculate_arbitrary_numerical_parameter( + arbitrary_numerical_parameter: FloatFieldIJ, +): + """Set a scaling factor, used primarially (but not exclusively) for mass flux calculations. + + Args: + arbitrary_numerical_parameter (FloatFieldIJ) + """ + with computation(FORWARD), interval(0, 1): + arbitrary_numerical_parameter = 0.1 + # approx xmb * timescale + # other options (from fortran): + # 0.1 * dtime*xmb_nm1(i) + # 100.*(p_cup(i,kbcon(i))-p_cup(i,kbcon(i)+1))/(g*dtime) + # 0.1*mbdt(i) + + +class Setup(NDSLRuntime): + """Setup the GF2020 cumulus parameterization core. Prefill locals with appropriate values based on + configuration, set plume dependent constants for the correct plume, and + + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cu_param_config = cumulus_parameterization_config + + # construct stencils and functions + self._set_plume_dependent_fields = stencil_factory.from_dims_halo( + func=set_plume_dependent_fields, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DT_MOIST": config.DT_MOIST}, + ) + + self._prefil_internal_fields = stencil_factory.from_dims_halo( + func=prefil_internal_fields, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "CAP_MAXS": cumulus_parameterization_config.CAP_MAXS, + "ENSEMBLE_MEMBERS": MAXENS1 * MAXENS2 * MAXENS3, + }, + ) + + self._compute_scale_dependence_factor = stencil_factory.from_dims_halo( + func=compute_scale_dependence_factor, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"USE_SCALE_DEP": cumulus_parameterization_config.USE_SCALE_DEP}, + ) + + self._get_random_number = stencil_factory.from_dims_halo( + func=get_random_number, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"USE_RANDOM_NUMBER": cumulus_parameterization_config.USE_RANDOM_NUMBER}, + ) + + self._initial_entrainment_detrainment = stencil_factory.from_dims_halo( + func=initial_entrainment_detrainment, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._epsilon_min_max = stencil_factory.from_dims_halo( + func=epsilon_min_max, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._calculate_arbitrary_numerical_parameter = stencil_factory.from_dims_halo( + func=calculate_arbitrary_numerical_parameter, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + error_code_2: Quantity, + error_code_3: Quantity, + maximum_updraft_origin_level: Quantity, + kstabm: Quantity, + t_excess: Quantity, + t_excess_local: Quantity, + vapor_excess: Quantity, + vapor_excess_local: Quantity, + ocean_fraction: Quantity, + ocean_fraction_local: Quantity, + t_old: Quantity, + t_new: Quantity, + t_new_pbl: Quantity, + vapor_old: Quantity, + vapor_forced: Quantity, + vapor_forced_pbl: Quantity, + downdraft_saturation_vapor_forced: Quantity, + grid_scale_forcing_t: Quantity, + grid_scale_forcing_vapor: Quantity, + subgrid_scale_forcing_t: Quantity, + subgrid_scale_forcing_vapor: Quantity, + dmoist_static_energydt: Quantity, + cloud_moist_static_energy_downdraft_forced: Quantity, + cloud_moist_static_energy_forced_transported: Quantity, + cap_max: Quantity, + cap_max_increment: Quantity, + geopotential_height: Quantity, + geopotential_height_local: Quantity, + geopotential_height_modified_local: Quantity, + cloud_workfunction_0: Quantity, + cloud_workfunction_1: Quantity, + cloud_workfunction_2: Quantity, + cloud_workfunction_3: Quantity, + cloud_workfunction_0_pbl: Quantity, + cloud_workfunction_1_pbl: Quantity, + cloud_workfunction_1_fa: Quantity, + cin_1: Quantity, + k_x_modified: Quantity, + epsilon_forced: Quantity, + epsilon_local: Quantity, + epsilon_min: Quantity, + epsilon_max: Quantity, + pbl_time_scale: Quantity, + t_wetbulb: Quantity, + vapor_wetbulb: Quantity, + cape_removal_time_scale: Quantity, + f_dicycle_modified: Quantity, + add_buoyancy: Quantity, + scale_dependence_factor: Quantity, + scale_dependence_factor_downdraft: Quantity, + c1d: Quantity, + evaporation_below_cloud_base: Quantity, + mass_flux_ensemble: Quantity, + precipitation_ensemble: Quantity, + precip: Quantity, + lightning_density: Quantity, + seed_convection: Quantity, + grid_length: Quantity, + random_number: Quantity, + lateral_entrainment_rate: Quantity, + entrainment_rate: Quantity, + detrainment_function_updraft: Quantity, + arbitrary_numerical_parameter: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + plume: str, + ): + plume_dependent_constants = set_constants(self.cu_param_config, plume_dependent_constants, plume) + + if plume_dependent_constants.ENABLE_PLUME == 1: + # compute/prefil the last few fields needed for the rest of the scheme + self._set_plume_dependent_fields( + t_excess=t_excess, + t_excess_local=t_excess_local, + vapor_excess=vapor_excess, + vapor_excess_local=vapor_excess_local, + ocean_fraction=ocean_fraction, + use_excess=plume_dependent_constants.USE_EXCESS, + t_old=t_old, + vapor_old=vapor_old, + grid_scale_forcing_t=grid_scale_forcing_t, + grid_scale_forcing_vapor=grid_scale_forcing_vapor, + subgrid_scale_forcing_t=subgrid_scale_forcing_t, + subgrid_scale_forcing_vapor=subgrid_scale_forcing_vapor, + t_new=t_new, + vapor_forced=vapor_forced, + t_new_pbl=t_new_pbl, + vapor_forced_pbl=vapor_forced_pbl, + dmoist_static_energydt=dmoist_static_energydt, + ) + + self._prefil_internal_fields( + plume=plume_dependent_constants.PLUME_INDEX, + maximum_updraft_origin_level=maximum_updraft_origin_level, + kstabm=kstabm, + ocean_fraction=ocean_fraction, + ocean_fraction_local=ocean_fraction_local, + cap_max=cap_max, + error_code_2=error_code_2, + error_code_3=error_code_3, + CAP_MAX_INC=plume_dependent_constants.CAP_MAX_INC, + cap_max_increment=cap_max_increment, + geopotential_height=geopotential_height, + geopotential_height_local=geopotential_height_local, + geopotential_height_modified_local=geopotential_height_modified_local, + cloud_workfunction_0=cloud_workfunction_0, + cloud_workfunction_1=cloud_workfunction_1, + cloud_workfunction_2=cloud_workfunction_2, + cloud_workfunction_3=cloud_workfunction_3, + cloud_workfunction_0_pbl=cloud_workfunction_0_pbl, + cloud_workfunction_1_pbl=cloud_workfunction_1_pbl, + cloud_workfunction_1_fa=cloud_workfunction_1_fa, + cin_1=cin_1, + k_x_modified=k_x_modified, + epsilon_forced=epsilon_forced, + epsilon_local=epsilon_local, + pbl_time_scale=pbl_time_scale, + t_wetbulb=t_wetbulb, + vapor_wetbulb=vapor_wetbulb, + cape_removal_time_scale=cape_removal_time_scale, + f_dicycle_modified=f_dicycle_modified, + add_buoyancy=add_buoyancy, + scale_dependence_factor=scale_dependence_factor, + scale_dependence_factor_downdraft=scale_dependence_factor_downdraft, + cloud_moist_static_energy_downdraft_forced=cloud_moist_static_energy_downdraft_forced, + downdraft_saturation_vapor_forced=downdraft_saturation_vapor_forced, + cloud_moist_static_energy_forced_transported=cloud_moist_static_energy_forced_transported, + c1d=c1d, + evaporation_below_cloud_base=evaporation_below_cloud_base, + mass_flux_ensemble=mass_flux_ensemble, + precipitation_ensemble=precipitation_ensemble, + precip=precip, + lightning_density=lightning_density, + ) + + # scale dependence factor (sig), version new + self._compute_scale_dependence_factor( + plume=plume_dependent_constants.PLUME_INDEX, + scale_dependence_factor=scale_dependence_factor, + seed_convection=seed_convection, + error_code=error_code, + grid_length=grid_length, + ) + + # create a real random number in the interval [-use_random_num, +use_random_num] + self._get_random_number( + plume=plume_dependent_constants.PLUME_INDEX, + random_number=random_number, + ) + + # define entrainment/detrainment profiles for updrafts + self._initial_entrainment_detrainment( + plume=plume_dependent_constants.PLUME_INDEX, + lateral_entrainment_rate=lateral_entrainment_rate, + current_plume_rate=plume_dependent_constants.ENTRAINMENT_RATE, + entrainment_rate=entrainment_rate, + detrainment_function_updraft=detrainment_function_updraft, + ) + + # max/min allowed value for epsilon (ratio downdraft base mass flux/updraft base mass flux + # note : to make the evaporation stronger => increase "epsilon_min" + self._epsilon_min_max( + ocean_fraction=ocean_fraction, + epsilon_min=epsilon_min, + epsilon_max=epsilon_max, + MINIMUM_EVAP_FRACTION_OCEAN=plume_dependent_constants.MINIMUM_EVAP_FRACTION_OCEAN, + MAXIMUM_EVAP_FRACTION_OCEAN=plume_dependent_constants.MAXIMUM_EVAP_FRACTION_OCEAN, + MINIMUM_EVAP_FRACTION_LAND=plume_dependent_constants.MINIMUM_EVAP_FRACTION_LAND, + MAXIMUM_EVAP_FRACTION_LAND=plume_dependent_constants.MAXIMUM_EVAP_FRACTION_LAND, + ) + + # calculate arbitrary numerical parameter + self._calculate_arbitrary_numerical_parameter( + arbitrary_numerical_parameter=arbitrary_numerical_parameter, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_functions.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_functions.py new file mode 100644 index 000000000..8fc2785e8 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_functions.py @@ -0,0 +1,227 @@ +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl.dsl.gt4py import function, log +from ndsl.stencils.column_operations import column_max +from pyMoist.shared_incloud_processes import ice_fraction + + +@function +def saturation_vapor_pressure(t): + """ + Compute saturation vapor pressure + + Args: + t (in): temperature in Kelvin + """ + # CAUTION: This function has not been verified!!! + t_celcius = t - 273.155 + if t_celcius < -20.0: + toot = 273.16 / t + toto = 1 / toot + eilog = ( + -9.09718 * (toot - 1) + - 3.56654 * (log(toot) / log(10.0)) + + 0.876793 * (1 - toto) + + (log(6.1071) / log(10.0)) + ) + saturation_vapor_pressure = 10**eilog + else: + tsot = 373.16 / t + ewlog = -7.90298 * (tsot - 1) + 5.02808 * (log(tsot) / log(10.0)) + ewlog2 = ewlog - 1.3816e-07 * (10 ** (11.344 * (1 - (1 / tsot))) - 1) + ewlog3 = ewlog2 + 0.0081328 * (10 ** (-3.49149 * (tsot - 1)) - 1) + ewlog4 = ewlog3 + (log(1013.246) / log(10.0)) + saturation_vapor_pressure = 10**ewlog4 + + return saturation_vapor_pressure + + +@function +def updraft_origin_conditions_perturbation(value, perturbation, start_index, end_index): + """ + Modify the value based by the maximum value of perturbation + + Args: + value + perturbation + start_index: starting index for max value search + end_index: ending index for max value search + """ + max_value, _ = column_max(perturbation, start_index, end_index) + return value + cumulus_parameterization_constants.CP * max_value + + +@function +def get_cloud_boundary_conditions( + field, + scalar_perturbation, + p, + updraft_origin_level, + ocean_fraction, + BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH, + k_end, + compute_perturbation, + perturbation_field, +): + """ + Get the conditions of a field at the origin level of a updraft by taking an average across layers. + + Dimensions of the average: + a) to pick the value at k22 level, instead of an average between + (k22-order_aver, ..., k22-1, k22) set order_aver=kts + b) to average between kts and k22 => set order_aver = k22 + + order_aver = 4: + BOUNDARY_CONDITION_METHOD == 0: average between updraft_origin_level, + updraft_origin_level - 1, updraft_origin_level - 2 + BOUNDARY_CONDITION_METHOD == 1: average between updraft_origin_level + 1, + updraft_origin_level, updraft_origin_level - 1 + + order_aver = 0: + average between updraft_origin_level, updraft_origin_level - 1, updraft_origin_level - 2 + + Args: + field: field from which value is extracted + scalar_perturbation: scalar perturbation to be applied to the value after extraction + p: pressure field + updraft_origin_level: origin level of updraft (LCL, LFC, surface, etc.) + ocean_fraction: fraction of surface in the grid cell which is ocean + BOUNDARY_CONDITION_METHOD: switch which controls function + AVERAGE_LAYER_DEPTH: average depth of a model layer + k_end: number of model layers + compute_perturbation: switch which controls if a full field perturbation is applied to the value + perturbation_field: perturbation field which may be applied to the value (must supply dummy field if + compute_perturbation is False) + """ + # internal constant + frac_ave_layer_ocean = 0.3 + + if BOUNDARY_CONDITION_METHOD == 0: + order_aver = 3 + local_order_aver = min(updraft_origin_level + 1, order_aver) + + source_parcel_value = 0.0 + count = 0 + while count < local_order_aver: + source_parcel_value = source_parcel_value + field.at(K=updraft_origin_level - count) + count += 1 + + source_parcel_value = source_parcel_value / local_order_aver + + elif BOUNDARY_CONDITION_METHOD == 1: + effec_frac = (1.0 - ocean_fraction) + ocean_fraction * frac_ave_layer_ocean + average_layer = AVERAGE_LAYER_DEPTH * effec_frac + + start_index = 0 + level = 0 + p_at_origin = p.at(K=updraft_origin_level) + 0.5 * average_layer + while level <= k_end - 1: + new = abs(p.at(K=level) - p_at_origin) + old = abs(p.at(K=start_index) - p_at_origin) + if new < old: + start_index = level + level += 1 + + end_index = 0 + level = 0 + p_at_origin = p.at(K=updraft_origin_level) - 0.5 * average_layer + while level <= k_end - 1: + new = abs(p.at(K=level) - p_at_origin) + old = abs(p.at(K=end_index) - p_at_origin) + if new < old: + end_index = level + level += 1 + + if start_index >= end_index: + source_parcel_value = field.at(K=updraft_origin_level) + dp_layer = 1.0 + level_c = start_index + + else: + dp_layer = 1.0e-06 + source_parcel_value = 0.0 + level_c = 0 + stop_computation = False + level = start_index + while level <= k_end - 1 and stop_computation == False: + dp = -(p.at(K=level + 1) - p.at(K=level)) + if dp_layer + dp <= average_layer: + dp_layer = dp_layer + dp + source_parcel_value = source_parcel_value + field.at(K=level) * dp + level += 1 + else: + dp = average_layer - dp_layer + dp_layer = dp_layer + dp + source_parcel_value = source_parcel_value + field.at(K=level) * dp + stop_computation = True + level += 1 + + source_parcel_value = source_parcel_value / dp_layer + level_c = max(start_index, level) + + # this perturbation is only used for moist static energy + if compute_perturbation == True: + source_parcel_value = updraft_origin_conditions_perturbation( + value=source_parcel_value, + perturbation=perturbation_field, + start_index=start_index, + end_index=level_c, + ) + + return source_parcel_value + scalar_perturbation + + +@function +def compute_dewpoint(p, vapor): + rr = vapor + 1e-8 + es = p * rr / (0.622 + rr) + esln = log(es) + return (35.86 * esln - 4947.2325) / (esln - 23.6837) + + +@function +def liquid_fraction( + t, + convection_fraction, + surface_type, + FRAC_MODIS, +): + + if FRAC_MODIS == 1: + liquid_fraction = 1.0 - ice_fraction(t, convection_fraction, surface_type) + else: + liquid_fraction = min( + 1.0, + ( + max(0.0, (t - cumulus_parameterization_constants.T_ICE)) + / (cumulus_parameterization_constants.T_0 - cumulus_parameterization_constants.T_ICE) + ) + ** 2, + ) + + return liquid_fraction + + +# @function +# def get_delmix( +# ocean_fraction, +# sub_cloud_level, +# p_forced, +# in_var, +# out_var, +# ): +# internal_var = out_var.at(K=0) +# x1 = 0.0 +# x2 = 0.0 +# level = 0 +# while level <= sub_cloud_level: +# dp = p_forced[0,0,1] - p_forced +# x1 = x1 + dp*in_var +# x2 = x2 + dp +# level += 1 + +# delta = abs(internal_var-x1/(x2 + 1.e-12)) + +# level = 0 +# while level <= sub_cloud_level: +# out_var = diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_stencils.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_stencils.py new file mode 100644 index 000000000..a7e7b9c3d --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/shared_stencils.py @@ -0,0 +1,277 @@ +from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, K, computation, interval, sqrt + +import pyMoist.constants as constants +from ndsl.dsl.typing import BoolFieldIJ, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume + + +def generic_find_level( + array: FloatField, + start_index: IntFieldIJ_Plume, + end_index: IntFieldIJ, + out_index: IntFieldIJ_Plume, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """Find the first level (searching from bottom to top of the atmosphere) where a field exceeds + a threshold. The potential search area is restricted by the start_index and end_index arguments. + + Args: + array (FloatField): field to be searched + start_index (IntFieldIJ_Plume): starting index for the search + end_index (IntFieldIJ): ending index for the search + out_index (IntFieldIJ_Plume): level of interest identified by the stencil + error_code (IntFieldIJ_Plume) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + out_index[0, 0][plume] = start_index[0, 0][plume] + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + x = array.at(K=start_index[0, 0][plume]) + stop_index = max(start_index[0, 0][plume] + 1, end_index) + + level = start_index[0, 0][plume] + 1 + while level >= start_index[0, 0][plume] + 1 and level <= stop_index: + if array.at(K=level) < x: + x = array.at(K=level) + out_index[0, 0][plume] = level + level += 1 + + +def updraft_vertical_velocity( + vertical_velocity_3d: FloatField, + vertical_velocity_2d: FloatFieldIJ, + convective_scale_velocity: FloatFieldIJ, + entrainment_rate: FloatField_Plume, + detrainment_function_updraft: FloatField, + geopotential_height_cloud_levels_forced: FloatField, + t_cloud_levels_forced: FloatField, + updraft_column_temperature_forced: FloatField, + cloud_total_water_after_entrainment_forced: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + vapor_forced: FloatField, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """Compute the vertical velocity within the updraft. Return this as a 3D field + and accumulate (considering layer thickness) into a 2D field. + + Args: + vertical_velocity_3d (FloatField) + vertical_velocity_2d (FloatFieldIJ) + convective_scale_velocity (FloatFieldIJ) + entrainment_rate (FloatField_Plume) + detrainment_function_updraft (FloatField) + geopotential_height_cloud_levels_forced (FloatField) + t_cloud_levels_forced (FloatField) + updraft_column_temperature_forced (FloatField) + cloud_total_water_after_entrainment_forced (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + vapor_forced (FloatField) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + error_code (IntFieldIJ_Plume) + plume (Int) + """ + from __externals__ import ZERO_DIFF, k_end + + # internal constants + with computation(FORWARD), interval(0, 1): + ctea: FloatFieldIJ = 1.0 / 3.0 + cteb: FloatFieldIJ = 2.0 + visc: FloatFieldIJ = 2000.0 + eps: FloatFieldIJ = 0.622 + f: FloatFieldIJ = 2.0 + C_d: FloatFieldIJ = 0.506 + gam: FloatFieldIJ = 0.5 + beta: FloatFieldIJ = 1.875 + smooth: BoolFieldIJ = True + n_smooth: IntFieldIJ = 1 + + if ZERO_DIFF == 1: + ftun1: FloatFieldIJ = 1.0 + ftun2: FloatFieldIJ = 0.5 + else: + ftun1: FloatFieldIJ = 0.25 + ftun2: FloatFieldIJ = 1.0 + + # initialize arrays to zero + with computation(PARALLEL), interval(...): + vertical_velocity_3d = 0.0 + + with computation(FORWARD), interval(0, 1): + vertical_velocity_2d = 0.0 + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K <= updraft_lfc_level[0, 0][plume]: + vertical_velocity_3d = max(1.0, convective_scale_velocity**2) + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if K >= updraft_lfc_level[0, 0][plume] and K <= cloud_top_level[0, 0][plume]: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) + + t_ve = 0.5 * ( + t_cloud_levels_forced * (1.0 + (vapor_forced / eps) / (1.0 + vapor_forced)) + + t_cloud_levels_forced[0, 0, 1] + * (1.0 + (vapor_forced[0, 0, 1] / eps) / (1.0 + vapor_forced[0, 0, 1])) + ) + + t_v = 0.5 * ( + updraft_column_temperature_forced + * ( + 1.0 + + (cloud_total_water_after_entrainment_forced / eps) + / (1.0 + cloud_total_water_after_entrainment_forced) + ) + + updraft_column_temperature_forced[0, 0, 1] + * ( + 1.0 + + (cloud_total_water_after_entrainment_forced[0, 0, 1] / eps) + / (1.0 + cloud_total_water_after_entrainment_forced[0, 0, 1]) + ) + ) + + bu = constants.MAPL_GRAV * ( + (t_v - t_ve) / t_ve + - ftun2 + * 0.50 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + ) + + dw1 = 2.0 / (f * (1.0 + gam)) * bu * dz + if ZERO_DIFF == 1: + kx = max(entrainment_rate[0, 0, 0][plume], detrainment_function_updraft) * dz + else: + kx = ( + (1.0 + beta * C_d) + * max(entrainment_rate[0, 0, 0][plume], detrainment_function_updraft) + * dz + * ftun1 + ) + + dw2 = (vertical_velocity_3d) - 2.0 * kx * (vertical_velocity_3d) + + vertical_velocity_3d[0, 0, 1] = (dw1 + dw2) / (1.0 + kx) + + if vertical_velocity_3d[0, 0, 1] < 0.0: + vertical_velocity_3d[0, 0, 1] = 0.5 * vertical_velocity_3d + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if smooth == True: # noqa + if ZERO_DIFF == 1: + if K <= cloud_top_level[0, 0][plume] - 2: + nvs: IntFieldIJ = 0 + vs: FloatFieldIJ = 0.0 + level_inner_loop = max(K - n_smooth, 0) + while level_inner_loop <= min(K + n_smooth, k_end - 1): + nvs = nvs + 1 + vs = vs + vertical_velocity_3d.at(K=level_inner_loop) + level_inner_loop += 1 + vertical_velocity_3d = vs / (1.0e-16 + nvs) + else: + if K <= cloud_top_level[0, 0][plume] + 1: + vs: FloatFieldIJ = 0.0 + dz1m: FloatFieldIJ = 0.0 + level_inner_loop = max(K - n_smooth, 0) + while level_inner_loop <= min(K + n_smooth, k_end - 1): + dz = geopotential_height_cloud_levels_forced.at( + K=level_inner_loop + 1 + ) - geopotential_height_cloud_levels_forced.at(K=level_inner_loop) + vs = vs + dz * vertical_velocity_3d.at(K=level_inner_loop) + dz1m = dz1m + dz + level_inner_loop += 1 + vertical_velocity_3d = vs / (1.0e-16 + dz1m) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + # convert to vertical velocity + vertical_velocity_3d = sqrt(max(0.1, vertical_velocity_3d)) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + max_val, _ = column_max(vertical_velocity_3d, 0, k_end) + if max_val < 1.0: + error_code[0, 0][plume] = 54 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 or error_code[0, 0][plume] == 54: + # sanity check + if vertical_velocity_3d < 1.0: + vertical_velocity_3d = 1.0 + if vertical_velocity_3d > 20.0: + vertical_velocity_3d = 20.0 + + if K >= cloud_top_level[0, 0][plume] + 1 and ZERO_DIFF == 0: + vertical_velocity_3d = 0.1 + + # get the column average vertical velocity + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 or error_code[0, 0][plume] == 54: + if K >= updraft_lfc_level[0, 0][plume] and K <= cloud_top_level[0, 0][plume]: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] - geopotential_height_cloud_levels_forced + ) + vertical_velocity_2d = vertical_velocity_2d + vertical_velocity_3d * dz + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 or error_code[0, 0][plume] == 54: + vertical_velocity_2d = vertical_velocity_2d / ( + geopotential_height_cloud_levels_forced.at(K=cloud_top_level[0, 0][plume] + 1) + - geopotential_height_cloud_levels_forced.at(K=updraft_lfc_level[0, 0][plume]) + + 1.0e-16 + ) + vertical_velocity_2d = max(1.0, vertical_velocity_2d) + + +def tridiag( + m: IntFieldIJ_Plume, + a: FloatField, + b: FloatField, + c: FloatField, + f: FloatField, + error_code: IntFieldIJ_Plume, + plume: Int, +): + """This routine solves the problem: aa*f(k-1,t+1) + bb*f(k,t+1) + cc*f(k+1,t+1) = dd + an updated "f" at time t+1 is the output + + Args: + m (IntFieldIJ_Plume) + a (FloatField) + b (FloatField) + c (FloatField) + f (FloatField) + error_code (IntFieldIJ_Plume) + plume (Int) + """ + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and K == m[0, 0][plume]: + c = 0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + q = -c / b + f = f / b + + with computation(FORWARD), interval(1, None): + if error_code[0, 0][plume] == 0 and K <= m[0, 0][plume]: + p = 1.0 / (b + a * q[0, 0, -1]) + q = -c * p + f = p * (f - a * f[0, 0, -1]) + + with computation(BACKWARD), interval(...): + if error_code[0, 0][plume] == 0 and K < m[0, 0][plume]: + f = f + q * f[0, 0, 1] diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/smoothing.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/smoothing.py new file mode 100644 index 000000000..4ea1efa44 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/smoothing.py @@ -0,0 +1,102 @@ +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Int +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import FloatField_Plume, IntFieldIJ_Plume + + +def smooth_tendencies( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + del_moist_static_energy_cloud_ensemble: FloatField, + del_vapor_cloud_ensemble: FloatField, + del_cloud_liquid_cloud_ensemble: FloatField, + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + plume: Int, +): + """Smooth computed tendencies before output and model state update. The number of levels considered + by the smoother (in both directions) is controlled by USE_SMOOTH_TENDENCIES. Smoothing is disabled + if this value is less than one. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + del_moist_static_energy_cloud_ensemble (FloatField) + del_vapor_cloud_ensemble (FloatField) + del_cloud_liquid_cloud_ensemble (FloatField) + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + plume (Int) + """ + from __externals__ import USE_SMOOTH_TENDENCIES + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0 and USE_SMOOTH_TENDENCIES >= 1: + # initialize 2d temporaries + internal_tendency_1_2d: FloatFieldIJ = 0.0 + internal_tendency_2_2d: FloatFieldIJ = 0.0 + internal_tendency_3_2d: FloatFieldIJ = 0.0 + internal_tendency_4_2d: FloatFieldIJ = 0.0 + internal_tendency_5_2d: FloatFieldIJ = 0.0 + rcount: FloatFieldIJ = 0.0 + dp: FloatFieldIJ = 0.0 + + with computation(FORWARD), interval(...): + # NOTE this entire block can be (and should be, for the sake of performance) rewritten with + # dynamic intervals, but this cannot be done until the feature is more stable. will revisit later + if error_code[0, 0][plume] == 0 and USE_SMOOTH_TENDENCIES >= 1 and K <= cloud_top_level[0, 0][plume]: + rcount = 1.0e-8 + internal_tendency_1_2d = 0 + internal_tendency_2_2d = 0 + internal_tendency_3_2d = 0 + internal_tendency_4_2d = 0 + internal_tendency_5_2d = 0 + inner_loop = max(0, K - USE_SMOOTH_TENDENCIES) + while inner_loop <= min(cloud_top_level[0, 0][plume], K + USE_SMOOTH_TENDENCIES): + dp = p_cloud_levels_forced.at(K=inner_loop, ddim=[plume]) - p_cloud_levels_forced.at( + K=inner_loop + 1, ddim=[plume] + ) + rcount = rcount + dp + internal_tendency_1_2d = ( + internal_tendency_1_2d + dp * del_moist_static_energy_cloud_ensemble.at(K=inner_loop) + ) + internal_tendency_2_2d = internal_tendency_2_2d + dp * del_vapor_cloud_ensemble.at( + K=inner_loop + ) + internal_tendency_3_2d = internal_tendency_3_2d + dp * del_cloud_liquid_cloud_ensemble.at( + K=inner_loop + ) + internal_tendency_4_2d = internal_tendency_4_2d + dp * del_u_cloud_ensemble.at(K=inner_loop) + internal_tendency_5_2d = internal_tendency_5_2d + dp * del_v_cloud_ensemble.at(K=inner_loop) + inner_loop += 1 + + internal_tendency_1_3d = internal_tendency_1_2d / rcount + internal_tendency_2_3d = internal_tendency_2_2d / rcount + internal_tendency_3_3d = internal_tendency_3_2d / rcount + internal_tendency_4_3d = internal_tendency_4_2d / rcount + internal_tendency_5_3d = internal_tendency_5_2d / rcount + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and USE_SMOOTH_TENDENCIES >= 1 and K <= cloud_top_level[0, 0][plume]: + del_moist_static_energy_cloud_ensemble = internal_tendency_1_3d + del_vapor_cloud_ensemble = internal_tendency_2_3d + del_cloud_liquid_cloud_ensemble = internal_tendency_3_3d + del_u_cloud_ensemble = internal_tendency_4_3d + del_v_cloud_ensemble = internal_tendency_5_3d + + +class MakeSmoother: + def __init__(self): + pass + + def __call__(self, *args, **kwds): + pass + + +class ApplySmoother: + def __init__(self): + pass + + def __call__(self, *args, **kwds): + pass diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/sounding.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/sounding.py new file mode 100644 index 000000000..85e583ffb --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/sounding.py @@ -0,0 +1,29 @@ +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig + + +class Sounding: + def __init__(self, cumulus_parameterization_config: GF2020CumulusParameterizationConfig): + if cumulus_parameterization_config.OUTPUT_SOUNDING == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->Sounding: Output Sounding capabilities have not" + "been implemented. You should have been caught before getting here by the config checker." + "Beware, something likely failing in the config checker as well - you may be unknowingly" + "calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass + + +class GATESounding: + def __init__(self, cumulus_parameterization_config: GF2020CumulusParameterizationConfig): + if cumulus_parameterization_config.WRTGRADS == 1: + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->GATESounding: GATE Sounding capabilities have not" + "been implemented. You should have been caught before getting here by the config checker." + "Beware, something likely failing in the config checker as well - you may be unknowingly" + "calling other untested/unimplemented sections." + ) + + def __call__(self, *args, **kwds): + pass diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/state.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/state.py new file mode 100644 index 000000000..0e37dda96 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/state.py @@ -0,0 +1,894 @@ +import dataclasses + +from ndsl import Quantity, State +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import Float, Int + + +@dataclasses.dataclass +class GF2020CumulusParameterizationState(State): + @dataclasses.dataclass + class Input: + t_excess: Quantity = dataclasses.field( + metadata={ + "name": "t_excess", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_excess: Quantity = dataclasses.field( + metadata={ + "name": "vapor_excess", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + grid_scale_forcing_t: Quantity = dataclasses.field( + metadata={ + "name": "grid_scale_forcing_t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + grid_scale_forcing_vapor: Quantity = dataclasses.field( + metadata={ + "name": "grid_scale_forcing_vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + subgrid_scale_forcing_t: Quantity = dataclasses.field( + metadata={ + "name": "subgrid_scale_forcing_t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + subgrid_scale_forcing_vapor: Quantity = dataclasses.field( + metadata={ + "name": "subgrid_scale_forcing_vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + seed_convection: Quantity = dataclasses.field( + metadata={ + "name": "seed_convection", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + saturation_water_vapor: Quantity = dataclasses.field( + metadata={ + "name": "saturation_water_vapor", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ocean_fraction: Quantity = dataclasses.field( + metadata={ + "name": "ocean_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convection_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + surface_type: Quantity = dataclasses.field( + metadata={ + "name": "surface_type", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + last_error_code: Quantity = dataclasses.field( + metadata={ + "name": "last_error_code", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + + @dataclasses.dataclass + class Output: + dtdt: Quantity = dataclasses.field( + metadata={ + "name": "t", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvapordt: Quantity = dataclasses.field( + metadata={ + "name": "dvapordt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dcloudicedt: Quantity = dataclasses.field( + metadata={ + "name": "dcloudicedt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dudt: Quantity = dataclasses.field( + metadata={ + "name": "dudt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvdt: Quantity = dataclasses.field( + metadata={ + "name": "dvdt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dnliquiddt: Quantity = dataclasses.field( + metadata={ + "name": "dnliquiddt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dnicedt: Quantity = dataclasses.field( + metadata={ + "name": "dnicedt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dbuoyancydt: Quantity = dataclasses.field( + metadata={ + "name": "dbuoyancydt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvectiveicedt: Quantity = dataclasses.field( + metadata={ + "name": "dconvectiveicedt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlargescaleicedt: Quantity = dataclasses.field( + metadata={ + "name": "dlargescaleicedt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvectiveliquiddt: Quantity = dataclasses.field( + metadata={ + "name": "dconvectiveliquiddt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlargescaleliquiddt: Quantity = dataclasses.field( + metadata={ + "name": "dlargescaleliquiddt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvectivecloudfractiondt: Quantity = dataclasses.field( + metadata={ + "name": "dconvectivecloudfractiondt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlargescalecloudfractiondt: Quantity = dataclasses.field( + metadata={ + "name": "dlargescalecloudfractiondt", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + error_code: Quantity = dataclasses.field( + metadata={ + "name": "error_code", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + downdraft_origin_level: Quantity = dataclasses.field( + metadata={ + "name": "downdraft_origin_level", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + lcl_level: Quantity = dataclasses.field( + metadata={ + "name": "lcl_level", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + updraft_origin_level: Quantity = dataclasses.field( + metadata={ + "name": "updraft_origin_level", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + updraft_lfc_level: Quantity = dataclasses.field( + metadata={ + "name": "updraft_lfc_level", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + cloud_top_level: Quantity = dataclasses.field( + metadata={ + "name": "cloud_top_level", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + kstabi: Quantity = dataclasses.field( + metadata={ + "name": "kstabi", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + kstabm: Quantity = dataclasses.field( + metadata={ + "name": "kstabm", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + precip: Quantity = dataclasses.field( + metadata={ + "name": "precip", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_base_mass_flux_modified: Quantity = dataclasses.field( + metadata={ + "name": "cloud_base_mass_flux_modified", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + epsilon_forced: Quantity = dataclasses.field( + metadata={ + "name": "epsilon_forced", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_normalized_integrated_condensate_forced: Quantity = dataclasses.field( + metadata={ + "name": "total_normalized_integrated_condensate_forced", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + scale_dependence_factor: Quantity = dataclasses.field( + metadata={ + "name": "scale_dependence_factor", + "dims": [X_DIM, Y_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_cloud_levels_forced: Quantity = dataclasses.field( + metadata={ + "name": "p_cloud_levels_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + entrainment_rate: Quantity = dataclasses.field( + metadata={ + "name": "entrainment_rate", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_updraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_updraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_entrainment_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "mass_entrainment_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_updraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_updraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_detrainment_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "mass_detrainment_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_updraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + normalized_massflux_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "normalized_massflux_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + condensate_to_fall_forced: Quantity = dataclasses.field( + metadata={ + "name": "condensate_to_fall_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporate_in_downdraft_forced: Quantity = dataclasses.field( + metadata={ + "name": "evaporate_in_downdraft_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_liquid_after_rain_forced: Quantity = dataclasses.field( + metadata={ + "name": "cloud_liquid_after_rain_forced", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_updraft: Quantity = dataclasses.field( + metadata={ + "name": "t_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convective_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_0: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_0", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_1", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_2: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_2", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_3: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_3", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_pbl: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_1_pbl", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_cin: Quantity = dataclasses.field( + metadata={ + "name": "cloud_workfunction_1_cin", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cape_removal_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "cape_removal_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "pbl_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lightning_density: Quantity = dataclasses.field( + metadata={ + "name": "lightning_density", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation_sublimation_tendency: Quantity = dataclasses.field( + metadata={ + "name": "evaporation_sublimation_tendency", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_precip_flux: Quantity = dataclasses.field( + metadata={ + "name": "convective_precip_flux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_perturbation: Quantity = dataclasses.field( + metadata={ + "name": "t_perturbation", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + + @dataclasses.dataclass + class InputOutput: + grid_length: Quantity = dataclasses.field( + metadata={ + "name": "grid_length", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_level: Quantity = dataclasses.field( + metadata={ + "name": "pbl_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + ccn: Quantity = dataclasses.field( + metadata={ + "name": "ccn", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + air_density: Quantity = dataclasses.field( + metadata={ + "name": "air_density", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + omega: Quantity = dataclasses.field( + metadata={ + "name": "omega", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + topography_height_no_negative: Quantity = dataclasses.field( + metadata={ + "name": "topography_height_no_negative", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sensible_heat_flux: Quantity = dataclasses.field( + metadata={ + "name": "sensible_heat_flux", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + latent_heat_flux: Quantity = dataclasses.field( + metadata={ + "name": "latent_heat_flux", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + longitude_degrees: Quantity = dataclasses.field( + metadata={ + "name": "longitude_degrees", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + latitude_degrees: Quantity = dataclasses.field( + metadata={ + "name": "latitude_degrees", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_old: Quantity = dataclasses.field( + metadata={ + "name": "t_old", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_old: Quantity = dataclasses.field( + metadata={ + "name": "vapor_old", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_modified_by_advection: Quantity = dataclasses.field( + metadata={ + "name": "t_modified_by_advection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_modified_by_advection: Quantity = dataclasses.field( + metadata={ + "name": "vapor_modified_by_advection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_forced: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_forced: Quantity = dataclasses.field( + metadata={ + "name": "p_forced", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_surface: Quantity = dataclasses.field( + metadata={ + "name": "p_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_surface: Quantity = dataclasses.field( + metadata={ + "name": "t_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u: Quantity = dataclasses.field( + metadata={ + "name": "u", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v: Quantity = dataclasses.field( + metadata={ + "name": "v", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + w: Quantity = dataclasses.field( + metadata={ + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass: Quantity = dataclasses.field( + metadata={ + "name": "mass", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_scale_velocity: Quantity = dataclasses.field( + metadata={ + "name": "convective_scale_velocity", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + buoyancy_excess: Quantity = dataclasses.field( + metadata={ + "name": "buoyancy_excess", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_ice: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_ice: Quantity = dataclasses.field( + metadata={ + "name": "convective_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_liquid: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_liquid: Quantity = dataclasses.field( + metadata={ + "name": "convective_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convective_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + chemistry_tracers_output: Quantity = dataclasses.field( + metadata={ + "name": "chemistry_tracers_output", + "dims": [X_DIM, Y_DIM, Z_DIM, "plumes", "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + + input: Input + output: Output + input_output: InputOutput diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/triggers.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/triggers.py new file mode 100644 index 000000000..fc70e1340 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/triggers.py @@ -0,0 +1,52 @@ +from ndsl.dsl.gt4py import FORWARD, computation, interval +from ndsl.dsl.typing import FloatFieldIJ, Int +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import IntFieldIJ_Plume +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) + + +def convection_trigger( + error_code: IntFieldIJ_Plume, + convective_scale_velosity: FloatFieldIJ, + cin_0: FloatFieldIJ, + plume: Int, +): + """Check if sufficient velocity exists to overcome present convective inhibition. + + Args: + error_code (IntFieldIJ_Plume) + convective_scale_velosity (FloatFieldIJ) + cin_0 (FloatFieldIJ) + plume (Int) + """ + from __externals__ import DICYCLE + + with computation(FORWARD), interval(...): + if DICYCLE > 1: + if error_code[0, 0][plume] == 0: + # think about including the grid scale vertical velocity at KE calculation + if cin_0 + 0.5 * convective_scale_velosity**2 < 0.0: + error_code[0, 0][plume] = 19 + + +class XieTriggerFunction: + def __init__( + self, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + def __call__(self, plume_dependent_constants: GF2020PlumeDependentConstants): + + if self.config.ADV_TRIGGER == 3 and (plume_dependent_constants.PLUME_INDEX in (1, 2)): + raise NotImplementedError( + "[NDSL] GF2020-->CumulusParameterization-->XieTriggerFunction this code" + "has not been impemented. You should have been caught before getting here by the config" + "checker. Beware, something likely failing in the config checker as well - you may be" + "unknowingly calling other untested/unimplemented sections." + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/updraft.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/updraft.py new file mode 100644 index 000000000..8ad53e6b2 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/updraft.py @@ -0,0 +1,1227 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, Field, GlobalTable, K, computation, interval +from ndsl.dsl.typing import BoolFieldIJ, Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max_ddim, column_min +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Ensemble, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_functions import get_cloud_boundary_conditions + + +# initialize constants and field type for UpdraftMassFlux stencil +_X_ALPHA = [ + 3.699999, + 3.699999, + 3.699999, + 3.699999, + 3.024999, + 2.559999, + 2.249999, + 2.028571, + 1.862500, + 1.733333, + 1.630000, + 1.545454, + 1.475000, + 1.415385, + 1.364286, + 1.320000, + 1.281250, + 1.247059, + 1.216667, + 1.189474, + 1.165000, + 1.142857, + 1.122727, + 1.104348, + 1.087500, + 1.075000, + 1.075000, + 1.075000, + 1.075000, + 1.075000, +] + +_G_ALPHA = [ + 4.1706450, + 4.1706450, + 4.1706450, + 4.1706450, + 2.0469250, + 1.3878370, + 1.1330030, + 1.012418, + 0.9494680, + 0.9153771, + 0.8972442, + 0.8885444, + 0.8856795, + 0.8865333, + 0.8897996, + 0.8946404, + 0.9005030, + 0.9070138, + 0.9139161, + 0.9210315, + 0.9282347, + 0.9354376, + 0.9425780, + 0.9496124, + 0.9565111, + 0.9619183, + 0.9619183, + 0.9619183, + 0.9619183, + 0.9619183, +] +_CONSTANTS_TABLES_TYPE: Field = GlobalTable[(Float, len(_X_ALPHA))] + + +def updraft_mass_flux( + error_code: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + pbl_level: IntFieldIJ, + updraft_lfc_level: IntFieldIJ_Plume, + lcl_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + p_surface: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + normalized_massflux_updraft: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_updraft_modified: FloatField, + random_number: FloatFieldIJ, + UPDRAFT_MAX_HEIGHT_LAND: Float, + UPDRAFT_MAX_HEIGHT_OCEAN: Float, + plume: Int, + X_ALPHA: _CONSTANTS_TABLES_TYPE, + G_ALPHA: _CONSTANTS_TABLES_TYPE, +): + """Handle mass fluxes in the updraft. This code has a number of potential paths depending on configuration + settings - some of these paths are not yet implemented. + + This code may execute for all plumes - shallow plume execution is currently not implemented. + + Args: + error_code (IntFieldIJ_Plume): _description_ + updraft_origin_level (IntFieldIJ_Plume): _description_ + cloud_top_level (IntFieldIJ_Plume): _description_ + pbl_level (IntFieldIJ): _description_ + updraft_lfc_level (IntFieldIJ_Plume): _description_ + lcl_level (IntFieldIJ_Plume): _description_ + p_cloud_levels_forced (FloatField_Plume): _description_ + p_surface (FloatFieldIJ): _description_ + ocean_fraction (FloatFieldIJ): _description_ + normalized_massflux_updraft (FloatField): _description_ + normalized_massflux_updraft_forced (FloatField_Plume): _description_ + normalized_massflux_updraft_modified (FloatField): _description_ + random_number (FloatFieldIJ): _description_ + UPDRAFT_MAX_HEIGHT_LAND (Float): _description_ + UPDRAFT_MAX_HEIGHT_OCEAN (Float): _description_ + plume (Int): _description_ + X_ALPHA (_CONSTANTS_TABLES_TYPE): _description_ + G_ALPHA (_CONSTANTS_TABLES_TYPE): _description_ + """ + from __externals__ import ZERO_DIFF, k_end + + with computation(FORWARD), interval(0, 1): + # initialize constants + PX: FloatFieldIJ = 45 / 120 # px sets the pressure level of max zu + BETA_DEEP: FloatFieldIJ = 1.25 + G_BETA_DEEP: FloatFieldIJ = 0.8974707 + execution_choice: IntFieldIJ = -999 + + # set up masks + stop_computation: BoolFieldIJ = False + stop_do_loop: BoolFieldIJ = False + + if ZERO_DIFF == 1: + if plume == cumulus_parameterization_constants.DEEP and ocean_fraction > 0.90: + # deep plume over ocean + execution_choice = 11 + if plume == cumulus_parameterization_constants.DEEP and ocean_fraction <= 0.90: + # deep plume over land + execution_choice = 12 + if plume == cumulus_parameterization_constants.MID: + # mid plume + execution_choice = 5 + else: + if plume == cumulus_parameterization_constants.DEEP: + # deep plume + execution_choice = 20 + if plume == cumulus_parameterization_constants.MID: + # mid plume + execution_choice = 20 + + with computation(PARALLEL), interval(...): + # fill momentum with 0 + normalized_massflux_updraft_forced[0, 0, 0][plume] = 0.0 + normalized_massflux_updraft_h = 0.0 + normalized_massflux_updraft_l = 0.0 + + # EXECUTION CHOICE 20 + with computation(FORWARD), interval(0, 1): + # land/ocean + if error_code[0, 0][plume] == 0: + if execution_choice == 20: + + height_updraft: FloatFieldIJ = ( + 1.0 - ocean_fraction + ) * UPDRAFT_MAX_HEIGHT_LAND + ocean_fraction * UPDRAFT_MAX_HEIGHT_OCEAN + # add a randomic perturbation + height_updraft = height_updraft + random_number + + # height_updraft parameter goes from 0 to 1 = rainfall decreases with height_updraft + p_max_normalized_massflux_updraft: FloatFieldIJ = (p_surface - 100.0) * ( + 1.0 - 0.5 * height_updraft + ) + 0.6 * p_cloud_levels_forced.at( + K=cloud_top_level[0, 0][plume], ddim=[plume] + ) * 0.5 * height_updraft + + # beta parameter: must be larger than 1 + # higher makes the profile sharper around the maximum zu + beta: FloatFieldIJ = max(1.1, 2.1 - 0.5 * height_updraft) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + if execution_choice == 20: + p_internal = abs(p_cloud_levels_forced[0, 0, 0][plume] - p_max_normalized_massflux_updraft) + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if execution_choice == 20: + _, _min_loc = column_min(p_internal, 0, cloud_top_level[0, 0][plume]) + updraft_origin_level_adj: IntFieldIJ = _min_loc + 1 + updraft_origin_level_adj = max( + updraft_origin_level[0, 0][plume] + 1, updraft_origin_level_adj + ) + updraft_origin_level_adj = min(updraft_origin_level_adj, cloud_top_level[0, 0][plume] + 1) + + # this alpha constrains the location of the maximun normalized_massflux_updraft_forced + # to be at "updraft_origin_level_adj" vertical level + alpha: FloatFieldIJ = 1.0 + ( + (beta - 1.0) + * ((updraft_origin_level_adj / (cloud_top_level[0, 0][plume] + 2))) + / (1.0 - ((updraft_origin_level_adj) / (cloud_top_level[0, 0][plume] + 2))) + ) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0: + if ( + execution_choice == 20 + and K >= lcl_level[0, 0][plume] - 1 + and K <= min(k_end, cloud_top_level[0, 0][plume]) + ): + ratio = float(K + 1) / (cloud_top_level[0, 0][plume] + 2) + normalized_massflux_updraft_forced[0, 0, 0][plume] = (ratio ** (alpha - 1.0)) * ( + (1.0 - ratio) ** (beta - 1.0) + ) + + with computation(BACKWARD), interval(1, -1): + if error_code[0, 0][plume] == 0: + if execution_choice == 20 and K <= lcl_level[0, 0][plume]: + # special treatment below updraft_origin_level/updraft_lcl_level + normalized_massflux_updraft_forced[0, 0, 0][plume] = ( + normalized_massflux_updraft_forced[0, 0, 1][plume] * 0.5 + ) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if execution_choice == 20: + normalized_massflux_updraft_forced[0, 0, 0][plume] = 0.0 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + _max_val, _ = column_max_ddim( + normalized_massflux_updraft_forced, + plume, + 0, + min(k_end, cloud_top_level[0, 0][plume] + 1), + ) + max_val: FloatFieldIJ = _max_val + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + if max_val <= 0: + normalized_massflux_updraft_forced[0, 0, 0][plume] = 0.0 + error_code[0, 0][plume] = 51 + else: + if K <= min(k_end, cloud_top_level[0, 0][plume] + 1): + normalized_massflux_updraft_forced[0, 0, 0][plume] = normalized_massflux_updraft_forced[ + 0, 0, 0 + ][plume] / (1.0e-9 + max_val) + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0: + normalized_massflux_updraft_modified = normalized_massflux_updraft_forced[0, 0, 0][plume] + normalized_massflux_updraft = normalized_massflux_updraft_forced[0, 0, 0][plume] + + +def updraft_moisture( + start_level: IntFieldIJ, + error_code: IntFieldIJ_Plume, + geopotential_height_cloud_levels_forced: FloatField, + cloud_total_water_after_entrainment_forced: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + condensate_to_fall_forced: FloatField_Plume, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + cloud_moist_static_energy_forced: FloatField, + updraft_column_temperature_forced: FloatField, + ocean_fraction: FloatFieldIJ, + convection_fraction: FloatFieldIJ, + surface_type: FloatFieldIJ, + p_forced: FloatField, + cloud_top_level: IntFieldIJ_Plume, + d_buoyancy_forced: FloatField, + cloud_liquid_before_rain_forced: FloatField, + t_cloud_levels: FloatField, + vapor_forced: FloatField, + gamma_cloud_levels_forced: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + environment_saturation_mixing_ratio_cloud_levels_forced: FloatField, + updraft_origin_level: IntFieldIJ_Plume, + vapor_cloud_levels_forced: FloatField, + vapor_excess: FloatFieldIJ, + ccn: FloatFieldIJ, + mass_entrainment_updraft: FloatField, + mass_detrainment_updraft: FloatField, + psum: FloatFieldIJ, + psumh: FloatFieldIJ, + c1d: FloatField, + add_buoyancy: FloatFieldIJ, + vertical_velocity_3d: FloatField, + C0: Float, + AVERAGE_LAYER_DEPTH: Float, + plume: Int, +): + """Compute moisture properties of the updraft and quantify and segregate precipitable water. + + Args: + start_level (IntFieldIJ) + error_code (IntFieldIJ_Plume) + geopotential_height_cloud_levels_forced (FloatField) + cloud_total_water_after_entrainment_forced (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + condensate_to_fall_forced (FloatField_Plume) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + cloud_moist_static_energy_forced (FloatField) + updraft_column_temperature_forced (FloatField) + ocean_fraction (FloatFieldIJ) + convection_fraction (FloatFieldIJ) + surface_type (FloatFieldIJ) + p_forced (FloatField) + cloud_top_level (IntFieldIJ_Plume) + d_buoyancy_forced (FloatField) + cloud_liquid_before_rain_forced (FloatField) + t_cloud_levels (FloatField) + vapor_forced (FloatField) + gamma_cloud_levels_forced (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + environment_saturation_mixing_ratio_cloud_levels_forced (FloatField) + updraft_origin_level (IntFieldIJ_Plume) + vapor_cloud_levels_forced (FloatField) + vapor_excess (FloatFieldIJ) + ccn (FloatFieldIJ) + mass_entrainment_updraft (FloatField) + mass_detrainment_updraft (FloatField) + psum (FloatFieldIJ) + psumh (FloatFieldIJ) + c1d (FloatField) + add_buoyancy (FloatFieldIJ) + vertical_velocity_3d (FloatField) + C0 (Float) + AVERAGE_LAYER_DEPTH (Float) + plume (Int) + """ + from __externals__ import ( + AUTOCONV, + BOUNDARY_CONDITION_METHOD, + CRITICAL_MIXING_RATIO_OVER_LAND, + CRITICAL_MIXING_RATIO_OVER_OCEAN, + USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + ZERO_DIFF, + k_end, + ) + + with computation(PARALLEL), interval(...): + # make garbage field so the get_cloud_boundary_conditions call does not break + # this is never touched so long as compute_perturbation=False + dummy_field_no_read = 0.0 + + with computation(FORWARD), interval(0, 1): + # internal constants + BDISPM: FloatFieldIJ = 0.366 # berry--size dispersion (maritime) + BDISPC: FloatFieldIJ = 0.146 # berry--size dispersion (continental) + T_BF: FloatFieldIJ = 268.16 + T_ICE_BF: FloatFieldIJ = 235.16 + RK: FloatFieldIJ = 3.0 + XEXP: FloatFieldIJ = 2.0 + + with computation(FORWARD), interval(0, 1): + # no precip for small clouds + total_normalized_integrated_condensate_forced[0, 0][plume] = 0.0 + psum = 0.0 + psumh = 0.0 + + with computation(PARALLEL), interval(0, -1): + condensate_to_fall_forced[0, 0, 0][plume] = 0.0 + updraft_column_temperature_forced = t_cloud_levels + cloud_liquid_before_rain_forced = 0.0 + cloud_liquid_after_rain_forced[0, 0, 0][plume] = 0.0 # liq/ice water + cloud_total_water_after_entrainment_forced = ( + vapor_cloud_levels_forced # total water: liq/ice = vapor water + ) + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + # get boundary condition for cloud_total_water_after_entrainment_forced + vapor_source: FloatFieldIJ = get_cloud_boundary_conditions( + field=vapor_cloud_levels_forced, + scalar_perturbation=0, + p=p_forced, + updraft_origin_level=updraft_origin_level[0, 0][plume], + ocean_fraction=ocean_fraction, + BOUNDARY_CONDITION_METHOD=BOUNDARY_CONDITION_METHOD, + AVERAGE_LAYER_DEPTH=AVERAGE_LAYER_DEPTH, + k_end=k_end, + compute_perturbation=False, + perturbation_field=dummy_field_no_read, + ) + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0 and K <= start_level: + # get boundary condition for cloud_total_water_after_entrainment_forced + cloud_total_water_after_entrainment_forced = ( + vapor_source + vapor_excess + 0.5 * add_buoyancy / cumulus_parameterization_constants.XLV + ) + cloud_liquid_after_rain_forced[0, 0, 0][plume] = 0.0 + + with computation(FORWARD), interval(0, 1): + if ( + error_code[0, 0][plume] == 0 + and USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES == 1 + and plume == cumulus_parameterization_constants.SHALLOW + ): # only for shallow plume + # option to produce linear fluxes in the sub-cloud layer + get_delmix_implementation_here = True + + with computation(PARALLEL), interval(...): + # initialize mask to stop computation in the next block + stop_current_index = False + + with computation(FORWARD), interval(1, None): + if error_code[0, 0][plume] == 0: + if K >= start_level + 1 and K <= cloud_top_level[0, 0][plume] + 1: + dz = ( + geopotential_height_cloud_levels_forced + - geopotential_height_cloud_levels_forced[0, 0, -1] + ) + # saturation in cloud, this is what is allowed to be in it + saturation_cloud_liquid = ( + environment_saturation_mixing_ratio_cloud_levels_forced + + (1.0 / cumulus_parameterization_constants.XLV) + * (gamma_cloud_levels_forced / (1.0 + gamma_cloud_levels_forced)) + * d_buoyancy_forced + ) + + # steady state plume equation, for what could be in cloud without condensation + denom = ( + normalized_massflux_updraft_forced[0, 0, -1][plume] + - 0.5 * mass_detrainment_updraft[0, 0, -1] + + mass_entrainment_updraft[0, 0, -1] + ) + + if denom > 0.0: + cloud_total_water_after_entrainment_forced = ( + cloud_total_water_after_entrainment_forced[0, 0, -1] + * normalized_massflux_updraft_forced[0, 0, -1][plume] + - 0.5 + * mass_detrainment_updraft[0, 0, -1] + * cloud_total_water_after_entrainment_forced[0, 0, -1] + + mass_entrainment_updraft[0, 0, -1] * vapor_forced[0, 0, -1] + ) / denom + + if K == start_level + 1: + cloud_total_water_after_entrainment_forced = ( + cloud_total_water_after_entrainment_forced + + vapor_excess * mass_entrainment_updraft[0, 0, -1] / denom + ) + # assuming no liq/ice water in the environment + cloud_liquid_after_rain_forced[0, 0, 0][plume] = ( + cloud_liquid_after_rain_forced[0, 0, -1][plume] + * normalized_massflux_updraft_forced[0, 0, -1][plume] + - 0.5 + * mass_detrainment_updraft[0, 0, -1] + * cloud_liquid_after_rain_forced[0, 0, -1][plume] + ) / denom + + else: + cloud_total_water_after_entrainment_forced = cloud_total_water_after_entrainment_forced[ + 0, 0, -1 + ] + cloud_liquid_after_rain_forced[0, 0, 0][plume] = cloud_liquid_after_rain_forced[0, 0, -1][ + plume + ] + + # updraft temp + updraft_column_temperature_forced = (1.0 / cumulus_parameterization_constants.CP) * ( + cloud_moist_static_energy_forced + - constants.MAPL_GRAV * geopotential_height_cloud_levels_forced + - cumulus_parameterization_constants.XLV * saturation_cloud_liquid + ) + + # total condensed water before rainout + cloud_liquid_before_rain_forced = max( + 0.0, cloud_total_water_after_entrainment_forced - saturation_cloud_liquid + ) + + cloud_liquid_after_rain_forced[0, 0, 0][plume] = min( + cloud_liquid_before_rain_forced, cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + + # production term => condensation/diffusional growth + cup = ( + max( + 0.0, + cloud_total_water_after_entrainment_forced + - saturation_cloud_liquid + - cloud_liquid_after_rain_forced[0, 0, 0][plume], + ) + / dz + ) + + if C0 < 1.0e-6: + cloud_liquid_after_rain_forced[0, 0, 0][plume] = cloud_liquid_before_rain_forced + cloud_total_water_after_entrainment_forced = cloud_liquid_after_rain_forced[0, 0, 0][ + plume + ] + min(cloud_total_water_after_entrainment_forced, saturation_cloud_liquid) + total_normalized_integrated_condensate_forced[0, 0][plume] = 0.0 + psum = ( + psum + + cloud_liquid_before_rain_forced + * normalized_massflux_updraft_forced[0, 0, 0][plume] + * dz + ) + + stop_current_index = True + + if stop_current_index == False: + if AUTOCONV == 1: + min_liq = ( + ocean_fraction * CRITICAL_MIXING_RATIO_OVER_OCEAN + + (1.0 - ocean_fraction) * CRITICAL_MIXING_RATIO_OVER_LAND + ) + cx0 = (c1d + C0) * dz + cloud_liquid_after_rain_forced[0, 0, 0][plume] = cloud_liquid_before_rain_forced / ( + 1.0 + cx0 + ) + condensate_to_fall_forced[0, 0, 0][plume] = cx0 * max( + 0.0, cloud_liquid_after_rain_forced[0, 0, 0][plume] - min_liq + ) # units kg[rain]/kg[air] + # convert precipitable_water_updraft_forced to + # normalized precipitable_water_updraft_forced + condensate_to_fall_forced[0, 0, 0][plume] = ( + condensate_to_fall_forced[0, 0, 0][plume] + * normalized_massflux_updraft_forced[0, 0, 0][plume] + ) + + # total water (vapor + condensed) in updraft after the rainout + cloud_total_water_after_entrainment_forced = cloud_liquid_after_rain_forced[0, 0, 0][ + plume + ] + min(cloud_total_water_after_entrainment_forced, saturation_cloud_liquid) + + # integrated normalized condensates + total_normalized_integrated_condensate_forced[0, 0][plume] = ( + total_normalized_integrated_condensate_forced[0, 0][plume] + + condensate_to_fall_forced[0, 0, 0][plume] + ) + psum = ( + psum + + cloud_liquid_before_rain_forced + * normalized_massflux_updraft_forced[0, 0, 0][plume] + * dz + ) + + if ZERO_DIFF == 0 and total_normalized_integrated_condensate_forced[0, 0][plume] < 0.0: + error_code[0, 0][plume] = 66 + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume] + 1: + # get back cloud_vapor_mixing_ratio_forced + cloud_total_water_after_entrainment_forced = ( + cloud_total_water_after_entrainment_forced - cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + + +def updraft_moist_static_energy_and_momentum_budget( + error_code: IntFieldIJ_Plume, + start_level: IntFieldIJ, + cloud_top_level: IntFieldIJ_Plume, + p_forced: FloatField, + environment_moist_static_energy: FloatField, + environment_moist_static_energy_forced: FloatField, + environment_moist_static_energy_cloud_levels: FloatField, + environment_moist_static_energy_cloud_levels_forced: FloatField, + environment_saturation_moist_static_energy_cloud_levels: FloatField, + environment_saturation_moist_static_energy_cloud_levels_forced: FloatField, + cloud_moist_static_energy: FloatField, + cloud_moist_static_energy_forced: FloatField, + normalized_massflux_updraft: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + mass_entrainment_updraft: FloatField, + mass_detrainment_updraft: FloatField, + mass_entrainment_u_updraft: FloatField, + mass_detrainment_u_updraft: FloatField, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + u: FloatField, + v: FloatField, + u_c: FloatField, + v_c: FloatField, + u_cloud_levels: FloatField, + v_cloud_levels: FloatField, + partition_liquid_ice: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + vapor_excess: FloatFieldIJ, + t_excess: FloatFieldIJ, + add_buoyancy: FloatFieldIJ, + plume: Int, +): + """Compute the impate of precipitation on the moist static energy, horizontal winds, + and massflux within the updraft. + + Args: + error_code (IntFieldIJ_Plume) + start_level (IntFieldIJ) + cloud_top_level (IntFieldIJ_Plume) + p_forced (FloatField) + environment_moist_static_energy (FloatField) + environment_moist_static_energy_forced (FloatField) + environment_moist_static_energy_cloud_levels (FloatField) + environment_moist_static_energy_cloud_levels_forced (FloatField) + environment_saturation_moist_static_energy_cloud_levels (FloatField) + environment_saturation_moist_static_energy_cloud_levels_forced (FloatField) + cloud_moist_static_energy (FloatField) + cloud_moist_static_energy_forced (FloatField) + normalized_massflux_updraft (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + mass_entrainment_updraft (FloatField) + mass_detrainment_updraft (FloatField) + mass_entrainment_u_updraft (FloatField) + mass_detrainment_u_updraft (FloatField) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + u (FloatField) + v (FloatField) + u_c (FloatField) + v_c (FloatField) + u_cloud_levels (FloatField) + v_cloud_levels (FloatField) + partition_liquid_ice (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + vapor_excess (FloatFieldIJ) + t_excess (FloatFieldIJ) + add_buoyancy (FloatFieldIJ) + plume (Int) + """ + from __externals__ import PRESSURE_GRADIENT_CONSTANT, USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES + + with computation(PARALLEL), interval(...): + if ( + error_code[0, 0][plume] == 0 + and plume == cumulus_parameterization_constants.SHALLOW + and USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES == 1 + ): + # only for shallow plume + get_delmix_implementation_here = True + + with computation(FORWARD), interval(1, None): + if error_code[0, 0][plume] == 0: + if K >= start_level + 1 and K <= cloud_top_level[0, 0][plume] + 1: + denom = ( + normalized_massflux_updraft[0, 0, -1] + - 0.5 * mass_detrainment_updraft[0, 0, -1] + + mass_entrainment_updraft[0, 0, -1] + ) + denom_u = ( + normalized_massflux_updraft[0, 0, -1] + - 0.5 * mass_detrainment_u_updraft[0, 0, -1] + + mass_entrainment_u_updraft[0, 0, -1] + ) + + if denom > 0.0 and denom_u > 0.0: + cloud_moist_static_energy = ( + cloud_moist_static_energy[0, 0, -1] * normalized_massflux_updraft[0, 0, -1] + - (0.5 * mass_detrainment_updraft[0, 0, -1]) * cloud_moist_static_energy[0, 0, -1] + + mass_entrainment_updraft[0, 0, -1] * environment_moist_static_energy[0, 0, -1] + ) / denom + + cloud_moist_static_energy_forced = ( + cloud_moist_static_energy_forced[0, 0, -1] + * normalized_massflux_updraft_forced[0, 0, -1][plume] + - 0.5 + * mass_detrainment_updraft_forced[0, 0, -1][plume] + * cloud_moist_static_energy_forced[0, 0, -1] + + mass_entrainment_updraft_forced[0, 0, -1][plume] + * environment_moist_static_energy_forced[0, 0, -1] + ) / denom + + if K == start_level + 1: + modification = ( + cumulus_parameterization_constants.XLV * vapor_excess + + cumulus_parameterization_constants.CP * t_excess + ) + add_buoyancy + cloud_moist_static_energy_forced = ( + cloud_moist_static_energy_forced + + modification * mass_entrainment_updraft_forced[0, 0, -1][plume] / denom + ) + cloud_moist_static_energy = ( + cloud_moist_static_energy + + modification * mass_entrainment_updraft[0, 0, -1] / denom + ) + + u_c = ( + u_c[0, 0, -1] * normalized_massflux_updraft[0, 0, -1] + - 0.5 * mass_detrainment_u_updraft[0, 0, -1] * u_c[0, 0, -1] + + mass_entrainment_u_updraft[0, 0, -1] * u[0, 0, -1] + - PRESSURE_GRADIENT_CONSTANT + * 0.5 + * (normalized_massflux_updraft + normalized_massflux_updraft[0, 0, -1]) + * (u_cloud_levels - u_cloud_levels[0, 0, -1]) + ) / denom_u + + v_c = ( + v_c[0, 0, -1] * normalized_massflux_updraft[0, 0, -1] + - 0.5 * mass_detrainment_u_updraft[0, 0, -1] * v_c[0, 0, -1] + + mass_entrainment_u_updraft[0, 0, -1] * v[0, 0, -1] + - PRESSURE_GRADIENT_CONSTANT + * 0.5 + * (normalized_massflux_updraft + normalized_massflux_updraft[0, 0, -1]) + * (v_cloud_levels - v_cloud_levels[0, 0, -1]) + ) / denom_u + + else: + cloud_moist_static_energy = cloud_moist_static_energy[0, 0, -1] + cloud_moist_static_energy_forced = cloud_moist_static_energy_forced[0, 0, -1] + u_c = u_c[0, 0, -1] + v_c = v_c[0, 0, -1] + + cloud_moist_static_energy = ( + cloud_moist_static_energy + + (1.0 - partition_liquid_ice) + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + * cumulus_parameterization_constants.XLF + ) + cloud_moist_static_energy_forced = ( + cloud_moist_static_energy_forced + + (1.0 - partition_liquid_ice) + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + * cumulus_parameterization_constants.XLF + ) + + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0 and K >= cloud_top_level[0, 0][plume] + 2: + cloud_moist_static_energy = environment_saturation_moist_static_energy_cloud_levels + u_c = u_cloud_levels + v_c = v_cloud_levels + cloud_moist_static_energy_forced = environment_saturation_moist_static_energy_cloud_levels_forced + normalized_massflux_updraft = 0.0 + normalized_massflux_updraft_forced[0, 0, 0][plume] = 0.0 + + +def updraft_temperature( + error_code: IntFieldIJ_Plume, + updraft_column_temperature_forced: FloatField, + cloud_moist_static_energy_forced: FloatField, + geopotential_height_cloud_levels_forced: FloatField, + cloud_total_water_after_entrainment_forced: FloatField, + t_cloud_levels_forced: FloatField, + plume: Int, +): + """Compute the temperature within the updraft. + + Args: + error_code (IntFieldIJ_Plume) + updraft_column_temperature_forced (FloatField) + cloud_moist_static_energy_forced (FloatField) + geopotential_height_cloud_levels_forced (FloatField) + cloud_total_water_after_entrainment_forced (FloatField) + t_cloud_levels_forced (FloatField) + plume (Int) + """ + with computation(PARALLEL), interval(0, -1): + if error_code[0, 0][plume] == 0: + updraft_column_temperature_forced = (1.0 / cumulus_parameterization_constants.CP) * ( + cloud_moist_static_energy_forced + - constants.MAPL_GRAV * geopotential_height_cloud_levels_forced + - cumulus_parameterization_constants.XLV * cloud_total_water_after_entrainment_forced + ) + + with computation(PARALLEL), interval(-1, None): + if error_code[0, 0][plume] == 0: + updraft_column_temperature_forced = t_cloud_levels_forced + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] != 0: + updraft_column_temperature_forced = t_cloud_levels_forced + + +def cloud_workfunction_aa0( + error_code: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + geopotential_height: FloatField, + normalized_massflux_updraft: FloatField, + d_buoyancy: FloatField, + gamma_cloud_levels: FloatField, + t_cloud_levels: FloatField, + workfunction: FloatFieldIJ, + mode: Int, + plume: Int, +): + """Compute cloud workfunction aa0 + + Stencil calculation is controlled my "mode" argument: + mode = 0: default. uses dynamic bounds (updraft_lfc_level --> cloud_top_level) + mode = 1: boundary layer. uses fixed lower bound (surface --> updraft_lfc_level - 1) + mode = 2: convective inhibition. uses alternate dynamic bounds + (updraft_origin_level --> updraft_lfc_level - 1) + + Args: + error_code (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + geopotential_height (FloatField) + normalized_massflux_updraft (FloatField) + d_buoyancy (FloatField) + gamma_cloud_levels (FloatField) + t_cloud_levels (FloatField) + workfunction (FloatFieldIJ) + mode (Int) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + # initialize workfunction to zero + workfunction = 0.0 + + # set up bounds for next block + if mode == 0: + lower_bound: IntFieldIJ = updraft_lfc_level[0, 0][plume] + upper_bound: IntFieldIJ = cloud_top_level[0, 0][plume] + elif mode == 1: + lower_bound: IntFieldIJ = 0 + upper_bound: IntFieldIJ = updraft_lfc_level[0, 0][plume] - 1 + elif mode == 2: + lower_bound: IntFieldIJ = updraft_origin_level[0, 0][plume] + upper_bound: IntFieldIJ = updraft_lfc_level[0, 0][plume] - 1 + + upper_bound = upper_bound + 1 + + with computation(FORWARD), interval(lower_bound, upper_bound): + if error_code[0, 0][plume] == 0: + dz = geopotential_height[0, 0, 1] - geopotential_height + workfunction_current_level = ( + normalized_massflux_updraft + * (constants.MAPL_GRAV / (cumulus_parameterization_constants.CP * t_cloud_levels)) + * d_buoyancy + / (1.0 + gamma_cloud_levels) + ) + workfunction_level_above = ( + normalized_massflux_updraft[0, 0, 1] + * (constants.MAPL_GRAV / (cumulus_parameterization_constants.CP * t_cloud_levels[0, 0, 1])) + * d_buoyancy[0, 0, 1] + / (1.0 + gamma_cloud_levels[0, 0, 1]) + ) + d_workfunction = 0.5 * (workfunction_current_level + workfunction_level_above) * dz + + workfunction = workfunction + d_workfunction + + +def check_cloud_workfunction_1( + error_code: IntFieldIJ_Plume, + cloud_workfunction_1: FloatFieldIJ, + plume: Int, +): + """Check the validity of cloud_workfunction_1. + + Args: + error_code (IntFieldIJ_Plume) + cloud_workfunction_1 (FloatFieldIJ) + plume (Int) + """ + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if cloud_workfunction_1 == 0.0: + error_code[0, 0][plume] = 17 + + +def compute_precipitation_ensemble( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + epsilon_forced: FloatFieldIJ_Plume, + precipitation_ensemble: FloatFieldIJ_Ensemble, + plume: Int, +): + from __externals__ import C0_MID, ENSEMBLE_MEMBERS + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + member = 0 + while member < ENSEMBLE_MEMBERS: + precipitation_ensemble[0, 0][member] = ( + precipitation_ensemble[0, 0][member] + + condensate_to_fall_forced[0, 0, 0][plume] + + epsilon_forced[0, 0][plume] * evaporate_in_downdraft_forced[0, 0, 0][plume] + ) + member += 1 + + with computation(FORWARD), interval(0, 1): + if error_code[0, 0][plume] == 0: + if precipitation_ensemble[0, 0][6] < 1.0e-6 and C0_MID > 0.0 and plume != 0: + error_code[0, 0][plume] = 18 + member = 0 + while member < ENSEMBLE_MEMBERS: + precipitation_ensemble[0, 0][member] = 0.0 + member += 1 + + member = 0 + while member < ENSEMBLE_MEMBERS: + if precipitation_ensemble[0, 0][plume] < 1.0e-5: + precipitation_ensemble[0, 0][plume] = 0.0 + member += 1 + + +class UpdraftMassFlux(NDSLRuntime): + """Handle mass fluxes in the updraft. This code has a number of potential paths depending on configuration + settings - some of these paths are not yet implemented. + + This code may execute for all plumes - shallow plume execution is currently not implemented. + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # add dimension to quantityfactory and create classes for constants + quantity_factory.add_data_dimensions({"UpdraftMassFlux_constants": len(_X_ALPHA)}) + + self._X_ALPHA: Local = quantity_factory.zeros(["UpdraftMassFlux_constants"], "n/a") + self._G_ALPHA: Local = quantity_factory.zeros(["UpdraftMassFlux_constants"], "n/a") + + self._X_ALPHA.field[:] = _X_ALPHA + self._G_ALPHA.field[:] = _G_ALPHA + + # construct stencil + self._updraft_mass_flux = stencil_factory.from_dims_halo( + func=updraft_mass_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "BETA_SHALLOW": cumulus_parameterization_config.BETA_SHALLOW, + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES": cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + }, + ) + + def __call__( + self, + error_code: Quantity, + updraft_origin_level: Quantity, + cloud_top_level: Quantity, + pbl_level: Quantity, + updraft_lfc_level: Quantity, + lcl_level: Quantity, + p_cloud_levels_forced: Quantity, + p_surface: Quantity, + ocean_fraction: Quantity, + normalized_massflux_updraft: Quantity, + normalized_massflux_updraft_forced: Quantity, + normalized_massflux_updraft_modified: Quantity, + random_number: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._updraft_mass_flux( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + cloud_top_level=cloud_top_level, + pbl_level=pbl_level, + updraft_lfc_level=updraft_lfc_level, + lcl_level=lcl_level, + p_cloud_levels_forced=p_cloud_levels_forced, + p_surface=p_surface, + ocean_fraction=ocean_fraction, + normalized_massflux_updraft=normalized_massflux_updraft, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_updraft_modified=normalized_massflux_updraft_modified, + random_number=random_number, + UPDRAFT_MAX_HEIGHT_LAND=plume_dependent_constants.UPDRAFT_MAX_HEIGHT_LAND, + UPDRAFT_MAX_HEIGHT_OCEAN=plume_dependent_constants.UPDRAFT_MAX_HEIGHT_OCEAN, + plume=plume_dependent_constants.PLUME_INDEX, + X_ALPHA=self._X_ALPHA, + G_ALPHA=self._G_ALPHA, + ) + + +class UpdraftInitialWorkfunctions(NDSLRuntime): + """Compute initial estimates for cloud_workfunction_0 and cloud_workfunction_1.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._cloud_workfunction_aa0 = stencil_factory.from_dims_halo( + func=cloud_workfunction_aa0, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._check_cloud_workfunction_1 = stencil_factory.from_dims_halo( + func=check_cloud_workfunction_1, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + updraft_origin_level: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + geopotential_height_cloud_levels: Quantity, + geopotential_height_cloud_levels_forced: Quantity, + normalized_massflux_updraft: Quantity, + normalized_massflux_updraft_forced: Quantity, + d_buoyancy: Quantity, + d_buoyancy_forced: Quantity, + gamma_cloud_levels: Quantity, + gamma_cloud_levels_forced: Quantity, + t_cloud_levels: Quantity, + t_cloud_levels_forced: Quantity, + cloud_workfunction_0: Quantity, + cloud_workfunction_1: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._cloud_workfunction_aa0( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height=geopotential_height_cloud_levels, + normalized_massflux_updraft=normalized_massflux_updraft, + d_buoyancy=d_buoyancy, + gamma_cloud_levels=gamma_cloud_levels, + t_cloud_levels=t_cloud_levels, + workfunction=cloud_workfunction_0, + mode=Int(0), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._cloud_workfunction_aa0( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height=geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + d_buoyancy=d_buoyancy_forced, + gamma_cloud_levels=gamma_cloud_levels_forced, + t_cloud_levels=t_cloud_levels_forced, + workfunction=cloud_workfunction_1, + mode=Int(0), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._check_cloud_workfunction_1( + error_code=error_code, + cloud_workfunction_1=cloud_workfunction_1, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class UpdraftCIN(NDSLRuntime): + """Compute initial estimate of updraft CIN""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._cloud_workfunction_aa0 = stencil_factory.from_dims_halo( + func=cloud_workfunction_aa0, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def __call__( + self, + error_code: Quantity, + updraft_origin_level: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + geopotential_height_cloud_levels: Quantity, + geopotential_height_cloud_levels_forced: Quantity, + normalized_massflux_updraft: Quantity, + normalized_massflux_updraft_forced: Quantity, + d_buoyancy: Quantity, + d_buoyancy_forced: Quantity, + gamma_cloud_levels: Quantity, + gamma_cloud_levels_forced: Quantity, + t_cloud_levels: Quantity, + t_cloud_levels_forced: Quantity, + cin_0: Quantity, + cin_1: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._cloud_workfunction_aa0( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height=geopotential_height_cloud_levels, + normalized_massflux_updraft=normalized_massflux_updraft, + d_buoyancy=d_buoyancy, + gamma_cloud_levels=gamma_cloud_levels, + t_cloud_levels=t_cloud_levels, + workfunction=cin_0, + mode=Int(2), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._cloud_workfunction_aa0( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height=geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + d_buoyancy=d_buoyancy_forced, + gamma_cloud_levels=gamma_cloud_levels_forced, + t_cloud_levels=t_cloud_levels_forced, + workfunction=cin_1, + mode=Int(2), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + +class UpdateWorkfunctionAndPrecipitationEnsemble(NDSLRuntime): + """Update cloud workfunctions and precipitation ensemble with the result of convection.""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # construct stencils and functions + self._cloud_workfunction_aa0 = stencil_factory.from_dims_halo( + func=cloud_workfunction_aa0, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._compute_precipitation_ensemble = stencil_factory.from_dims_halo( + func=compute_precipitation_ensemble, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ENSEMBLE_MEMBERS": cumulus_parameterization_constants.MAXENS1 + * cumulus_parameterization_constants.MAXENS2 + * cumulus_parameterization_constants.MAXENS3, + "C0_MID": cumulus_parameterization_config.C0_MID, + }, + ) + + def __call__( + self, + error_code: Quantity, + updraft_origin_level: Quantity, + updraft_lfc_level: Quantity, + cloud_top_level: Quantity, + geopotential_height_cloud_levels_modified: Quantity, + normalized_massflux_updraft_modified: Quantity, + d_buoyancy_modified: Quantity, + gamma_cloud_levels: Quantity, + t_cloud_levels_modified: Quantity, + cloud_workfunction_0_modified: Quantity, + condensate_to_fall_forced: Quantity, + evaporate_in_downdraft_forced: Quantity, + epsilon_forced: Quantity, + precipitation_ensemble: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._cloud_workfunction_aa0( + error_code=error_code, + updraft_origin_level=updraft_origin_level, + updraft_lfc_level=updraft_lfc_level, + cloud_top_level=cloud_top_level, + geopotential_height=geopotential_height_cloud_levels_modified, + normalized_massflux_updraft=normalized_massflux_updraft_modified, + d_buoyancy=d_buoyancy_modified, + gamma_cloud_levels=gamma_cloud_levels, + t_cloud_levels=t_cloud_levels_modified, + workfunction=cloud_workfunction_0_modified, + mode=Int(0), + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._compute_precipitation_ensemble( + error_code=error_code, + cloud_top_level=cloud_top_level, + condensate_to_fall_forced=condensate_to_fall_forced, + evaporate_in_downdraft_forced=evaporate_in_downdraft_forced, + epsilon_forced=epsilon_forced, + precipitation_ensemble=precipitation_ensemble, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/vertical_discretization.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/vertical_discretization.py new file mode 100644 index 000000000..67a7cdd9f --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/cumulus_parameterization/vertical_discretization.py @@ -0,0 +1,836 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import Local, NDSLRuntime, Quantity, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import FORWARD, PARALLEL, K, computation, interval +from ndsl.dsl.typing import FloatField, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_max +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_Plume, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import tridiag + + +def zero_tendencies( + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + del_moist_static_energy_cloud_ensemble: FloatField, + del_t_cloud_ensemble: FloatField, + del_vapor_cloud_ensemble: FloatField, + del_cloud_liquid_cloud_ensemble: FloatField, + del_buoyancy_cloud_ensemble: FloatField, + moist_static_energy_tendency_from_environmental_subsidence: FloatField, + vapor_tendency_from_environmental_subsidence: FloatField, + t_tendency_from_environmental_subsidence: FloatField, +): + """Ensure that all tendencies are zero, and there is no lingering data from the previous call. + + Args: + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + del_moist_static_energy_cloud_ensemble (FloatField) + del_t_cloud_ensemble (FloatField) + del_vapor_cloud_ensemble (FloatField) + del_cloud_liquid_cloud_ensemble (FloatField) + del_buoyancy_cloud_ensemble (FloatField) + moist_static_energy_tendency_from_environmental_subsidence (FloatField) + vapor_tendency_from_environmental_subsidence (FloatField) + t_tendency_from_environmental_subsidence (FloatField) + """ + with computation(PARALLEL), interval(...): + del_u_cloud_ensemble = 0.0 + del_v_cloud_ensemble = 0.0 + del_moist_static_energy_cloud_ensemble = 0.0 + del_t_cloud_ensemble = 0.0 + del_vapor_cloud_ensemble = 0.0 + del_cloud_liquid_cloud_ensemble = 0.0 + del_buoyancy_cloud_ensemble = 0.0 + moist_static_energy_tendency_from_environmental_subsidence = 0.0 + vapor_tendency_from_environmental_subsidence = 0.0 + t_tendency_from_environmental_subsidence = 0.0 + + +def convective_transport_of_momentum( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + environment_massflux: FloatField, + u: FloatField, + v: FloatField, + u_cloud_levels: FloatField, + v_cloud_levels: FloatField, + u_c: FloatField, + v_c: FloatField, + u_c_downdraft: FloatField, + v_c_downdraft: FloatField, + del_u_cloud_ensemble: FloatField, + del_v_cloud_ensemble: FloatField, + epsilon_forced: FloatFieldIJ_Plume, + fp: FloatField, + fm: FloatField, + aa: FloatField, + bb: FloatField, + cc: FloatField, + ddu: FloatField, + ddv: FloatField, + plume: Int, +): + """Compute u & v tendnecies - the effect of convection on the state wind field. These tencencies are not + directly used to update the overarching model state, but are the first step in that process. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + environment_massflux (FloatField) + u (FloatField) + v (FloatField) + u_cloud_levels (FloatField) + v_cloud_levels (FloatField) + u_c (FloatField) + v_c (FloatField) + u_c_downdraft (FloatField) + v_c_downdraft (FloatField) + del_u_cloud_ensemble (FloatField) + del_v_cloud_ensemble (FloatField) + epsilon_forced (FloatFieldIJ_Plume) + fp (FloatField) + fm (FloatField) + aa (FloatField) + bb (FloatField) + cc (FloatField) + ddu (FloatField) + ddv (FloatField) + plume (Int) + """ + from __externals__ import ALP1, DTIME, VERTICAL_DISCRETIZATION_OPTION + + with computation(FORWARD), interval(0, 1): + # prepare bounds for subsequent computation + upper_bound: IntFieldIJ = cloud_top_level[0, 0][plume] + 1 + upper_bound_up_1: IntFieldIJ = upper_bound + 1 + + with computation(FORWARD), interval(0, upper_bound): + if ( + error_code[0, 0][plume] == 0 and VERTICAL_DISCRETIZATION_OPTION == 1 and ALP1 == 0.0 + ): # fully time explicit + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + + del_u_cloud_ensemble = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * (u_c[0, 0, 1] - u_cloud_levels[0, 0, 1]) + - normalized_massflux_updraft_forced[0, 0, 0][plume] * (u_c - u_cloud_levels) + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (u_c_downdraft[0, 0, 1] - u_cloud_levels[0, 0, 1]) + - normalized_massflux_downdraft_forced[0, 0, 0][plume] * (u_c_downdraft - u_cloud_levels) + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + + del_v_cloud_ensemble = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * (v_c_downdraft[0, 0, 1] - v_cloud_levels[0, 0, 1]) + - normalized_massflux_updraft_forced[0, 0, 0][plume] * (v_c_downdraft - v_cloud_levels) + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (v_c_downdraft[0, 0, 1] - v_cloud_levels[0, 0, 1]) + - normalized_massflux_downdraft_forced[0, 0, 0][plume] * (v_c_downdraft - v_cloud_levels) + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + + with computation(PARALLEL), interval(0, upper_bound_up_1): + if ( + error_code[0, 0][plume] == 0 and VERTICAL_DISCRETIZATION_OPTION == 1 and ALP1 > 0.0 + ): # time alp0*explict + alp1*implicit + upstream + alp0 = 1.0 - ALP1 + fp = 0.5 * (environment_massflux + abs(environment_massflux)) + fm = 0.5 * (environment_massflux - abs(environment_massflux)) + + with computation(FORWARD), interval(0, upper_bound): + if error_code[0, 0][plume] == 0 and VERTICAL_DISCRETIZATION_OPTION == 1 and ALP1 > 0.0: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + + beta1 = DTIME * constants.MAPL_GRAV / dp + aa = ALP1 * beta1 * fm + bb = 1.0 + ALP1 * beta1 * (fp - fm[0, 0, 1]) + cc = -ALP1 * beta1 * fp[0, 0, 1] + + ddu = ( + u + - ( + normalized_massflux_updraft_forced[0, 0, 1][plume] * u_c[0, 0, 1] + - normalized_massflux_updraft_forced[0, 0, 0][plume] * u_c[0, 0, 0] + ) + * beta1 + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] * u_c_downdraft[0, 0, 1] + - normalized_massflux_downdraft_forced[0, 0, 1][plume] * u_c_downdraft + ) + * beta1 + * epsilon_forced[0, 0][plume] + ) + + _, max_index = column_max(u, 0, K - 1) + ddu = ddu + alp0 * beta1 * ( + -fm * u.at(K=max_index) + (fm[0, 0, 1] - fp) * u + fp[0, 0, 1] * u[0, 0, 1] + ) + + ddv = ( + v + - ( + normalized_massflux_updraft_forced[0, 0, 1][plume] * v_c[0, 0, 1] + - normalized_massflux_updraft_forced[0, 0, 0][plume] * v_c + ) + * beta1 + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] * v_c_downdraft[0, 0, 1] + - normalized_massflux_downdraft_forced[0, 0, 0][plume] * v_c_downdraft + ) + * beta1 + * epsilon_forced[0, 0][plume] + ) + + _, max_index = column_max(u, 0, K - 1) + ddv = ddv + alp0 * beta1 * ( + -fm * v.at(K=max_index) + (fm[0, 0, 1] - fp) * v + fp[0, 0, 1] * v[0, 0, 1] + ) + + +def update_after_tridiag( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + in_field: FloatField, + out_field: FloatField, + wind: FloatField, + plume: Int, +): + """Update a field with the output of the tridiagonal solver. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + in_field (FloatField) + out_field (FloatField) + wind (FloatField) + plume (Int) + """ + from __externals__ import DTIME + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + out_field = (in_field - wind) / DTIME + + +def convective_transport_of_mse( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + cloud_moist_static_energy_forced: FloatField, + cloud_moist_static_energy_downdraft_forced: FloatField, + environment_moist_static_energy_cloud_levels_forced: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + melting: FloatField, + partition_liquid_ice: FloatField, + epsilon_forced: FloatFieldIJ_Plume, + del_moist_static_energy_cloud_ensemble: FloatField, + moist_static_energy_tendency_from_environmental_subsidence: FloatField, + plume: Int, +): + """Compute moist static energy tendency - the effect of convection and environmental + subsidence on environmental moist static energy independently. This tencency is not + directly used to update the overarching model state, but is the first step in that process. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + cloud_moist_static_energy_forced (FloatField) + cloud_moist_static_energy_downdraft_forced (FloatField) + environment_moist_static_energy_cloud_levels_forced (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + melting (FloatField) + partition_liquid_ice (FloatField) + epsilon_forced (FloatFieldIJ_Plume) + del_moist_static_energy_cloud_ensemble (FloatField) + moist_static_energy_tendency_from_environmental_subsidence (FloatField) + plume (Int) + """ + from __externals__ import USE_FCT + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0: + # moist static energy : flux form + source/sink terms + time explicit + if USE_FCT == 0: + if K <= cloud_top_level[0, 0][plume]: + dp = 100.0 * ( + p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume] + ) + + del_moist_static_energy_cloud_ensemble = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * ( + cloud_moist_static_energy_forced[0, 0, 1] + - environment_moist_static_energy_cloud_levels_forced[0, 0, 1] + ) + - normalized_massflux_updraft_forced[0, 0, 0][plume] + * ( + cloud_moist_static_energy_forced + - environment_moist_static_energy_cloud_levels_forced + ) + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * ( + cloud_moist_static_energy_downdraft_forced[0, 0, 1] + - environment_moist_static_energy_cloud_levels_forced[0, 0, 1] + ) + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + * ( + cloud_moist_static_energy_downdraft_forced + - environment_moist_static_energy_cloud_levels_forced + ) + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + + del_moist_static_energy_cloud_ensemble = ( + del_moist_static_energy_cloud_ensemble + + cumulus_parameterization_constants.XLF + * ( + (1.0 - partition_liquid_ice) + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + - melting + ) + * constants.MAPL_GRAV + / dp + ) + + # for output only + moist_static_energy_tendency_from_environmental_subsidence = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * (-environment_moist_static_energy_cloud_levels_forced[0, 0, 1]) + - normalized_massflux_updraft_forced[0, 0, 0][plume] + * (-environment_moist_static_energy_cloud_levels_forced) + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (-environment_moist_static_energy_cloud_levels_forced[0, 0, 1]) + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + * (-environment_moist_static_energy_cloud_levels_forced) + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + + +def convective_transport_of_water_vapor_and_condensates( + error_code: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + geopotential_height_cloud_levels_forced: FloatField, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + c1d: FloatField, + vapor_cloud_levels_forced: FloatField, + cloud_total_water_after_entrainment_forced: FloatField, + cloud_total_water_after_entrainment_downdraft_forced: FloatField, + cloud_liquid_after_rain_forced: FloatField_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + epsilon_forced: FloatFieldIJ_Plume, + d_buoyancy_downdraft_forced: FloatField, + del_cloud_liquid_cloud_ensemble: FloatField, + del_vapor_cloud_ensemble: FloatField, + vapor_tendency_from_environmental_subsidence: FloatField, + plume: Int, +): + """Compute cloud liquid and vapor tendencies - the effect of convection on environmental water, + and the contribution from environmental subsidence. These tencencies are not + directly used to update the overarching model state, but are the first step in that process. + + Args: + error_code (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + geopotential_height_cloud_levels_forced (FloatField) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + c1d (FloatField) + vapor_cloud_levels_forced (FloatField) + cloud_total_water_after_entrainment_forced (FloatField) + cloud_total_water_after_entrainment_downdraft_forced (FloatField) + cloud_liquid_after_rain_forced (FloatField_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + epsilon_forced (FloatFieldIJ_Plume) + d_buoyancy_downdraft_forced (FloatField) + del_cloud_liquid_cloud_ensemble (FloatField) + del_vapor_cloud_ensemble (FloatField) + vapor_tendency_from_environmental_subsidence (FloatField) + plume (Int) + """ + from __externals__ import C1, USE_FCT + + with computation(FORWARD), interval(0, -1): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + + # take out cloud liquid/ice water for detrainment + if ( + plume == cumulus_parameterization_constants.SHALLOW + or plume == cumulus_parameterization_constants.MID + ): # shallow or mid plume + del_cloud_liquid_cloud_ensemble = ( + mass_detrainment_updraft_forced[0, 0, 0][plume] + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + ) + elif plume == cumulus_parameterization_constants.DEEP: # deep plume + if not (abs(C1) > 0): + del_cloud_liquid_cloud_ensemble = ( + mass_detrainment_updraft_forced[0, 0, 0][plume] + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + ) + elif C1 > 0.0: + if K == cloud_top_level[0, 0][plume]: + del_cloud_liquid_cloud_ensemble = ( + mass_detrainment_updraft_forced[0, 0, 0][plume] + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + ) + else: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] + - geopotential_height_cloud_levels_forced + ) + del_cloud_liquid_cloud_ensemble = ( + normalized_massflux_updraft_forced[0, 0, 0][plume] + * c1d + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + * dz + / dp + * constants.MAPL_GRAV + ) + else: + if K == cloud_top_level[0, 0][plume]: + del_cloud_liquid_cloud_ensemble = ( + mass_detrainment_updraft_forced[0, 0, 0][plume] + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + ) + else: + dz = ( + geopotential_height_cloud_levels_forced[0, 0, 1] + - geopotential_height_cloud_levels_forced + ) + del_cloud_liquid_cloud_ensemble = ( + normalized_massflux_updraft_forced[0, 0, 0][plume] + * c1d + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + * dz + / dp + * constants.MAPL_GRAV + + mass_detrainment_updraft_forced[0, 0, 0][plume] + * 0.5 + * ( + cloud_liquid_after_rain_forced[0, 0, 1][plume] + + cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + ) * 0.5 + + g_rain = ( + 0.5 + * (condensate_to_fall_forced[0, 0, 0][plume] + condensate_to_fall_forced[0, 0, 1][plume]) + * constants.MAPL_GRAV + / dp + ) + e_dn = ( + -0.5 + * ( + evaporate_in_downdraft_forced[0, 0, 0][plume] + + evaporate_in_downdraft_forced[0, 0, 1][plume] + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) # pwdo < 0 and E_dn must > 0 + # condensation source term = detrained + flux divergence of + # cloud liquid/ice water (cloud_liquid_after_rain_forced) + converted to rain + c_up = ( + del_cloud_liquid_cloud_ensemble + + ( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * cloud_liquid_after_rain_forced[0, 0, 1][plume] + - normalized_massflux_updraft_forced[0, 0, 0][plume] + * cloud_liquid_after_rain_forced[0, 0, 0][plume] + ) + * constants.MAPL_GRAV + / dp + + g_rain + ) + + # water vapor budget + # = flux divergence z*(Q_c - Q_env)_up_and_down - condensation term + evaporation + del_vapor_cloud_ensemble = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] + * cloud_total_water_after_entrainment_forced[0, 0, 1] + - normalized_massflux_updraft_forced[0, 0, 0][plume] + * cloud_total_water_after_entrainment_forced + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * cloud_total_water_after_entrainment_downdraft_forced[0, 0, 1] + - normalized_massflux_downdraft_forced[0, 0, 0][plume] + * cloud_total_water_after_entrainment_downdraft_forced + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + - c_up + + e_dn + ) + + del_buoyancy_cloud_ensemble = ( + epsilon_forced[0, 0][plume] + * mass_detrainment_downdraft_forced[0, 0, 0][plume] + * 0.5 + * (d_buoyancy_downdraft_forced[0, 0, 1] + d_buoyancy_downdraft_forced) + * constants.MAPL_GRAV + / dp + ) + + with computation(FORWARD), interval(...): + if error_code[0, 0][plume] == 0 and USE_FCT == 0 and K <= cloud_top_level[0, 0][plume]: + dp = 100.0 * (p_cloud_levels_forced[0, 0, 0][plume] - p_cloud_levels_forced[0, 0, 1][plume]) + subsidence_tendency = ( + -( + normalized_massflux_updraft_forced[0, 0, 1][plume] * (-vapor_cloud_levels_forced[0, 0, 1]) + - normalized_massflux_updraft_forced[0, 0, 0][plume] * (-vapor_cloud_levels_forced) + ) + * constants.MAPL_GRAV + / dp + + ( + normalized_massflux_downdraft_forced[0, 0, 1][plume] + * (-vapor_cloud_levels_forced[0, 0, 1]) + - normalized_massflux_downdraft_forced[0, 0, 0][plume] * (-vapor_cloud_levels_forced) + ) + * constants.MAPL_GRAV + / dp + * epsilon_forced[0, 0][plume] + ) + + with computation(PARALLEL), interval(...): + if error_code[0, 0][plume] == 0 and K <= cloud_top_level[0, 0][plume]: + # add the contribuition from the environmental subsidence + del_vapor_cloud_ensemble = del_vapor_cloud_ensemble + subsidence_tendency + + # for output only + vapor_tendency_from_environmental_subsidence = subsidence_tendency + + +class VerticalDiscretization(NDSLRuntime): + """Connect the effect of convection on the environment state. Tendencies computed within this module + ARE NOT used outside of the cumulus parameterization core, but are the "first guess" for a value + which will eventually get fed back into the rest of the model. + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + ): + # init NDSLRuntime + super().__init__(stencil_factory) + + # make configuration visible at runtime + self.config = config + self.cumulus_parameterization_config = cumulus_parameterization_config + + # initialize local fields + self._fp: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._fm: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._aa: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._bb: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._cc: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._ddu: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + self._ddv: Local = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + + self._zero_tendencies = stencil_factory.from_dims_halo( + func=zero_tendencies, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._convective_transport_of_momentum = stencil_factory.from_dims_halo( + func=convective_transport_of_momentum, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "VERTICAL_DISCRETIZATION_OPTION": cumulus_parameterization_config.VERTICAL_DISCRETIZATION_OPTION, + "ALP1": cumulus_parameterization_config.ALP1, + "DTIME": cumulus_parameterization_config.DTIME, + }, + ) + + self._tridiag = stencil_factory.from_dims_halo( + func=tridiag, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._update_after_tridiag = stencil_factory.from_dims_halo( + func=update_after_tridiag, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DTIME": cumulus_parameterization_config.DTIME}, + ) + + self._convective_transport_of_mse_and_liquid_water = stencil_factory.from_dims_halo( + func=convective_transport_of_mse, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"USE_FCT": cumulus_parameterization_config.USE_FCT}, + ) + + self._convective_transport_of_water_vapor_and_condensates = stencil_factory.from_dims_halo( + func=convective_transport_of_water_vapor_and_condensates, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"C1": config.C1, "USE_FCT": cumulus_parameterization_config.USE_FCT}, + ) + + def __call__( + self, + error_code: Quantity, + cloud_top_level: Quantity, + p_cloud_levels_forced: Quantity, + geopotential_height_cloud_levels_forced: Quantity, + normalized_massflux_updraft_forced: Quantity, + normalized_massflux_downdraft_forced: Quantity, + environment_massflux: Quantity, + mass_detrainment_updraft_forced: Quantity, + mass_detrainment_downdraft_forced: Quantity, + c1d: Quantity, + u: Quantity, + v: Quantity, + u_cloud_levels: Quantity, + v_cloud_levels: Quantity, + u_c: Quantity, + v_c: Quantity, + u_c_downdraft: Quantity, + v_c_downdraft: Quantity, + cloud_moist_static_energy_forced: Quantity, + cloud_moist_static_energy_downdraft_forced: Quantity, + environment_moist_static_energy_cloud_levels_forced: Quantity, + vapor_cloud_levels_forced: Quantity, + cloud_total_water_after_entrainment_forced: Quantity, + cloud_total_water_after_entrainment_downdraft_forced: Quantity, + cloud_liquid_after_rain_forced: Quantity, + condensate_to_fall_forced: Quantity, + evaporate_in_downdraft_forced: Quantity, + melting: Quantity, + partition_liquid_ice: Quantity, + epsilon_forced: Quantity, + d_buoyancy_downdraft_forced: Quantity, + del_u_cloud_ensemble: Quantity, + del_v_cloud_ensemble: Quantity, + del_moist_static_energy_cloud_ensemble: Quantity, + del_t_cloud_ensemble: Quantity, + del_vapor_cloud_ensemble: Quantity, + del_cloud_liquid_cloud_ensemble: Quantity, + del_buoyancy_cloud_ensemble: Quantity, + moist_static_energy_tendency_from_environmental_subsidence: Quantity, + vapor_tendency_from_environmental_subsidence: Quantity, + t_tendency_from_environmental_subsidence: Quantity, + plume_dependent_constants: GF2020PlumeDependentConstants, + ): + self._zero_tendencies( + del_u_cloud_ensemble=del_u_cloud_ensemble, + del_v_cloud_ensemble=del_v_cloud_ensemble, + del_moist_static_energy_cloud_ensemble=del_moist_static_energy_cloud_ensemble, + del_t_cloud_ensemble=del_t_cloud_ensemble, + del_vapor_cloud_ensemble=del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=del_cloud_liquid_cloud_ensemble, + del_buoyancy_cloud_ensemble=del_buoyancy_cloud_ensemble, + moist_static_energy_tendency_from_environmental_subsidence=moist_static_energy_tendency_from_environmental_subsidence, + vapor_tendency_from_environmental_subsidence=vapor_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=t_tendency_from_environmental_subsidence, + ) + + if self.cumulus_parameterization_config.VERTICAL_DISCRETIZATION_OPTION in (0, 1): + self._convective_transport_of_momentum( + error_code=error_code, + cloud_top_level=cloud_top_level, + p_cloud_levels_forced=p_cloud_levels_forced, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + environment_massflux=environment_massflux, + u=u, + v=v, + u_cloud_levels=u_cloud_levels, + v_cloud_levels=v_cloud_levels, + u_c=u_c, + v_c=v_c, + u_c_downdraft=u_c_downdraft, + v_c_downdraft=v_c_downdraft, + del_u_cloud_ensemble=del_u_cloud_ensemble, + del_v_cloud_ensemble=del_v_cloud_ensemble, + epsilon_forced=epsilon_forced, + fp=self._fp, + fm=self._fm, + aa=self._aa, + bb=self._bb, + cc=self._cc, + ddu=self._ddu, + ddv=self._ddv, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + if self.cumulus_parameterization_config.VERTICAL_DISCRETIZATION_OPTION == 1: + self._tridiag( + m=cloud_top_level, + a=self._aa, + b=self._bb, + c=self._cc, + f=self._ddu, + error_code=error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._update_after_tridiag( + error_code=error_code, + cloud_top_level=cloud_top_level, + in_field=self._ddu, + out_field=del_u_cloud_ensemble, + wind=u, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._tridiag( + m=cloud_top_level, + a=self._aa, + b=self._bb, + c=self._cc, + f=self._ddv, + error_code=error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._update_after_tridiag( + error_code=error_code, + cloud_top_level=cloud_top_level, + in_field=self._ddv, + out_field=del_v_cloud_ensemble, + wind=v, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._convective_transport_of_mse_and_liquid_water( + error_code=error_code, + cloud_top_level=cloud_top_level, + p_cloud_levels_forced=p_cloud_levels_forced, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + cloud_moist_static_energy_forced=cloud_moist_static_energy_forced, + cloud_moist_static_energy_downdraft_forced=cloud_moist_static_energy_downdraft_forced, + environment_moist_static_energy_cloud_levels_forced=environment_moist_static_energy_cloud_levels_forced, + cloud_liquid_after_rain_forced=cloud_liquid_after_rain_forced, + melting=melting, + partition_liquid_ice=partition_liquid_ice, + epsilon_forced=epsilon_forced, + del_moist_static_energy_cloud_ensemble=del_moist_static_energy_cloud_ensemble, + moist_static_energy_tendency_from_environmental_subsidence=moist_static_energy_tendency_from_environmental_subsidence, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + self._convective_transport_of_water_vapor_and_condensates( + error_code=error_code, + cloud_top_level=cloud_top_level, + p_cloud_levels_forced=p_cloud_levels_forced, + geopotential_height_cloud_levels_forced=geopotential_height_cloud_levels_forced, + normalized_massflux_updraft_forced=normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=normalized_massflux_downdraft_forced, + mass_detrainment_updraft_forced=mass_detrainment_updraft_forced, + mass_detrainment_downdraft_forced=mass_detrainment_downdraft_forced, + c1d=c1d, + vapor_cloud_levels_forced=vapor_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=cloud_total_water_after_entrainment_forced, + cloud_total_water_after_entrainment_downdraft_forced=cloud_total_water_after_entrainment_downdraft_forced, + cloud_liquid_after_rain_forced=cloud_liquid_after_rain_forced, + condensate_to_fall_forced=condensate_to_fall_forced, + evaporate_in_downdraft_forced=evaporate_in_downdraft_forced, + epsilon_forced=epsilon_forced, + d_buoyancy_downdraft_forced=d_buoyancy_downdraft_forced, + del_cloud_liquid_cloud_ensemble=del_cloud_liquid_cloud_ensemble, + del_vapor_cloud_ensemble=del_vapor_cloud_ensemble, + vapor_tendency_from_environmental_subsidence=vapor_tendency_from_environmental_subsidence, + plume=plume_dependent_constants.PLUME_INDEX, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/finalize.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/finalize.py new file mode 100644 index 000000000..a89960f51 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/finalize.py @@ -0,0 +1,1465 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import NDSLRuntime, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, K, abs, computation, interval, max, min +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ +from ndsl.stencils.column_operations import column_min +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_ConvectionTracers, + FloatField_ConvectionTracers_Plume, + FloatField_Plume, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.locals import GF2020Locals +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables import ( + GlobalTable_saturation_tables, + saturation_specific_humidity, + saturation_specific_humidity_liquid_surface, +) +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable +from pyMoist.shared_incloud_processes import ice_fraction + + +def copy_from_cumulus_parameterization_state( + pbl_time_scale_from_cumulus_parameterization: FloatFieldIJ, + pbl_time_scale: FloatFieldIJ, + cape_removal_time_scale_from_cumulus_parameterization: FloatFieldIJ, + cape_removal_time_scale: FloatFieldIJ, + cloud_workfunction_0_from_cumulus_parameterization: FloatFieldIJ, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + cloud_workfunction_1_from_cumulus_parameterization: FloatFieldIJ, +): + """Copy fields from the cumulus parameterization state to the overarching model state. + + Args: + pbl_time_scale_from_cumulus_parameterization (FloatFieldIJ) + pbl_time_scale (FloatFieldIJ) + cape_removal_time_scale_from_cumulus_parameterization (FloatFieldIJ) + cape_removal_time_scale (FloatFieldIJ) + cloud_workfunction_0_from_cumulus_parameterization (FloatFieldIJ) + cloud_workfunction_0 (FloatFieldIJ) + cloud_workfunction_1 (FloatFieldIJ) + cloud_workfunction_1_from_cumulus_parameterization (FloatFieldIJ) + """ + with computation(FORWARD), interval(0, 1): + pbl_time_scale = pbl_time_scale_from_cumulus_parameterization + cape_removal_time_scale = cape_removal_time_scale_from_cumulus_parameterization + + cloud_workfunction_0 = cloud_workfunction_0_from_cumulus_parameterization + cloud_workfunction_1 = cloud_workfunction_1_from_cumulus_parameterization + + +def flag_computed_plumes_and_columns( + error_code: IntFieldIJ_Plume, + do_this_column: IntFieldIJ, +): + """Flag which plumes were computed, and which columns made it through the entire scheme + (error code = 0) for one or more columns + + Args: + error_code (IntFieldIJ_Plume): contains error codes from cumulus parameterization core + do_this_column (IntFieldIJ): mask which shows "successful" columns + """ + from __externals__ import ENABLE_DEEP, ENABLE_MID, ENABLE_SHALLOW + + with computation(FORWARD), interval(0, 1): + if ENABLE_SHALLOW == 0: + error_code[0, 0][cumulus_parameterization_constants.SHALLOW] = -99 + if ENABLE_MID == 0: + error_code[0, 0][cumulus_parameterization_constants.MID] = -99 + if ENABLE_DEEP == 0: + error_code[0, 0][cumulus_parameterization_constants.DEEP] = -99 + + # generate a mask (which has been initialized to zero in GF2020Setup) which flags only columns + # which made it through the entire cumulus parameterization scheme at one or more plumes + with computation(FORWARD), interval(0, 1): + plume = 0 + while plume < cumulus_parameterization_constants.NUMBER_OF_PLUMES: + if error_code[0, 0][plume] == 0: + do_this_column = 1 + plume += 1 + + +def check_vapor_mixing_ratio( + vapor_current: FloatField, + dvapordt: FloatField_Plume, + t_tendency_from_vapor: FloatField, + fix_out_vapor: FloatFieldIJ, + do_this_column: IntFieldIJ, +): + """Ensure that output water vapor mixing ratio is a reasonable value + + Args: + vapor_current (FloatField): vapor mixing ratio before GF2020CumulusParameterization call + dvapordt (FloatField_Plume): vapor tendency output from GF2020CumulusParameterization call + t_tendency_from_vapor (FloatField): temperature tendency from water vapor + fix_out_vapor (FloatFieldIJ): modification to "fix" water vapor mixing ratio + do_this_column (IntFieldIJ): mask which shows "successful" columns + """ + from __externals__ import DT_MOIST, k_end + + with computation(PARALLEL), interval(0, -1): + if do_this_column != 0: + t_tendency_from_vapor = ( + dvapordt[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dvapordt[0, 0, 0][cumulus_parameterization_constants.MID] + + dvapordt[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) + + distance = vapor_current + t_tendency_from_vapor * DT_MOIST + + with computation(FORWARD), interval(0, 1): + # prefill + fix_out_vapor = 1.0 + + if do_this_column != 0: + min_value, min_index = column_min(distance, 0, k_end - 1) + if min_value < 0.0: + if abs(t_tendency_from_vapor.at(K=min_index) * DT_MOIST) < constants.FLOAT_TINY: + fix_out_vapor = 0.999999 + else: + fix_out_vapor = ( + cumulus_parameterization_constants.smaller_qv - vapor_current.at(K=min_index) + ) / (t_tendency_from_vapor.at(K=min_index) * DT_MOIST) + fix_out_vapor = max(0.0, min(fix_out_vapor, 1.0)) + + +def feedback( + fix_out_vapor: FloatFieldIJ, + precip: FloatFieldIJ, + precip_from_cumulus_parameterization: FloatFieldIJ_Plume, + evaporation_sublimation_tendency: FloatField, + evaporation_sublimation_tendency_from_cumulus_parameterization: FloatField, + convective_precip_flux: FloatField, + convective_precip_flux_from_cumulus_parameterization: FloatField, + dtdt: FloatField, + dtdt_from_cumulus_parameterization: FloatField_Plume, + dvapordt: FloatField, + dvapordt_from_cumulus_parameterization: FloatField_Plume, + dcloudicedt: FloatField, + dcloudicedt_from_cumulus_parameterization: FloatField_Plume, + dudt: FloatField, + dudt_from_cumulus_parameterization: FloatField_Plume, + dvdt: FloatField, + dvdt_from_cumulus_parameterization: FloatField_Plume, + dlarge_scale_icedt: FloatField, + dlarge_scale_icedt_from_cumulus_parameterization: FloatField_Plume, + dconvective_icedt: FloatField, + dconvective_icedt_from_cumulus_parameterization: FloatField_Plume, + dlarge_scale_liquiddt: FloatField, + dlarge_scale_liquiddt_from_cumulus_parameterization: FloatField_Plume, + dconvective_liquiddt: FloatField, + dconvective_liquiddt_from_cumulus_parameterization: FloatField_Plume, + dlarge_scale_cloud_fractiondt: FloatField, + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization: FloatField_Plume, + dconvective_cloud_fractiondt: FloatField, + dconvective_cloud_fractiondt_from_cumulus_parameterization: FloatField_Plume, + dbuoyancydt_from_cumulus_parameterization: FloatField_Plume, + dbuoyancydt: FloatField, + do_this_column: IntFieldIJ, +): + """Feedback output from the cumulus parameterization core to the local state. + This cannot be fed straight into the model as the local state and model state have opposing k-axis + directions, and further calculations may be performed on one or more of these outputs before this flip + and final exchange is performed. + + These could be condensed, but have been kept separate for readibility and clarity. + + Args: + fix_out_vapor (FloatFieldIJ) + evaporation_sublimation_tendency (FloatField) + evaporation_sublimation_tendency_from_cumulus_parameterization (FloatField) + convective_precip_flux (FloatField) + convective_precip_flux_from_cumulus_parameterization (FloatField_Plume) + dtdt (FloatField) + dtdt_from_cumulus_parameterization (FloatField_Plume) + dvapordt (FloatField) + dvapordt_from_cumulus_parameterization (FloatField_Plume) + dcloudicedt (FloatField) + dcloudicedt_from_cumulus_parameterization (FloatField_Plume) + dudt (FloatField) + dudt_from_cumulus_parameterization (FloatField_Plume) + dvdt (FloatField) + dvdt_from_cumulus_parameterization (FloatField_Plume) + dlarge_scale_icedt (FloatField) + dlarge_scale_icedt_from_cumulus_parameterization (FloatField_Plume) + dconvective_icedt (FloatField) + dconvective_icedt_from_cumulus_parameterization (FloatField_Plume) + dlarge_scale_liquiddt (FloatField) + dlarge_scale_liquiddt_from_cumulus_parameterization (FloatField_Plume) + dconvective_liquiddt (FloatField) + dconvective_liquiddt_from_cumulus_parameterization (FloatField_Plume) + dlarge_scale_cloud_fractiondt (FloatField) + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization (FloatField_Plume) + dconvective_cloud_fractiondt (FloatField) + dconvective_cloud_fractiondt_from_cumulus_parameterization (FloatField_Plume) + dbuoyancydt_from_cumulus_parameterization (FloatField_Plume) + dbuoyancydt (FloatField) + do_this_column (IntFieldIJ) + """ + from __externals__ import APPLY_SUBSIDENCE_MICROPHYSICS, CONVECTION_TRACER, USE_MOMENTUM_TRANSPORT + + with computation(FORWARD), interval(0, 1): + if do_this_column != 0: + precip = ( + precip_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.SHALLOW] + + precip_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.MID] + + precip_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + + with computation(PARALLEL), interval(...): + # combining effects of shallow + mid + deep convection + if do_this_column != 0: + # feedback the tendencies from convection + dtdt = ( + dtdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dtdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dtdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + dvapordt = ( + dvapordt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dvapordt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dvapordt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + dcloudicedt = ( + dcloudicedt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dcloudicedt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dcloudicedt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + evaporation_sublimation_tendency = ( + evaporation_sublimation_tendency_from_cumulus_parameterization * fix_out_vapor + ) # already contains deep and mid amounts. + + # precip flux is only computed for deep plume + convective_precip_flux = ( + convective_precip_flux_from_cumulus_parameterization * fix_out_vapor + ) # ice/liquid precip flux of the deep plume + + if USE_MOMENTUM_TRANSPORT > 0: + dudt = ( + dudt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dudt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dudt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + dvdt = ( + dvdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dvdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dvdt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + + if APPLY_SUBSIDENCE_MICROPHYSICS == 1: + dlarge_scale_icedt = ( + dlarge_scale_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dlarge_scale_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dlarge_scale_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + dconvective_icedt = ( + dconvective_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dconvective_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dconvective_icedt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + dlarge_scale_liquiddt = ( + dlarge_scale_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dlarge_scale_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dlarge_scale_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + dconvective_liquiddt = ( + dconvective_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dconvective_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dconvective_liquiddt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + dlarge_scale_cloud_fractiondt = ( + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + dconvective_cloud_fractiondt = ( + dconvective_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + dconvective_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID + ] + + dconvective_cloud_fractiondt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP + ] + ) * fix_out_vapor + + with computation(PARALLEL), interval(...): + if do_this_column != 0 and CONVECTION_TRACER == 1: + dbuoyancydt = ( + dbuoyancydt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.SHALLOW] + + dbuoyancydt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.MID] + + dbuoyancydt_from_cumulus_parameterization[0, 0, 0][cumulus_parameterization_constants.DEEP] + ) * fix_out_vapor + + +def feedback_tracers( + tracer: Int, + fix_out_vapor: FloatFieldIJ, + dconvection_tracersdt: FloatField_ConvectionTracers, + dconvection_tracersdt_from_cumulus_parameterization: FloatField_ConvectionTracers_Plume, + chemistry_tracers_from_cumulus_parameterization: FloatField_ConvectionTracers, + do_this_column: IntFieldIJ, +): + """See documentation for "feedback". This is a separate stencil to accomodate the tracer data dimension. + + Args: + tracer (Int) + fix_out_vapor (FloatFieldIJ) + dconvection_tracersdt (FloatField_ConvectionTracers) + dconvection_tracersdt_from_cumulus_parameterization (FloatField_ConvectionTracers) + chemistry_tracers_from_cumulus_parameterization (FloatField_ConvectionTracers) + do_this_column (IntFieldIJ) + """ + from __externals__ import DT_MOIST, USE_TRACER_TRANSPORT, k_end + + with computation(PARALLEL), interval(...): + if do_this_column != 0 and USE_TRACER_TRANSPORT == 1: + dconvection_tracersdt[0, 0, 0][tracer] = ( + dconvection_tracersdt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.SHALLOW, tracer + ] + + dconvection_tracersdt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.MID, tracer + ] + + dconvection_tracersdt_from_cumulus_parameterization[0, 0, 0][ + cumulus_parameterization_constants.DEEP, tracer + ] + ) * fix_out_vapor + + with computation(PARALLEL), interval(0, -1): + if do_this_column != 0 and USE_TRACER_TRANSPORT == 1: + # constrain positivity for tracers + distance = ( + chemistry_tracers_from_cumulus_parameterization[0, 0, 0][tracer] + + dconvection_tracersdt[0, 0, 0][tracer] * DT_MOIST + ) + + with computation(FORWARD), interval(0, 1): + # ensure temporary is initialized properly + fix_tracers: FloatFieldIJ = 0.0 + + if do_this_column != 0 and USE_TRACER_TRANSPORT == 1: + # fixer for mass of tracer + min_value, min_index = column_min(distance, 0, k_end - 1) + if min_value < 0.0: + if ( + abs(dconvection_tracersdt.at(K=min_index, ddim=[tracer]) * DT_MOIST) + < constants.FLOAT_TINY + ): + fix_tracers = 0.999999 + else: + fix_tracers = ( + constants.FLOAT_TINY + - chemistry_tracers_from_cumulus_parameterization.at(K=min_index, ddim=[tracer]) + ) / (dconvection_tracersdt.at(K=min_index, ddim=[tracer]) * DT_MOIST) + if fix_tracers > 1.0 or fix_tracers < 0.0: + fix_tracers = 0.0 + + with computation(PARALLEL), interval(0, -1): + if do_this_column != 0 and USE_TRACER_TRANSPORT == 1: + # apply fixer + dconvection_tracersdt[0, 0, 0][tracer] = fix_tracers * dconvection_tracersdt[0, 0, 0][tracer] + + +def cloud_workfunction_output( + precip: FloatFieldIJ_Plume, + fix_out_vapor: FloatFieldIJ, + cloud_workfunction_2: FloatFieldIJ, + cloud_workfunction_3: FloatFieldIJ, +): + """Fill cloud workfunction 2 and 3 with data from the cumulus parameterization core + + Args: + precip (FloatFieldIJ_Plume) + fix_out_vapor (FloatFieldIJ) + cloud_workfunction_2 (FloatFieldIJ) + cloud_workfunction_3 (FloatFieldIJ) + """ + with computation(FORWARD), interval(0, 1): + cloud_workfunction_2 = precip[0, 0][cumulus_parameterization_constants.MID] * fix_out_vapor + cloud_workfunction_3 = precip[0, 0][cumulus_parameterization_constants.DEEP] * fix_out_vapor + + +def prefill_entrainment( + lateral_entrainment_rate_shallow: FloatField, + lateral_entrainment_rate_mid: FloatField, + lateral_entrainment_rate_deep: FloatField, +): + with computation(PARALLEL), interval(...): + lateral_entrainment_rate_shallow = constants.MAPL_UNDEF + lateral_entrainment_rate_mid = constants.MAPL_UNDEF + lateral_entrainment_rate_deep = constants.MAPL_UNDEF + + +def feed_3d_model( + do_this_column: IntFieldIJ, + precip: FloatFieldIJ, + convective_precipitation_GF: FloatFieldIJ, + evaporation_sublimation_tendency_from_cumulus_parameterization: FloatField, + evaporation_sublimation_tendency: FloatField, + convective_precip_flux_from_cumulus_parameterization: FloatField, + convective_precip_flux: FloatField, + dconvection_tracersdt: FloatField_ConvectionTracers, + convection_tracers: FloatField_ConvectionTracers, +): + """Update the model state with feedback from the plume-independent cumulus parameterization core fields. + + Args: + do_this_column (IntFieldIJ) + precip (FloatFieldIJ) + convective_precipitation_GF (FloatFieldIJ) + evaporation_sublimation_tendency_from_cumulus_parameterization (FloatField) + evaporation_sublimation_tendency (FloatField) + convective_precip_flux_from_cumulus_parameterization (FloatField) + convective_precip_flux (FloatField) + dconvection_tracersdt (FloatField_ConvectionTracers) + convection_tracers (FloatField_ConvectionTracers) + """ + from __externals__ import DT_MOIST, FEED_3D_MODEL, ITEST, USE_TRACER_TRANSPORT, k_end + + with computation(FORWARD), interval(0, 1): + if FEED_3D_MODEL == True and do_this_column != 0: + # convective precip rate: mm/s = kg m-2 s-1 + if ITEST == 0: + convective_precipitation_GF = 0.0 + else: + convective_precipitation_GF = precip + + with computation(PARALLEL), interval(...): + if FEED_3D_MODEL == True and do_this_column != 0: + # sublimation/evaporation tendencies (kg/kg/s) + evaporation_sublimation_tendency = ( + evaporation_sublimation_tendency_from_cumulus_parameterization.at(K=k_end - K) + ) + # preciptation fluxes (kg/kg/s) + convective_precip_flux = convective_precip_flux_from_cumulus_parameterization.at(K=k_end - K) + + if USE_TRACER_TRANSPORT == 1: + # update tracer mass mixing ratios + tracer = 0 + while tracer < constants.NUMBER_OF_TRACERS: + convection_tracers[0, 0, 0][tracer] = convection_tracers[0, 0, 0][ + tracer + ] + DT_MOIST * dconvection_tracersdt.at(K=k_end - K, ddim=[tracer]) + + # final check for negative tracer mass mixing ratio + convection_tracers[0, 0, 0][tracer] = max( + convection_tracers[0, 0, 0][tracer], constants.FLOAT_TINY + ) + tracer += 1 + + +def feed_3d_model_from_plumes( + plume: Int, + do_this_column: IntFieldIJ, + dz: FloatField, + air_density: FloatField, + p_flipped: FloatField, + vapor_flipped: FloatField, + dcloudicedt: FloatField, + lateral_entrainment_rate_shallow: FloatField, + lateral_entrainment_rate_mid: FloatField, + lateral_entrainment_rate_deep: FloatField, + entrainment_rate_from_cumulus_parameterization: FloatField_Plume, + error_code_from_cumulus_parameterization: IntFieldIJ_Plume, + cloud_top_level_from_cumulus_parameterization: IntFieldIJ_Plume, + t_updraft_from_cumulus_parameterization: FloatField_Plume, + mass_detrainment_updraft_forced_from_cumulus_parameterization: FloatField_Plume, + mass_entrainment_updraft_forced_from_cumulus_parameterization: FloatField_Plume, + normalized_massflux_updraft_forced_from_cumulus_parameterization: FloatField_Plume, + cloud_liquid_after_rain_forced_from_cumulus_parameterization: FloatField_Plume, + condensate_to_fall_forced_from_cumulus_parameterization: FloatField_Plume, + epsilon_forced_from_cumulus_parameterization: FloatFieldIJ_Plume, + evaporate_in_downdraft_forced_from_cumulus_parameterization: FloatField_Plume, + total_water_flux_deep_convection: FloatField, + convective_condensate_source: FloatField, + mass_flux_cloud_base: FloatField, + mass_flux_deep_updraft_detrained: FloatField, + mass_flux_deep_updraft_interface: FloatField, + entrainment_parameter: FloatField, + updraft_vertical_velocity: FloatField, + convective_condensate_grid_mean: FloatField, + convective_precipitation_RAS: FloatField, + updraft_areal_fraction: FloatField, + esw: GlobalTable_saturation_tables, + estlqu: Float, +): + """Update the model state with feedback from the cumulus parameterization core. + + Args: + plume (Float) + do_this_column (IntFieldIJ) + dz (FloatField) + air_density (FloatField) + p_flipped (FloatField) + vapor_flipped (FloatField) + dcloudicedt (FloatField) + lateral_entrainment_rate_shallow (FloatField) + lateral_entrainment_rate_mid (FloatField) + lateral_entrainment_rate_deep (FloatField) + error_code_from_cumulus_parameterization (FloatFieldIJ_Plume) + cloud_top_level_from_cumulus_parameterization (FloatFieldIJ_Plume) + lateral_entrainment_rate_from_cumulus_parameterization (FloatField_Plume) + t_updraft_from_cumulus_parameterization (FloatField_Plume) + mass_detrainment_updraft_forced_from_cumulus_parameterization (FloatField_Plume) + mass_entrainment_updraft_forced_from_cumulus_parameterization (FloatField_Plume) + normalized_massflux_updraft_forced_from_cumulus_parameterization (FloatField_Plume) + cloud_liquid_after_rain_forced_from_cumulus_parameterization (FloatField_Plume) + condensate_to_fall_forced_from_cumulus_parameterization (FloatField_Plume) + epsilon_forced_from_cumulus_parameterization (FloatFieldIJ_Plume) + evaporate_in_downdraft_forced_from_cumulus_parameterization (FloatField_Plume) + total_water_flux_deep_convection (FloatField) + convective_condensate_source (FloatField) + mass_flux_cloud_base (FloatField) + mass_flux_deep_updraft_detrained (FloatField) + mass_flux_deep_updraft_interface (FloatField) + entrainment_parameter (FloatField) + updraft_vertical_velocity (FloatField) + convective_condensate_grid_mean (FloatField) + convective_precipitation_RAS (FloatField) + esw (GlobalTable_saturation_tables) + estlqu (Float) + """ + from __externals__ import DT_MOIST, FEED_3D_MODEL, k_end + + with computation(FORWARD), interval(...): + if ( + FEED_3D_MODEL == True + and K >= k_end - cloud_top_level_from_cumulus_parameterization[0, 0][plume] - 1 + and error_code_from_cumulus_parameterization[0, 0][plume] == 0 + ): + # deep convective total water flux + # assumes .033 fractional area + saturation_specific_humidity_updraft, _ = saturation_specific_humidity_liquid_surface( + esw=esw, + lqu=estlqu, + t=t_updraft_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]), + p=p_flipped.at(K=k_end - K), + pressure_correction=True, + ) + + saturation_specific_humidity_updraft = ( + saturation_specific_humidity_updraft + + cloud_liquid_after_rain_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + / 0.033 + ) + total_water_flux_deep_convection[0, 0, 1] = total_water_flux_deep_convection[ + 0, 0, 1 + ] + normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) * ( + saturation_specific_humidity_updraft - vapor_flipped.at(K=k_end - K) + ) + + with computation(PARALLEL), interval(...): + if ( + FEED_3D_MODEL == True + and K >= k_end - cloud_top_level_from_cumulus_parameterization[0, 0][plume] - 1 + and error_code_from_cumulus_parameterization[0, 0][plume] == 0 + ): + if plume == cumulus_parameterization_constants.SHALLOW: + # export entrainment rates used by GF + lateral_entrainment_rate_shallow = entrainment_rate_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + if plume == cumulus_parameterization_constants.MID: + # export entrainment rates used by GF + lateral_entrainment_rate_mid = entrainment_rate_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + if plume == cumulus_parameterization_constants.DEEP: + # export entrainment rates used by GF + lateral_entrainment_rate_deep = entrainment_rate_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + # special treatment for convective_condensate_source + # units = 'kg m-2 s-1', + # dcloudicedt contains contributions from all plumes, so no need to accumulate across levels + convective_condensate_source = dcloudicedt.at(K=k_end - K) * dz * air_density + + # detraining_mass_flux + # units = 'kg m-2 s-1' + mass_flux_deep_updraft_detrained = ( + mass_flux_deep_updraft_detrained + + mass_detrainment_updraft_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + ) + + # cloud_base_mass_flux + # units = 'kg m-2 s-1' + mass_flux_cloud_base = ( + mass_flux_cloud_base + + normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + ) + + if ( + normalized_massflux_updraft_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + > 1.0e-6 + ): + # entrainment parameter + # units ='m-1', + entrainment_parameter = entrainment_parameter + ( + mass_entrainment_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + / ( + dz + * normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + ) + ) + + # # updraft_vertical_velocity + # # units = 'hPa s-1', + updraft_vertical_velocity = -0.2 # hPa/s => 4 m/s + + # convective_condensate_grid_mean + # units ='kg kg-1' + convective_condensate_grid_mean = ( + convective_condensate_grid_mean + + cloud_liquid_after_rain_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + ) + + # not using progno-cloud to calculate the precip from the convective column + # if CNV_PRC3 will be send to progno-cloud, set CNPCPRATE = zero + # 'convective_precipitation_from_GF',UNITS = 'kg m-2 s-1', + # JAN/17/2017 : the units above are wrong. The correct are kg[precip water]/kg[air] + convective_precipitation_RAS = convective_precipitation_RAS + ( + condensate_to_fall_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + + epsilon_forced_from_cumulus_parameterization[0, 0][plume] + * evaporate_in_downdraft_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + ) * DT_MOIST / (dz * air_density) + + # updraft_area_fraction + if ( + normalized_massflux_updraft_forced_from_cumulus_parameterization.at(K=k_end - K, ddim=[plume]) + > 1.0e-6 + ): + updraft_areal_fraction = 0.033 + + with computation(BACKWARD), interval(...): + # this must be done in a separate computation because the offset write is incompatable with PARALLEL + if FEED_3D_MODEL == True: + if ( + K >= k_end - cloud_top_level_from_cumulus_parameterization[0, 0][plume] - 1 + and error_code_from_cumulus_parameterization[0, 0][plume] == 0 + ): + # convective mass flux - only updraft + # units = 'kg m-2 s-1' + mass_flux_deep_updraft_interface[0, 0, 1] = mass_flux_deep_updraft_interface[ + 0, 0, 1 + ] + normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[plume] + ) + + +def update_convection_tracer( + land_fraction: FloatFieldIJ, + dbuoyancydt: FloatField, + convection_tracer: FloatField, +): + """Update the convection tracer field with output from the cumulus parameterizaiton core. + + Args: + land_fraction (FloatFieldIJ) + dbuoyancydt (FloatField) + convection_tracer (FloatField) + """ + from __externals__ import CONVECTION_TRACER, DT_MOIST, k_end + + # cold pool/"convection tracer" + with computation(PARALLEL), interval(...): + if CONVECTION_TRACER == 1: + cold_pool_timescale = land_fraction * (6.0 / 3600.0) + (1 - land_fraction) * (6.0 / 3600.0) + + # sink term (exp decay 1h) + sink = DT_MOIST * abs(convection_tracer) / cold_pool_timescale + # source term + # downdraft detrainment of buoyancy [ J/kg s^{-1}] + # negative sign => source for updraft lifting + source = DT_MOIST * min(0.0, dbuoyancydt.at(K=k_end - K)) + + # 'continuity' equation = ADV + SRC - SINK + convection_tracer = convection_tracer + source - sink + + +def update_outputs( + p_flipped: FloatField, + error_code_from_cumulus_parameterization: IntFieldIJ_Plume, + cloud_top_level_from_cumulus_parameterization: IntFieldIJ_Plume, + cloud_base_mass_flux_modified_from_cumulus_parameterization: FloatFieldIJ_Plume, + normalized_massflux_updraft_forced_from_cumulus_parameterization: FloatField_Plume, + normalized_massflux_downdraft_forced_from_cumulus_parameterization: FloatField_Plume, + epsilon_forced_from_cumulus_parameterization: FloatFieldIJ_Plume, + scale_dependence_factor_from_cumulus_parameterizaiton: FloatFieldIJ_Plume, + pressure_shallow_convective_cloud_top: FloatFieldIJ, + pressure_mid_convective_cloud_top: FloatFieldIJ, + pressure_deep_convective_cloud_top: FloatFieldIJ, + mass_flux_cloud_base_shallow: FloatFieldIJ, + mass_flux_cloud_base_mid: FloatFieldIJ, + mass_flux_cloud_base_deep: FloatFieldIJ, + mass_flux_shallow: FloatField, + mass_flux_mid: FloatField, + mass_flux_deep_updraft: FloatField, + mass_flux_deep_downdraft: FloatField, + sigma_mid: FloatFieldIJ, + sigma_deep: FloatFieldIJ, +): + """Update various outputs in preparation for a push back to the model state. + + Args: + p_flipped (FloatField) + error_code_from_cumulus_parameterization (IntFieldIJ_Plume) + cloud_top_level_from_cumulus_parameterization (IntFieldIJ_Plume) + cloud_base_mass_flux_modified_from_cumulus_parameterization (FloatFieldIJ_Plume) + normalized_massflux_updraft_forced_from_cumulus_parameterization (FloatField_Plume) + normalized_massflux_downdraft_forced_from_cumulus_parameterization (FloatField_Plume) + epsilon_forced_from_cumulus_parameterization (FloatFieldIJ_Plume) + scale_dependence_factor_from_cumulus_parameterizaiton (FloatFieldIJ) + pressure_shallow_convective_cloud_top (FloatFieldIJ) + pressure_mid_convective_cloud_top (FloatFieldIJ) + pressure_deep_convective_cloud_top (FloatFieldIJ) + mass_flux_cloud_base_shallow (FloatFieldIJ) + mass_flux_cloud_base_mid (FloatFieldIJ) + mass_flux_cloud_base_deep (FloatFieldIJ) + mass_flux_shallow (FloatField) + mass_flux_mid (FloatField) + mass_flux_deep_updraft (FloatField) + mass_flux_deep_downdraft (FloatField) + sigma_mid (FloatFieldIJ) + sigma_deep (FloatFieldIJ) + """ + from __externals__ import ENABLE_DEEP, ENABLE_MID, ENABLE_SHALLOW, k_end + + with computation(FORWARD), interval(0, 1): + pressure_shallow_convective_cloud_top = constants.MAPL_UNDEF + pressure_mid_convective_cloud_top = constants.MAPL_UNDEF + pressure_deep_convective_cloud_top = constants.MAPL_UNDEF + + with computation(FORWARD), interval(0, 1): + if ( + ENABLE_SHALLOW == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.SHALLOW] + == 0 + ): + pressure_shallow_convective_cloud_top = p_flipped.at( + K=cloud_top_level_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + ) + mass_flux_cloud_base_shallow = cloud_base_mass_flux_modified_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + + if ( + ENABLE_MID == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.MID] == 0 + ): + pressure_mid_convective_cloud_top = p_flipped.at( + K=cloud_top_level_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.MID] + ) + mass_flux_cloud_base_mid = cloud_base_mass_flux_modified_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.MID + ] + sigma_mid = scale_dependence_factor_from_cumulus_parameterizaiton[0, 0][ + cumulus_parameterization_constants.MID + ] + + if ( + ENABLE_DEEP == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.DEEP] == 0 + ): + pressure_deep_convective_cloud_top = p_flipped.at( + K=cloud_top_level_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.DEEP] + ) + mass_flux_cloud_base_deep = cloud_base_mass_flux_modified_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.DEEP + ] + sigma_deep = scale_dependence_factor_from_cumulus_parameterizaiton[0, 0][ + cumulus_parameterization_constants.DEEP + ] + + with computation(PARALLEL), interval(...): + if ( + ENABLE_SHALLOW == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.SHALLOW] + == 0 + ): + mass_flux_shallow = normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[cumulus_parameterization_constants.SHALLOW] + ) + + if ( + ENABLE_MID == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.MID] == 0 + ): + mass_flux_mid = normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[cumulus_parameterization_constants.MID] + ) + + if ( + ENABLE_DEEP == True + and error_code_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.DEEP] == 0 + ): + mass_flux_deep_updraft = normalized_massflux_updraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[cumulus_parameterization_constants.DEEP] + ) + mass_flux_deep_downdraft = ( + normalized_massflux_downdraft_forced_from_cumulus_parameterization.at( + K=k_end - K, ddim=[cumulus_parameterization_constants.DEEP] + ) + * epsilon_forced_from_cumulus_parameterization[0, 0][cumulus_parameterization_constants.DEEP] + ) + + with computation(FORWARD), interval(0, 1): + # for output purposes, error_code=0 (convection scheme runs without error) will be changed to 1 + # and all other internal errors are forced to zero + plume = 0 + while plume < cumulus_parameterization_constants.NUMBER_OF_PLUMES: + if error_code_from_cumulus_parameterization[0, 0][plume] == 0: + error_code_from_cumulus_parameterization[0, 0][plume] = 1 + if error_code_from_cumulus_parameterization[0, 0][plume] > 1: + error_code_from_cumulus_parameterization[0, 0][plume] = 0 + plume += 1 + + +def update_tendencies( + dtdt: FloatField, + dvapordt: FloatField, + dudt: FloatField, + dvdt: FloatField, + dtdt_deep_convection: FloatField, + dvapordt_deep_convection: FloatField, + dudt_deep_convection: FloatField, + dvdt_deep_convection: FloatField, +): + """Push tendency values computed in previous stencils (in the finalize class) back to the model state. + + Args: + dtdt (FloatField) + dvapordt (FloatField) + dudt (FloatField) + dvdt (FloatField) + dtdt_deep_convection (FloatField) + dvapordt_deep_convection (FloatField) + dudt_deep_convection (FloatField) + dvdt_deep_convection (FloatField) + """ + from __externals__ import k_end + + with computation(PARALLEL), interval(...): + # tendencies + dtdt_deep_convection = dtdt.at(K=k_end - K) + dvapordt_deep_convection = dvapordt.at(K=k_end - K) + dudt_deep_convection = dudt.at(K=k_end - K) + dvdt_deep_convection = dvdt.at(K=k_end - K) + + +def update_convection_codes( + error_code_from_cumulus_parameterization: IntFieldIJ_Plume, + convection_code_shallow: FloatFieldIJ, + convection_code_mid: FloatFieldIJ, + convection_code_deep: FloatFieldIJ, +): + with computation(FORWARD), interval(0, 1): + # error codes + convection_code_shallow = error_code_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.SHALLOW + ] + convection_code_mid = error_code_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.MID + ] + convection_code_deep = error_code_from_cumulus_parameterization[0, 0][ + cumulus_parameterization_constants.DEEP + ] + + +def update_state_with_tendencies( + convection_fraction: FloatFieldIJ, + surface_type: FloatFieldIJ, + u: FloatField, + v: FloatField, + vapor: FloatField, + t: FloatField, + p: FloatField, + p_kappa: FloatField, + mass: FloatField, + mass_flux_deep_updraft_detrained: FloatField, + mass_flux_deep_updraft_interface: FloatField, + total_cumulative_mass_flux_interface: FloatField, + total_detraining_mass_flux: FloatField, + dudt_deep_convection: FloatField, + dvdt_deep_convection: FloatField, + dvapordt_deep_convection: FloatField, + dtdt_deep_convection: FloatField, + dliquiddt_deep_convection: FloatField, + dicedt_deep_convection: FloatField, + dcloudfractiondt_deep_convection: FloatField, + convective_condensate_source: FloatField, + evaporation_sublimation_tendency: FloatField, + convective_precip_flux: FloatField, + sublimation_of_convective_precipitation: FloatField, + evaporation_of_convective_precipitation: FloatField, + ice_fraction_in_convective_tower: FloatField, + ice_precip_flux_interface: FloatField, + liquid_precip_flux_interface: FloatField, + convective_liquid: FloatField, + convective_ice: FloatField, + convective_cloud_fraction: FloatField, + convective_rainwater_source: FloatField, + convective_precipitation_RAS: FloatField, + ese: GlobalTable_saturation_tables, + esx: GlobalTable_saturation_tables, +): + """Update the model state (excluding the few fields which have already been updated in earlier + stencils) with the output from the cumulus parameterization core. + + Containts a call to saturation_specific_humidity, which is techincally a port of the GEOS_QSAT function. + In fortran + + Args: + convection_fraction (FloatFieldIJ) + surface_type (FloatFieldIJ) + u (FloatField) + v (FloatField) + vapor (FloatField) + t (FloatField) + p (FloatField) + p_kappa (FloatField) + mass (FloatField) + mass_flux_deep_updraft_detrained (FloatField) + mass_flux_deep_updraft_interface (FloatField) + total_cumulative_mass_flux_interface (FloatField) + total_detraining_mass_flux (FloatField) + dudt_deep_convection (FloatField) + dvdt_deep_convection (FloatField) + dvapordt_deep_convection (FloatField) + dtdt_deep_convection (FloatField) + dliquiddt_deep_convection (FloatField) + dicedt_deep_convection (FloatField) + dcloudfractiondt_deep_convection (FloatField) + convective_condensate_source (FloatField) + evaporation_sublimation_tendency (FloatField) + convective_precip_flux (FloatField) + sublimation_of_convective_precipitation (FloatField) + evaporation_of_convective_precipitation (FloatField) + ice_fraction_in_convective_tower (FloatField) + ice_precip_flux_interface (FloatField) + liquid_precip_flux_interface (FloatField) + convective_liquid (FloatField) + convective_ice (FloatField) + convective_cloud_fraction (FloatField) + convective_rainwater_source (FloatField) + convective_precipitation_RAS (FloatField) + ese (GlobalTable_saturation_tables) + esx (GlobalTable_saturation_tables) + """ + from __externals__ import DT_MOIST, FIX_CONVECTIVE_CLOUD, SCLM_DEEP + + with computation(PARALLEL), interval(...): + u = u + dudt_deep_convection * DT_MOIST + v = v + dvdt_deep_convection * DT_MOIST + vapor = vapor + dvapordt_deep_convection * DT_MOIST + t = t + dtdt_deep_convection * DT_MOIST + + # update deep cumulus liquid/ice/cloud fraction tendencies + fraction_ice = ice_fraction(t, convection_fraction, surface_type) + condensate_per_mass = convective_condensate_source / mass + dliquiddt_deep_convection = (1.0 - fraction_ice) * condensate_per_mass + dicedt_deep_convection = fraction_ice * condensate_per_mass + dcloudfractiondt_deep_convection = mass_flux_deep_updraft_detrained * SCLM_DEEP / mass + + # sublimation/evaporation tendencies (kg/kg/s) + sublimation_of_convective_precipitation = evaporation_sublimation_tendency * fraction_ice + evaporation_of_convective_precipitation = evaporation_sublimation_tendency * (1.0 - fraction_ice) + + with computation(FORWARD), interval(...): + # preciptation fluxes (kg/kg/s) + ice_precip_flux_interface[0, 0, 1] = convective_precip_flux * fraction_ice + liquid_precip_flux_interface[0, 0, 1] = convective_precip_flux * (1.0 - fraction_ice) + + with computation(PARALLEL), interval(...): + # add liquid/ice/cloud fraction tendencies + convective_liquid = convective_liquid + dliquiddt_deep_convection * DT_MOIST + convective_ice = convective_ice + dicedt_deep_convection * DT_MOIST + convective_cloud_fraction = max( + min(convective_cloud_fraction + dcloudfractiondt_deep_convection * DT_MOIST, 1.0), 0.0 + ) + + ice_fraction_in_convective_tower = fraction_ice + + # fix convective cloud fraction + if FIX_CONVECTIVE_CLOUD == True: + saturation_humidity, _ = saturation_specific_humidity(t, p, ese, esx) + + if convective_cloud_fraction < 1.0: + modification = (vapor - saturation_humidity * convective_cloud_fraction) / ( + 1.0 - convective_cloud_fraction + ) + min_saturation_humidity = 0.001 + if ( + modification - min_saturation_humidity * saturation_humidity + ) < 0.0 and convective_cloud_fraction > 0.0: + convective_cloud_fraction = (vapor - min_saturation_humidity * saturation_humidity) / ( + saturation_humidity * (1.0 - min_saturation_humidity) + ) + # if a suitable environment relative humidity cannot be made then destroy anvil + if convective_cloud_fraction < 0.0: + convective_cloud_fraction = 0.0 + dliquiddt_deep_convection = dliquiddt_deep_convection - (convective_liquid) / DT_MOIST + dicedt_deep_convection = dicedt_deep_convection - (convective_ice) / DT_MOIST + dvapordt_deep_convection = ( + dvapordt_deep_convection + (convective_liquid + convective_ice) / DT_MOIST + ) + vapor = vapor + (convective_liquid + convective_ice) + modification = ( + constants.MAPL_ALHL * convective_liquid + constants.MAPL_ALHS * convective_ice + ) / constants.MAPL_CP + dtdt_deep_convection = dtdt_deep_convection - modification / DT_MOIST + t = t - modification + convective_liquid = 0.0 + convective_ice = 0.0 + + total_cumulative_mass_flux_interface = ( + total_cumulative_mass_flux_interface + mass_flux_deep_updraft_interface + ) + total_detraining_mass_flux = total_detraining_mass_flux + mass_flux_deep_updraft_detrained + convective_rainwater_source = convective_precipitation_RAS / DT_MOIST + + +class GF2020Finalize(NDSLRuntime): + """This class performs the entire finalization sequence for the GF2020 convection parameterization scheme + + In the source Fortran codee, this code is split across three subroutines nested as follows: + + - GF_Run + - GF2020_INTERFACE + - GF2020_DRV beginning at the end (outside) of the of the plume loop (after the call to CUP_gf) + + This python implementation simplifies this structure by bringing all setup calculations to the same level. + An effort has been made to reduce duplicate/unnecessary locals where possible, but some have been retained + for the sake of readibility. + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + cumulus_parameterization_config: GF2020CumulusParameterizationConfig, + saturation_tables: SaturationVaporPressureTable, + ): + super().__init__(stencil_factory) + + # make status of plumes visible at runtime + self._plume_status = [ + cumulus_parameterization_config.ENABLE_SHALLOW, + cumulus_parameterization_config.ENABLE_MID, + cumulus_parameterization_config.ENABLE_DEEP, + ] + + # make saturation table data visible at runtime + # NOTE: this is an orchestration workaround. Direct call to + # `self.tables.X` fails closure capture for + # argument reconstruction at call time + self._ese = saturation_tables.ese + self._esw = saturation_tables.esw + self._esx = saturation_tables.esx + self._estfrz = saturation_tables.frz + self._estlqu = saturation_tables.lqu + + # construct stencils + self._copy_from_cumulus_parameterization_state = stencil_factory.from_dims_halo( + func=copy_from_cumulus_parameterization_state, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._flag_computed_plumes_and_columns = stencil_factory.from_dims_halo( + func=flag_computed_plumes_and_columns, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ENABLE_SHALLOW": cumulus_parameterization_config.ENABLE_SHALLOW, + "ENABLE_MID": cumulus_parameterization_config.ENABLE_MID, + "ENABLE_DEEP": cumulus_parameterization_config.ENABLE_DEEP, + }, + ) + + self._check_vapor_mixing_ratio = stencil_factory.from_dims_halo( + func=check_vapor_mixing_ratio, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DT_MOIST": config.DT_MOIST}, + ) + + self._feedback = stencil_factory.from_dims_halo( + func=feedback, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_MOMENTUM_TRANSPORT": config.USE_MOMENTUM_TRANSPORT, + "APPLY_SUBSIDENCE_MICROPHYSICS": config.APPLY_SUBSIDENCE_MICROPHYSICS, + "CONVECTION_TRACER": config.CONVECTION_TRACER, + }, + ) + + self._feedback_tracers = stencil_factory.from_dims_halo( + func=feedback_tracers, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_TRACER_TRANSPORT": config.USE_TRACER_TRANSPORT, + "DT_MOIST": config.DT_MOIST, + }, + ) + + self._cloud_workfunction_output = stencil_factory.from_dims_halo( + func=cloud_workfunction_output, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._prefill_entrainment = stencil_factory.from_dims_halo( + func=prefill_entrainment, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._feed_3d_model = stencil_factory.from_dims_halo( + func=feed_3d_model, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "FEED_3D_MODEL": cumulus_parameterization_config.FEED_3D_MODEL, + "ITEST": cumulus_parameterization_config.ITEST, + "USE_TRACER_TRANSPORT": config.USE_TRACER_TRANSPORT, + "DT_MOIST": config.DT_MOIST, + }, + ) + + self._feed_3d_model_from_plumes = stencil_factory.from_dims_halo( + func=feed_3d_model_from_plumes, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "FEED_3D_MODEL": cumulus_parameterization_config.FEED_3D_MODEL, + "ITEST": cumulus_parameterization_config.ITEST, + "USE_TRACER_TRANSPORT": config.USE_TRACER_TRANSPORT, + "DT_MOIST": config.DT_MOIST, + }, + ) + + self._update_convection_tracer = stencil_factory.from_dims_halo( + func=update_convection_tracer, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"CONVECTION_TRACER": config.CONVECTION_TRACER, "DT_MOIST": config.DT_MOIST}, + ) + + self._update_outputs = stencil_factory.from_dims_halo( + func=update_outputs, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ENABLE_SHALLOW": cumulus_parameterization_config.ENABLE_SHALLOW, + "ENABLE_MID": cumulus_parameterization_config.ENABLE_MID, + "ENABLE_DEEP": cumulus_parameterization_config.ENABLE_DEEP, + }, + ) + + self._update_tendencies = stencil_factory.from_dims_halo( + func=update_tendencies, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._update_convection_codes = stencil_factory.from_dims_halo( + func=update_convection_codes, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._update_state_with_tendencies = stencil_factory.from_dims_halo( + func=update_state_with_tendencies, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SCLM_DEEP": config.SCLM_DEEP, + "DT_MOIST": config.DT_MOIST, + "FIX_CONVECTIVE_CLOUD": config.FIX_CONVECTIVE_CLOUD, + }, + ) + + def __call__( + self, + state: GF2020State, + locals: GF2020Locals, + cumulus_parameterization_state: GF2020CumulusParameterizationState, + convection_tracers: ConvectionTracers, + ): + """Finish the GF2020 Cumulus Parameterization scheme + + Args: + state (GF2020State): NDSL State containing all model fields required for GF2020. + locals (GF2020Locals): NDSL LocalState containing all locals for GF2020. + cumulus_parameterization_state (GF2020CumulusParameterizationState): NDSL State containing all + fields required for the CumulusParameterization. + convection_tracers (ConvectionTracers): Collection of tracers from the rest of the model which + will be updated within convection. These may come from a variaty of sources, and need to be + collected into the expected ConvectionTracers data type before being passed down. + """ + + self._copy_from_cumulus_parameterization_state( + pbl_time_scale_from_cumulus_parameterization=cumulus_parameterization_state.output.pbl_time_scale, + pbl_time_scale=state.pbl_time_scale, + cape_removal_time_scale_from_cumulus_parameterization=cumulus_parameterization_state.output.cape_removal_time_scale, + cape_removal_time_scale=state.cape_removal_time_scale, + cloud_workfunction_0_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_workfunction_0, + cloud_workfunction_0=state.cloud_workfunction_0, + cloud_workfunction_1_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_workfunction_1, + cloud_workfunction_1=state.cloud_workfunction_1, + ) + + self._flag_computed_plumes_and_columns( + error_code=cumulus_parameterization_state.output.error_code, + do_this_column=locals.do_this_column, + ) + + self._check_vapor_mixing_ratio( + do_this_column=locals.do_this_column, + vapor_current=locals.flipped_copy.vapor_current, + dvapordt=cumulus_parameterization_state.output.dvapordt, + t_tendency_from_vapor=locals.t_tendency_from_vapor, + fix_out_vapor=locals.fix_out_vapor, + ) + + self._feedback( + fix_out_vapor=locals.fix_out_vapor, + precip=locals.precip, + precip_from_cumulus_parameterization=cumulus_parameterization_state.output.precip, + evaporation_sublimation_tendency=locals.evaporation_sublimation_tendency, + evaporation_sublimation_tendency_from_cumulus_parameterization=cumulus_parameterization_state.output.evaporation_sublimation_tendency, + convective_precip_flux=locals.convective_precip_flux, + convective_precip_flux_from_cumulus_parameterization=cumulus_parameterization_state.output.convective_precip_flux, + dtdt=locals.dtdt, + dtdt_from_cumulus_parameterization=cumulus_parameterization_state.output.dtdt, + dvapordt=locals.dvapordt, + dvapordt_from_cumulus_parameterization=cumulus_parameterization_state.output.dvapordt, + dcloudicedt=locals.dcloudicedt, + dcloudicedt_from_cumulus_parameterization=cumulus_parameterization_state.output.dcloudicedt, + dudt=locals.dudt, + dudt_from_cumulus_parameterization=cumulus_parameterization_state.output.dudt, + dvdt=locals.dvdt, + dvdt_from_cumulus_parameterization=cumulus_parameterization_state.output.dvdt, + dlarge_scale_icedt=locals.dlarge_scale_icedt, + dlarge_scale_icedt_from_cumulus_parameterization=cumulus_parameterization_state.output.dlargescaleicedt, + dconvective_icedt=locals.dconvective_icedt, + dconvective_icedt_from_cumulus_parameterization=cumulus_parameterization_state.output.dconvectiveicedt, + dlarge_scale_liquiddt=locals.dlarge_scale_liquiddt, + dlarge_scale_liquiddt_from_cumulus_parameterization=cumulus_parameterization_state.output.dlargescaleliquiddt, + dconvective_liquiddt=locals.dconvective_liquiddt, + dconvective_liquiddt_from_cumulus_parameterization=cumulus_parameterization_state.output.dconvectiveliquiddt, + dlarge_scale_cloud_fractiondt=locals.dlarge_scale_cloud_fractiondt, + dlarge_scale_cloud_fractiondt_from_cumulus_parameterization=cumulus_parameterization_state.output.dlargescalecloudfractiondt, + dconvective_cloud_fractiondt=locals.dconvective_cloud_fractiondt, + dconvective_cloud_fractiondt_from_cumulus_parameterization=cumulus_parameterization_state.output.dconvectivecloudfractiondt, + dbuoyancydt_from_cumulus_parameterization=cumulus_parameterization_state.output.dbuoyancydt, + dbuoyancydt=locals.dbuoyancydt, + do_this_column=locals.do_this_column, + ) + + for tracer in range(constants.NUMBER_OF_TRACERS): + self._feedback_tracers( + tracer=Int(tracer), + fix_out_vapor=locals.fix_out_vapor, + dconvection_tracersdt=locals.dconvection_tracersdt, + dconvection_tracersdt_from_cumulus_parameterization=cumulus_parameterization_state.input_output.chemistry_tracers_output, + chemistry_tracers_from_cumulus_parameterization=cumulus_parameterization_state.input_output.chemistry_tracers, + do_this_column=locals.do_this_column, + ) + + self._cloud_workfunction_output( + precip=cumulus_parameterization_state.output.precip, + fix_out_vapor=locals.fix_out_vapor, + cloud_workfunction_2=state.cloud_workfunction_2, + cloud_workfunction_3=state.cloud_workfunction_3, + ) + + self._prefill_entrainment( + lateral_entrainment_rate_shallow=state.lateral_entrainment_rate_shallow, + lateral_entrainment_rate_mid=state.lateral_entrainment_rate_mid, + lateral_entrainment_rate_deep=state.lateral_entrainment_rate_deep, + ) + + self._feed_3d_model( + do_this_column=locals.do_this_column, + precip=locals.precip, + convective_precipitation_GF=state.convective_precipitation_GF, + evaporation_sublimation_tendency_from_cumulus_parameterization=cumulus_parameterization_state.output.evaporation_sublimation_tendency, + evaporation_sublimation_tendency=locals.evaporation_sublimation_tendency, + convective_precip_flux_from_cumulus_parameterization=cumulus_parameterization_state.output.convective_precip_flux, + convective_precip_flux=locals.convective_precip_flux, + dconvection_tracersdt=locals.dconvection_tracersdt, + convection_tracers=convection_tracers.tracers, + ) + + for plume in range(cumulus_parameterization_constants.NUMBER_OF_PLUMES): + # Only call the function if the current plume is enabled self._plume_status is fixed to + # [shallow, mid, deep] order. This conditional ensures the correct is checked in plume_status + # and the correct plume is written in the stencil even if the overarching data order changes + if plume == cumulus_parameterization_constants.SHALLOW: + index = 0 + elif plume == cumulus_parameterization_constants.MID: + index = 1 + elif plume == cumulus_parameterization_constants.DEEP: + index = 2 + if self._plume_status[index] == 1: + self._feed_3d_model_from_plumes( + plume=Int(index), + do_this_column=locals.do_this_column, + dz=locals.derived_state.dz, + air_density=locals.derived_state.air_density, + p_flipped=locals.flipped_copy.p, + vapor_flipped=locals.flipped_copy.vapor, + dcloudicedt=locals.dcloudicedt, + lateral_entrainment_rate_shallow=state.lateral_entrainment_rate_shallow, + lateral_entrainment_rate_mid=state.lateral_entrainment_rate_mid, + lateral_entrainment_rate_deep=state.lateral_entrainment_rate_deep, + entrainment_rate_from_cumulus_parameterization=cumulus_parameterization_state.output.entrainment_rate, + error_code_from_cumulus_parameterization=cumulus_parameterization_state.output.error_code, + cloud_top_level_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_top_level, + t_updraft_from_cumulus_parameterization=cumulus_parameterization_state.output.t_updraft, + mass_detrainment_updraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.mass_entrainment_updraft_forced, + normalized_massflux_updraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.normalized_massflux_updraft_forced, + cloud_liquid_after_rain_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_liquid_after_rain_forced, + condensate_to_fall_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.condensate_to_fall_forced, + epsilon_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.epsilon_forced, + evaporate_in_downdraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.evaporate_in_downdraft_forced, + total_water_flux_deep_convection=state.total_water_flux_deep_convection_interface, + convective_condensate_source=state.convective_condensate_source, + mass_flux_cloud_base=state.mass_flux_cloud_base, + mass_flux_deep_updraft_detrained=state.mass_flux_deep_updraft_detrained, + mass_flux_deep_updraft_interface=state.mass_flux_deep_updraft_interface, + entrainment_parameter=state.entrainment_parameter, + updraft_vertical_velocity=state.updraft_vertical_velocity, + convective_condensate_grid_mean=state.convective_condensate_grid_mean, + convective_precipitation_RAS=state.convective_precipitation_RAS, + updraft_areal_fraction=state.updraft_areal_fraction, + esw=self._esw, + estlqu=self._estlqu, + ) + + self._update_convection_tracer( + land_fraction=state.land_fraction, + dbuoyancydt=locals.dbuoyancydt, + convection_tracer=state.convection_tracer, + ) + + self._update_outputs( + p_flipped=locals.flipped_copy.p, + error_code_from_cumulus_parameterization=cumulus_parameterization_state.output.error_code, + cloud_top_level_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_top_level, + cloud_base_mass_flux_modified_from_cumulus_parameterization=cumulus_parameterization_state.output.cloud_base_mass_flux_modified, + normalized_massflux_updraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.normalized_massflux_downdraft_forced, + epsilon_forced_from_cumulus_parameterization=cumulus_parameterization_state.output.epsilon_forced, + scale_dependence_factor_from_cumulus_parameterizaiton=cumulus_parameterization_state.output.scale_dependence_factor, + pressure_shallow_convective_cloud_top=state.pressure_shallow_convective_cloud_top, + pressure_mid_convective_cloud_top=state.pressure_mid_convective_cloud_top, + pressure_deep_convective_cloud_top=state.pressure_deep_convective_cloud_top, + mass_flux_cloud_base_shallow=state.mass_flux_cloud_base_shallow, + mass_flux_cloud_base_mid=state.mass_flux_cloud_base_mid, + mass_flux_cloud_base_deep=state.mass_flux_cloud_base_deep, + mass_flux_shallow=state.mass_flux_shallow, + mass_flux_mid=state.mass_flux_mid, + mass_flux_deep_updraft=state.mass_flux_deep_updraft, + mass_flux_deep_downdraft=state.mass_flux_deep_downdraft, + sigma_mid=state.sigma_mid, + sigma_deep=state.sigma_deep, + ) + + self._update_tendencies( + dtdt=locals.dtdt, + dvapordt=locals.dvapordt, + dudt=locals.dudt, + dvdt=locals.dvdt, + dtdt_deep_convection=state.dtdt_deep_convection, + dvapordt_deep_convection=state.dvapordt_deep_convection, + dudt_deep_convection=state.dudt_deep_convection, + dvdt_deep_convection=state.dvdt_deep_convection, + ) + + self._update_convection_codes( + error_code_from_cumulus_parameterization=cumulus_parameterization_state.output.error_code, + convection_code_shallow=state.convection_code_shallow, + convection_code_mid=state.convection_code_mid, + convection_code_deep=state.convection_code_deep, + ) + + self._update_state_with_tendencies( + convection_fraction=state.convection_fraction, + surface_type=state.surface_type, + u=state.u, + v=state.v, + vapor=state.vapor, + t=state.t, + p=locals.derived_state.p, + p_kappa=locals.derived_state.p_kappa, + mass=locals.derived_state.mass, + mass_flux_deep_updraft_detrained=state.mass_flux_deep_updraft_detrained, + mass_flux_deep_updraft_interface=state.mass_flux_deep_updraft_interface, + total_cumulative_mass_flux_interface=state.total_cumulative_mass_flux_interface, + total_detraining_mass_flux=state.total_detraining_mass_flux, + dudt_deep_convection=state.dudt_deep_convection, + dvdt_deep_convection=state.dvdt_deep_convection, + dvapordt_deep_convection=state.dvapordt_deep_convection, + dtdt_deep_convection=state.dtdt_deep_convection, + dliquiddt_deep_convection=state.dliquiddt_deep_convection, + dicedt_deep_convection=state.dicedt_deep_convection, + dcloudfractiondt_deep_convection=state.dcloudfractiondt_deep_convection, + convective_condensate_source=state.convective_condensate_source, + evaporation_sublimation_tendency=locals.evaporation_sublimation_tendency, + convective_precip_flux=locals.convective_precip_flux, + sublimation_of_convective_precipitation=state.sublimation_of_convective_precipitation, + evaporation_of_convective_precipitation=state.evaporation_of_convective_precipitation, + ice_fraction_in_convective_tower=state.ice_fraction_in_convective_tower, + ice_precip_flux_interface=state.ice_precip_flux_interface, + liquid_precip_flux_interface=state.liquid_precip_flux_interface, + convective_liquid=state.convective_liquid, + convective_ice=state.convective_ice, + convective_cloud_fraction=state.convective_cloud_fraction, + convective_rainwater_source=state.convective_rainwater_source, + convective_precipitation_RAS=state.convective_precipitation_RAS, + ese=self._ese, + esx=self._esx, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/locals.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/locals.py new file mode 100644 index 000000000..aa0130365 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/locals.py @@ -0,0 +1,657 @@ +import dataclasses + +from ndsl import Quantity, State +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM +from ndsl.dsl.typing import Float, Int + + +@dataclasses.dataclass +class GF2020Locals(State): + @dataclasses.dataclass + class DerivedState: + edge_height_above_surface: Quantity = dataclasses.field( + metadata={ + "name": "edge_height_above_surface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + layer_height_above_surface: Quantity = dataclasses.field( + metadata={ + "name": "layer_height_above_surface", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p: Quantity = dataclasses.field( + metadata={ + "name": "p", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_kappa: Quantity = dataclasses.field( + metadata={ + "name": "p_kappa", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + th: Quantity = dataclasses.field( + metadata={ + "name": "th", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass: Quantity = dataclasses.field( + metadata={ + "name": "mass", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vertical_velocity: Quantity = dataclasses.field( + metadata={ + "name": "vertical_velocity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ice_fraciton: Quantity = dataclasses.field( + metadata={ + "name": "ice_fraciton", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + saturation_specific_humidity: Quantity = dataclasses.field( + metadata={ + "name": "saturation_specific_humidity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + seed_convection: Quantity = dataclasses.field( + metadata={ + "name": "seed_convection", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + modified_area: Quantity = dataclasses.field( + metadata={ + "name": "modified_area", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_perturbation_horizontal: Quantity = dataclasses.field( + metadata={ + "name": "t_perturbation_horizontal", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_perturbation_vertical: Quantity = dataclasses.field( + metadata={ + "name": "t_perturbation_vertical", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dz: Quantity = dataclasses.field( + metadata={ + "name": "dz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + air_density: Quantity = dataclasses.field( + metadata={ + "name": "air_density", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + scalar_diffusivity: Quantity = dataclasses.field( + metadata={ + "name": "scalar_diffusivity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + + @dataclasses.dataclass + class FlippedCopy: + t_2m: Quantity = dataclasses.field( + metadata={ + "name": "t_2m", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t: Quantity = dataclasses.field( + metadata={ + "name": "t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_surface: Quantity = dataclasses.field( + metadata={ + "name": "t_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p: Quantity = dataclasses.field( + metadata={ + "name": "p", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_surface: Quantity = dataclasses.field( + metadata={ + "name": "p_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor: Quantity = dataclasses.field( + metadata={ + "name": "vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_current: Quantity = dataclasses.field( + metadata={ + "name": "vapor_current", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u: Quantity = dataclasses.field( + metadata={ + "name": "u", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v: Quantity = dataclasses.field( + metadata={ + "name": "v", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + w: Quantity = dataclasses.field( + metadata={ + "name": "vertical_velocity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + layer_height_above_surface: Quantity = dataclasses.field( + metadata={ + "name": "layer_height_above_surface", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + edge_height_above_surface: Quantity = dataclasses.field( + metadata={ + "name": "edge_height_above_surface", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass: Quantity = dataclasses.field( + metadata={ + "name": "mass", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + scalar_diffusivity: Quantity = dataclasses.field( + metadata={ + "name": "scalar_diffusivity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_liquid: Quantity = dataclasses.field( + metadata={ + "name": "convective_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_ice: Quantity = dataclasses.field( + metadata={ + "name": "convective_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convective_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_liquid: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_ice: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_level: Quantity = dataclasses.field( + metadata={ + "name": "pbl_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + + do_this_column: Quantity = dataclasses.field( + metadata={ + "name": "do_this_column", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Int, + } + ) + fix_out_vapor: Quantity = dataclasses.field( + metadata={ + "name": "fix_out_vapor", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_tendency_from_vapor: Quantity = dataclasses.field( + metadata={ + "name": "t_tendency_from_vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + rtgt: Quantity = dataclasses.field( + metadata={ + "name": "rtgt", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + aot500: Quantity = dataclasses.field( + metadata={ + "name": "aot500", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation: Quantity = dataclasses.field( + metadata={ + "name": "evaporation", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sensible_heat_flux: Quantity = dataclasses.field( + metadata={ + "name": "sensible_heat_flux", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + topography_height: Quantity = dataclasses.field( + metadata={ + "name": "topography_height", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ocean_fraction: Quantity = dataclasses.field( + metadata={ + "name": "ocean_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + grid_length: Quantity = dataclasses.field( + metadata={ + "name": "grid_length", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + buoyancy_excess: Quantity = dataclasses.field( + metadata={ + "name": "buoyancy_excess", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + grid_scale_forcing_t: Quantity = dataclasses.field( + metadata={ + "name": "grid_scale_forcing_t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + grid_scale_forcing_vapor: Quantity = dataclasses.field( + metadata={ + "name": "grid_scale_forcing_vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + subgrid_scale_forcing_t: Quantity = dataclasses.field( + metadata={ + "name": "subgrid_scale_forcing_t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + subgrid_scale_forcing_vapor: Quantity = dataclasses.field( + metadata={ + "name": "subgrid_scale_forcing_vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + advective_forcing_t: Quantity = dataclasses.field( + metadata={ + "name": "advective_forcing_t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + saturation_water_vapor: Quantity = dataclasses.field( + metadata={ + "name": "saturation_water_vapor", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt: Quantity = dataclasses.field( + metadata={ + "name": "dtdt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvapordt: Quantity = dataclasses.field( + metadata={ + "name": "dvapordt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dcloudicedt: Quantity = dataclasses.field( + metadata={ + "name": "dcloudicedt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dudt: Quantity = dataclasses.field( + metadata={ + "name": "dudt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvdt: Quantity = dataclasses.field( + metadata={ + "name": "dvdt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlarge_scale_icedt: Quantity = dataclasses.field( + metadata={ + "name": "dlarge_scale_icedt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvective_icedt: Quantity = dataclasses.field( + metadata={ + "name": "dconvective_icedt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlarge_scale_liquiddt: Quantity = dataclasses.field( + metadata={ + "name": "dlarge_scale_liquiddt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvective_liquiddt: Quantity = dataclasses.field( + metadata={ + "name": "dconvective_liquiddt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dlarge_scale_cloud_fractiondt: Quantity = dataclasses.field( + metadata={ + "name": "dlarge_scale_cloud_fractiondt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvective_cloud_fractiondt: Quantity = dataclasses.field( + metadata={ + "name": "dconvective_cloud_fractiondt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dbuoyancydt: Quantity = dataclasses.field( + metadata={ + "name": "dbuoyancydt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dconvection_tracersdt: Quantity = dataclasses.field( + metadata={ + "name": "dconvection_tracersdt", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation_sublimation_tendency: Quantity = dataclasses.field( + metadata={ + "name": "evaporation_sublimation_tendency", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_precip_flux: Quantity = dataclasses.field( + metadata={ + "name": "convective_precip_flux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + precip: Quantity = dataclasses.field( + metadata={ + "name": "precip", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + + derived_state: DerivedState + flipped_copy: FlippedCopy diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/setup.py new file mode 100644 index 000000000..cc24667ab --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/setup.py @@ -0,0 +1,1660 @@ +import pyMoist.constants as constants +import pyMoist.convection.GF_2020.cumulus_parameterization.constants as cumulus_parameterization_constants +from ndsl import NDSLRuntime, QuantityFactory, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM +from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, K, abs, computation, floor, interval, max, min, sqrt +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.field_types import ( + FloatField_ConvectionTracers, + FloatField_ConvectionTracers_Plume, + FloatField_Plume, + FloatFieldIJ_Plume, + IntFieldIJ_Plume, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.locals import GF2020Locals +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.qsat_functions import saturation_specific_humidity +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable +from pyMoist.saturation_tables.types import GlobalTable_saturation_tables + + +def compute_extra_inputs_from_state( + p_interface: FloatField, + p: FloatField, + p_kappa: FloatField, + edge_height_above_surface: FloatField, + layer_height_above_surface: FloatField, + geopotential_height_interface: FloatField, + t: FloatField, + th: FloatField, + vapor: FloatField, + mass: FloatField, + w: FloatField, + omega: FloatField, + vertical_motion: FloatField, + tpwi: FloatFieldIJ, + tpwi_star: FloatFieldIJ, + seed_convection: FloatFieldIJ, + area: FloatFieldIJ, + modified_area: FloatFieldIJ, + convection_fraction: FloatFieldIJ, + ese: GlobalTable_saturation_tables, + esx: GlobalTable_saturation_tables, +): + """ + Performs initial setup for the GF 2020 convection scheme: + - Compute derived states + - initialize stochastic variability for convection + - Modify area (m^2) here so GF scale dependence has a convection_fraction dependence + + This stencil MUST be built using Z_INTERFACE_DIM to function properly. + + Args: + p_interface (FloatField) + p (FloatField) + p_kappa (FloatField) + edge_height_above_surface (FloatField) + layer_height_above_surface (FloatField) + geopotential_height_interface (FloatField) + t (FloatField) + th (FloatField) + vapor (FloatField) + mass (FloatField) + w (FloatField) + omega (FloatField) + vertical_motion (FloatField) + tpwi (FloatFieldIJ) + tpwi_star (FloatFieldIJ) + seed_convection (FloatFieldIJ) + area (FloatFieldIJ) + modified_area (FloatFieldIJ) + convection_fraction (FloatFieldIJ) + ese (GlobalTable_saturation_tables) + esx (GlobalTable_saturation_tables) + """ + from __externals__ import GF_MIN_AREA, LHYDROSTATIC, STOCH_BOT, STOCH_TOP, STOCHASTIC_CONVECTION, k_end + + # compute derived states + with computation(PARALLEL), interval(...): + edge_height_above_surface = geopotential_height_interface - geopotential_height_interface.at(K=k_end) + + with computation(PARALLEL), interval(0, -1): + p = 0.5 * (p_interface + p_interface[0, 0, 1]) + p_kappa = (p / constants.MAPL_P00) ** (constants.MAPL_KAPPA) + layer_height_above_surface = 0.5 * (edge_height_above_surface + edge_height_above_surface[0, 0, 1]) + th = t / p_kappa + mass = (p_interface[0, 0, 1] - p_interface) / constants.MAPL_GRAV + + if LHYDROSTATIC: + vertical_motion = ( + -1 + * omega + / ( + constants.MAPL_GRAV + * p + / (constants.MAPL_RDRY * t * (1.0 + constants.MAPL_VIREPS * vapor)) + ) + ) + else: + vertical_motion = w + + with computation(FORWARD), interval(0, 1): + tpwi = vapor * mass + qsat, _ = saturation_specific_humidity(t, p, ese, esx) + tpwi_star = qsat * mass + + with computation(FORWARD), interval(1, -1): + tpwi = tpwi + vapor * mass + qsat, _ = saturation_specific_humidity(t, p, ese, esx) + tpwi_star = tpwi_star + qsat * mass + + # initialize stochastic variability for convection + if STOCHASTIC_CONVECTION == True: # noqa + # Create bit-processor-reproducible random white noise for convection [0:1] + seedini = 1000000 * (100 * t.at(K=k_end) - floor(100 * t.at(K=k_end))) + seed_convection = sqrt(max(min(seedini / 1000000.0, 1.0), 0.0)) + # Create stochastic variability to GF sigma + seed_convection = sqrt(1.0 - (1.0 - seed_convection)) * (STOCH_TOP - STOCH_BOT) + STOCH_BOT + else: + seed_convection = 1.0 + + # Modify area (m^2) here so GF scale dependence has a convection_fraction dependence + if GF_MIN_AREA > 0: + if area > GF_MIN_AREA: + modified_area = GF_MIN_AREA * convection_fraction + area * (1.0 - convection_fraction) + else: + modified_area = area + elif GF_MIN_AREA < 0: + if area > abs(GF_MIN_AREA): + modified_area = area * convection_fraction + abs(GF_MIN_AREA) * (1.0 - convection_fraction) + else: + modified_area = area + else: + modified_area = area + + +def zero_state( + dvapordt_deep_convection: FloatField, + dtdt_deep_convection: FloatField, + dudt_deep_convection: FloatField, + dvdt_deep_convection: FloatField, + sigma_deep: FloatFieldIJ, + sigma_mid: FloatFieldIJ, + mass_flux_shalow: FloatField, + mass_flux_mid: FloatField, + mass_flux_deep_updraft: FloatField, + mass_flux_deep_updraft_interface: FloatField, + mass_flux_deep_updraft_detrained: FloatField, + mass_flux_deep_downdraft: FloatField, + mass_flux_cloud_base: FloatField, + mass_flux_cloud_base_shallow: FloatFieldIJ, + mass_flux_cloud_base_mid: FloatFieldIJ, + mass_flux_cloud_base_deep: FloatFieldIJ, + convection_code_shallow: FloatFieldIJ, + convection_code_mid: FloatFieldIJ, + convection_code_deep: FloatFieldIJ, + cloud_workfunction_0: FloatFieldIJ, + cloud_workfunction_1: FloatFieldIJ, + cloud_workfunction_2: FloatFieldIJ, + cloud_workfunction_3: FloatFieldIJ, + cloud_workfunction_1_pbl: FloatFieldIJ, + cloud_workfunction_1_cin: FloatFieldIJ, + convective_precipitation_RAS: FloatField, + convective_precipitation_GF: FloatFieldIJ, + convective_condensate_source: FloatField, + convective_condensate_grid_mean: FloatField, + total_water_flux_deep_convection_interface: FloatField, + updraft_area_fraction: FloatField, + updraft_vertical_velocity: FloatField, + entrainment_parameter: FloatField, + lightning_density: FloatFieldIJ, + pbl_time_scale: FloatFieldIJ, + cape_removal_time_scale: FloatFieldIJ, +): + """ + Zero fields from the GEOS model state. + + All of these fields are outputs of GF2020 which may be set in the CumulusParameterization core routine, + so we are really just clearing any data from the previous timestep to ensure nothing leaks through. + + Must be built with Z_INTERFACE_DIM. + + Args: + dvapordt_deep_convection (FloatField) + dtdt_deep_convection (FloatField) + dudt_deep_convection (FloatField) + dvdt_deep_convection (FloatField) + sigma_deep (FloatFieldIJ) + sigma_mid (FloatFieldIJ) + mass_flux_shalow (FloatField) + mass_flux_mid (FloatField) + mass_flux_deep_updraft (FloatField) + mass_flux_deep_updraft_interface (FloatField) + mass_flux_deep_updraft_detrained (FloatField) + mass_flux_deep_downdraft (FloatField) + mass_flux_cloud_base (FloatField) + mass_flux_cloud_base_shallow (FloatFieldIJ) + mass_flux_cloud_base_mid (FloatFieldIJ) + mass_flux_cloud_base_deep (FloatFieldIJ) + convection_code_shallow (FloatFieldIJ) + convection_code_mid (FloatFieldIJ) + convection_code_deep (FloatFieldIJ) + cloud_work_function_0 (FloatFieldIJ) + cloud_work_function_1 (FloatFieldIJ) + cloud_work_function_2 (FloatFieldIJ) + cloud_work_function_3 (FloatFieldIJ) + cloud_work_function_1_pbl (FloatFieldIJ) + cloud_work_function_1_cin (FloatFieldIJ) + convective_precipitation_RAS (FloatField) + convective_precipitation_GF (FloatFieldIJ) + convective_condensate_source (FloatField) + convective_condensate_grid_mean (FloatField) + total_water_flux_deep_convection_interface (FloatField) + updraft_area_fraction (FloatField) + updraft_vertical_velocity (FloatField) + entrainment_parameter (FloatField) + lightning_density (FloatFieldIJ) + pbl_time_scale (FloatFieldIJ) + cape_removal_time_scale (FloatFieldIJ) + """ + with computation(PARALLEL), interval(0, -1): + dvapordt_deep_convection = 0.0 + dtdt_deep_convection = 0.0 + dudt_deep_convection = 0.0 + dvdt_deep_convection = 0.0 + + with computation(FORWARD), interval(0, 1): + sigma_deep = 0.0 + sigma_mid = 0.0 + + with computation(PARALLEL), interval(0, -1): + mass_flux_shalow = 0.0 + mass_flux_mid = 0.0 + mass_flux_deep_updraft = 0.0 + + with computation(PARALLEL), interval(...): + mass_flux_deep_updraft_interface = 0.0 + + with computation(PARALLEL), interval(0, -1): + mass_flux_deep_updraft_detrained = 0.0 + mass_flux_deep_downdraft = 0.0 + mass_flux_cloud_base = 0.0 + + with computation(FORWARD), interval(0, 1): + mass_flux_cloud_base_shallow = 0.0 + mass_flux_cloud_base_mid = 0.0 + mass_flux_cloud_base_deep = 0.0 + + convection_code_shallow = 0.0 + convection_code_mid = 0.0 + convection_code_deep = 0.0 + + cloud_workfunction_0 = 0.0 + cloud_workfunction_1 = 0.0 + cloud_workfunction_2 = 0.0 + cloud_workfunction_3 = 0.0 + cloud_workfunction_1_pbl = 0.0 + cloud_workfunction_1_cin = 0.0 + + with computation(PARALLEL), interval(0, -1): + convective_precipitation_RAS = 0.0 + + with computation(FORWARD), interval(0, 1): + convective_precipitation_GF = 0.0 + + with computation(PARALLEL), interval(0, -1): + convective_condensate_source = 0.0 + convective_condensate_grid_mean = 0.0 + + with computation(PARALLEL), interval(...): + total_water_flux_deep_convection_interface = 0.0 + + with computation(PARALLEL), interval(0, -1): + updraft_area_fraction = 0.0 + updraft_vertical_velocity = 0.0 + entrainment_parameter = 0.0 + + with computation(FORWARD), interval(0, 1): + lightning_density = 0.0 + pbl_time_scale = 0.0 + cape_removal_time_scale = 0.0 + + +def prefill_cumulus_parameterization_state( + error_code: IntFieldIJ_Plume, + downdraft_origin_level: IntFieldIJ_Plume, + lcl_level: IntFieldIJ_Plume, + updraft_origin_level: IntFieldIJ_Plume, + updraft_lfc_level: IntFieldIJ_Plume, + cloud_top_level: IntFieldIJ_Plume, + kstabi: IntFieldIJ_Plume, + kstabm: IntFieldIJ_Plume, + precip: FloatFieldIJ_Plume, + cloud_base_mass_flux_modified: FloatFieldIJ_Plume, + epsilon_forced: FloatFieldIJ_Plume, + total_normalized_integrated_condensate_forced: FloatFieldIJ_Plume, + scale_dependence_factor: FloatFieldIJ_Plume, + p_cloud_levels_forced: FloatField_Plume, + entrainment_rate: FloatField_Plume, + mass_entrainment_updraft_forced: FloatField_Plume, + mass_entrainment_downdraft_forced: FloatField_Plume, + mass_detrainment_updraft_forced: FloatField_Plume, + mass_detrainment_downdraft_forced: FloatField_Plume, + normalized_massflux_updraft_forced: FloatField_Plume, + normalized_massflux_downdraft_forced: FloatField_Plume, + condensate_to_fall_forced: FloatField_Plume, + evaporate_in_downdraft_forced: FloatField_Plume, + cloud_liquid_after_rain_forced: FloatField_Plume, + t_updraft: FloatField_Plume, + convective_cloud_fraction_output: FloatField_Plume, + dtdt: FloatField_Plume, + dudt: FloatField_Plume, + dvdt: FloatField_Plume, + dvapordt: FloatField_Plume, + dcloudicedt: FloatField_Plume, + dnicedt: FloatField_Plume, + dnliquiddt: FloatField_Plume, + dbuoyancydt: FloatField_Plume, + chemistry_tracers_output: FloatField_ConvectionTracers_Plume, + evaporation_sublimation_tendency: FloatField, + convective_precip_flux: FloatField, + t_perturbation: FloatField, + omega: FloatField, + large_scale_ice: FloatField, + convective_ice: FloatField, + large_scale_liquid: FloatField, + convective_liquid: FloatField, + large_scale_cloud_fraction: FloatField, + convective_cloud_fraction: FloatField, + lightning_density: FloatFieldIJ, + t_excess: FloatFieldIJ, + vapor_excess: FloatFieldIJ, + last_error_code: IntFieldIJ, +): + """ + Zero fields from the CumulusParameterization state. + + All of these fields may be set in the CumulusParameterization core routine, + so they need to be reset to ensure no lingering data from the previous + timestep makes it through. + + Args: + error_code (IntFieldIJ_Plume) + downdraft_origin_level (IntFieldIJ_Plume) + lcl_level (IntFieldIJ_Plume) + updraft_origin_level (IntFieldIJ_Plume) + updraft_lfc_level (IntFieldIJ_Plume) + cloud_top_level (IntFieldIJ_Plume) + kstabi (IntFieldIJ_Plume) + kstabm (IntFieldIJ_Plume) + precip (FloatFieldIJ_Plume) + cloud_base_mass_flux_modified (FloatFieldIJ_Plume) + epsilon_forced (FloatFieldIJ_Plume) + total_normalized_integrated_condensate_forced (FloatFieldIJ_Plume) + scale_dependence_factor (FloatFieldIJ_Plume) + p_cloud_levels_forced (FloatField_Plume) + entrainment_rate (FloatField_Plume) + mass_entrainment_updraft_forced (FloatField_Plume) + mass_entrainment_downdraft_forced (FloatField_Plume) + mass_detrainment_updraft_forced (FloatField_Plume) + mass_detrainment_downdraft_forced (FloatField_Plume) + normalized_massflux_updraft_forced (FloatField_Plume) + normalized_massflux_downdraft_forced (FloatField_Plume) + condensate_to_fall_forced (FloatField_Plume) + evaporate_in_downdraft_forced (FloatField_Plume) + cloud_liquid_after_rain_forced (FloatField_Plume) + t_updraft (FloatField_Plume) + convective_cloud_fraction_output (FloatField_Plume) + dtdt (FloatField_Plume) + dudt (FloatField_Plume) + dvdt (FloatField_Plume) + dvapordt (FloatField_Plume) + dcloudicedt (FloatField_Plume) + dnicedt (FloatField_Plume) + dnliquiddt (FloatField_Plume) + dbuoyancydt (FloatField_Plume) + chemistry_tracers_output (FloatField_ConvectionTracers_Plume) + evaporation_sublimation_tendency (FloatField) + convective_precip_flux (FloatField) + t_perturbation (FloatField) + omega (FloatField) + large_scale_ice (FloatField) + convective_ice (FloatField) + large_scale_liquid (FloatField) + convective_liquid (FloatField) + large_scale_cloud_fraction (FloatField) + convective_cloud_fraction (FloatField) + lightning_density (FloatFieldIJ) + t_excess (FloatFieldIJ) + vapor_excess (FloatFieldIJ) + last_error_code (IntFieldIJ) + """ + from __externals__ import APPLY_SUBSIDENCE_MICROPHYSICS, NUMBER_OF_PLUMES + + with computation(FORWARD), interval(0, 1): + plume = 0 + while plume < NUMBER_OF_PLUMES: + error_code[0, 0][plume] = 0 + downdraft_origin_level[0, 0][plume] = 0 + lcl_level[0, 0][plume] = 0 + updraft_origin_level[0, 0][plume] = 0 + updraft_lfc_level[0, 0][plume] = 0 + cloud_top_level[0, 0][plume] = 0 + kstabi[0, 0][plume] = 0 + kstabm[0, 0][plume] = 0 + precip[0, 0][plume] = 0.0 + cloud_base_mass_flux_modified[0, 0][plume] = 0.0 + epsilon_forced[0, 0][plume] = 0.0 + total_normalized_integrated_condensate_forced[0, 0][plume] = 0.0 + scale_dependence_factor[0, 0][plume] = 0.0 + plume += 1 + + with computation(PARALLEL), interval(...): + plume = 0 + while plume < NUMBER_OF_PLUMES: + p_cloud_levels_forced[0, 0, 0][plume] = 0.0 + entrainment_rate[0, 0, 0][plume] = 0.0 + mass_entrainment_updraft_forced[0, 0, 0][plume] = 0.0 + mass_entrainment_downdraft_forced[0, 0, 0][plume] = 0.0 + mass_detrainment_updraft_forced[0, 0, 0][plume] = 0.0 + mass_detrainment_downdraft_forced[0, 0, 0][plume] = 0.0 + normalized_massflux_updraft_forced[0, 0, 0][plume] = 0.0 + normalized_massflux_downdraft_forced[0, 0, 0][plume] = 0.0 + condensate_to_fall_forced[0, 0, 0][plume] = 0.0 + evaporate_in_downdraft_forced[0, 0, 0][plume] = 0.0 + cloud_liquid_after_rain_forced[0, 0, 0][plume] = 0.0 + t_updraft[0, 0, 0][plume] = 0.0 + convective_cloud_fraction_output[0, 0, 0][plume] = 0.0 + + dtdt[0, 0, 0][plume] = 0.0 + dudt[0, 0, 0][plume] = 0.0 + dvdt[0, 0, 0][plume] = 0.0 + dvapordt[0, 0, 0][plume] = 0.0 + dcloudicedt[0, 0, 0][plume] = 0.0 + dnicedt[0, 0, 0][plume] = 0.0 + dnliquiddt[0, 0, 0][plume] = 0.0 + dbuoyancydt[0, 0, 0][plume] = 0.0 + + tracer = 0 + while tracer < constants.NUMBER_OF_TRACERS: + chemistry_tracers_output[0, 0, 0][plume, tracer] = 0.0 + tracer += 1 + + plume += 1 + + evaporation_sublimation_tendency = 0.0 + convective_precip_flux = 0.0 + t_perturbation = 0.0 + omega = 0.0 + + if APPLY_SUBSIDENCE_MICROPHYSICS == 1: + large_scale_ice = 0.0 + convective_ice = 0.0 + large_scale_liquid = 0.0 + convective_liquid = 0.0 + large_scale_cloud_fraction = 0.0 + convective_cloud_fraction = 0.0 + + with computation(FORWARD), interval(0, 1): + lightning_density = 0.0 + t_excess = 0.0 + vapor_excess = 0.0 + last_error_code = -999 + + +def prefill_locals( + rtgt: FloatFieldIJ, + precip: FloatFieldIJ, + t_tendency_from_vapor: FloatField, + dtdt: FloatField, + dvapordt: FloatField, + dcloudicedt: FloatField, + dudt: FloatField, + dvdt: FloatField, + dlarge_scale_icedt: FloatField, + dconvective_icedt: FloatField, + dlarge_scale_liquiddt: FloatField, + dconvective_liquiddt: FloatField, + dlarge_scale_cloud_fractiondt: FloatField, + dconvective_cloud_fractiondt: FloatField, + dbuoyancydt: FloatField, + evaporation_sublimation_tendency: FloatField, + convective_precip_flux: FloatField, +): + """ + Zero local fields which are conditionally written to ensure no data remains from the previous timestep. + + Args: + rtgt (FloatFieldIJ) + precip (FloatFieldIJ) + t_tendency_from_vapor (FloatField) + dtdt (FloatField) + dvapordt (FloatField) + dcloudicedt (FloatField) + dudt (FloatField) + dvdt (FloatField) + dlarge_scale_icedt (FloatField) + dconvective_icedt (FloatField) + dlarge_scale_liquiddt (FloatField) + dconvective_liquiddt (FloatField) + dlarge_scale_cloud_fractiondt (FloatField) + dconvective_cloud_fractiondt (FloatField) + dbuoyancydt (FloatField) + evaporation_sublimation_tendency (FloatField) + convective_precip_flux (FloatField) + """ + with computation(FORWARD), interval(0, 1): + rtgt = 1.0 + precip = 0.0 + + with computation(PARALLEL), interval(...): + t_tendency_from_vapor = 0.0 + dtdt = 0.0 + dvapordt = 0.0 + dcloudicedt = 0.0 + dudt = 0.0 + dvdt = 0.0 + dlarge_scale_icedt = 0.0 + dconvective_icedt = 0.0 + dlarge_scale_liquiddt = 0.0 + dconvective_liquiddt = 0.0 + dlarge_scale_cloud_fractiondt = 0.0 + dconvective_cloud_fractiondt = 0.0 + dbuoyancydt = 0.0 + evaporation_sublimation_tendency = 0.0 + convective_precip_flux = 0.0 + + +def set_2d_fields( + aot500: FloatFieldIJ, + t: FloatField, + t_2m_max: Float, + t_2m: FloatFieldIJ, + t_2m_local: FloatFieldIJ, + evaporation: FloatFieldIJ, + evaporation_local: FloatFieldIJ, + sensible_heat_flux: FloatFieldIJ, + sensible_heat_flux_local: FloatFieldIJ, + p_interface: FloatField, + vapor: FloatField, + geopotential_height_surface: FloatFieldIJ, + topography_height: FloatFieldIJ, + land_fraction: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + area: FloatFieldIJ, + grid_length: FloatFieldIJ, + pbl_level: FloatFieldIJ, + pbl_level_flipped: IntFieldIJ, +): + """ + Compute 2-dimensional inputs for the CumulusParameterization core routine. + + Args: + aot500 (FloatFieldIJ) + t (FloatField) + t_2m_max (Float) + t_2m (FloatFieldIJ) + t_2m_local (FloatFieldIJ) + evaporation (FloatFieldIJ) + evaporation_local (FloatFieldIJ) + sensible_heat_flux (FloatFieldIJ) + sensible_heat_flux_local (FloatFieldIJ) + p_interface (FloatField) + vapor (FloatField) + geopotential_height_surface (FloatFieldIJ) + topography_height (FloatFieldIJ) + land_fraction (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + area (FloatFieldIJ) + grid_length (FloatFieldIJ) + pbl_level (FloatFieldIJ) + pbl_level_flipped (IntFieldIJ) + """ + from __externals__ import SIZE_I_DIM, SIZE_J_DIM, k_end + + with computation(FORWARD), interval(0, 1): + aot500 = 0.1 + + if t_2m_max < 1.0e-6: + # in case convection is called before surface data is computed + t_2m_local = t.at(K=k_end) # kelvin + else: + t_2m_local = t_2m # kelvin + + # moisture flux from sfc + evaporation_local = evaporation # kg m-2 s-1 + + # sensible–heat_flux comes in W m-2, below it is converted to K m s-1 + sensible_heat_flux_local = sensible_heat_flux / ( + 1004.0 * p_interface.at(K=k_end) / (287.04 * t.at(K=k_end) * (1.0 + 0.608 * vapor.at(K=k_end))) + ) # K m s-1 + + # topography height (m) + topography_height = geopotential_height_surface / constants.MAPL_GRAV + + # land/ocean fraction: land if < 1 ,ocean if = 1 + ocean_fraction = 1.0 - land_fraction + + # grid length for the scale awareness (m) + grid_length = sqrt(area) + # special setting for SCM runs + if SIZE_I_DIM == 1 and SIZE_J_DIM == 1: + grid_length = 100000.0 + + # flip pbl_level + if pbl_level != -1.0: + pbl_level_flipped = k_end - int(round(pbl_level)) + else: + pbl_level_flipped = 0 + + +def choose_environment_and_flip_k_axis( + dz: FloatField, + air_density: FloatField, + geopotential_height_interface: FloatField, + t: FloatField, + t_flipped: FloatField, + t_timestep_start: FloatField, + p: FloatField, + p_interface: FloatField, + p_interface_timestep_start: FloatField, + p_flipped: FloatField, + p_surface_flipped: FloatFieldIJ, + vapor: FloatField, + vapor_timestep_start: FloatField, + vapor_flipped: FloatField, + vapor_current_flipped: FloatField, + u: FloatField, + u_timestep_start: FloatField, + u_flipped: FloatField, + v: FloatField, + v_timestep_start: FloatField, + v_flipped: FloatField, + w: FloatField, + w_flipped: FloatField, + layer_height_above_surface: FloatField, + layer_height_above_surface_flipped: FloatField, + edge_height_above_surface: FloatField, + edge_height_above_surface_flipped: FloatField, + mass: FloatField, + mass_flipped: FloatField, + scalar_diffusivity: FloatField, + scalar_diffusivity_flipped: FloatField, + lateral_entrainment_rate: FloatField, + lateral_entrainment_rate_flipped: FloatField, + buoyancy: FloatField, + buoyancy_excess: FloatField, + dtdt_shortwave: FloatField, + dtdt_longwave: FloatField, + dtdt_from_dynamics: FloatField, + dtdt_pbl: FloatField, + dvapordt_from_dynamics: FloatField, + dspecific_humiditydt_pbl: FloatField, + grid_scale_forcing_t: FloatField, + grid_scale_forcing_vapor: FloatField, + subgrid_scale_forcing_t: FloatField, + subgrid_scale_forcing_vapor: FloatField, + advective_forcing_t: FloatField, + convective_liquid: FloatField, + convective_liquid_flipped: FloatField, + convective_ice: FloatField, + convective_ice_flipped: FloatField, + convective_cloud_fraction: FloatField, + convective_cloud_fraction_flipped: FloatField, + large_scale_liquid: FloatField, + large_scale_liquid_flipped: FloatField, + large_scale_ice: FloatField, + large_scale_ice_flipped: FloatField, + large_scale_cloud_fraction: FloatField, + large_scale_cloud_fraction_flipped: FloatField, + convection_tracer: FloatField, + total_precipitable_water_initial: FloatFieldIJ, + saturation_total_precipitable_water_initial: FloatFieldIJ, + saturation_water_vapor: FloatFieldIJ, +): + """ + Get the desired state for the convection scheme, controlled by external GF_ENV_SETTING. + + GF_ENV_SETTING = 0: + use state updated by during current timestep (dynamics and physics) + GF_ENV_SETTING = 1: + use state from start of timestep (pre-dynamics) + + Args: + dz (FloatField) + air_density (FloatField) + geopotential_height_interface (FloatField) + t (FloatField) + t_flipped (FloatField) + t_timestep_start (FloatField) + p (FloatField) + p_interface (FloatField) + p_interface_timestep_start (FloatField) + p_flipped (FloatField) + p_surface_flipped (FloatFieldIJ) + vapor (FloatField) + vapor_timestep_start (FloatField) + vapor_flipped (FloatField) + vapor_current_flipped (FloatField) + u (FloatField) + u_timestep_start (FloatField) + u_flipped (FloatField) + v (FloatField) + v_timestep_start (FloatField) + v_flipped (FloatField) + w (FloatField) + w_flipped (FloatField) + layer_height_above_surface (FloatField) + layer_height_above_surface_flipped (FloatField) + edge_height_above_surface (FloatField) + edge_height_above_surface_flipped (FloatField) + mass (FloatField) + mass_flipped (FloatField) + scalar_diffusivity (FloatField) + scalar_diffusivity_flipped (FloatField) + lateral_entrainment_rate (FloatField) + lateral_entrainment_rate_flipped (FloatField) + buoyancy (FloatField) + buoyancy_excess (FloatField) + dtdt_shortwave (FloatField) + dtdt_longwave (FloatField) + dtdt_from_dynamics (FloatField) + dtdt_pbl (FloatField) + dvapordt_from_dynamics (FloatField) + dspecific_humiditydt_pbl (FloatField) + grid_scale_forcing_t (FloatField) + grid_scale_forcing_vapor (FloatField) + subgrid_scale_forcing_t (FloatField) + subgrid_scale_forcing_vapor (FloatField) + advective_forcing_t (FloatField) + convective_liquid (FloatField) + convective_liquid_flipped (FloatField) + convective_ice (FloatField) + convective_ice_flipped (FloatField) + convective_cloud_fraction (FloatField) + convective_cloud_fraction_flipped (FloatField) + large_scale_liquid (FloatField) + large_scale_liquid_flipped (FloatField) + large_scale_ice (FloatField) + large_scale_ice_flipped (FloatField) + large_scale_cloud_fraction (FloatField) + large_scale_cloud_fraction_flipped (FloatField) + convection_tracer (FloatField) + total_precipitable_water_initial (FloatFieldIJ) + saturation_total_precipitable_water_initial (FloatFieldIJ) + saturation_water_vapor (FloatFieldIJ) + """ + from __externals__ import CONVECTION_TRACER, ENTRVERSION, GF_ENV_SETTING, k_end + + # 1st setting: enviromental state is the one already modified by dyn + physics + with computation(PARALLEL), interval(...): + if GF_ENV_SETTING == 0: + dz = -(geopotential_height_interface[0, 0, 1] - geopotential_height_interface) + air_density = p / (287.04 * t * (1.0 + 0.608 * vapor)) + t_flipped = t.at(K=k_end - K) + p_flipped = p.at(K=k_end - K) + vapor_flipped = vapor.at(K=k_end - K) + vapor_current_flipped = vapor.at(K=k_end - K) + u_local = u.at(K=k_end - K) + v_local = v.at(K=k_end - K) + w_flipped = w.at(K=k_end - K) + layer_height_above_surface_flipped = layer_height_above_surface.at(K=k_end - K) + edge_height_above_surface_flipped = edge_height_above_surface.at(K=k_end - K) + mass_flipped = mass.at(K=k_end - K) + scalar_diffusivity_flipped = scalar_diffusivity.at(K=k_end - K) + + # Grid and sub-grid scale forcings for convection + grid_scale_forcing_t = 0.0 + grid_scale_forcing_vapor = 0.0 + subgrid_scale_forcing_t = 0.0 + subgrid_scale_forcing_vapor = 0.0 + advective_forcing_t = 0.0 + + with computation(FORWARD), interval(0, 1): + if GF_ENV_SETTING == 0: + p_surface_flipped = p_interface.at(K=k_end + 1) + + # 2nd setting: environmental state is that one before any tendency + # is applied (i.e, at begin of each time step). + # Get back the model state, heights and others variables at time N + # (or at the beggining of current time step) + # In physics, the state vars (t,u,v,PLE) are untouched and represent the + # model state after dynamics phase 1. But, "Q" is modified by physics, so + # depending on what was called before this subroutine, "Q" may be already + # changed from what it was just after dynamics phase 1. To solve this issue, + # "Q" just after dynamics is saved in the var named "QV_DYN_IN" in "GEOS_AgcmGridComp.F90". + with computation(PARALLEL), interval(...): + if GF_ENV_SETTING == 1: + mass_n = (p_interface_timestep_start[0, 0, 1] - p_interface_timestep_start) * ( + 1.0 / constants.MAPL_GRAV + ) + p_n = 0.5 * (p_interface_timestep_start + p_interface_timestep_start[0, 0, 1]) + p_kappa_interface_n = (p_interface_timestep_start / constants.MAPL_P00) ** ( + constants.MAPL_RGAS / constants.MAPL_CP + ) + if K == k_end: + p_kappa_surface_n = (p_interface_timestep_start[0, 0, 1] / constants.MAPL_P00) ** ( + constants.MAPL_RGAS / constants.MAPL_CP + ) + p_kappa_n = (p_n / constants.MAPL_P00) ** (constants.MAPL_RGAS / constants.MAPL_CP) + edge_height_above_surface_n = (t_timestep_start / p_kappa_n) * ( + 1.0 + constants.MAPL_VIREPS * vapor_timestep_start + ) + + with computation(BACKWARD), interval(...): + if GF_ENV_SETTING == 1: + if K == k_end: + layer_height_above_surface_n = ( + 0 + + (constants.MAPL_CP / constants.MAPL_GRAV) + * (p_kappa_surface_n - p_kappa_n) + * edge_height_above_surface_n + ) + else: + layer_height_above_surface_n = ( + edge_height_above_surface_n[0, 0, 1] + + (constants.MAPL_CP / constants.MAPL_GRAV) + * (p_kappa_interface_n[0, 0, 1] - p_kappa_n) + * edge_height_above_surface_n + ) + edge_height_above_surface_n = ( + layer_height_above_surface_n + + (constants.MAPL_CP / constants.MAPL_GRAV) + * (p_kappa_n - p_kappa_interface_n) + * edge_height_above_surface_n + ) + + with computation(PARALLEL), interval(...): + if GF_ENV_SETTING == 1: + if K == k_end: + dz = -(0 - edge_height_above_surface_n) + else: + dz = -(edge_height_above_surface_n[0, 0, 1] - edge_height_above_surface_n) + air_density = p_n / (287.04 * t_timestep_start * (1.0 + 0.608 * vapor_timestep_start)) + t_flipped = t_timestep_start.at(K=k_end - K) + p_flipped = p_n.at(K=k_end - K) + vapor_flipped = vapor_timestep_start.at(K=k_end - K) + vapor_current_flipped = vapor.at(K=k_end - K) + u_flipped = u_timestep_start.at(K=k_end - K) + v_flipped = v_timestep_start.at(K=k_end - K) + w_flipped = w.at(K=k_end - K) + layer_height_above_surface_flipped = layer_height_above_surface_n.at(K=k_end - K) + if K == 0: + edge_height_above_surface_flipped = 0 + else: + edge_height_above_surface_flipped = edge_height_above_surface_n.at(K=k_end - K + 1) + mass_flipped = mass_n.at(K=k_end - K) + scalar_diffusivity_flipped = scalar_diffusivity.at(K=k_end - K) + + # Grid and sub-grid scale forcings for convection + grid_scale_forcing_t = ( + dtdt_from_dynamics.at(K=k_end - K) + + dtdt_shortwave.at(K=k_end - K) + + dtdt_longwave.at(K=k_end - K) + ) + grid_scale_forcing_vapor = dvapordt_from_dynamics.at(K=k_end - K) + subgrid_scale_forcing_t = dtdt_pbl.at(K=k_end - K) + subgrid_scale_forcing_vapor = dspecific_humiditydt_pbl.at(K=k_end - K) + advective_forcing_t = dtdt_from_dynamics.at(K=k_end - K) + + with computation(FORWARD), interval(0, 1): + if GF_ENV_SETTING == 1: + p_surface_flipped = p_interface_timestep_start.at(K=k_end + 1) + + # remainder is the same for both settings + with computation(PARALLEL), interval(...): + convective_liquid_flipped = convective_liquid.at(K=k_end - K) + convective_ice_flipped = convective_ice.at(K=k_end - K) + convective_cloud_fraction_flipped = convective_cloud_fraction.at(K=k_end - K) + large_scale_liquid_flipped = large_scale_liquid.at(K=k_end - K) + large_scale_ice_flipped = large_scale_ice.at(K=k_end - K) + large_scale_cloud_fraction_flipped = large_scale_cloud_fraction.at(K=k_end - K) + + if ENTRVERSION == 0: + # eq 6 of https://doi.org/10.1029/2021JD034881 + lateral_entrainment_rate_flipped = ( + 0.71 * max(0.5, w.at(K=k_end - K)) ** (-1.17) * max(0.1, buoyancy.at(K=k_end - K)) ** (-0.36) + ) + else: + lateral_entrainment_rate_flipped = 1.0 + + with computation(PARALLEL), interval(...): + # must be in separate computation to ensure lateral_entrainment_rate_local is written before read + lateral_entrainment_rate = lateral_entrainment_rate_flipped.at(K=k_end - K) + + if CONVECTION_TRACER == 1: + buoyancy_excess = convection_tracer.at(K=k_end - K) + else: + buoyancy_excess = 0.0 + + with computation(FORWARD), interval(0, 1): + # saturation column_water_vapor + if CONVECTION_TRACER == 1: + saturation_water_vapor = total_precipitable_water_initial / ( + 1.0e-6 + saturation_total_precipitable_water_initial + ) + saturation_water_vapor = min(1.0, max(0.0, saturation_water_vapor)) + + +def copy_into_cumulus_parameterization_state( + grid_length_local: FloatFieldIJ, + grid_length: FloatFieldIJ, + saturation_water_vapor_local: FloatFieldIJ, + saturation_water_vapor: FloatFieldIJ, + seed_convection_model_state: FloatFieldIJ, + seed_convection: FloatFieldIJ, + convection_fraction_model_state: FloatFieldIJ, + convection_fraction: FloatFieldIJ, + surface_type_model_state: FloatFieldIJ, + surface_type: FloatFieldIJ, + grid_scale_forcing_t_local: FloatField, + grid_scale_forcing_t: FloatField, + grid_scale_forcing_vapor_local: FloatField, + grid_scale_forcing_vapor: FloatField, + subgrid_scale_forcing_t_local: FloatField, + subgrid_scale_forcing_t: FloatField, + subgrid_scale_forcing_vapor_local: FloatField, + subgrid_scale_forcing_vapor: FloatField, + lateral_entrainment_rate_flipped: FloatField, + lateral_entrainment_rate: FloatField, +): + """ + One-to-one copy fields into the CumulusParameterization state. + + This division has been inforced for readibility. Some of the earlier functions could + write directly to the CumulusParameterization state, but local copies have been + introduced for consistency, so that the data is being written to only one state at a time. + + Args: + grid_length_local (FloatFieldIJ) + grid_length (FloatFieldIJ) + saturation_water_vapor_local (FloatFieldIJ) + saturation_water_vapor (FloatFieldIJ) + seed_convection_model_state (FloatFieldIJ) + seed_convection (FloatFieldIJ) + convection_fraction_model_state (FloatFieldIJ) + convection_fraction (FloatFieldIJ) + surface_type_model_state (FloatFieldIJ) + surface_type (FloatFieldIJ) + grid_scale_forcing_t_local (FloatField) + grid_scale_forcing_t (FloatField) + grid_scale_forcing_vapor_local (FloatField) + grid_scale_forcing_vapor (FloatField) + subgrid_scale_forcing_t_local (FloatField) + subgrid_scale_forcing_t (FloatField) + subgrid_scale_forcing_vapor_local (FloatField) + subgrid_scale_forcing_vapor (FloatField) + lateral_entrainment_rate_flipped (FloatField) + lateral_entrainment_rate (FloatField) + """ + with computation(FORWARD), interval(0, 1): + grid_length = grid_length_local + saturation_water_vapor = saturation_water_vapor_local + seed_convection = seed_convection_model_state + convection_fraction = convection_fraction_model_state + surface_type = surface_type_model_state + + with computation(PARALLEL), interval(...): + grid_scale_forcing_t = grid_scale_forcing_t_local + grid_scale_forcing_vapor = grid_scale_forcing_vapor_local + subgrid_scale_forcing_t = subgrid_scale_forcing_t_local + subgrid_scale_forcing_vapor = subgrid_scale_forcing_vapor_local + lateral_entrainment_rate = lateral_entrainment_rate_flipped + + +def prepare_cumulus_paramaterization_state( + aot500: FloatFieldIJ, + ccn: FloatFieldIJ, + ocean_fraction_local: FloatFieldIJ, + ocean_fraction: FloatFieldIJ, + p_surface_flipped: FloatFieldIJ, + p_surface: FloatFieldIJ, + t_2m_flipped: FloatFieldIJ, + t_surface: FloatFieldIJ, + topography_height: FloatFieldIJ, + topography_height_no_negative: FloatFieldIJ, + pbl_level_flipped: IntFieldIJ, + pbl_level: IntFieldIJ, + latitude_model_state: FloatFieldIJ, + latitude: FloatFieldIJ, + longitude_model_state: FloatFieldIJ, + longitude: FloatFieldIJ, + rtgt: FloatFieldIJ, + geopotential_height_forced: FloatField, + layer_height_above_surface_flipped: FloatField, + edge_height_above_surface_flipped: FloatField, + p_flipped: FloatField, + p_forced: FloatField, + t_flipped: FloatField, + t_old: FloatField, + vapor_flipped: FloatField, + vapor_old: FloatField, + air_density: FloatField, + u_flipped: FloatField, + u: FloatField, + v_flipped: FloatField, + v: FloatField, + w_flipped: FloatField, + w: FloatField, + mass_flipped: FloatField, + mass: FloatField, + omega: FloatField, + buoyancy_excess_local: FloatField, + buoyancy_excess: FloatField, + advective_forcing_t: FloatField, + t_modified_by_advection: FloatField, + grid_scale_forcing_vapor: FloatField, + vapor_modified_by_advection: FloatField, + convective_liquid_flipped: FloatField, + convective_liquid: FloatField, + convective_ice_flipped: FloatField, + convective_ice: FloatField, + convective_cloud_fraction_flipped: FloatField, + convective_cloud_fraction: FloatField, + large_scale_liquid_flipped: FloatField, + large_scale_liquid: FloatField, + large_scale_ice_flipped: FloatField, + large_scale_ice: FloatField, + large_scale_cloud_fraction_flipped: FloatField, + large_scale_cloud_fraction: FloatField, + convection_tracers: FloatField_ConvectionTracers, + chemistry_tracers: FloatField_ConvectionTracers, + sensible_heat_flux_local: FloatFieldIJ, + sensible_heat_flux: FloatFieldIJ, + evaporation_local: FloatFieldIJ, + latent_heat_flux: FloatFieldIJ, + convective_scale_velocity: FloatFieldIJ, + t_excess: FloatFieldIJ, + vapor_excess: FloatFieldIJ, +): + """ + Compute inputs for the cumulus parameterization and fill the rest of the state with non-one-to-one-copies. + + Args: + aot500 (FloatFieldIJ) + ccn (FloatFieldIJ) + ocean_fraction_local (FloatFieldIJ) + ocean_fraction (FloatFieldIJ) + p_surface_flipped (FloatFieldIJ) + p_surface (FloatFieldIJ) + t_2m_flipped (FloatFieldIJ) + t_surface (FloatFieldIJ) + topography_height (FloatFieldIJ) + topography_height_no_negative (FloatFieldIJ) + pbl_level_flipped (IntFieldIJ) + pbl_level (IntFieldIJ) + latitude_model_state (FloatFieldIJ) + latitude (FloatFieldIJ) + longitude_model_state (FloatFieldIJ) + longitude (FloatFieldIJ) + rtgt (FloatFieldIJ) + geopotential_height_forced (FloatField) + layer_height_above_surface_flipped (FloatField) + edge_height_above_surface_flipped (FloatField) + p_flipped (FloatField) + p_forced (FloatField) + t_flipped (FloatField) + t_old (FloatField) + vapor_flipped (FloatField) + vapor_old (FloatField) + air_density (FloatField) + u_flipped (FloatField) + u (FloatField) + v_flipped (FloatField) + v (FloatField) + w_flipped (FloatField) + w (FloatField) + mass_flipped (FloatField) + mass (FloatField) + omega (FloatField) + buoyancy_excess_local (FloatField) + buoyancy_excess (FloatField) + advective_forcing_t (FloatField) + t_modified_by_advection (FloatField) + grid_scale_forcing_vapor (FloatField) + vapor_modified_by_advection (FloatField) + convective_liquid_flipped (FloatField) + convective_liquid (FloatField) + convective_ice_flipped (FloatField) + convective_ice (FloatField) + convective_cloud_fraction_flipped (FloatField) + convective_cloud_fraction (FloatField) + large_scale_liquid_flipped (FloatField) + large_scale_liquid (FloatField) + large_scale_ice_flipped (FloatField) + large_scale_ice (FloatField) + large_scale_cloud_fraction_flipped (FloatField) + large_scale_cloud_fraction (FloatField) + convection_tracers (FloatField_ConvectionTracers) + chemistry_tracers (FloatField_ConvectionTracers) + sensible_heat_flux_local (FloatFieldIJ) + sensible_heat_flux (FloatFieldIJ) + evaporation_local (FloatFieldIJ) + latent_heat_flux (FloatFieldIJ) + convective_scale_velocity (FloatFieldIJ) + t_excess (FloatFieldIJ) + vapor_excess (FloatFieldIJ) + """ + from __externals__ import APPLY_SUBSIDENCE_MICROPHYSICS, AUTOCONV, DT_MOIST, USE_TRACER_TRANSPORT, k_end + + with computation(FORWARD), interval(0, 1): + if AUTOCONV == 2: + ccn = max(100.0, (370.37 * (0.01 + max(0.0, aot500))) ** 1.555) + else: + ccn = 100.0 + + ocean_fraction = ocean_fraction_local + p_surface = p_surface_flipped * 1.0e-2 # mbar + t_surface = t_2m_flipped + topography_height_no_negative = max(0.0, topography_height) + pbl_level = pbl_level_flipped + latitude = latitude_model_state * 180.0 / 3.14159 + longitude = longitude_model_state * 180.0 / 3.14159 + + with computation(PARALLEL), interval(0, -1): + # heights, current pressure, temperature and water vapor mixing ratio + geopotential_height_forced = layer_height_above_surface_flipped * rtgt + topography_height + p_forced = p_flipped * 1.0e-2 # mbar + t_old = t_flipped + vapor_old = vapor_flipped # at beginning of the timestep + + # air density, TKE and cloud liquid water mixing ratio + air_density = 1.0e2 * p_forced / (287.04 * t_old * (1.0 + 0.608 * vapor_old)) + + # wind velocities + u = u_flipped + v = v_flipped + w = w_flipped + mass = mass_flipped + omega = -constants.MAPL_GRAV * air_density * w + + # buoyancy excess + buoyancy_excess = buoyancy_excess_local + + # temperature/water vapor modified only by advection + t_modified_by_advection = t_old + advective_forcing_t * DT_MOIST + vapor_modified_by_advection = vapor_old + grid_scale_forcing_vapor * DT_MOIST + + if APPLY_SUBSIDENCE_MICROPHYSICS == 1: + # microphysics ice and liquid mixing ratio, and cloud fraction of the host model + # (only subsidence is applied) + convective_liquid = convective_liquid_flipped + convective_ice = convective_ice_flipped + convective_cloud_fraction = convective_cloud_fraction_flipped + large_scale_liquid = large_scale_liquid_flipped + large_scale_ice = large_scale_ice_flipped + large_scale_cloud_fraction = large_scale_cloud_fraction_flipped + + with computation(PARALLEL), interval(...): + if USE_TRACER_TRANSPORT == 1: + tracer = 0 + while tracer < constants.NUMBER_OF_TRACERS: + chemistry_tracers[0, 0, 0][tracer] = max( + convection_tracers.at(K=k_end - K, ddim=[tracer]), constants.FLOAT_TINY + ) + tracer += 1 + + with computation(FORWARD), interval(0, 1): + pbl_internal: FloatFieldIJ = geopotential_height_forced.at(K=pbl_level) - topography_height + + # NOTE variables in this section have unintelligible names + # could not desipher during Fortran --> Python port, so Fortran names remain + with computation(FORWARD), interval(0, 1): + # get execess temperature and water vapor for source air parcels + pten = t_old.at(K=0) + pqen = vapor_old.at(K=0) + paph = 100.0 * p_surface + zrho = paph / (287.04 * (t_old.at(K=0) * (1.0 + 0.608 * vapor_old.at(K=0)))) + + # sensible and latent sfc fluxes for the heat-engine closure + sensible_heat_flux = zrho * cumulus_parameterization_constants.CP * sensible_heat_flux_local # W/m^2 + latent_heat_flux = zrho * cumulus_parameterization_constants.XLV * evaporation_local # W/m^2 + + # local le and h fluxes for W* + pahfs = -sensible_heat_flux_local * zrho * 1004.64 # W/m^2 + pqhfl = -evaporation_local # kg/m^2/s + + # buoyancy flux (h+le) + zkhvfl = (pahfs / 1004.64 + 0.608 * pten * pqhfl) / zrho # K m s-1 + + # depth of 1st model layer + # geopotential_height_forced.at(K=0) - topography_height is ~ 1/2 of the depth + # of 1st model layer, so multiply by 2 + pgeoh = ( + 2.0 * (geopotential_height_forced.at(K=0) - topography_height) * constants.MAPL_GRAV + ) # m+2 s-2 + + # convective-scale velocity w* + # in the future, change 0.001 by ustar^3 + convective_scale_velocity = max(0.0, 0.001 - 1.5 * 0.41 * zkhvfl * pgeoh / pten) # m+3 s-3 + + if convective_scale_velocity > constants.FLOAT_TINY: + # convective-scale velocity w* + convective_scale_velocity = 1.2 * convective_scale_velocity**0.3333 + + # temperature excess + t_excess = max(0.0, -1.5 * pahfs / (zrho * convective_scale_velocity * 1004.64)) # K + + # moisture excess + vapor_excess = max(0.0, -1.5 * pqhfl / (zrho * convective_scale_velocity)) # kg kg-1 + + # convective_scale_velocity for shallow convection closure (Grant 2001) + + # depth of the pbl + pgeoh = pbl_level * constants.MAPL_GRAV + + # convective_scale_velocity W* (m/s) + convective_scale_velocity = max(0.0, 0.001 - 1.5 * 0.41 * zkhvfl * pgeoh / pten) + convective_scale_velocity = 1.2 * convective_scale_velocity**0.3333 + + +class GF2020Setup(NDSLRuntime): + """ + This class performs the entire setup sequence for the GF2020 convection parameterization scheme + + In the source Fortran codee, this code is split across three subroutines nested as follows: + + - GF_Run + - GF2020_INTERFACE + - GF2020_DRV up to "------ CALL CUMULUS PARAMETERIZATION" + + This python implementation simplifies this structure by bringing all setup calculations to the same level. + An effort has been made to reduce duplicate/unnecessary locals where possible, but some have been retained + for the sake of readibility. + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: GF2020Config, + saturation_tables: SaturationVaporPressureTable, + ): + super().__init__(stencil_factory) + + # make inputs visible at runtime + self.stencil_factory = stencil_factory + self.config = config + self.saturation_tables = saturation_tables + + # check config for unimplemented paths + if config.ADV_TRIGGER == 2: + raise NotImplementedError( + "[NDSL] GF2020-->Setup initialized with shallow plume enabled. This requires" + "an umplemented portion of ensemble_output_and_feedback. Please impelment, then disable this" + "error manually to proceed." + ) + + # Construct stencils + self._compute_extra_inputs_from_state = stencil_factory.from_dims_halo( + func=compute_extra_inputs_from_state, + compute_dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], + externals={ + "STOCHASTIC_CONVECTION": config.STOCHASTIC_CNV, + "STOCH_TOP": config.STOCH_TOP, + "STOCH_BOT": config.STOCH_BOT, + "GF_MIN_AREA": config.GF_MIN_AREA, + "LHYDROSTATIC": config.LHYDROSTATIC, + }, + ) + + self._zero_state = stencil_factory.from_dims_halo( + func=zero_state, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._prefill_cumulus_parameterization_state = stencil_factory.from_dims_halo( + func=prefill_cumulus_parameterization_state, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "NUMBER_OF_PLUMES": cumulus_parameterization_constants.NUMBER_OF_PLUMES, + "APPLY_SUBSIDENCE_MICROPHYSICS": config.APPLY_SUBSIDENCE_MICROPHYSICS, + }, + ) + + self._prefill_locals = stencil_factory.from_dims_halo( + func=prefill_locals, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._set_2d_fields = stencil_factory.from_dims_halo( + func=set_2d_fields, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SIZE_I_DIM": stencil_factory.grid_indexing.get_shape([X_DIM])[0], + "SIZE_J_DIM": stencil_factory.grid_indexing.get_shape([Y_DIM])[0], + }, + ) + + self._choose_environment_and_flip_k_axis = stencil_factory.from_dims_halo( + func=choose_environment_and_flip_k_axis, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "GF_ENV_SETTING": config.GF_ENV_SETTING, + "ENTRVERSION": config.ENTRVERSION, + "CONVECTION_TRACER": config.CONVECTION_TRACER, + }, + ) + + self._copy_into_cumulus_parameterization_state = stencil_factory.from_dims_halo( + func=copy_into_cumulus_parameterization_state, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._prepare_cumulus_paramaterization_state = stencil_factory.from_dims_halo( + func=prepare_cumulus_paramaterization_state, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "AUTOCONV": config.AUTOCONV, + "DT_MOIST": config.DT_MOIST, + "APPLY_SUBSIDENCE_MICROPHYSICS": config.APPLY_SUBSIDENCE_MICROPHYSICS, + "USE_TRACER_TRANSPORT": config.USE_TRACER_TRANSPORT, + }, + ) + + def __call__( + self, + state: GF2020State, + locals: GF2020Locals, + cumulus_parameterization_state: GF2020CumulusParameterizationState, + convection_tracers: ConvectionTracers, + scm_stop: bool, + ): + """ + Perform setup calculations + + Args: + state (GF2020State): NDSL State containing all model fields required for GF2020. + locals (GF2020Locals): NDSL LocalState containing all locals for GF2020. + cumulus_parameterization_state (GF2020CumulusParameterizationState): NDSL State containing all + fields required for the CumulusParameterization. + convection_tracers (ConvectionTracers): Collection of tracers from the rest of the model which + will be updated within convection. These may come from a variaty of sources, and need to be + collected into the expected ConvectionTracers data type before being passed down. + scm_stop (bool): flag which can stop the execution of GF2020 + + """ + self._compute_extra_inputs_from_state( + p_interface=state.p_interface, + p=locals.derived_state.p, + p_kappa=locals.derived_state.p_kappa, + edge_height_above_surface=locals.derived_state.edge_height_above_surface, + layer_height_above_surface=locals.derived_state.layer_height_above_surface, + geopotential_height_interface=state.geopotential_height_interface, + t=state.t, + th=locals.derived_state.th, + vapor=state.vapor, + mass=locals.derived_state.mass, + w=state.w, + omega=state.omega, + vertical_motion=locals.derived_state.vertical_velocity, + tpwi=state.total_precipitable_water_initial, + tpwi_star=state.saturation_total_precipitable_water_initial, + seed_convection=locals.derived_state.seed_convection, + area=state.area, + modified_area=locals.derived_state.modified_area, + convection_fraction=state.convection_fraction, + ese=self.saturation_tables.ese, + esx=self.saturation_tables.esx, + ) + + self._zero_state( + dvapordt_deep_convection=state.dvapordt_deep_convection, + dtdt_deep_convection=state.dtdt_deep_convection, + dudt_deep_convection=state.dudt_deep_convection, + dvdt_deep_convection=state.dvdt_deep_convection, + sigma_deep=state.sigma_deep, + sigma_mid=state.sigma_mid, + mass_flux_shalow=state.mass_flux_shallow, + mass_flux_mid=state.mass_flux_mid, + mass_flux_deep_updraft=state.mass_flux_deep_updraft, + mass_flux_deep_updraft_interface=state.mass_flux_deep_updraft_interface, + mass_flux_deep_updraft_detrained=state.mass_flux_deep_updraft_detrained, + mass_flux_deep_downdraft=state.mass_flux_deep_downdraft, + mass_flux_cloud_base=state.mass_flux_cloud_base, + mass_flux_cloud_base_shallow=state.mass_flux_cloud_base_shallow, + mass_flux_cloud_base_mid=state.mass_flux_cloud_base_mid, + mass_flux_cloud_base_deep=state.mass_flux_cloud_base_deep, + convection_code_shallow=state.convection_code_shallow, + convection_code_mid=state.convection_code_mid, + convection_code_deep=state.convection_code_deep, + cloud_workfunction_0=state.cloud_workfunction_0, + cloud_workfunction_1=state.cloud_workfunction_1, + cloud_workfunction_2=state.cloud_workfunction_2, + cloud_workfunction_3=state.cloud_workfunction_3, + cloud_workfunction_1_pbl=state.cloud_workfunction_1_pbl, + cloud_workfunction_1_cin=state.cloud_workfunction_1_cin, + convective_precipitation_RAS=state.convective_precipitation_RAS, + convective_precipitation_GF=state.convective_precipitation_GF, + convective_condensate_source=state.convective_condensate_source, + convective_condensate_grid_mean=state.convective_condensate_grid_mean, + total_water_flux_deep_convection_interface=state.total_water_flux_deep_convection_interface, + updraft_area_fraction=state.updraft_areal_fraction, + updraft_vertical_velocity=state.updraft_vertical_velocity, + entrainment_parameter=state.entrainment_parameter, + lightning_density=state.lightning_density, + pbl_time_scale=state.pbl_time_scale, + cape_removal_time_scale=state.cape_removal_time_scale, + ) + + self._prefill_cumulus_parameterization_state( + error_code=cumulus_parameterization_state.output.error_code, + downdraft_origin_level=cumulus_parameterization_state.output.downdraft_origin_level, + lcl_level=cumulus_parameterization_state.output.lcl_level, + updraft_origin_level=cumulus_parameterization_state.output.updraft_origin_level, + updraft_lfc_level=cumulus_parameterization_state.output.updraft_lfc_level, + cloud_top_level=cumulus_parameterization_state.output.cloud_top_level, + kstabi=cumulus_parameterization_state.output.kstabi, + kstabm=cumulus_parameterization_state.output.kstabm, + precip=cumulus_parameterization_state.output.precip, + cloud_base_mass_flux_modified=cumulus_parameterization_state.output.cloud_base_mass_flux_modified, + epsilon_forced=cumulus_parameterization_state.output.epsilon_forced, + total_normalized_integrated_condensate_forced=cumulus_parameterization_state.output.total_normalized_integrated_condensate_forced, + scale_dependence_factor=cumulus_parameterization_state.output.scale_dependence_factor, + p_cloud_levels_forced=cumulus_parameterization_state.output.p_cloud_levels_forced, + entrainment_rate=cumulus_parameterization_state.output.entrainment_rate, + mass_entrainment_updraft_forced=cumulus_parameterization_state.output.mass_entrainment_updraft_forced, + mass_entrainment_downdraft_forced=cumulus_parameterization_state.output.mass_entrainment_downdraft_forced, + mass_detrainment_updraft_forced=cumulus_parameterization_state.output.mass_detrainment_updraft_forced, + mass_detrainment_downdraft_forced=cumulus_parameterization_state.output.mass_detrainment_downdraft_forced, + normalized_massflux_updraft_forced=cumulus_parameterization_state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=cumulus_parameterization_state.output.normalized_massflux_downdraft_forced, + condensate_to_fall_forced=cumulus_parameterization_state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=cumulus_parameterization_state.output.evaporate_in_downdraft_forced, + cloud_liquid_after_rain_forced=cumulus_parameterization_state.output.cloud_liquid_after_rain_forced, + t_updraft=cumulus_parameterization_state.output.t_updraft, + convective_cloud_fraction_output=cumulus_parameterization_state.output.convective_cloud_fraction, + dtdt=cumulus_parameterization_state.output.dtdt, + dudt=cumulus_parameterization_state.output.dudt, + dvdt=cumulus_parameterization_state.output.dvdt, + dvapordt=cumulus_parameterization_state.output.dvapordt, + dcloudicedt=cumulus_parameterization_state.output.dcloudicedt, + dnicedt=cumulus_parameterization_state.output.dnicedt, + dnliquiddt=cumulus_parameterization_state.output.dnliquiddt, + dbuoyancydt=cumulus_parameterization_state.output.dbuoyancydt, + chemistry_tracers_output=cumulus_parameterization_state.input_output.chemistry_tracers_output, + evaporation_sublimation_tendency=cumulus_parameterization_state.output.evaporation_sublimation_tendency, + convective_precip_flux=cumulus_parameterization_state.output.convective_precip_flux, + t_perturbation=cumulus_parameterization_state.output.t_perturbation, + omega=cumulus_parameterization_state.input_output.omega, + large_scale_ice=cumulus_parameterization_state.input_output.large_scale_ice, + convective_ice=cumulus_parameterization_state.input_output.convective_ice, + large_scale_liquid=cumulus_parameterization_state.input_output.large_scale_liquid, + convective_liquid=cumulus_parameterization_state.input_output.convective_liquid, + large_scale_cloud_fraction=cumulus_parameterization_state.input_output.large_scale_cloud_fraction, + convective_cloud_fraction=cumulus_parameterization_state.input_output.convective_cloud_fraction, + lightning_density=cumulus_parameterization_state.output.lightning_density, + t_excess=cumulus_parameterization_state.input.t_excess, + vapor_excess=cumulus_parameterization_state.input.vapor_excess, + last_error_code=cumulus_parameterization_state.input.last_error_code, + ) + + self._prefill_locals( + rtgt=locals.rtgt, + precip=locals.precip, + t_tendency_from_vapor=locals.t_tendency_from_vapor, + dtdt=locals.dtdt, + dvapordt=locals.dvapordt, + dcloudicedt=locals.dcloudicedt, + dudt=locals.dudt, + dvdt=locals.dvdt, + dlarge_scale_icedt=locals.dlarge_scale_icedt, + dconvective_icedt=locals.dconvective_icedt, + dlarge_scale_liquiddt=locals.dlarge_scale_liquiddt, + dconvective_liquiddt=locals.dconvective_liquiddt, + dlarge_scale_cloud_fractiondt=locals.dlarge_scale_cloud_fractiondt, + dconvective_cloud_fractiondt=locals.dconvective_cloud_fractiondt, + dbuoyancydt=locals.dbuoyancydt, + evaporation_sublimation_tendency=locals.evaporation_sublimation_tendency, + convective_precip_flux=locals.convective_precip_flux, + ) + + # workaround because max of full field cannot be determined inside a stencil + t_2m_max = state.t_2m.field.max() + + # # if surface temperature is not yet set in single column mode, stop the entire convection scheme + if self.stencil_factory.grid_indexing.get_shape([X_DIM, Y_DIM]) == (1, 1) and t_2m_max < 1.0e-6: + scm_stop = True + return + + self._set_2d_fields( + aot500=locals.aot500, + t=state.t, + t_2m_max=t_2m_max, + t_2m=state.t_2m, + t_2m_local=locals.flipped_copy.t_2m, + evaporation=state.evaporation, + evaporation_local=locals.evaporation, + sensible_heat_flux=state.sensible_heat_flux, + sensible_heat_flux_local=locals.sensible_heat_flux, + p_interface=state.p_interface, + vapor=state.vapor, + geopotential_height_surface=state.geopotential_height_surface, + topography_height=locals.topography_height, + land_fraction=state.land_fraction, + ocean_fraction=locals.ocean_fraction, + area=state.area, + grid_length=locals.grid_length, + pbl_level=state.pbl_level, + pbl_level_flipped=locals.flipped_copy.pbl_level, + ) + + self._choose_environment_and_flip_k_axis( + dz=locals.derived_state.dz, + air_density=locals.derived_state.air_density, + geopotential_height_interface=state.geopotential_height_interface, + t=state.t, + t_flipped=locals.flipped_copy.t, + t_timestep_start=state.t_timestep_start, + p=locals.derived_state.p, + p_interface=state.p_interface, + p_interface_timestep_start=state.p_interface_timestep_start, + p_flipped=locals.flipped_copy.p, + p_surface_flipped=locals.flipped_copy.p_surface, + vapor=state.vapor, + vapor_timestep_start=state.vapor_timestep_start, + vapor_flipped=locals.flipped_copy.vapor, + vapor_current_flipped=locals.flipped_copy.vapor_current, + u=state.u, + u_timestep_start=state.u_timestep_start, + u_flipped=locals.flipped_copy.u, + v=state.v, + v_timestep_start=state.v_timestep_start, + v_flipped=locals.flipped_copy.v, + w=state.w, + w_flipped=locals.flipped_copy.w, + layer_height_above_surface=locals.derived_state.layer_height_above_surface, + layer_height_above_surface_flipped=locals.flipped_copy.layer_height_above_surface, + edge_height_above_surface=locals.derived_state.edge_height_above_surface, + edge_height_above_surface_flipped=locals.flipped_copy.edge_height_above_surface, + mass=locals.derived_state.mass, + mass_flipped=locals.flipped_copy.mass, + scalar_diffusivity=state.scalar_diffusivity, + scalar_diffusivity_flipped=locals.flipped_copy.scalar_diffusivity, + lateral_entrainment_rate=state.lateral_entrainment_rate, + lateral_entrainment_rate_flipped=locals.flipped_copy.lateral_entrainment_rate, + buoyancy=state.buoyancy, + buoyancy_excess=locals.buoyancy_excess, + dtdt_shortwave=state.dtdt_shortwave, + dtdt_longwave=state.dtdt_longwave, + dtdt_from_dynamics=state.dtdt_from_dynamics, + dtdt_pbl=state.dtdt_pbl, + dvapordt_from_dynamics=state.dvapordt_from_dynamics, + dspecific_humiditydt_pbl=state.dspecific_humiditydt_pbl, + grid_scale_forcing_t=locals.grid_scale_forcing_t, + grid_scale_forcing_vapor=locals.grid_scale_forcing_vapor, + subgrid_scale_forcing_t=locals.subgrid_scale_forcing_t, + subgrid_scale_forcing_vapor=locals.subgrid_scale_forcing_vapor, + advective_forcing_t=locals.advective_forcing_t, + convective_liquid=state.convective_liquid, + convective_liquid_flipped=locals.flipped_copy.convective_liquid, + convective_ice=state.convective_ice, + convective_ice_flipped=locals.flipped_copy.convective_ice, + convective_cloud_fraction=state.convective_cloud_fraction, + convective_cloud_fraction_flipped=locals.flipped_copy.convective_cloud_fraction, + large_scale_liquid=state.large_scale_liquid, + large_scale_liquid_flipped=locals.flipped_copy.large_scale_liquid, + large_scale_ice=state.large_scale_ice, + large_scale_ice_flipped=locals.flipped_copy.large_scale_ice, + large_scale_cloud_fraction=state.large_scale_cloud_fraction, + large_scale_cloud_fraction_flipped=locals.flipped_copy.large_scale_cloud_fraction, + convection_tracer=state.convection_tracer, + total_precipitable_water_initial=state.total_precipitable_water_initial, + saturation_total_precipitable_water_initial=state.saturation_total_precipitable_water_initial, + saturation_water_vapor=locals.saturation_water_vapor, + ) + + if self.config.ADV_TRIGGER == 2: + raise NotImplementedError("option not implemented, should have been caught at initialization") + + self._copy_into_cumulus_parameterization_state( + grid_length_local=locals.grid_length, + grid_length=cumulus_parameterization_state.input_output.grid_length, + saturation_water_vapor_local=locals.saturation_water_vapor, + saturation_water_vapor=cumulus_parameterization_state.input.saturation_water_vapor, + seed_convection_model_state=state.seed_convection, + seed_convection=cumulus_parameterization_state.input.seed_convection, + convection_fraction_model_state=state.convection_fraction, + convection_fraction=cumulus_parameterization_state.input.convection_fraction, + surface_type_model_state=state.surface_type, + surface_type=cumulus_parameterization_state.input.surface_type, + grid_scale_forcing_t_local=locals.grid_scale_forcing_t, + grid_scale_forcing_t=cumulus_parameterization_state.input.grid_scale_forcing_t, + grid_scale_forcing_vapor_local=locals.grid_scale_forcing_vapor, + grid_scale_forcing_vapor=cumulus_parameterization_state.input.grid_scale_forcing_vapor, + subgrid_scale_forcing_t_local=locals.subgrid_scale_forcing_t, + subgrid_scale_forcing_t=cumulus_parameterization_state.input.subgrid_scale_forcing_t, + subgrid_scale_forcing_vapor_local=locals.subgrid_scale_forcing_vapor, + subgrid_scale_forcing_vapor=cumulus_parameterization_state.input.subgrid_scale_forcing_vapor, + lateral_entrainment_rate_flipped=locals.flipped_copy.lateral_entrainment_rate, + lateral_entrainment_rate=cumulus_parameterization_state.input.lateral_entrainment_rate, + ) + + self._prepare_cumulus_paramaterization_state( + aot500=locals.aot500, + ccn=cumulus_parameterization_state.input_output.ccn, + ocean_fraction_local=locals.ocean_fraction, + ocean_fraction=cumulus_parameterization_state.input.ocean_fraction, + p_surface_flipped=locals.flipped_copy.p_surface, + p_surface=cumulus_parameterization_state.input_output.p_surface, + t_2m_flipped=locals.flipped_copy.t_2m, + t_surface=cumulus_parameterization_state.input_output.t_surface, + topography_height=locals.topography_height, + topography_height_no_negative=cumulus_parameterization_state.input_output.topography_height_no_negative, + pbl_level_flipped=locals.flipped_copy.pbl_level, + pbl_level=cumulus_parameterization_state.input_output.pbl_level, + latitude_model_state=state.latitude, + latitude=cumulus_parameterization_state.input_output.latitude_degrees, + longitude_model_state=state.longitude, + longitude=cumulus_parameterization_state.input_output.longitude_degrees, + rtgt=locals.rtgt, + geopotential_height_forced=cumulus_parameterization_state.input_output.geopotential_height_forced, + layer_height_above_surface_flipped=locals.flipped_copy.layer_height_above_surface, + edge_height_above_surface_flipped=locals.flipped_copy.edge_height_above_surface, + p_flipped=locals.flipped_copy.p, + p_forced=cumulus_parameterization_state.input_output.p_forced, + t_flipped=locals.flipped_copy.t, + t_old=cumulus_parameterization_state.input_output.t_old, + vapor_flipped=locals.flipped_copy.vapor, + vapor_old=cumulus_parameterization_state.input_output.vapor_old, + air_density=cumulus_parameterization_state.input_output.air_density, + u_flipped=locals.flipped_copy.u, + u=cumulus_parameterization_state.input_output.u, + v_flipped=locals.flipped_copy.v, + v=cumulus_parameterization_state.input_output.v, + w_flipped=locals.flipped_copy.w, + w=cumulus_parameterization_state.input_output.w, + mass_flipped=locals.flipped_copy.mass, + mass=cumulus_parameterization_state.input_output.mass, + omega=cumulus_parameterization_state.input_output.omega, + buoyancy_excess_local=locals.buoyancy_excess, + buoyancy_excess=cumulus_parameterization_state.input_output.buoyancy_excess, + advective_forcing_t=locals.advective_forcing_t, + t_modified_by_advection=cumulus_parameterization_state.input_output.t_modified_by_advection, + grid_scale_forcing_vapor=cumulus_parameterization_state.input.grid_scale_forcing_vapor, + vapor_modified_by_advection=cumulus_parameterization_state.input_output.vapor_modified_by_advection, + convective_liquid_flipped=locals.flipped_copy.convective_liquid, + convective_liquid=cumulus_parameterization_state.input_output.convective_liquid, + convective_ice_flipped=locals.flipped_copy.convective_ice, + convective_ice=cumulus_parameterization_state.input_output.convective_ice, + convective_cloud_fraction_flipped=locals.flipped_copy.convective_cloud_fraction, + convective_cloud_fraction=cumulus_parameterization_state.input_output.convective_cloud_fraction, + large_scale_liquid_flipped=locals.flipped_copy.large_scale_liquid, + large_scale_liquid=cumulus_parameterization_state.input_output.large_scale_liquid, + large_scale_ice_flipped=locals.flipped_copy.large_scale_ice, + large_scale_ice=cumulus_parameterization_state.input_output.large_scale_ice, + large_scale_cloud_fraction_flipped=locals.flipped_copy.large_scale_cloud_fraction, + large_scale_cloud_fraction=cumulus_parameterization_state.input_output.large_scale_cloud_fraction, + convection_tracers=convection_tracers.tracers, + chemistry_tracers=cumulus_parameterization_state.input_output.chemistry_tracers, + sensible_heat_flux_local=locals.sensible_heat_flux, + sensible_heat_flux=cumulus_parameterization_state.input_output.sensible_heat_flux, + evaporation_local=locals.evaporation, + latent_heat_flux=cumulus_parameterization_state.input_output.latent_heat_flux, + convective_scale_velocity=cumulus_parameterization_state.input_output.convective_scale_velocity, + t_excess=cumulus_parameterization_state.input.t_excess, + vapor_excess=cumulus_parameterization_state.input.vapor_excess, + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/state.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/state.py new file mode 100644 index 000000000..a2ab29a62 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/state.py @@ -0,0 +1,900 @@ +import dataclasses + +from ndsl import Quantity, State +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM +from ndsl.dsl.typing import Float + + +@dataclasses.dataclass +class GF2020State(State): + latitude: Quantity = dataclasses.field( + metadata={ + "name": "latitude", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + longitude: Quantity = dataclasses.field( + metadata={ + "name": "longitude", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_interface: Quantity = dataclasses.field( + metadata={ + "name": "p_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t: Quantity = dataclasses.field( + metadata={ + "name": "t", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u: Quantity = dataclasses.field( + metadata={ + "name": "u", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v: Quantity = dataclasses.field( + metadata={ + "name": "v", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + w: Quantity = dataclasses.field( + metadata={ + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + omega: Quantity = dataclasses.field( + metadata={ + "name": "omega", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_2m: Quantity = dataclasses.field( + metadata={ + "name": "t_2m", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + specific_humidity_2m: Quantity = dataclasses.field( + metadata={ + "name": "specific_humidity_2m", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_surface: Quantity = dataclasses.field( + metadata={ + "name": "t_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + specific_humidity_surface: Quantity = dataclasses.field( + metadata={ + "name": "specific_humidity_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor: Quantity = dataclasses.field( + metadata={ + "name": "vapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_liquid: Quantity = dataclasses.field( + metadata={ + "name": "convective_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_ice: Quantity = dataclasses.field( + metadata={ + "name": "convective_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convective_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_liquid: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_liquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_ice: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_ice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + large_scale_cloud_fraction: Quantity = dataclasses.field( + metadata={ + "name": "large_scale_cloud_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ice_fraction_in_convective_tower: Quantity = dataclasses.field( + metadata={ + "name": "ice_fraction_in_convective_tower", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + p_interface_timestep_start: Quantity = dataclasses.field( + metadata={ + "name": "p_interface_timestep_start", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + t_timestep_start: Quantity = dataclasses.field( + metadata={ + "name": "t_timestep_start", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + u_timestep_start: Quantity = dataclasses.field( + metadata={ + "name": "u_timestep_start", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + v_timestep_start: Quantity = dataclasses.field( + metadata={ + "name": "v_timestep_start", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vapor_timestep_start: Quantity = dataclasses.field( + metadata={ + "name": "vapor_timestep_start", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_interface: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + geopotential_height_surface: Quantity = dataclasses.field( + metadata={ + "name": "geopotential_height_surface", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + area: Quantity = dataclasses.field( + metadata={ + "name": "area", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_level: Quantity = dataclasses.field( + metadata={ + "name": "pbl_level", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_fraction: Quantity = dataclasses.field( + metadata={ + "name": "convection_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + surface_type: Quantity = dataclasses.field( + metadata={ + "name": "surface_type", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + seed_convection: Quantity = dataclasses.field( + metadata={ + "name": "seed_convection", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + land_fraction: Quantity = dataclasses.field( + metadata={ + "name": "land_fraction", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + scalar_diffusivity: Quantity = dataclasses.field( + metadata={ + "name": "scalar_diffusivity", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + buoyancy: Quantity = dataclasses.field( + metadata={ + "name": "buoyancy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_precipitation_GF: Quantity = dataclasses.field( + metadata={ + "name": "convective_precipitation_GF", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_precipitation_RAS: Quantity = dataclasses.field( + metadata={ + "name": "convective_precipitation_RAS", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_rainwater_source: Quantity = dataclasses.field( + metadata={ + "name": "convective_rainwater_source", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sensible_heat_flux: Quantity = dataclasses.field( + metadata={ + "name": "sensible_heat_flux", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_water_flux_deep_convection_interface: Quantity = dataclasses.field( + metadata={ + "name": "total_water_flux_deep_convection_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sublimation_of_convective_precipitation: Quantity = dataclasses.field( + metadata={ + "name": "sublimation_of_convective_precipitation", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation_of_convective_precipitation: Quantity = dataclasses.field( + metadata={ + "name": "evaporation_of_convective_precipitation", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ice_precip_flux_interface: Quantity = dataclasses.field( + metadata={ + "name": "ice_precip_flux_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + liquid_precip_flux_interface: Quantity = dataclasses.field( + metadata={ + "name": "liquid_precip_flux_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + evaporation: Quantity = dataclasses.field( + metadata={ + "name": "evaporation", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_condensate_source: Quantity = dataclasses.field( + metadata={ + "name": "convective_condensate_source", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convective_condensate_grid_mean: Quantity = dataclasses.field( + metadata={ + "name": "convective_condensate_grid_mean", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + entrainment_parameter: Quantity = dataclasses.field( + metadata={ + "name": "entrainment_parameter", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate_shallow: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate_shallow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate_mid: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate_mid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lateral_entrainment_rate_deep: Quantity = dataclasses.field( + metadata={ + "name": "lateral_entrainment_rate_deep", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + updraft_areal_fraction: Quantity = dataclasses.field( + metadata={ + "name": "updraft_areal_fraction", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + updraft_vertical_velocity: Quantity = dataclasses.field( + metadata={ + "name": "updraft_vertical_velocity", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt_shortwave: Quantity = dataclasses.field( + metadata={ + "name": "dtdt_shortwave", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt_longwave: Quantity = dataclasses.field( + metadata={ + "name": "dtdt_longwave", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dspecific_humiditydt_pbl: Quantity = dataclasses.field( + metadata={ + "name": "dspecific_humiditydt_pbl", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt_pbl: Quantity = dataclasses.field( + metadata={ + "name": "dtdt_pbl", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt_from_dynamics: Quantity = dataclasses.field( + metadata={ + "name": "dtdt_from_dynamics", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvapordt_from_dynamics: Quantity = dataclasses.field( + metadata={ + "name": "dvapordt_from_dynamics", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sigma_mid: Quantity = dataclasses.field( + metadata={ + "name": "sigma_mid", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + sigma_deep: Quantity = dataclasses.field( + metadata={ + "name": "sigma_deep", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_precipitable_water_initial: Quantity = dataclasses.field( + metadata={ + "name": "total_precipitable_water_initial", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + saturation_total_precipitable_water_initial: Quantity = dataclasses.field( + metadata={ + "name": "saturation_total_precipitable_water_initial", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvapordt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dvapordt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dtdt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dtdt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dudt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dudt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dvdt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dvdt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dliquiddt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dliquiddt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dicedt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dicedt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + dcloudfractiondt_deep_convection: Quantity = dataclasses.field( + metadata={ + "name": "dcloudfractiondt_deep_convection", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pressure_shallow_convective_cloud_top: Quantity = dataclasses.field( + metadata={ + "name": "pressure_shallow_convective_cloud_top", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pressure_mid_convective_cloud_top: Quantity = dataclasses.field( + metadata={ + "name": "pressure_mid_convective_cloud_top", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pressure_deep_convective_cloud_top: Quantity = dataclasses.field( + metadata={ + "name": "pressure_deep_convective_cloud_top", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_shallow: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_shallow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_mid: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_mid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_deep_updraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_deep_updraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_deep_updraft_interface: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_deep_updraft_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_deep_updraft_detrained: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_deep_updraft_detrained", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_deep_downdraft: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_deep_downdraft", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_cloud_base: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_cloud_base", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_cloud_base_shallow: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_cloud_base_shallow", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_cloud_base_mid: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_cloud_base_mid", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + mass_flux_cloud_base_deep: Quantity = dataclasses.field( + metadata={ + "name": "mass_flux_cloud_base_deep", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_cumulative_mass_flux_interface: Quantity = dataclasses.field( + metadata={ + "name": "total_cumulative_mass_flux_interface", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + total_detraining_mass_flux: Quantity = dataclasses.field( + metadata={ + "name": "total_detraining_mass_flux", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_code_shallow: Quantity = dataclasses.field( + metadata={ + "name": "convection_code_shallow", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_code_mid: Quantity = dataclasses.field( + metadata={ + "name": "convection_code_mid", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_code_deep: Quantity = dataclasses.field( + metadata={ + "name": "convection_code_deep", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_0: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_0", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_1", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_2: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_2", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_3: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_3", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_pbl: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_1_pbl", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cloud_workfunction_1_cin: Quantity = dataclasses.field( + metadata={ + "name": "cloud_work_function_1_cin", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + pbl_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "pbl_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + cape_removal_time_scale: Quantity = dataclasses.field( + metadata={ + "name": "cape_removal_time_scale", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + lightning_density: Quantity = dataclasses.field( + metadata={ + "name": "lighting_density", + "dims": [X_DIM, Y_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convection_tracer: Quantity = dataclasses.field( + metadata={ + "name": "convection_tracer", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection_tracers.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection_tracers.py new file mode 100644 index 000000000..e10d8c1a9 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection_tracers.py @@ -0,0 +1,136 @@ +import dataclasses + +from ndsl import Quantity, State +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import Bool, Float + + +@dataclasses.dataclass +class ConvectionTracers(State): + """ + Dataclass of Convection Tracers, contains both the numerical data of the tracers + (stored in the "tracer" field) and metadata, each stored in its off-grid field + + Must be initialized with the following extra dimensions: + "convection_tracers": number of convective tracers, must be defined prior to initalization + "size_three_dimension": fixed dimension of size three for metadata + "size_four_dimension": fixed dimension of size four for metadata + """ + + tracers: Quantity = dataclasses.field( + metadata={ + "name": "tracers", + "dims": [X_DIM, Y_DIM, Z_DIM, "convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + fscav: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + vect_hcts: Quantity = dataclasses.field( + metadata={ + "name": "vect_hcts", + "dims": ["convection_tracers", "size_four_dimension"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + kc_scal: Quantity = dataclasses.field( + metadata={ + "name": "kc_scal", + "dims": ["convection_tracers", "size_three_dimension"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + convfaci2g: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + retfactor: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + liq_and_gas: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + online_cldliq: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + online_vud: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + ftemp_threshold: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Float, + } + ) + use_gcc_washout: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Bool, + } + ) + use_gocart: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Bool, + } + ) + is_wetdep: Quantity = dataclasses.field( + metadata={ + "name": "fscav", + "dims": ["convection_tracers"], + "units": "?", + "intent": "?", + "dtype": Bool, + } + ) diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/field_types.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/field_types.py index 6c09ad617..1e68c2d94 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/field_types.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/field_types.py @@ -1,10 +1,12 @@ -import gt4py.cartesian.gtscript as gtscript - -from ndsl.dsl.gt4py import IJK, Field -from ndsl.dsl.typing import Float +from ndsl.dsl.gt4py import IJ, IJK, Field, GlobalTable +from ndsl.dsl.typing import Bool, Float from pyMoist.constants import N_MODES, NCNST FloatField_NModes = Field[IJK, (Float, (N_MODES))] -FloatField_NTracers = gtscript.Field[gtscript.IJK, (Float, (int(NCNST)))] -FloatFieldIJ_NTracers = gtscript.Field[gtscript.IJ, (Float, (int(NCNST)))] +FloatField_NTracers = Field[IJK, (Float, (int(NCNST)))] +FloatFieldIJ_NTracers = Field[IJ, (Float, (int(NCNST)))] +ConvectionTracerMetaDataTable_Float = GlobalTable[(Float, int(NCNST))] +ConvectionTracerMetaDataTable_Bool = GlobalTable[(Bool, int(NCNST))] +ConvectionTracerMetaDataTable_x3 = GlobalTable[(Float, (int(NCNST), 3))] +ConvectionTracerMetaDataTable_x4 = GlobalTable[(Float, (int(NCNST), 4))] diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/shared_incloud_processes.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/shared_incloud_processes.py index 4cf9a4c2c..bf83c01df 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/shared_incloud_processes.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/shared_incloud_processes.py @@ -3,7 +3,20 @@ processes/quantities.""" import pyMoist.constants as constants -from ndsl.dsl.gt4py import PARALLEL, computation, exp, function, interval, log10, sin +from ndsl.dsl.gt4py import ( + PARALLEL, + GlobalTable, + computation, + exp, + float32, + float64, + floor, + function, + interval, + log10, + round_away_from_zero, + sin, +) from ndsl.dsl.typing import Float, FloatField from pyMoist.shared_generic_math import air_density @@ -263,3 +276,232 @@ def fix_up_clouds( large_scale_cloud_fraction = 0.0 large_scale_liquid = 0.0 large_scale_ice = 0.0 + + +# able of lookup values of radiative effective radius of ice crystals as a function of temperature from +# -94C to 0C for make_ice_number. Taken from WRF RRTMG radiation code where it is attributed to +# Jon Egill Kristjansson and coauthors. This must be built into a custom shape off-grid quantity, +# passed into the stencil which calls make_ice_number with a custom field type (defined below), then +# passed into the function make_ice_number and accessed with .A[index] to function properly +RADIATIVE_EFFECTIVE_RADIUS = [ + 5.92779, + 6.26422, + 6.61973, + 6.99539, + 7.39234, + 7.81177, + 8.25496, + 8.72323, + 9.21800, + 9.74075, + 10.2930, + 10.8765, + 11.4929, + 12.1440, + 12.8317, + 13.5581, + 14.2319, + 15.0351, + 15.8799, + 16.7674, + 17.6986, + 18.6744, + 19.6955, + 20.7623, + 21.8757, + 23.0364, + 24.2452, + 25.5034, + 26.8125, + 27.7895, + 28.6450, + 29.4167, + 30.1088, + 30.7306, + 31.2943, + 31.8151, + 32.3077, + 32.7870, + 33.2657, + 33.7540, + 34.2601, + 34.7892, + 35.3442, + 35.9255, + 36.5316, + 37.1602, + 37.8078, + 38.4720, + 39.1508, + 39.8442, + 40.5552, + 41.2912, + 42.0635, + 42.8876, + 43.7863, + 44.7853, + 45.9170, + 47.2165, + 48.7221, + 50.4710, + 52.4980, + 54.8315, + 57.4898, + 60.4785, + 63.7898, + 65.5604, + 71.2885, + 75.4113, + 79.7368, + 84.2351, + 88.8833, + 93.6658, + 98.5739, + 103.603, + 108.752, + 114.025, + 119.424, + 124.954, + 130.630, + 136.457, + 142.446, + 148.608, + 154.956, + 161.503, + 168.262, + 175.248, + 182.473, + 189.952, + 197.699, + 205.728, + 214.055, + 222.694, + 231.661, + 240.971, + 250.639, +] +RADIATIVE_EFFECTIVE_RADIUS_Table_Type = GlobalTable[(Float, len(RADIATIVE_EFFECTIVE_RADIUS))] + + +@function +def make_ice_number( + cloud_ice_mixing_ratio, + t, + RADIATIVE_EFFECTIVE_RADIUS, +): + """ + Get the ice crystal number given cloud ice mixing ratio and temperature. + Returns the number of droplets per kg per m3. + + Args: + cloud_ice_mixing_ratio (in): units kg/m3 + t (in): units K + RADIATIVE_EFFECTIVE_RADIUS: table used for calculations + + Returns: + crystal_number: units number/(kg*m3) + + Developed by H. Barnes @ NOAA/OAR/ESRL/GSL Earth Prediction Advancement Division + """ + + # DEBUG + crystal_number = 0.0 + # internal constant + ice_density = 890.0 + + if cloud_ice_mixing_ratio == 0.0: + crystal_number = 0.0 + + else: + # From the model 3D temperature field, subtract 180K for which + # index value of RADIATIVE_EFFECTIVE_RADIUS as a start. Value of corr is for + # interpolating between neighboring values in the table. + + idx_rei = int(t - 180.0) + idx_rei = min(max(idx_rei, 0), 93) + corr = t - floor(t) + reice = ( + RADIATIVE_EFFECTIVE_RADIUS.A[idx_rei] * (1.0 - corr) + + RADIATIVE_EFFECTIVE_RADIUS.A[idx_rei + 1] * corr + ) + deice = 2.0 * reice * 1.0e-6 + + internal_lambda = float64(3.0 / deice) + + # value of the dispersion parameter according to Heymsfield et al 2002, Table3. + t_celcius = t - 273.15 + + t_celcius = min(max(t_celcius, -70.0), -15.0) + + if t_celcius > -27.0: + lambdai = 6.8 * exp(-0.096 * t_celcius) + else: + lambdai = 24.8 * exp(-0.049 * t_celcius) + + mui = (0.13 * (lambdai**0.64)) - 2.0 + + k = (mui + 3) * (mui * 3) / (mui + 2) / (mui + 1) + + crystal_number = ( + k + * cloud_ice_mixing_ratio + * internal_lambda + * internal_lambda + * internal_lambda + / (constants.MAPL_PI * ice_density) + ) + + return crystal_number + + +# table of constants for make_droplet_number, this must be built into a custom shape off-grid quantity, +# passed into the stencil which calls make_droplet_number with a custom field type (defined below), then +# passed into the function make_droplet_number and accessed with .A[index] to function properly +G_RATIO = [24, 60, 120, 210, 336, 504, 720, 990, 1320, 1716, 2184, 2730, 3360, 4080, 4896] +G_RATIO_Table_Type = GlobalTable[(Float, len(G_RATIO))] + + +@function +def make_droplet_number( + cloud_water_mixing_ratio, + num_water_friendly_aerosols, + G_RATIO, +): + """ + Get the droplet number given cloud water mixing ratio and number of water-friendly aerosols. + Returns the number of droplets per kg per m3. + + Args: + cloud_water_mixing_ratio (in): units kg/m3 + num_water_friendly_aerosols (in): units number/kg + G_RATIO: table used for calculations + + Returns: + droplet_number: units number/(kg*m3) + + Developed by H. Barnes @ NOAA/OAR/ESRL/GSL Earth Prediction Advancement Division + """ + am_r = constants.MAPL_PI * 1000.0 / 6.0 + + if cloud_water_mixing_ratio <= 0.0: + droplet_number = 0.0 + else: + internal_num_water_friendly_aerosols = max(99.0e6, min(num_water_friendly_aerosols, 5.0e10)) + nu_c = max(2, min(round_away_from_zero(2.5e10 / internal_num_water_friendly_aerosols), 15)) + + x1 = max(1.0, min(internal_num_water_friendly_aerosols * 1.0e-9, 10.0)) - 1.0 + xDc = (30.0 - x1 * 20.0 / 9.0) * 1.0e-6 + + internal_lambda = (float64(4.0) + nu_c) / xDc + + qnc = ( + cloud_water_mixing_ratio + / G_RATIO.A[int(nu_c - 1)] + * internal_lambda + * internal_lambda + * internal_lambda + / am_r + ) + droplet_number = float32(qnc) + + return droplet_number diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/__init__.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/__init__.py index 3d64ea923..388aeb9eb 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/__init__.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/__init__.py @@ -1,3 +1,305 @@ +from .convection.GF_2020.cumulus_parameterization.air_density.translate_GF2020_CumulusParameterization_HydrostaticAirDensity import ( + TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_deep, + TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_mid, + TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_shallow, +) +from .convection.GF_2020.cumulus_parameterization.buoyancy.translate_GF2020_CumulusParameterization_GetBuoyancy_1 import ( + TranslateGF2020_CumulusParameterization_GetBuoyancy_1_deep, + TranslateGF2020_CumulusParameterization_GetBuoyancy_1_mid, + TranslateGF2020_CumulusParameterization_GetBuoyancy_1_shallow, +) +from .convection.GF_2020.cumulus_parameterization.buoyancy.translate_GF2020_CumulusParameterization_GetBuoyancy_2 import ( + TranslateGF2020_CumulusParameterization_GetBuoyancy_2_deep, + TranslateGF2020_CumulusParameterization_GetBuoyancy_2_mid, + TranslateGF2020_CumulusParameterization_GetBuoyancy_2_shallow, +) +from .convection.GF_2020.cumulus_parameterization.buoyancy.translate_GF2020_CumulusParameterization_GetBuoyancy_3 import ( + TranslateGF2020_CumulusParameterization_GetBuoyancy_3_deep, + TranslateGF2020_CumulusParameterization_GetBuoyancy_3_mid, + TranslateGF2020_CumulusParameterization_GetBuoyancy_3_shallow, +) +from .convection.GF_2020.cumulus_parameterization.convection_tracers.translate_GF2020_Cumulus_Parameterization_AtmosphericComposition import ( + TranslateGF2020_CumulusParameterization_AtmosphericComposition_deep, + TranslateGF2020_CumulusParameterization_AtmosphericComposition_mid, + TranslateGF2020_CumulusParameterization_AtmosphericComposition_shallow, +) +from .convection.GF_2020.cumulus_parameterization.diurnal_cycle.translate_GF2020_CumulusParameterization_DiurnalCycle import ( + TranslateGF2020_CumulusParameterization_DiurnalCycle_deep, + TranslateGF2020_CumulusParameterization_DiurnalCycle_mid, + TranslateGF2020_CumulusParameterization_DiurnalCycle_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftLateralMassFlux import ( + TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_deep, + TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_mid, + TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftMassFlux import ( + TranslateGF2020_CumulusParameterization_DowndraftMassFlux_deep, + TranslateGF2020_CumulusParameterization_DowndraftMassFlux_mid, + TranslateGF2020_CumulusParameterization_DowndraftMassFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftMoisture import ( + TranslateGF2020_CumulusParameterization_DowndraftMoisture_deep, + TranslateGF2020_CumulusParameterization_DowndraftMoisture_mid, + TranslateGF2020_CumulusParameterization_DowndraftMoisture_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftMSEAndBuoyancy import ( + TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_deep, + TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_mid, + TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftOriginLevel import ( + TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_deep, + TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_mid, + TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftTemperature import ( + TranslateGF2020_CumulusParameterization_DowndraftTemperature_deep, + TranslateGF2020_CumulusParameterization_DowndraftTemperature_mid, + TranslateGF2020_CumulusParameterization_DowndraftTemperature_shallow, +) +from .convection.GF_2020.cumulus_parameterization.downdraft.translate_GF2020_CumulusParameterization_DowndraftWindShear import ( + TranslateGF2020_CumulusParameterization_DowndraftWindShear_deep, + TranslateGF2020_CumulusParameterization_DowndraftWindShear_mid, + TranslateGF2020_CumulusParameterization_DowndraftWindShear_shallow, +) +from .convection.GF_2020.cumulus_parameterization.entrainment.translate_GF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment import ( + TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_deep, + TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_mid, + TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_shallow, +) +from .convection.GF_2020.cumulus_parameterization.entrainment.translate_GF2020_CumulusParameterization_DowndraftEntrainmentProfiles import ( + TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_deep, + TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_mid, + TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_shallow, +) +from .convection.GF_2020.cumulus_parameterization.entrainment.translate_GF2020_CumulusParameterization_EntrainmentRates import ( + TranslateGF2020_CumulusParameterization_EntrainmentRates_deep, + TranslateGF2020_CumulusParameterization_EntrainmentRates_mid, + TranslateGF2020_CumulusParameterization_EntrainmentRates_shallow, +) +from .convection.GF_2020.cumulus_parameterization.entrainment.translate_GF2020_CumulusParameterization_StableDetrainment import ( + TranslateGF2020_CumulusParameterization_StableDetrainment_deep, + TranslateGF2020_CumulusParameterization_StableDetrainment_mid, + TranslateGF2020_CumulusParameterization_StableDetrainment_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_1 import ( + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_deep, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_mid, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_2 import ( + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_deep, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_mid, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_3 import ( + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_deep, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_mid, + TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentConditions_1 import ( + TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_deep, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_mid, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentConditions_2 import ( + TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_deep, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_mid, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentConditions_3 import ( + TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_deep, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_mid, + TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_EnvironmentMassFlux import ( + TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_deep, + TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_mid, + TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.environment.translate_GF2020_CumulusParameterization_ModifyEnvironmentProfiles import ( + TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_deep, + TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_mid, + TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_shallow, +) +from .convection.GF_2020.cumulus_parameterization.get_levels.translate_GF2020_CumulusParameterization_CloudTop import ( + TranslateGF2020_CumulusParameterization_CloudTop_deep, + TranslateGF2020_CumulusParameterization_CloudTop_mid, + TranslateGF2020_CumulusParameterization_CloudTop_shallow, +) +from .convection.GF_2020.cumulus_parameterization.get_levels.translate_GF2020_CumulusParameterization_ConvectiveCloudBaseLevel import ( + TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_deep, + TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_mid, + TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_shallow, +) +from .convection.GF_2020.cumulus_parameterization.get_levels.translate_GF2020_CumulusParameterization_GetLCL import ( + TranslateGF2020_CumulusParameterization_GetLCL_deep, + TranslateGF2020_CumulusParameterization_GetLCL_mid, + TranslateGF2020_CumulusParameterization_GetLCL_shallow, +) +from .convection.GF_2020.cumulus_parameterization.get_levels.translate_GF2020_CumulusParameterization_HighestMoistStaticEnergyLevel import ( + TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_deep, + TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_mid, + TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_shallow, +) +from .convection.GF_2020.cumulus_parameterization.kinetic_energy_to_heating.translate_GF2020_CumulusParameterization_KineticEnergyToHeating import ( + TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_deep, + TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_mid, + TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_shallow, +) +from .convection.GF_2020.cumulus_parameterization.large_scale_forcing.translate_GF2020_CumulusParameterization_LargeScaleForcing import ( + TranslateGF2020_CumulusParameterization_LargeScaleForcing_deep, + TranslateGF2020_CumulusParameterization_LargeScaleForcing_mid, + TranslateGF2020_CumulusParameterization_LargeScaleForcing_shallow, +) +from .convection.GF_2020.cumulus_parameterization.moist_static_energy.translate_GF2020_CumulusParameterization_FirstGuessMoistStaticEnergy import ( + TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_deep, + TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_mid, + TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_shallow, +) +from .convection.GF_2020.cumulus_parameterization.moist_static_energy.translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_1 import ( + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_deep, + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_mid, + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_shallow, +) +from .convection.GF_2020.cumulus_parameterization.moist_static_energy.translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_2 import ( + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_deep, + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_mid, + TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_shallow, +) +from .convection.GF_2020.cumulus_parameterization.moist_static_energy.translate_GF2020_CumulusParameterization_StaticControl import ( + TranslateGF2020_CumulusParameterization_StaticControl_deep, + TranslateGF2020_CumulusParameterization_StaticControl_mid, + TranslateGF2020_CumulusParameterization_StaticControl_shallow, +) +from .convection.GF_2020.cumulus_parameterization.precip.translate_GF2020_CumulusParameterization_CloudDissipation import ( + TranslateGF2020_CumulusParameterization_CloudDissipation_deep, + TranslateGF2020_CumulusParameterization_CloudDissipation_mid, + TranslateGF2020_CumulusParameterization_CloudDissipation_shallow, +) +from .convection.GF_2020.cumulus_parameterization.precip.translate_GF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels import ( + TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_deep, + TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_mid, + TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_shallow, +) +from .convection.GF_2020.cumulus_parameterization.precip.translate_GF2020_CumulusParameterization_PrecipitationFlux import ( + TranslateGF2020_CumulusParameterization_PrecipitationFlux_deep, + TranslateGF2020_CumulusParameterization_PrecipitationFlux_mid, + TranslateGF2020_CumulusParameterization_PrecipitationFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.precip.translate_GF2020_CumulusParameterization_RainEvaporationBelowCloudBase import ( + TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_deep, + TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_mid, + TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_DeepPrecipitationOutput import ( + TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_deep, + TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_mid, + TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_EnsembleOutputAndFeedback import ( + TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_deep, + TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_mid, + TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_OutputUpdraftTemperature import ( + TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_deep, + TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_mid, + TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations import ( + TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_deep, + TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_mid, + TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_PrepareOutput import ( + TranslateGF2020_CumulusParameterization_PrepareOutput_deep, + TranslateGF2020_CumulusParameterization_PrepareOutput_mid, + TranslateGF2020_CumulusParameterization_PrepareOutput_shallow, +) +from .convection.GF_2020.cumulus_parameterization.prepare_output.translate_GF2020_CumulusParameterization_TotalEvaporationFlux import ( + TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_deep, + TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_mid, + TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.profiles.translate_GF2020_CumulusParameterization_C1DProfile import ( + TranslateGF2020_CumulusParameterization_C1DProfile_deep, + TranslateGF2020_CumulusParameterization_C1DProfile_mid, + TranslateGF2020_CumulusParameterization_C1DProfile_shallow, +) +from .convection.GF_2020.cumulus_parameterization.profiles.translate_GF2020_CumulusParameterization_MeltingProfile import ( + TranslateGF2020_CumulusParameterization_MeltingProfile_deep, + TranslateGF2020_CumulusParameterization_MeltingProfile_mid, + TranslateGF2020_CumulusParameterization_MeltingProfile_shallow, +) +from .convection.GF_2020.cumulus_parameterization.setup.translate_GF2020_CumulusParameterization_Setup import ( + TranslateGF2020_CumulusParameterization_Setup_deep, + TranslateGF2020_CumulusParameterization_Setup_mid, + TranslateGF2020_CumulusParameterization_Setup_shallow, +) +from .convection.GF_2020.cumulus_parameterization.smoothing.translate_GF2020_CumulusParameterization_SmoothTendencies import ( + TranslateGF2020_CumulusParameterization_SmoothTendencies_deep, + TranslateGF2020_CumulusParameterization_SmoothTendencies_mid, + TranslateGF2020_CumulusParameterization_SmoothTendencies_shallow, +) +from .convection.GF_2020.cumulus_parameterization.translate_GF2020_CumulusParameterization import ( + TranslateGF2020_CumulusParameterization, +) +from .convection.GF_2020.cumulus_parameterization.triggers.translate_GF2020_CumulusParameterization_ConvectionTrigger import ( + TranslateGF2020_CumulusParameterization_ConvectionTrigger_deep, + TranslateGF2020_CumulusParameterization_ConvectionTrigger_mid, + TranslateGF2020_CumulusParameterization_ConvectionTrigger_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble import ( + TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_deep, + TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_mid, + TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftCIN import ( + TranslateGF2020_CumulusParameterization_UpdraftCIN_deep, + TranslateGF2020_CumulusParameterization_UpdraftCIN_mid, + TranslateGF2020_CumulusParameterization_UpdraftCIN_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftInitialWorkfunctions import ( + TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_deep, + TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_mid, + TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftMassFlux import ( + TranslateGF2020_CumulusParameterization_UpdraftMassFlux_deep, + TranslateGF2020_CumulusParameterization_UpdraftMassFlux_mid, + TranslateGF2020_CumulusParameterization_UpdraftMassFlux_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget import ( + TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_deep, + TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_mid, + TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftMoisture import ( + TranslateGF2020_CumulusParameterization_UpdraftMoisture_deep, + TranslateGF2020_CumulusParameterization_UpdraftMoisture_mid, + TranslateGF2020_CumulusParameterization_UpdraftMoisture_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftTemperature import ( + TranslateGF2020_CumulusParameterization_UpdraftTemperature_deep, + TranslateGF2020_CumulusParameterization_UpdraftTemperature_mid, + TranslateGF2020_CumulusParameterization_UpdraftTemperature_shallow, +) +from .convection.GF_2020.cumulus_parameterization.updraft.translate_GF2020_CumulusParameterization_UpdraftVerticalVelocity import ( + TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_deep, + TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_mid, + TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_shallow, +) +from .convection.GF_2020.cumulus_parameterization.vertical_discretization.translate_GF2020_CumulusParameterization_VerticalDiscretization import ( + TranslateGF2020_CumulusParameterization_VerticalDiscretization_deep, + TranslateGF2020_CumulusParameterization_VerticalDiscretization_mid, + TranslateGF2020_CumulusParameterization_VerticalDiscretization_shallow, +) + +from .convection.GF_2020.translate_GF2020 import TranslateGF2020 +from .convection.GF_2020.translate_GF2020_Finalize import TranslateGF2020_Finalize +from .convection.GF_2020.translate_GF2020_Setup import TranslateGF2020_Setup from .GFDL_1M.driver.translate_GFDL_1M_Driver import TranslateGFDL_1M_Driver from .GFDL_1M.driver.translate_GFDL_1M_DriverFinish import TranslateGFDL_1M_DriverFinish from .GFDL_1M.driver.translate_GFDL_1M_DriverSetup import TranslateGFDL_1M_DriverSetup diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/air_density/translate_GF2020_CumulusParameterization_HydrostaticAirDensity.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/air_density/translate_GF2020_CumulusParameterization_HydrostaticAirDensity.py new file mode 100644 index 000000000..4e6798143 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/air_density/translate_GF2020_CumulusParameterization_HydrostaticAirDensity.py @@ -0,0 +1,179 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.air_density import hydrostatic_air_density +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_geopotential_height_cloud_levels_forced": {}, + "p_cloud_levels_forced": {}, + "error_code": {}, + "local_hydrostatic_air_density": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.hydrostatic_air_density.data[:] = inputs["local_hydrostatic_air_density"] + + code = self.stencil_factory.from_dims_halo( + func=hydrostatic_air_density, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.output.p_cloud_levels_forced, + geopotential_height=locals.geopotential_height_cloud_levels_forced, + error_code=state.output.error_code, + air_density=locals.hydrostatic_air_density, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_hydrostatic_air_density": locals.hydrostatic_air_density.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_HydrostaticAirDensity_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_1.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_1.py new file mode 100644 index 000000000..f3cf72b31 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_1.py @@ -0,0 +1,213 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.buoyancy import get_buoyancy +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "lcl_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_d_buoyancy_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=get_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=locals.cloud_moist_static_energy_forced, + environment_moist_static_energy=locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + d_buoyancy=locals.d_buoyancy_forced, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_d_buoyancy_forced": locals.d_buoyancy_forced.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_1_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_1_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_1_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_2.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_2.py new file mode 100644 index 000000000..ecd080f63 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_2.py @@ -0,0 +1,209 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.buoyancy import get_buoyancy +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "lcl_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_cloud_moist_static_energy": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "local_env_saturation_moist_static_energy_cloud_levels": {}, + "local_d_buoyancy": {}, + "local_geopotential_height_cloud_levels": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.cloud_moist_static_energy.data[:] = inputs["local_cloud_moist_static_energy"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + locals.environment_saturation_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels" + ] + locals.d_buoyancy.data[:] = inputs["local_d_buoyancy"] + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=get_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=locals.cloud_moist_static_energy, + environment_moist_static_energy=locals.environment_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_cloud_levels, + d_buoyancy=locals.d_buoyancy, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_cloud_moist_static_energy": locals.cloud_moist_static_energy.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels": locals.environment_saturation_moist_static_energy_cloud_levels.field[ + : + ], + "local_d_buoyancy": locals.d_buoyancy.field[:], + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_2_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_2_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_2_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_3.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_3.py new file mode 100644 index 000000000..99c83a75d --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/buoyancy/translate_GF2020_CumulusParameterization_GetBuoyancy_3.py @@ -0,0 +1,213 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.buoyancy import get_buoyancy +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "lcl_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_d_buoyancy_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=get_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy=locals.cloud_moist_static_energy_forced, + environment_moist_static_energy=locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + d_buoyancy=locals.d_buoyancy_forced, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_d_buoyancy_forced": locals.d_buoyancy_forced.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_3_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_3_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetBuoyancy_3_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/convection_tracers/translate_GF2020_Cumulus_Parameterization_AtmosphericComposition.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/convection_tracers/translate_GF2020_Cumulus_Parameterization_AtmosphericComposition.py new file mode 100644 index 000000000..e59eeea75 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/convection_tracers/translate_GF2020_Cumulus_Parameterization_AtmosphericComposition.py @@ -0,0 +1,415 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.convective_tracers import AtmosphericComposition +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection_tracers import ConvectionTracers + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "updraft_origin_level": {}, + "downdraft_origin_level": {}, + "ocean_fraction": {}, + "p_forced": {}, + "p_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels": {}, + "local_environment_massflux": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "mass_entrainment_updraft_forced": {}, + "mass_detrainment_updraft_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_vertical_velocity_3d": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_total_normalized_integrated_evaporate_forced": {}, + "evaporate_in_downdraft_forced": {}, + "epsilon_forced": {}, + # # THESE BREAK BECAUSE TRANSLATE TEST CANNOT HANDLE 4D FIELDS + # "chemistry_tracers": {}, + # "chemistry_tracers_output": {}, + # "local_chemistry_tracers_cloud_levels": {}, + # "local_chemistry_tracers_sc_updraft": {}, + # "local_chemistry_tracers_sc_downdraft": {}, + # "local_chemistry_tracers_pw_updraft": {}, + # "local_chemistry_tracers_pw_downdraft": {}, + # "local_chemistry_tracers_total_pw_updraft": {}, + # "local_chemistry_tracers_total_pw_downdraft": {}, + } + + out_vars.update(in_vars["data_vars"]) + out_vars.update( + { + "chemistry_tracers": {}, + "chemistry_tracers_output": {}, + "local_chemistry_tracers_cloud_levels": {}, + "local_chemistry_tracers_sc_updraft": {}, + "local_chemistry_tracers_sc_downdraft": {}, + "local_chemistry_tracers_pw_updraft": {}, + "local_chemistry_tracers_pw_downdraft": {}, + "local_chemistry_tracers_total_pw_updraft": {}, + "local_chemistry_tracers_total_pw_downdraft": {}, + } + ) + + def __call__( + self, + constants: dict, + cu_param_constants: dict, + convection_tracers_input: dict, + plume: str, + ddim_fields: dict, + **inputs, + ): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize convection tracers + convection_tracers = ConvectionTracers.ones( + self.quantity_factory, + data_dimensions={ + "convection_tracers": config.NUMBER_OF_TRACERS, + "size_three_dimension": 3, + "size_four_dimension": 4, + }, + ) + + convection_tracers.tracers.field[:] = np.moveaxis(convection_tracers_input["tracers"], 0, 3) + convection_tracers.vect_hcts.field[:] = convection_tracers_input["vect_hcts"] + convection_tracers.kc_scal.field[:] = convection_tracers_input["kc_scal"] + convection_tracers.fscav.field[:] = convection_tracers_input["fscav"] + convection_tracers.convfaci2g.field[:] = convection_tracers_input["convfaci2g"] + convection_tracers.retfactor.field[:] = convection_tracers_input["retfactor"] + convection_tracers.liq_and_gas.field[:] = convection_tracers_input["liq_and_gas"] + convection_tracers.online_cldliq.field[:] = convection_tracers_input["online_cldliq"] + convection_tracers.online_vud.field[:] = convection_tracers_input["online_vud"] + convection_tracers.ftemp_threshold.field[:] = convection_tracers_input["ftemp_threshold"] + convection_tracers.use_gcc_washout.field[:] = convection_tracers_input["use_gcc_washout"] + convection_tracers.use_gocart.field[:] = convection_tracers_input["use_gocart"] + convection_tracers.is_wetdep.field[:] = convection_tracers_input["is_wetdep"] + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + locals.environment_massflux.data[:] = inputs["local_environment_massflux"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_entrainment_downdraft_forced"] + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.vertical_velocity_3d.data[:] = inputs["local_vertical_velocity_3d"] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.total_normalized_integrated_evaporate_forced.data[:] = inputs[ + "local_total_normalized_integrated_evaporate_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + state.input_output.chemistry_tracers.field[:] = ddim_fields["chemistry_tracers"] + state.input_output.chemistry_tracers_output.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX, : + ] = ddim_fields["chemistry_tracers_output"] + locals.chemistry_tracers_cloud_levels.field[:] = ddim_fields["local_chemistry_tracers_cloud_levels"] + locals.chemistry_tracers_sc_updraft.field[:] = ddim_fields["local_chemistry_tracers_sc_updraft"] + locals.chemistry_tracers_sc_downdraft.field[:] = ddim_fields["local_chemistry_tracers_sc_downdraft"] + locals.chemistry_tracers_pw_updraft.field[:] = ddim_fields["local_chemistry_tracers_pw_updraft"] + locals.chemistry_tracers_pw_downdraft.field[:] = ddim_fields["local_chemistry_tracers_pw_downdraft"] + locals.chemistry_tracers_total_pw_updraft.field[:] = ddim_fields[ + "local_chemistry_tracers_total_pw_updraft" + ] + locals.chemistry_tracers_total_pw_downdraft.field[:] = ddim_fields[ + "local_chemistry_tracers_total_pw_downdraft" + ] + + code = AtmosphericComposition( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + updraft_origin_level=state.output.updraft_origin_level, + downdraft_origin_level=state.output.downdraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + p_forced=state.input_output.p_forced, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + environment_massflux=locals.environment_massflux, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + vertical_velocity_3d=locals.vertical_velocity_3d, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=locals.total_normalized_integrated_evaporate_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + epsilon_forced=state.output.epsilon_forced, + chemistry_tracers=state.input_output.chemistry_tracers, + chemistry_tracers_output=state.input_output.chemistry_tracers_output, + chemistry_tracers_cloud_levels=locals.chemistry_tracers_cloud_levels, + chemistry_tracers_sc_updraft=locals.chemistry_tracers_sc_updraft, + chemistry_tracers_sc_downdraft=locals.chemistry_tracers_sc_downdraft, + chemistry_tracers_pw_updraft=locals.chemistry_tracers_pw_updraft, + chemistry_tracers_pw_downdraft=locals.chemistry_tracers_pw_downdraft, + chemistry_tracers_total_pw_updraft=locals.chemistry_tracers_total_pw_updraft, + chemistry_tracers_total_pw_downdraft=locals.chemistry_tracers_total_pw_downdraft, + convection_tracers=convection_tracers, + plume_dependent_constants=plume_dependent_constants, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "downdraft_origin_level": state.output.downdraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "ocean_fraction": state.input.ocean_fraction.field[:], + "p_forced": state.input_output.p_forced.field[:], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.field[:], + "local_environment_massflux": locals.environment_massflux.field[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vertical_velocity_3d": locals.vertical_velocity_3d.field[:], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_total_normalized_integrated_evaporate_forced": locals.total_normalized_integrated_evaporate_forced.field[ + : + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "chemistry_tracers": state.input_output.chemistry_tracers.field[:], + "chemistry_tracers_output": state.input_output.chemistry_tracers_output.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX, : + ], + "local_chemistry_tracers_cloud_levels": locals.chemistry_tracers_cloud_levels.field[:], + "local_chemistry_tracers_sc_updraft": locals.chemistry_tracers_sc_updraft.field[:], + "local_chemistry_tracers_sc_downdraft": locals.chemistry_tracers_sc_downdraft.field[:], + "local_chemistry_tracers_pw_updraft": locals.chemistry_tracers_pw_updraft.field[:], + "local_chemistry_tracers_pw_downdraft": locals.chemistry_tracers_pw_downdraft.field[:], + "local_chemistry_tracers_total_pw_updraft": locals.chemistry_tracers_total_pw_updraft.field[:], + "local_chemistry_tracers_total_pw_downdraft": locals.chemistry_tracers_total_pw_downdraft.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_AtmosphericComposition_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers = data_loader.load("GF2020_ConvectionTracers") + self.ddim_fields = data_loader.load( + "GF2020_CumulusParameterization_AtmosphericComposition_shallow-In" + ) + + def compute_func(self, **inputs): + outputs = self.test_core( + self.constants, + self.cu_param_constants, + self.convection_tracers, + "shallow", + ddim_fields=self.ddim_fields, + **inputs, + ) + + return outputs + + +class TranslateGF2020_CumulusParameterization_AtmosphericComposition_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers = data_loader.load("GF2020_ConvectionTracers") + self.ddim_fields = data_loader.load("GF2020_CumulusParameterization_AtmosphericComposition_mid-In") + + def compute_func(self, **inputs): + outputs = self.test_core( + self.constants, + self.cu_param_constants, + self.convection_tracers, + "mid", + ddim_fields=self.ddim_fields, + **inputs, + ) + + return outputs + + +class TranslateGF2020_CumulusParameterization_AtmosphericComposition_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers = data_loader.load("GF2020_ConvectionTracers") + self.ddim_fields = data_loader.load("GF2020_CumulusParameterization_AtmosphericComposition_deep-In") + + def compute_func(self, **inputs): + outputs = self.test_core( + self.constants, + self.cu_param_constants, + self.convection_tracers, + "deep", + ddim_fields=self.ddim_fields, + **inputs, + ) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/diurnal_cycle/translate_GF2020_CumulusParameterization_DiurnalCycle.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/diurnal_cycle/translate_GF2020_CumulusParameterization_DiurnalCycle.py new file mode 100644 index 000000000..515c367b2 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/diurnal_cycle/translate_GF2020_CumulusParameterization_DiurnalCycle.py @@ -0,0 +1,259 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.diurnal_cycle import DiurnalCycle +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "pbl_level": {}, + "grid_length": {}, + "ocean_fraction": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "topography_height_no_negative": {}, + "t_old": {}, + "local_t_new": {}, + "local_t_cloud_levels_forced": {}, + "vapor_old": {}, + "local_vapor_forced": {}, + "u": {}, + "v": {}, + "local_vertical_velocity_2d": {}, + "local_cape_removal_time_scale": {}, + "cape_removal_time_scale": {}, + "local_pbl_time_scale": {}, + "pbl_time_scale": {}, + "local_cloud_work_function_1_pbl": {}, + "local_cloud_work_function_1_fa": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input_output.pbl_level.data[:] = inputs["pbl_level"] - 1 + state.input_output.grid_length.data[:] = inputs["grid_length"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + state.input_output.t_old.data[:] = inputs["t_old"] + locals.t_new.data[:] = inputs["local_t_new"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + state.input_output.vapor_old.data[:] = inputs["vapor_old"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.vertical_velocity_2d.data[:] = inputs["local_vertical_velocity_2d"] + locals.cape_removal_time_scale.data[:] = inputs["local_cape_removal_time_scale"] + state.output.cape_removal_time_scale.data[:] = inputs["cape_removal_time_scale"] + locals.pbl_time_scale.data[:] = inputs["local_pbl_time_scale"] + state.output.pbl_time_scale.data[:] = inputs["pbl_time_scale"] + locals.cloud_workfunction_1_pbl.data[:] = inputs["local_cloud_work_function_1_pbl"] + locals.cloud_workfunction_1_fa.data[:] = inputs["local_cloud_work_function_1_fa"] + + # initialize test code + code = DiurnalCycle( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + grid_length=state.input_output.grid_length, + ocean_fraction=state.input.ocean_fraction, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + t_old=state.input_output.t_old, + t_new=locals.t_new, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + vapor_old=state.input_output.vapor_old, + vapor_forced=locals.vapor_forced, + u=state.input_output.u, + v=state.input_output.v, + vertical_velocity_2d=locals.vertical_velocity_2d, + cape_removal_time_scale=locals.cape_removal_time_scale, + cape_removal_time_scale_from_state=state.output.cape_removal_time_scale, + pbl_time_scale=locals.pbl_time_scale, + pbl_time_scale_from_state=state.output.pbl_time_scale, + cloud_work_function_1_pbl=locals.cloud_workfunction_1_pbl, + cloud_work_function_1_fa=locals.cloud_workfunction_1_fa, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "pbl_level": state.input_output.pbl_level.field[:] + 1, + "grid_length": state.input_output.grid_length.field[:], + "ocean_fraction": state.input.ocean_fraction.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "t_old": state.input_output.t_old.field[:], + "local_t_new": locals.t_new.field[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + "vapor_old": state.input_output.vapor_old.field[:], + "local_vapor_forced": locals.vapor_forced.field[:], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_vertical_velocity_2d": locals.vertical_velocity_2d.field[:], + "local_cape_removal_time_scale": locals.cape_removal_time_scale.field[:], + "cape_removal_time_scale": state.output.cape_removal_time_scale.field[:], + "local_pbl_time_scale": locals.pbl_time_scale.field[:], + "pbl_time_scale": state.output.pbl_time_scale.field[:], + "local_cloud_work_function_1_pbl": locals.cloud_workfunction_1_pbl.field[:], + "local_cloud_work_function_1_fa": locals.cloud_workfunction_1_fa.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DiurnalCycle_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DiurnalCycle_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DiurnalCycle_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftLateralMassFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftLateralMassFlux.py new file mode 100644 index 000000000..3a6f3e8d8 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftLateralMassFlux.py @@ -0,0 +1,241 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import downdraft_lateral_massflux +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "downdraft_origin_level": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_normalized_massflux_downdraft": {}, + "normalized_massflux_downdraft_forced": {}, + "local_normalized_massflux_downdraft_modified": {}, + "local_detrainment_function_downdraft": {}, + "local_entrainment_rate_downdraft": {}, + "local_mass_entrainment_downdraft": {}, + "local_mass_detrainment_downdraft": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_mass_entrainment_u_downdraft": {}, + "local_mass_detrainment_u_downdraft": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.normalized_massflux_downdraft.data[:] = inputs["local_normalized_massflux_downdraft"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + locals.normalized_massflux_downdraft_modified.data[:] = inputs[ + "local_normalized_massflux_downdraft_modified" + ] + locals.detrainment_function_downdraft.data[:] = inputs["local_detrainment_function_downdraft"] + locals.entrainment_rate_downdraft.data[:] = inputs["local_entrainment_rate_downdraft"] + locals.mass_entrainment_downdraft.data[:] = inputs["local_mass_entrainment_downdraft"] + locals.mass_detrainment_downdraft.data[:] = inputs["local_mass_detrainment_downdraft"] + state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_entrainment_downdraft_forced"] + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.mass_entrainment_u_downdraft.data[:] = inputs["local_mass_entrainment_u_downdraft"] + locals.mass_detrainment_u_downdraft.data[:] = inputs["local_mass_detrainment_u_downdraft"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=downdraft_lateral_massflux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_downdraft=locals.normalized_massflux_downdraft, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + normalized_massflux_downdraft_modified=locals.normalized_massflux_downdraft_modified, + detrainment_function_downdraft=locals.detrainment_function_downdraft, + entrainment_rate_downdraft=locals.entrainment_rate_downdraft, + mass_entrainment_downdraft=locals.mass_entrainment_downdraft, + mass_detrainment_downdraft=locals.mass_detrainment_downdraft, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + mass_entrainment_u_downdraft=locals.mass_entrainment_u_downdraft, + mass_detrainment_u_downdraft=locals.mass_detrainment_u_downdraft, + LAMBDA_DOWN=plume_dependent_constants.LAMBDA_DOWN, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "downdraft_origin_level": state.output.downdraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "local_normalized_massflux_downdraft": locals.normalized_massflux_downdraft.field[:], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_normalized_massflux_downdraft_modified": locals.normalized_massflux_downdraft_modified.field[ + : + ], + "local_detrainment_function_downdraft": locals.detrainment_function_downdraft.field[:], + "local_entrainment_rate_downdraft": locals.entrainment_rate_downdraft.field[:], + "local_mass_entrainment_downdraft": locals.mass_entrainment_downdraft.field[:], + "local_mass_detrainment_downdraft": locals.mass_detrainment_downdraft.field[:], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_mass_entrainment_u_downdraft": locals.mass_entrainment_u_downdraft.field[:], + "local_mass_detrainment_u_downdraft": locals.mass_detrainment_u_downdraft.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftLateralMassFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMSEAndBuoyancy.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMSEAndBuoyancy.py new file mode 100644 index 000000000..c27f82088 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMSEAndBuoyancy.py @@ -0,0 +1,278 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import ( + downdraft_moist_static_energy_and_buoyancy, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "downdraft_origin_level": {}, + "u": {}, + "local_u_cloud_levels": {}, + "local_u_c_downdraft": {}, + "v": {}, + "local_v_cloud_levels": {}, + "local_v_c_downdraft": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_cloud_moist_static_energy": {}, + "local_cloud_moist_static_energy_downdraft_forced": {}, + "local_d_buoyancy_downdraft_forced": {}, + "local_t_wetbulb": {}, + "local_vapor_wetbulb": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_mass_entrainment_u_downdraft": {}, + "local_mass_detrainment_u_downdraft": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + state.input_output.u.data[:] = inputs["u"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.u_c_downdraft.data[:] = inputs["local_u_c_downdraft"] + state.input_output.v.data[:] = inputs["v"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.v_c_downdraft.data[:] = inputs["local_v_c_downdraft"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.cloud_moist_static_energy.data[:] = inputs["local_cloud_moist_static_energy"] + locals.cloud_moist_static_energy_downdraft_forced.data[:] = inputs[ + "local_cloud_moist_static_energy_downdraft_forced" + ] + locals.d_buoyancy_downdraft_forced.data[:] = inputs["local_d_buoyancy_downdraft_forced"] + locals.t_wetbulb.data[:] = inputs["local_t_wetbulb"] + locals.vapor_wetbulb.data[:] = inputs["local_vapor_wetbulb"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_entrainment_downdraft_forced"] + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.mass_entrainment_u_downdraft.data[:] = inputs["local_mass_entrainment_u_downdraft"] + locals.mass_detrainment_u_downdraft.data[:] = inputs["local_mass_detrainment_u_downdraft"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=downdraft_moist_static_energy_and_buoyancy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_WETBULB": cumulus_parameterization_config.USE_WETBULB, + "PRESSURE_GRADIENT_CONSTANT": cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT, + }, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + u=state.input_output.u, + u_cloud_levels=locals.u_cloud_levels, + u_c_downdraft=locals.u_c_downdraft, + v=state.input_output.v, + v_cloud_levels=locals.v_cloud_levels, + v_c_downdraft=locals.v_c_downdraft, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy=locals.cloud_moist_static_energy, + cloud_moist_static_energy_downdraft_forced=locals.cloud_moist_static_energy_downdraft_forced, + buoyancy_downdraft_forced=locals.d_buoyancy_downdraft_forced, + t_wetbulb=locals.t_wetbulb, + vapor_wetbulb=locals.vapor_wetbulb, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + mass_entrainment_u_downdraft=locals.mass_entrainment_u_downdraft, + mass_detrainment_u_downdraft=locals.mass_detrainment_u_downdraft, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "downdraft_origin_level": state.output.downdraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "u": state.input_output.u.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_u_c_downdraft": locals.u_c_downdraft.field[:], + "v": state.input_output.v.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_v_c_downdraft": locals.v_c_downdraft.field[:], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_cloud_moist_static_energy": locals.cloud_moist_static_energy.field[:], + "local_cloud_moist_static_energy_downdraft_forced": locals.cloud_moist_static_energy_downdraft_forced.field[ + : + ], + "local_d_buoyancy_downdraft_forced": locals.d_buoyancy_downdraft_forced.field[:], + "local_t_wetbulb": locals.t_wetbulb.field[:], + "local_vapor_wetbulb": locals.vapor_wetbulb.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_mass_entrainment_u_downdraft": locals.mass_entrainment_u_downdraft.field[:], + "local_mass_detrainment_u_downdraft": locals.mass_detrainment_u_downdraft.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMSEAndBuoyancy_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMassFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMassFlux.py new file mode 100644 index 000000000..e7c0acb0b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMassFlux.py @@ -0,0 +1,237 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import downdraft_mass_flux +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_detrainment_start_level": {}, + "downdraft_origin_level": {}, + "pbl_level": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "lcl_level": {}, + "p_cloud_levels_forced": {}, + "p_surface": {}, + "local_normalized_massflux_downdraft": {}, + "normalized_massflux_downdraft_forced": {}, + "ocean_fraction": {}, + "local_random_number": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.detrainment_start_level.data[:] = inputs["local_detrainment_start_level"] - 1 + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + state.input_output.pbl_level.data[:] = inputs["pbl_level"] - 1 + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.input_output.p_surface.data[:] = inputs["p_surface"] + locals.normalized_massflux_downdraft.data[:] = inputs["local_normalized_massflux_downdraft"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + locals.random_number.data[:] = inputs["local_random_number"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=downdraft_mass_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF}, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + detrainment_start_level=locals.detrainment_start_level, + downdraft_origin_level=state.output.downdraft_origin_level, + pbl_level=state.input_output.pbl_level, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + lcl_level=state.output.lcl_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + normalized_massflux_downdraft=locals.normalized_massflux_downdraft, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + ocean_fraction=state.input.ocean_fraction, + random_number=locals.random_number, + DOWNDRAFT_MAX_HEIGHT_LAND=plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_LAND, + DOWNDRAFT_MAX_HEIGHT_OCEAN=plume_dependent_constants.DOWNDRAFT_MAX_HEIGHT_OCEAN, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "local_detrainment_start_level": locals.detrainment_start_level.data[:] + 1, + "downdraft_origin_level": state.output.downdraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "pbl_level": state.input_output.pbl_level.data[:] + 1, + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "lcl_level": state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "p_surface": state.input_output.p_surface.data[:], + "local_normalized_massflux_downdraft": locals.normalized_massflux_downdraft.data[:], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "ocean_fraction": state.input.ocean_fraction.data[:], + "local_random_number": locals.random_number.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMassFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMassFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMassFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMoisture.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMoisture.py new file mode 100644 index 000000000..6aa7cca79 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftMoisture.py @@ -0,0 +1,309 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import downdraft_moisture +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "downdraft_origin_level": {}, + "local_t_cloud_levels_forced": {}, + "local_t_wetbulb": {}, + "local_vapor_forced": {}, + "local_vapor_cloud_levels_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + "local_cloud_total_water_after_entrainment_downdraft_forced": {}, + "local_downdraft_saturation_vapor_forced": {}, + "local_vapor_wetbulb": {}, + "normalized_massflux_downdraft_forced": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_cloud_moist_static_energy_downdraft_forced": {}, + "evaporate_in_downdraft_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_gamma_cloud_levels_forced": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_total_normalized_integrated_evaporate_forced": {}, + "local_buoyancy": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + locals.t_wetbulb.data[:] = inputs["local_t_wetbulb"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + locals.cloud_total_water_after_entrainment_downdraft_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_downdraft_forced" + ] + locals.downdraft_saturation_vapor_forced.data[:] = inputs["local_downdraft_saturation_vapor_forced"] + locals.vapor_wetbulb.data[:] = inputs["local_vapor_wetbulb"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.cloud_moist_static_energy_downdraft_forced.data[:] = inputs[ + "local_cloud_moist_static_energy_downdraft_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_entrainment_downdraft_forced"] + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.total_normalized_integrated_evaporate_forced.data[:,] = inputs[ + "local_total_normalized_integrated_evaporate_forced" + ] + locals.buoyancy.data[:] = inputs["local_buoyancy"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=downdraft_moisture, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_WETBULB": cumulus_parameterization_config.USE_WETBULB, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "EVAP_FIX": cumulus_parameterization_config.EVAP_FIX, + }, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + downdraft_origin_level=state.output.downdraft_origin_level, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + t_wetbulb=locals.t_wetbulb, + vapor_forced=locals.vapor_forced, + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=locals.cloud_total_water_after_entrainment_forced, + cloud_total_water_after_entrainment_downdraft_forced=locals.cloud_total_water_after_entrainment_downdraft_forced, + downdraft_saturation_vapor_forced=locals.downdraft_saturation_vapor_forced, + vapor_wetbulb=locals.vapor_wetbulb, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy_downdraft_forced=locals.cloud_moist_static_energy_downdraft_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + gamma_cloud_levels_forced=locals.gamma_cloud_levels_forced, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=locals.total_normalized_integrated_evaporate_forced, + buoyancy=locals.buoyancy, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "downdraft_origin_level": state.output.downdraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.data[:], + "local_t_wetbulb": locals.t_wetbulb.data[:], + "local_vapor_forced": locals.vapor_forced.data[:], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.data[:], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[ + : + ], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.data[ + : + ], + "local_cloud_total_water_after_entrainment_downdraft_forced": locals.cloud_total_water_after_entrainment_downdraft_forced.data[ + : + ], + "local_downdraft_saturation_vapor_forced": locals.downdraft_saturation_vapor_forced.data[:], + "local_vapor_wetbulb": locals.vapor_wetbulb.data[:], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.data[:], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[ + : + ], + "local_cloud_moist_static_energy_downdraft_forced": locals.cloud_moist_static_energy_downdraft_forced.data[ + : + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.data[ + : + ], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.data[:], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_total_normalized_integrated_evaporate_forced": locals.total_normalized_integrated_evaporate_forced.data[ + : + ], + "local_buoyancy": locals.buoyancy.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMoisture_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMoisture_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftMoisture_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftOriginLevel.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftOriginLevel.py new file mode 100644 index 000000000..262a48982 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftOriginLevel.py @@ -0,0 +1,227 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import DowndraftOriginLevel +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "updraft_origin_level": {}, + "downdraft_origin_level": {}, + "local_detrainment_start_level": {}, + "updraft_lfc_level": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "topography_height_no_negative": {}, + "local_melting_layer": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.downdraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["downdraft_origin_level"] - 1 + ) + locals.detrainment_start_level.data[:] = inputs["local_detrainment_start_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + locals.melting_layer.data[:] = inputs["local_melting_layer"] + + # initialize test code + code = DowndraftOriginLevel( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + updraft_origin_level=state.output.updraft_origin_level, + downdraft_origin_level=state.output.downdraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + detrainment_start_level=locals.detrainment_start_level, + melting_layer=locals.melting_layer, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "downdraft_origin_level": state.output.downdraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_detrainment_start_level": locals.detrainment_start_level.field[:] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "local_melting_layer": locals.melting_layer.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftOriginLevel_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftTemperature.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftTemperature.py new file mode 100644 index 000000000..9866ee592 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftTemperature.py @@ -0,0 +1,196 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import downdraft_temperature +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_downdraft_column_temperature_forced": {}, + "local_cloud_moist_static_energy_downdraft_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_cloud_total_water_after_entrainment_downdraft_forced": {}, + "local_t_cloud_levels_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.downdraft_column_temperature_forced.data[:] = inputs[ + "local_downdraft_column_temperature_forced" + ] + locals.cloud_moist_static_energy_downdraft_forced.data[:] = inputs[ + "local_cloud_moist_static_energy_downdraft_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.cloud_total_water_after_entrainment_downdraft_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_downdraft_forced" + ] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=downdraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + downdraft_column_temperature_forced=locals.downdraft_column_temperature_forced, + cloud_moist_static_energy_downdraft_forced=locals.cloud_moist_static_energy_downdraft_forced, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_downdraft_forced=locals.cloud_total_water_after_entrainment_downdraft_forced, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "local_downdraft_column_temperature_forced": locals.downdraft_column_temperature_forced.data[:], + "local_cloud_moist_static_energy_downdraft_forced": locals.cloud_moist_static_energy_downdraft_forced.data[ + : + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.data[ + : + ], + "local_cloud_total_water_after_entrainment_downdraft_forced": locals.cloud_total_water_after_entrainment_downdraft_forced.data[ + : + ], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftTemperature_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftTemperature_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftTemperature_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftWindShear.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftWindShear.py new file mode 100644 index 000000000..2c4d0d2cf --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/downdraft/translate_GF2020_CumulusParameterization_DowndraftWindShear.py @@ -0,0 +1,255 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.downdraft import DowndraftWindShear +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "geopotential_height_forced": {}, + "p_forced": {}, + "u": {}, + "v": {}, + "ccn": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_total_normalized_integrated_evaporate_forced": {}, + "local_psum": {}, + "local_psumh": {}, + "local_scale_dependence_factor_downdraft": {}, + "local_epsilon": {}, + "local_epsilon_min": {}, + "local_epsilon_max": {}, + "local_epsilon_computed": {}, + "epsilon_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + state.input_output.ccn.data[:] = inputs["ccn"] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.total_normalized_integrated_evaporate_forced.data[:,] = inputs[ + "local_total_normalized_integrated_evaporate_forced" + ] + locals.psum.data[:] = inputs["local_psum"] + locals.psumh.data[:] = inputs["local_psumh"] + locals.scale_dependence_factor_downdraft.data[:] = inputs["local_scale_dependence_factor_downdraft"] + locals.epsilon.data[:] = inputs["local_epsilon"] + locals.epsilon_min.data[:] = inputs["local_epsilon_min"] + locals.epsilon_max.data[:] = inputs["local_epsilon_max"] + size_epsilon_computed = len(inputs["local_epsilon_computed"].shape) + if size_epsilon_computed == 2: + import numpy as np + + locals.epsilon_computed.data[:] = inputs["local_epsilon_computed"][:, :, np.newaxis] + else: + locals.epsilon_computed.data[:] = inputs["local_epsilon_computed"] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + + # initialize test code + code = DowndraftWindShear( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_forced=state.input_output.geopotential_height_forced, + p_forced=state.input_output.p_forced, + u=state.input_output.u, + v=state.input_output.v, + ccn=state.input_output.ccn, + psum=locals.psum, + psumh=locals.psumh, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=locals.total_normalized_integrated_evaporate_forced, + scale_dependence_factor_downdraft=locals.scale_dependence_factor_downdraft, + epsilon=locals.epsilon, + epsilon_min=locals.epsilon_min, + epsilon_max=locals.epsilon_max, + epsilon_computed=locals.epsilon_computed, + epsilon_forced=state.output.epsilon_forced, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "p_forced": state.input_output.p_forced.field[:], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "ccn": state.input_output.ccn.field[:], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_total_normalized_integrated_evaporate_forced": locals.total_normalized_integrated_evaporate_forced.field[ + :, + ], + "local_psum": locals.psum.field[:], + "local_psumh": locals.psumh.field[:], + "local_scale_dependence_factor_downdraft": locals.scale_dependence_factor_downdraft.field[:], + "local_epsilon": locals.epsilon.field[:], + "local_epsilon_min": locals.epsilon_min.field[:], + "local_epsilon_max": locals.epsilon_max.field[:], + "local_epsilon_computed": locals.epsilon_computed.field[:], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftWindShear_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftWindShear_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftWindShear_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment.py new file mode 100644 index 000000000..ea4d56327 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment.py @@ -0,0 +1,329 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.entrainment import ( + compute_lateral_massflux, + compute_uc_vc, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "local_detrainment_function_updraft": {}, + "entrainment_rate": {}, + "p_cloud_levels_forced": {}, + "mass_entrainment_updraft_forced": {}, + "mass_detrainment_updraft_forced": {}, + "local_mass_entrainment_updraft": {}, + "local_mass_detrainment_updraft": {}, + "updraft_lfc_level": {}, + "updraft_origin_level": {}, + "pbl_level": {}, + "local_mass_entrainment_u_updraft": {}, + "local_mass_detrainment_u_updraft": {}, + "local_start_level": {}, + "ocean_fraction": {}, + "p_forced": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_u_c": {}, + "local_v_c": {}, + "local_moist_static_energy_origin_level": {}, + "local_moist_static_energy_origin_level_forced": {}, + "local_cloud_moist_static_energy": {}, + "local_cloud_moist_static_energy_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_top_level" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.detrainment_function_updraft.data[:] = inputs["local_detrainment_function_updraft"] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + locals.mass_entrainment_updraft.data[:] = inputs["local_mass_entrainment_updraft"] + locals.mass_detrainment_updraft.data[:] = inputs["local_mass_detrainment_updraft"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "updraft_lfc_level" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.input_output.pbl_level.data[:] = inputs["pbl_level"] + locals.mass_entrainment_u_updraft.data[:] = inputs["local_mass_entrainment_u_updraft"] + locals.mass_detrainment_u_updraft.data[:] = inputs["local_mass_detrainment_u_updraft"] + locals.start_level.data[:] = inputs["local_start_level"] - 1 + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.u_c.data[:] = inputs["local_u_c"] + locals.v_c.data[:] = inputs["local_v_c"] + locals.moist_static_energy_origin_level.data[:] = inputs["local_moist_static_energy_origin_level"] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + locals.cloud_moist_static_energy.data[:] = inputs["local_cloud_moist_static_energy"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + + code_part_1 = self.stencil_factory.from_dims_halo( + func=compute_lateral_massflux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + code_part_2 = self.stencil_factory.from_dims_halo( + func=compute_uc_vc, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code_part_1( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + geopotential_height=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=state.output.normalized_massflux_updraft_forced, + detrainment_function_updraft=locals.detrainment_function_updraft, + entrainment_rate=state.output.entrainment_rate, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft=locals.mass_entrainment_updraft, + mass_detrainment_updraft=locals.mass_detrainment_updraft, + updraft_lfc_level=state.output.updraft_lfc_level, + updraft_origin_level=state.output.updraft_origin_level, + pbl_level=state.input_output.pbl_level, + mass_entrainment_u_updraft=locals.mass_entrainment_u_updraft, + mass_detrainment_u_updraft=locals.mass_detrainment_u_updraft, + LAMBDA_DEEP=plume_dependent_constants.LAMBDA_DEEP, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + code_part_2( + u_c=locals.u_c, + v_c=locals.v_c, + cloud_moist_static_energy=locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + error_code=state.output.error_code, + start_level=locals.start_level, + moist_static_energy_origin_level=locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + p=state.input_output.p_forced, + updraft_origin_level=state.output.updraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_detrainment_function_updraft": locals.detrainment_function_updraft.field[:], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_mass_entrainment_updraft": locals.mass_entrainment_updraft.field[:], + "local_mass_detrainment_updraft": locals.mass_detrainment_updraft.field[:], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "pbl_level": state.input_output.pbl_level.field[:], + "local_mass_entrainment_u_updraft": locals.mass_entrainment_u_updraft.field[:], + "local_mass_detrainment_u_updraft": locals.mass_detrainment_u_updraft.field[:], + "local_start_level": locals.start_level.field[:] + 1, + "ocean_fraction": state.input.ocean_fraction.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_u_c": locals.u_c.field[:], + "local_v_c": locals.v_c.field[:], + "local_moist_static_energy_origin_level": locals.moist_static_energy_origin_level.field[:], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.field[ + : + ], + "local_cloud_moist_static_energy": locals.cloud_moist_static_energy.field[:], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_shallow( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_mid( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CalculateMassEntrainmentDetrainment_deep( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_DowndraftEntrainmentProfiles.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_DowndraftEntrainmentProfiles.py new file mode 100644 index 000000000..945b6d646 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_DowndraftEntrainmentProfiles.py @@ -0,0 +1,172 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.entrainment import downdraft_entrainment_profiles +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "lateral_entrainment_rate": {}, + "local_entrainment_rate_downdraft": {}, + "local_detrainment_function_downdraft": {}, + "local_scale_dependence_factor_downdraft": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.input.lateral_entrainment_rate.data[:] = inputs["lateral_entrainment_rate"] + locals.entrainment_rate_downdraft.data[:] = inputs["local_entrainment_rate_downdraft"] + locals.detrainment_function_downdraft.data[:] = inputs["local_detrainment_function_downdraft"] + locals.scale_dependence_factor_downdraft.data[:] = inputs["local_scale_dependence_factor_downdraft"] + + code = self.stencil_factory.from_dims_halo( + func=downdraft_entrainment_profiles, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DOWNDRAFT": cumulus_parameterization_config.DOWNDRAFT}, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + lateral_entrainment_rate=state.input.lateral_entrainment_rate, + entrainment_rate_downdraft=locals.entrainment_rate_downdraft, + detrainment_function_downdraft=locals.detrainment_function_downdraft, + scale_dependence_factor_downdraft=locals.scale_dependence_factor_downdraft, + plume_entrainment_rate=plume_dependent_constants.ENTRAINMENT_RATE, + ) + + outputs = { + "lateral_entrainment_rate": state.input.lateral_entrainment_rate.field[:], + "local_entrainment_rate_downdraft": locals.entrainment_rate_downdraft.field[:], + "local_detrainment_function_downdraft": locals.detrainment_function_downdraft.field[:], + "local_scale_dependence_factor_downdraft": locals.scale_dependence_factor_downdraft.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DowndraftEntrainmentProfiles_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_EntrainmentRates.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_EntrainmentRates.py new file mode 100644 index 000000000..73b9bf213 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_EntrainmentRates.py @@ -0,0 +1,191 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.entrainment import entrainment_rates +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_vapor_cloud_levels_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "lcl_level": {}, + "entrainment_rate": {}, + "local_updraft_detrainment_function": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + locals.detrainment_function_updraft.data[:] = inputs["local_updraft_detrainment_function"] + + code = self.stencil_factory.from_dims_halo( + func=entrainment_rates, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "MIN_ENTRAINMENT_RATE": cumulus_parameterization_config.MIN_ENTRAINMENT_RATE, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + lcl_level=state.output.lcl_level, + error_code=state.output.error_code, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=locals.detrainment_function_updraft, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.field[ + : + ], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_updraft_detrainment_function": locals.detrainment_function_updraft.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EntrainmentRates_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EntrainmentRates_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EntrainmentRates_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_StableDetrainment.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_StableDetrainment.py new file mode 100644 index 000000000..dcfa1b7e6 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/entrainment/translate_GF2020_CumulusParameterization_StableDetrainment.py @@ -0,0 +1,184 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import generic_find_level +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "updraft_lfc_level": {}, + "kstabm": {}, + "kstabi": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.kstabm.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["kstabm"] - 1 + state.output.kstabi.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["kstabi"] - 1 + + code = self.stencil_factory.from_dims_halo( + func=generic_find_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + array=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + start_index=state.output.updraft_lfc_level, + end_index=state.output.kstabm.field[:, :, plume_dependent_constants.PLUME_INDEX], + out_index=state.output.kstabi, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "kstabm": state.output.kstabm.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "kstabi": state.output.kstabi.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_StableDetrainment_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_StableDetrainment_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_StableDetrainment_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_1.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_1.py new file mode 100644 index 000000000..77aa1e7b8 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_1.py @@ -0,0 +1,264 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_cloud_levels +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "t_old": {}, + "local_env_saturation_mixing_ratio": {}, + "vapor_old": {}, + "local_env_moist_static_energy": {}, + "local_env_saturation_moist_static_energy": {}, + "local_geopotential_height": {}, + "p_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels": {}, + "local_vapor_cloud_levels": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "u": {}, + "v": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_env_saturation_moist_static_energy_cloud_levels": {}, + "local_geopotential_height_cloud_levels": {}, + "local_p_cloud_levels": {}, + "local_gamma_cloud_levels": {}, + "local_t_cloud_levels": {}, + "p_surface": {}, + "t_surface": {}, + "error_code": {}, + "topography_height_no_negative": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.input_output.t_old.data[:] = inputs["t_old"] + locals.environment_saturation_mixing_ratio.data[:] = inputs["local_env_saturation_mixing_ratio"] + state.input_output.vapor_old.data[:] = inputs["vapor_old"] + locals.environment_moist_static_energy.data[:] = inputs["local_env_moist_static_energy"] + locals.environment_saturation_moist_static_energy.data[:] = inputs[ + "local_env_saturation_moist_static_energy" + ] + locals.geopotential_height.data[:] = inputs["local_geopotential_height"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels" + ] + locals.vapor_cloud_levels.data[:] = inputs["local_vapor_cloud_levels"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.environment_saturation_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels" + ] + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + locals.p_cloud_levels.data[:] = inputs["local_p_cloud_levels"] + locals.gamma_cloud_levels.data[:] = inputs["local_gamma_cloud_levels"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.input_output.t_surface.data[:] = inputs["t_surface"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + + code = self.stencil_factory.from_dims_halo( + func=environment_cloud_levels, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"CLOUD_LEVEL_GRID": cumulus_parameterization_config.CLOUD_LEVEL_GRID}, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=locals.p_cloud_levels, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=locals.geopotential_height, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + t=state.input_output.t_old, + t_surface=state.input_output.t_surface, + t_cloud_levels=locals.t_cloud_levels, + vapor=state.input_output.vapor_old, + vapor_cloud_levels=locals.vapor_cloud_levels, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + environment_saturation_mixing_ratio=locals.environment_saturation_mixing_ratio, + environment_saturation_mixing_ratio_cloud_levels=locals.environment_saturation_mixing_ratio_cloud_levels, + environment_moist_static_energy=locals.environment_moist_static_energy, + environment_moist_static_energy_cloud_levels=locals.environment_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy, + environment_saturation_moist_static_energy_cloud_levels=locals.environment_saturation_moist_static_energy_cloud_levels, + gamma_cloud_levels=locals.gamma_cloud_levels, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "t_old": state.input_output.t_old.field[:], + "local_env_saturation_mixing_ratio": locals.environment_saturation_mixing_ratio.field[:], + "vapor_old": state.input_output.vapor_old.field[:], + "local_env_moist_static_energy": locals.environment_moist_static_energy.field[:], + "local_env_saturation_moist_static_energy": locals.environment_saturation_moist_static_energy.field[ + : + ], + "local_geopotential_height": locals.geopotential_height.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_env_saturation_mixing_ratio_cloud_levels": locals.environment_saturation_mixing_ratio_cloud_levels.field[ + : + ], + "local_vapor_cloud_levels": locals.vapor_cloud_levels.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_env_saturation_moist_static_energy_cloud_levels": locals.environment_saturation_moist_static_energy_cloud_levels.field[ + : + ], + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.field[:], + "local_p_cloud_levels": locals.p_cloud_levels.field[:], + "local_gamma_cloud_levels": locals.gamma_cloud_levels.field[:], + "local_t_cloud_levels": locals.t_cloud_levels.field[:], + "p_surface": state.input_output.p_surface.field[:], + "t_surface": state.input_output.t_surface.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_1_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_2.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_2.py new file mode 100644 index 000000000..a29bbb210 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_2.py @@ -0,0 +1,278 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_cloud_levels +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_t_new": {}, + "local_env_saturation_mixing_ratio_forced": {}, + "local_vapor_forced_env_cond": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_saturation_moist_static_energy_forced": {}, + "geopotential_height_forced": {}, + "p_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "local_vapor_cloud_levels_forced": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "u": {}, + "v": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "p_cloud_levels_forced": {}, + "local_gamma_cloud_levels_forced": {}, + "local_t_cloud_levels_forced": {}, + "p_surface": {}, + "t_surface": {}, + "error_code": {}, + "topography_height_no_negative": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.t_new.data[:] = inputs["local_t_new"] + locals.environment_saturation_mixing_ratio_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_forced" + ] + locals.vapor_forced.data[:] = inputs["local_vapor_forced_env_cond"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_saturation_moist_static_energy_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_forced" + ] + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.input_output.t_surface.data[:] = inputs["t_surface"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + + code = self.stencil_factory.from_dims_halo( + func=environment_cloud_levels, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"CLOUD_LEVEL_GRID": cumulus_parameterization_config.CLOUD_LEVEL_GRID}, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + p_3d = state.output.p_cloud_levels_forced.field[:, :, :, plume_dependent_constants.PLUME_INDEX] + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=p_3d, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=state.input_output.geopotential_height_forced, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels_forced, + t=locals.t_new, + t_surface=state.input_output.t_surface, + t_cloud_levels=locals.t_cloud_levels_forced, + vapor=locals.vapor_forced, + vapor_cloud_levels=locals.vapor_cloud_levels_forced, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + environment_saturation_mixing_ratio=locals.environment_saturation_mixing_ratio_forced, + environment_saturation_mixing_ratio_cloud_levels=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + environment_moist_static_energy=locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + gamma_cloud_levels=locals.gamma_cloud_levels_forced, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + state.output.p_cloud_levels_forced.field[:, :, :, plume_dependent_constants.PLUME_INDEX] = p_3d + + outputs = { + "local_t_new": locals.t_new.field[:], + "local_env_saturation_mixing_ratio_forced": locals.environment_saturation_mixing_ratio_forced.field[ + : + ], + "local_vapor_forced_env_cond": locals.vapor_forced.field[:], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_env_saturation_moist_static_energy_forced": locals.environment_saturation_moist_static_energy_forced.field[ + : + ], + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.field[ + : + ], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.field[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + "p_surface": state.input_output.p_surface.field[:], + "t_surface": state.input_output.t_surface.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_2_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_3.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_3.py new file mode 100644 index 000000000..f2d292811 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentCloudLevels_3.py @@ -0,0 +1,282 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_cloud_levels +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "p_forced": {}, + "p_cloud_levels_forced": {}, + "p_surface": {}, + "local_t_modified": {}, + "local_t_cloud_levels_modified": {}, + "t_surface": {}, + "local_vapor_modified": {}, + "local_vapor_cloud_levels_modified": {}, + "local_env_saturation_mixing_ratio_modified": {}, + "local_env_saturation_mixing_ratio_cloud_levels_modified": {}, + "u": {}, + "v": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_env_moist_static_energy_modified": {}, + "local_env_moist_static_energy_cloud_levels_modified": {}, + "local_env_saturation_moist_static_energy_modified": {}, + "local_env_saturation_moist_static_energy_cloud_levels_modified": {}, + "local_geopotential_height_modified": {}, + "local_geopotential_height_cloud_levels_modified": {}, + "topography_height_no_negative": {}, + "local_gamma_cloud_levels": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.input_output.p_surface.data[:] = inputs["p_surface"] + locals.t_modified.data[:] = inputs["local_t_modified"] + locals.t_cloud_levels_modified.data[:] = inputs["local_t_cloud_levels_modified"] + state.input_output.t_surface.data[:] = inputs["t_surface"] + locals.vapor_modified.data[:] = inputs["local_vapor_modified"] + locals.vapor_cloud_levels_modified.data[:] = inputs["local_vapor_cloud_levels_modified"] + locals.environment_saturation_mixing_ratio_modified.data[:] = inputs[ + "local_env_saturation_mixing_ratio_modified" + ] + locals.environment_saturation_mixing_ratio_cloud_levels_modified.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_modified" + ] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.environment_moist_static_energy_modified.data[:] = inputs[ + "local_env_moist_static_energy_modified" + ] + locals.environment_moist_static_energy_cloud_levels_modified.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_modified" + ] + locals.environment_saturation_moist_static_energy_modified.data[:] = inputs[ + "local_env_saturation_moist_static_energy_modified" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_modified.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_modified" + ] + locals.geopotential_height_modified.data[:] = inputs["local_geopotential_height_modified"] + locals.geopotential_height_cloud_levels_modified.data[:] = inputs[ + "local_geopotential_height_cloud_levels_modified" + ] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + locals.gamma_cloud_levels.data[:] = inputs["local_gamma_cloud_levels"] + + code = self.stencil_factory.from_dims_halo( + func=environment_cloud_levels, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"CLOUD_LEVEL_GRID": cumulus_parameterization_config.CLOUD_LEVEL_GRID}, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + p_3d = state.output.p_cloud_levels_forced.field[:, :, :, plume_dependent_constants.PLUME_INDEX] + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + p_cloud_levels=p_3d, + topography_height_no_negative=state.input_output.topography_height_no_negative, + geopotential_height=locals.geopotential_height_modified, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels_modified, + t=locals.t_modified, + t_surface=state.input_output.t_surface, + t_cloud_levels=locals.t_cloud_levels_modified, + vapor=locals.vapor_modified, + vapor_cloud_levels=locals.vapor_cloud_levels_modified, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + environment_saturation_mixing_ratio=locals.environment_saturation_mixing_ratio_modified, + environment_saturation_mixing_ratio_cloud_levels=locals.environment_saturation_mixing_ratio_cloud_levels_modified, + environment_moist_static_energy=locals.environment_moist_static_energy_modified, + environment_moist_static_energy_cloud_levels=locals.environment_moist_static_energy_cloud_levels_modified, + environment_saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_modified, + environment_saturation_moist_static_energy_cloud_levels=locals.environment_saturation_moist_static_energy_cloud_levels_modified, + gamma_cloud_levels=locals.gamma_cloud_levels, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + state.output.p_cloud_levels_forced.field[:, :, :, plume_dependent_constants.PLUME_INDEX] = p_3d + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "p_forced": state.input_output.p_forced.field[:], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "p_surface": state.input_output.p_surface.field[:], + "local_t_modified": locals.t_modified.field[:], + "local_t_cloud_levels_modified": locals.t_cloud_levels_modified.field[:], + "t_surface": state.input_output.t_surface.field[:], + "local_vapor_modified": locals.vapor_modified.field[:], + "local_vapor_cloud_levels_modified": locals.vapor_cloud_levels_modified.field[:], + "local_env_saturation_mixing_ratio_modified": locals.environment_saturation_mixing_ratio_modified.field[ + : + ], + "local_env_saturation_mixing_ratio_cloud_levels_modified": locals.environment_saturation_mixing_ratio_cloud_levels_modified.field[ + : + ], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_env_moist_static_energy_modified": locals.environment_moist_static_energy_modified.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_modified": locals.environment_moist_static_energy_cloud_levels_modified.field[ + : + ], + "local_env_saturation_moist_static_energy_modified": locals.environment_saturation_moist_static_energy_modified.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_modified": locals.environment_saturation_moist_static_energy_cloud_levels_modified.field[ + : + ], + "local_geopotential_height_modified": locals.geopotential_height_modified.field[:], + "local_geopotential_height_cloud_levels_modified": locals.geopotential_height_cloud_levels_modified.field[ + : + ], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "local_gamma_cloud_levels": locals.gamma_cloud_levels.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentCloudLevels_3_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_1.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_1.py new file mode 100644 index 000000000..3fcafb81c --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_1.py @@ -0,0 +1,203 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_conditions +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_geopotential_height": {}, + "local_env_saturation_mixing_ratio": {}, + "local_env_moist_static_energy": {}, + "local_env_saturation_moist_static_energy": {}, + "t_old": {}, + "vapor_old": {}, + "p_forced": {}, + "topography_height_no_negative": {}, + "p_surface": {}, + "error_code": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.geopotential_height.data[:] = inputs["local_geopotential_height"] + locals.environment_saturation_mixing_ratio.data[:] = inputs["local_env_saturation_mixing_ratio"] + locals.environment_moist_static_energy.data[:] = inputs["local_env_moist_static_energy"] + locals.environment_saturation_moist_static_energy.data[:] = inputs[ + "local_env_saturation_moist_static_energy" + ] + state.input_output.t_old.data[:] = inputs["t_old"] + state.input_output.vapor_old.data[:] = inputs["vapor_old"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + + code = self.stencil_factory.from_dims_halo( + func=environment_conditions, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SATURATION_CALCULATION_CHOICE": cumulus_parameterization_config.SATURATION_CALCULATION_CHOICE + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=state.input_output.t_old, + vapor=state.input_output.vapor_old, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=locals.environment_moist_static_energy, + saturation_moist_static_energy=locals.environment_saturation_moist_static_energy, + saturation_mixing_ratio=locals.environment_saturation_mixing_ratio, + geopotential_height=locals.geopotential_height, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + # state fields + "local_geopotential_height": locals.geopotential_height.field[:], + "local_env_saturation_mixing_ratio": locals.environment_saturation_mixing_ratio.field[:], + "local_env_moist_static_energy": locals.environment_moist_static_energy.field[:], + "local_env_saturation_moist_static_energy": locals.environment_saturation_moist_static_energy.field[ + : + ], + "t_old": state.input_output.t_old.field[:], + "vapor_old": state.input_output.vapor_old.field[:], + "p_forced": state.input_output.p_forced.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "p_surface": state.input_output.p_surface.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_1_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_2.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_2.py new file mode 100644 index 000000000..c2d7fa6d1 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_2.py @@ -0,0 +1,206 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_conditions +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "geopotential_height_forced": {}, + "local_env_saturation_mixing_ratio_forced": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_saturation_moist_static_energy_forced": {}, + "local_t_new": {}, + "local_vapor_forced": {}, + "p_forced": {}, + "topography_height_no_negative": {}, + "p_surface": {}, + "error_code": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + locals.environment_saturation_mixing_ratio_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_forced" + ] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_saturation_moist_static_energy_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_forced" + ] + locals.t_new.data[:] = inputs["local_t_new"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + + code = self.stencil_factory.from_dims_halo( + func=environment_conditions, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SATURATION_CALCULATION_CHOICE": cumulus_parameterization_config.SATURATION_CALCULATION_CHOICE + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=locals.t_new, + vapor=locals.vapor_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=locals.environment_moist_static_energy_forced, + saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_forced, + saturation_mixing_ratio=locals.environment_saturation_mixing_ratio_forced, + geopotential_height=state.input_output.geopotential_height_forced, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "local_env_saturation_mixing_ratio_forced": locals.environment_saturation_mixing_ratio_forced.field[ + : + ], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_env_saturation_moist_static_energy_forced": locals.environment_saturation_moist_static_energy_forced.field[ + : + ], + "local_t_new": locals.t_new.field[:], + "local_vapor_forced": locals.vapor_forced.field[:], + "p_forced": state.input_output.p_forced.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "p_surface": state.input_output.p_surface.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_2_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_3.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_3.py new file mode 100644 index 000000000..0e051256b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentConditions_3.py @@ -0,0 +1,210 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_conditions +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_geopotential_height_modified": {}, + "local_env_saturation_mixing_ratio_modified": {}, + "local_env_moist_static_energy_modified": {}, + "local_env_saturation_moist_static_energy_modified": {}, + "local_t_modified": {}, + "local_vapor_modified": {}, + "p_forced": {}, + "topography_height_no_negative": {}, + "p_surface": {}, + "error_code": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.geopotential_height_modified.data[:] = inputs["local_geopotential_height_modified"] + locals.environment_saturation_mixing_ratio_modified.data[:] = inputs[ + "local_env_saturation_mixing_ratio_modified" + ] + locals.environment_moist_static_energy_modified.data[:] = inputs[ + "local_env_moist_static_energy_modified" + ] + locals.environment_saturation_moist_static_energy_modified.data[:] = inputs[ + "local_env_saturation_moist_static_energy_modified" + ] + locals.t_modified.data[:] = inputs["local_t_modified"] + locals.vapor_modified.data[:] = inputs["local_vapor_modified"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + + code = self.stencil_factory.from_dims_halo( + func=environment_conditions, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "SATURATION_CALCULATION_CHOICE": cumulus_parameterization_config.SATURATION_CALCULATION_CHOICE + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.input_output.p_forced, + p_surface=state.input_output.p_surface, + t=locals.t_modified, + vapor=locals.vapor_modified, + topography_height_no_negative=state.input_output.topography_height_no_negative, + moist_static_energy=locals.environment_moist_static_energy_modified, + saturation_moist_static_energy=locals.environment_saturation_moist_static_energy_modified, + saturation_mixing_ratio=locals.environment_saturation_mixing_ratio_modified, + geopotential_height=locals.geopotential_height_modified, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "local_geopotential_height_modified": locals.geopotential_height_modified.field[:], + "local_env_saturation_mixing_ratio_modified": locals.environment_saturation_mixing_ratio_modified.field[ + : + ], + "local_env_moist_static_energy_modified": locals.environment_moist_static_energy_modified.field[ + : + ], + "local_env_saturation_moist_static_energy_modified": locals.environment_saturation_moist_static_energy_modified.field[ + : + ], + "local_t_modified": locals.t_modified.field[:], + "local_vapor_modified": locals.vapor_modified.field[:], + "p_forced": state.input_output.p_forced.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "p_surface": state.input_output.p_surface.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentConditions_3_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentMassFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentMassFlux.py new file mode 100644 index 000000000..4df7b302b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_EnvironmentMassFlux.py @@ -0,0 +1,185 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import environment_mass_flux +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "epsilon_forced": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "local_environment_massflux": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + locals.environment_massflux.data[:] = inputs["local_environment_massflux"] + + code = self.stencil_factory.from_dims_halo( + func=environment_mass_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + epsilon_forced=state.output.epsilon_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_massflux=locals.environment_massflux, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "epsilon_forced": state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_environment_massflux": locals.environment_massflux.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnvironmentMassFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_ModifyEnvironmentProfiles.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_ModifyEnvironmentProfiles.py new file mode 100644 index 000000000..4a0a1063a --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/environment/translate_GF2020_CumulusParameterization_ModifyEnvironmentProfiles.py @@ -0,0 +1,290 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.environment import modify_environment_profiles +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "updraft_origin_level": {}, + "ocean_fraction": {}, + "p_forced": {}, + "local_t_new": {}, + "local_t_modified": {}, + "local_vapor_forced": {}, + "local_vapor_modified": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_moist_static_energy_modified": {}, + "local_moist_static_energy_origin_level_forced": {}, + "local_moist_static_energy_origin_level_modified": {}, + "local_partition_liquid_ice": {}, + "local_del_moist_static_energy_cloud_ensemble": {}, + "local_del_t_cloud_ensemble": {}, + "local_del_vapor_cloud_ensemble": {}, + "local_del_cloud_liquid_cloud_ensemble": {}, + "local_del_u_cloud_ensemble": {}, + "local_del_v_cloud_ensemble": {}, + "local_moist_static_energy_tendency_from_environmental_subsidence": {}, + "local_vapor_tendency_from_environmental_subsidence": {}, + "local_t_tendency_from_environmental_subsidence": {}, + "local_arbitrary_numerical_parameter": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.t_new.data[:] = inputs["local_t_new"] + locals.t_modified.data[:] = inputs["local_t_modified"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + locals.vapor_modified.data[:] = inputs["local_vapor_modified"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_moist_static_energy_modified.data[:] = inputs[ + "local_env_moist_static_energy_modified" + ] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + locals.moist_static_energy_origin_level_modified.data[:] = inputs[ + "local_moist_static_energy_origin_level_modified" + ] + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + locals.del_moist_static_energy_cloud_ensemble.data[:] = inputs[ + "local_del_moist_static_energy_cloud_ensemble" + ] + locals.del_t_cloud_ensemble.data[:] = inputs["local_del_t_cloud_ensemble"] + locals.del_vapor_cloud_ensemble.data[:] = inputs["local_del_vapor_cloud_ensemble"] + locals.del_cloud_liquid_cloud_ensemble.data[:] = inputs["local_del_cloud_liquid_cloud_ensemble"] + locals.del_u_cloud_ensemble.data[:] = inputs["local_del_u_cloud_ensemble"] + locals.del_v_cloud_ensemble.data[:] = inputs["local_del_v_cloud_ensemble"] + locals.moist_static_energy_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_moist_static_energy_tendency_from_environmental_subsidence" + ] + locals.vapor_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_vapor_tendency_from_environmental_subsidence" + ] + locals.t_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_t_tendency_from_environmental_subsidence" + ] + locals.arbitrary_numerical_parameter.data[:] = inputs["local_arbitrary_numerical_parameter"] + + code = self.stencil_factory.from_dims_halo( + func=modify_environment_profiles, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "COUPLE_MICROPHYSICS": cumulus_parameterization_config.COUPLE_MICROPHYSICS, + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + updraft_origin_level=state.output.updraft_origin_level, + ocean_fraction=state.input.ocean_fraction, + p_forced=state.input_output.p_forced, + t_new=locals.t_new, + t_modified=locals.t_modified, + vapor_forced=locals.vapor_forced, + vapor_modified=locals.vapor_modified, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_moist_static_energy_modified=locals.environment_moist_static_energy_modified, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + moist_static_energy_origin_level_modified=locals.moist_static_energy_origin_level_modified, + partition_liquid_ice=locals.partition_liquid_ice, + del_moist_static_energy_cloud_ensemble=locals.del_moist_static_energy_cloud_ensemble, + del_t_cloud_ensemble=locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=locals.del_cloud_liquid_cloud_ensemble, + del_u_cloud_ensemble=locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=locals.del_v_cloud_ensemble, + moist_static_energy_tendency_from_environmental_subsidence=locals.moist_static_energy_tendency_from_environmental_subsidence, + vapor_tendency_from_environmental_subsidence=locals.vapor_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=locals.t_tendency_from_environmental_subsidence, + arbitrary_numerical_parameter=locals.arbitrary_numerical_parameter, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "ocean_fraction": state.input.ocean_fraction.data[:], + "p_forced": state.input_output.p_forced.data[:], + "local_t_new": locals.t_new.data[:], + "local_t_modified": locals.t_modified.data[:], + "local_vapor_forced": locals.vapor_forced.data[:], + "local_vapor_modified": locals.vapor_modified.data[:], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.data[:], + "local_env_moist_static_energy_modified": locals.environment_moist_static_energy_modified.data[:], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.data[ + : + ], + "local_moist_static_energy_origin_level_modified": locals.moist_static_energy_origin_level_modified.data[ + : + ], + "local_partition_liquid_ice": locals.partition_liquid_ice.data[:], + "local_del_moist_static_energy_cloud_ensemble": locals.del_moist_static_energy_cloud_ensemble.data[ + : + ], + "local_del_t_cloud_ensemble": locals.del_t_cloud_ensemble.data[:], + "local_del_vapor_cloud_ensemble": locals.del_vapor_cloud_ensemble.data[:], + "local_del_cloud_liquid_cloud_ensemble": locals.del_cloud_liquid_cloud_ensemble.data[:], + "local_del_u_cloud_ensemble": locals.del_u_cloud_ensemble.data[:], + "local_del_v_cloud_ensemble": locals.del_v_cloud_ensemble.data[:], + "local_moist_static_energy_tendency_from_environmental_subsidence": locals.moist_static_energy_tendency_from_environmental_subsidence.data[ + : + ], + "local_vapor_tendency_from_environmental_subsidence": locals.vapor_tendency_from_environmental_subsidence.data[ + : + ], + "local_t_tendency_from_environmental_subsidence": locals.t_tendency_from_environmental_subsidence.data[ + : + ], + "local_arbitrary_numerical_parameter": locals.arbitrary_numerical_parameter.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ModifyEnvironmentProfiles_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_CloudTop.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_CloudTop.py new file mode 100644 index 000000000..438849f65 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_CloudTop.py @@ -0,0 +1,252 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import cloud_top_checks, get_cloud_top +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "entrainment_rate": {}, + "local_moist_static_energy_origin_level_forced": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "updraft_lfc_level": {}, + "local_cloud_moist_static_energy_forced_transported": {}, + "p_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels": {}, + "last_error_code": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_top_level" + ] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + locals.cloud_moist_static_energy_forced_transported.data[:] = inputs[ + "local_cloud_moist_static_energy_forced_transported" + ] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + state.input.last_error_code.data[:] = inputs["last_error_code"] + + code_part_1 = self.stencil_factory.from_dims_halo( + func=get_cloud_top, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"OVERSHOOT": cumulus_parameterization_config.OVERSHOOT}, + ) + + code_part_2 = self.stencil_factory.from_dims_halo( + func=cloud_top_checks, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code_part_1( + entrainment_rate=state.output.entrainment_rate, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + cloud_moist_static_energy_forced_transported=locals.cloud_moist_static_energy_forced_transported, + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + code_part_2( + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + error_code=state.output.error_code, + last_error_code=state.input.last_error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + MINIMUM_DEPTH=plume_dependent_constants.MINIMUM_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + state.output.cloud_top_level.field[:] = state.output.cloud_top_level.field[:] + 1 + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.field[ + : + ], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_cloud_moist_static_energy_forced_transported": locals.cloud_moist_static_energy_forced_transported.field[ + : + ], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.field[:], + "last_error_code": state.input.last_error_code.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudTop_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudTop_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudTop_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_ConvectiveCloudBaseLevel.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_ConvectiveCloudBaseLevel.py new file mode 100644 index 000000000..0065a019b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_ConvectiveCloudBaseLevel.py @@ -0,0 +1,339 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import ( + get_convective_cloud_base_level, + set_start_level, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_cap_max_increment": {}, + "local_cap_max": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_vapor_cloud_levels_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "p_forced": {}, + "p_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_env_moist_static_energy_forced": {}, + "local_moist_static_energy_origin_level_forced": {}, + "local_vapor_forced": {}, + "local_env_saturation_mixing_ratio_forced": {}, + "entrainment_rate": {}, + "local_cloud_moist_static_energy_forced_transported": {}, + "updraft_origin_level": {}, + "local_maximum_updraft_origin_level": {}, + "lcl_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_negative_buoyancy_depth": {}, + "local_frh_lfc": {}, + "t_perturbation": {}, + "local_start_level": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoyancy": {}, + "ocean_fraction": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.cap_max_increment.data[:] = inputs["local_cap_max_increment"] + locals.cap_max.data[:] = inputs["local_cap_max"] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + locals.environment_saturation_mixing_ratio_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_forced" + ] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + locals.cloud_moist_static_energy_forced_transported.data[:] = inputs[ + "local_cloud_moist_static_energy_forced_transported" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + locals.maximum_updraft_origin_level.data[:] = inputs["local_maximum_updraft_origin_level"] - 1 + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.negative_buoyancy_depth.data[:] = inputs["local_negative_buoyancy_depth"] + locals.frh_lfc.data[:] = inputs["local_frh_lfc"] + state.output.t_perturbation.data[:] = inputs["t_perturbation"] + locals.start_level.data[:] = inputs["local_start_level"] - 1 + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + locals.ocean_fraction.data[:] = inputs["ocean_fraction"] + + code_part_1 = self.stencil_factory.from_dims_halo( + func=set_start_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + code_part_2 = self.stencil_factory.from_dims_halo( + func=get_convective_cloud_base_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "OVERSHOOT": cumulus_parameterization_config.OVERSHOOT, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + "MOIST_TRIGGER": cumulus_parameterization_config.MOIST_TRIGGER, + "USE_MEMORY": cumulus_parameterization_config.USE_MEMORY, + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code_part_1( + lcl_level=state.output.lcl_level, + start_level=locals.start_level, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + code_part_2( + error_code=state.output.error_code, + lcl_level=state.output.lcl_level, + cloud_moist_static_energy_forced_transported=locals.cloud_moist_static_energy_forced_transported, + cap_max=locals.cap_max, + updraft_origin_level=state.output.updraft_origin_level, + start_level=locals.start_level, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + negative_buoyancy_depth=locals.negative_buoyancy_depth, + frh_lfc=locals.frh_lfc, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + entrainment_rate=state.output.entrainment_rate, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels_forced=locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + t_excess=locals.t_excess, + vapor_excess=locals.vapor_excess, + add_buoyancy=locals.add_buoyancy, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + vapor_forced=locals.vapor_forced, + environment_saturation_mixing_ratio_forced=locals.environment_saturation_mixing_ratio_forced, + ocean_fraction=locals.ocean_fraction, + cap_max_increment=locals.cap_max_increment, + t_perturbation=state.output.t_perturbation, + p_forced=state.input_output.p_forced, + cloud_top_level=state.output.cloud_top_level, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_cap_max_increment": locals.cap_max_increment.field[:], + "local_cap_max": locals.cap_max.field[:], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.field[ + : + ], + "p_forced": state.input_output.p_forced.field[:], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.field[ + : + ], + "local_vapor_forced": locals.vapor_forced.field[:], + "local_env_saturation_mixing_ratio_forced": locals.environment_saturation_mixing_ratio_forced.field[ + : + ], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_cloud_moist_static_energy_forced_transported": locals.cloud_moist_static_energy_forced_transported.field[ + : + ], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_maximum_updraft_origin_level": locals.maximum_updraft_origin_level.field[:] + 1, + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_negative_buoyancy_depth": locals.negative_buoyancy_depth.field[:], + "local_frh_lfc": locals.frh_lfc.field[:], + "t_perturbation": state.output.t_perturbation.field[:], + "local_start_level": locals.start_level.field[:] + 1, + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_add_buoyancy": locals.add_buoyancy.field[:], + "ocean_fraction": locals.ocean_fraction.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectiveCloudBaseLevel_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_GetLCL.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_GetLCL.py new file mode 100644 index 000000000..3267f4759 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_GetLCL.py @@ -0,0 +1,229 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import find_lcl +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "p_forced": {}, + "local_p_cloud_levels": {}, + "local_t_excess": {}, + "local_t_cloud_levels": {}, + "t_perturbation": {}, + "local_vapor_excess": {}, + "local_vapor_cloud_levels_forced": {}, + "omega": {}, + "air_density": {}, + "local_geopotential_height_cloud_levels": {}, + "topography_height_no_negative": {}, + "ocean_fraction": {}, + "updraft_origin_level": {}, + "grid_length": {}, + "lcl_level": {}, + "error_code": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.p_cloud_levels.data[:] = inputs["local_p_cloud_levels"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + state.output.t_perturbation.data[:] = inputs["t_perturbation"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.vapor_cloud_levels.data[:] = inputs["local_vapor_cloud_levels_forced"] + state.input_output.omega.data[:] = inputs["omega"] + state.input_output.air_density.data[:] = inputs["air_density"] + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.input_output.grid_length.data[:] = inputs["grid_length"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + + code = self.stencil_factory.from_dims_halo( + func=find_lcl, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + "ADV_TRIGGER": config.ADV_TRIGGER, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + p=state.input_output.p_forced, + p_cloud_levels=locals.p_cloud_levels, + t_excess=locals.t_excess, + t_cloud_levels_forced=locals.t_cloud_levels, + t_perturbation=state.output.t_perturbation, + vapor_excess=locals.vapor_excess, + vapor_cloud_levels_forced=locals.vapor_cloud_levels, + omega=state.input_output.omega, + air_density=state.input_output.air_density, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + topography_height_no_negative=state.input_output.topography_height_no_negative, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + grid_length=state.input_output.grid_length, + lcl_level=state.output.lcl_level, + error_code=state.output.error_code, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "p_forced": state.input_output.p_forced.field[:], + "local_p_cloud_levels": locals.p_cloud_levels.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_t_cloud_levels": locals.t_cloud_levels.field[:], + "t_perturbation": state.output.t_perturbation.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels.field[:], + "omega": state.input_output.omega.field[:], + "air_density": state.input_output.air_density.field[:], + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "ocean_fraction": state.input.ocean_fraction.field[:], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "grid_length": state.input_output.grid_length.field[:], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetLCL_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetLCL_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_GetLCL_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_HighestMoistStaticEnergyLevel.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_HighestMoistStaticEnergyLevel.py new file mode 100644 index 000000000..2edc420ee --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/get_levels/translate_GF2020_CumulusParameterization_HighestMoistStaticEnergyLevel.py @@ -0,0 +1,182 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import ( + find_highest_moist_static_energy_level, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "updraft_origin_level": {}, + "error_code": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_maximum_updraft_origin_level": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.maximum_updraft_origin_level.data[:] = inputs["local_maximum_updraft_origin_level"] - 1 + + code = self.stencil_factory.from_dims_halo( + func=find_highest_moist_static_energy_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + moist_static_energy=locals.environment_moist_static_energy_cloud_levels_forced, + error_code=state.output.error_code, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + updraft_origin_level=state.output.updraft_origin_level, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_maximum_updraft_origin_level": locals.maximum_updraft_origin_level.field[:] + 1, + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_HighestMoistStaticEnergyLevel_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating/translate_GF2020_CumulusParameterization_KineticEnergyToHeating.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating/translate_GF2020_CumulusParameterization_KineticEnergyToHeating.py new file mode 100644 index 000000000..adff3564a --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/kinetic_energy_to_heating/translate_GF2020_CumulusParameterization_KineticEnergyToHeating.py @@ -0,0 +1,196 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.kinetic_energy_to_heating import ( + kinetic_energy_to_heating, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "p_cloud_levels_forced": {}, + "u": {}, + "v": {}, + "local_del_u_cloud_ensemble": {}, + "cloud_top_level": {}, + "local_del_v_cloud_ensemble": {}, + "local_del_t_cloud_ensemble": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.del_u_cloud_ensemble.data[:] = inputs["local_del_u_cloud_ensemble"] + locals.del_v_cloud_ensemble.data[:] = inputs["local_del_v_cloud_ensemble"] + locals.del_t_cloud_ensemble.data[:] = inputs["local_del_t_cloud_ensemble"] + + code = self.stencil_factory.from_dims_halo( + func=kinetic_energy_to_heating, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + u=state.input_output.u, + v=state.input_output.v, + del_u_cloud_ensemble=locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=locals.del_v_cloud_ensemble, + del_t_cloud_ensemble=locals.del_t_cloud_ensemble, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "u": state.input_output.u.data[:], + "v": state.input_output.v.data[:], + "local_del_u_cloud_ensemble": locals.del_u_cloud_ensemble.data[:], + "local_del_v_cloud_ensemble": locals.del_v_cloud_ensemble.data[:], + "local_del_t_cloud_ensemble": locals.del_t_cloud_ensemble.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_KineticEnergyToHeating_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/large_scale_forcing/translate_GF2020_CumulusParameterization_LargeScaleForcing.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/large_scale_forcing/translate_GF2020_CumulusParameterization_LargeScaleForcing.py new file mode 100644 index 000000000..a4100993e --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/large_scale_forcing/translate_GF2020_CumulusParameterization_LargeScaleForcing.py @@ -0,0 +1,339 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.large_scale_forcing import LargeScaleForcing +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_error_code_2": {}, + "local_error_code_3": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "pbl_level": {}, + "local_ocean_fraction": {}, + "p_cloud_levels_forced": {}, + "local_vapor_forced": {}, + "condensate_to_fall_forced": {}, + "local_effective_condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "omega": {}, + "convective_scale_velocity": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "local_cloud_moist_static_energy": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_dmoist_static_energydt": {}, + "local_cloud_workfunction_0": {}, + "local_cloud_workfunction_0_modified": {}, + "local_cloud_workfunction_1": {}, + "local_cloud_workfunction_1_pbl": {}, + "local_arbitrary_numerical_parameter": {}, + "local_f_dicycle_modified": {}, + "local_cape_removal_time_scale": {}, + "epsilon_forced": {}, + "local_k_x_modified": {}, + "local_mass_flux_ensemble": {}, + "local_precipitation_ensemble": {}, + "local_xff_mid": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.error_code_2.data[:] = inputs["local_error_code_2"] + locals.error_code_3.data[:] = inputs["local_error_code_3"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input_output.pbl_level.data[:] = inputs["pbl_level"] - 1 + locals.ocean_fraction.data[:] = inputs["local_ocean_fraction"] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + locals.effective_condensate_to_fall_forced.data[:, :, :] = inputs[ + "local_effective_condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + state.input_output.omega.data[:] = inputs["omega"] + state.input_output.convective_scale_velocity.data[:] = inputs["convective_scale_velocity"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + locals.cloud_moist_static_energy.data[:] = inputs["local_cloud_moist_static_energy"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.dmoist_static_energydt.data[:] = inputs["local_dmoist_static_energydt"] + locals.cloud_workfunction_0.data[:] = inputs["local_cloud_workfunction_0"] + locals.cloud_workfunction_0_modified.data[:] = inputs["local_cloud_workfunction_0_modified"] + locals.cloud_workfunction_1.data[:] = inputs["local_cloud_workfunction_1"] + locals.cloud_workfunction_1_pbl.data[:] = inputs["local_cloud_workfunction_1_pbl"] + locals.arbitrary_numerical_parameter.data[:] = inputs["local_arbitrary_numerical_parameter"] + locals.f_dicycle_modified.data[:] = inputs["local_f_dicycle_modified"] + locals.cape_removal_time_scale.data[:] = inputs["local_cape_removal_time_scale"] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + locals.k_x_modified.data[:] = inputs["local_k_x_modified"] + locals.mass_flux_ensemble.data[:] = inputs["local_mass_flux_ensemble"][:, :, 0:16] + locals.precipitation_ensemble.data[:] = inputs["local_precipitation_ensemble"][:, :, 0:16] + locals.xff_mid.data[:] = inputs["local_xff_mid"][:, :, 0:16] + + code = LargeScaleForcing( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + error_code_2=locals.error_code_2, + error_code_3=locals.error_code_3, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + ocean_fraction=locals.ocean_fraction, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + vapor_forced=locals.vapor_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + effective_condensate_to_fall_forced=locals.effective_condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + omega=state.input_output.omega, + convective_scale_velocity=state.input_output.convective_scale_velocity, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + cloud_moist_static_energy=locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=locals.environment_moist_static_energy_cloud_levels, + environment_moist_static_energy_cloud_levels_forced=locals.environment_moist_static_energy_cloud_levels_forced, + dmoist_static_energydt=locals.dmoist_static_energydt, + cloud_workfunction_0=locals.cloud_workfunction_0, + cloud_workfunction_0_modified=locals.cloud_workfunction_0_modified, + cloud_workfunction_1=locals.cloud_workfunction_1, + cloud_workfunction_1_pbl=locals.cloud_workfunction_1_pbl, + arbitrary_numerical_parameter=locals.arbitrary_numerical_parameter, + f_dicycle_modified=locals.f_dicycle_modified, + cape_removal_time_scale=locals.cape_removal_time_scale, + epsilon_forced=state.output.epsilon_forced, + k_x_modified=locals.k_x_modified, + mass_flux_ensemble=locals.mass_flux_ensemble, + precipitation_ensemble=locals.precipitation_ensemble, + xff_mid=locals.xff_mid, + plume_dependent_constants=plume_dependent_constants, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_error_code_2": locals.error_code_2.field[:], + "local_error_code_3": locals.error_code_3.field[:], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "pbl_level": state.input_output.pbl_level.field[:] + 1, + "local_ocean_fraction": locals.ocean_fraction.field[:], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vapor_forced": locals.vapor_forced.field[:], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_effective_condensate_to_fall_forced": locals.effective_condensate_to_fall_forced.field[ + :, :, : + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "omega": state.input_output.omega.field[:], + "convective_scale_velocity": state.input_output.convective_scale_velocity.field[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_cloud_moist_static_energy": locals.cloud_moist_static_energy.field[:], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_dmoist_static_energydt": locals.dmoist_static_energydt.field[:], + "local_cloud_workfunction_0": locals.cloud_workfunction_0.field[:], + "local_cloud_workfunction_0_modified": locals.cloud_workfunction_0_modified.field[:], + "local_cloud_workfunction_1": locals.cloud_workfunction_1.field[:], + "local_cloud_workfunction_1_pbl": locals.cloud_workfunction_1_pbl.field[:], + "local_arbitrary_numerical_parameter": locals.arbitrary_numerical_parameter.field[:], + "local_f_dicycle_modified": locals.f_dicycle_modified.field[:], + "local_cape_removal_time_scale": locals.cape_removal_time_scale.field[:], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_k_x_modified": locals.k_x_modified.field[:], + "local_mass_flux_ensemble": locals.mass_flux_ensemble.field[:], + "local_precipitation_ensemble": locals.precipitation_ensemble.field[:], + "local_xff_mid": locals.xff_mid.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_LargeScaleForcing_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_LargeScaleForcing_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_LargeScaleForcing_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_FirstGuessMoistStaticEnergy.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_FirstGuessMoistStaticEnergy.py new file mode 100644 index 000000000..040c70647 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_FirstGuessMoistStaticEnergy.py @@ -0,0 +1,231 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.moist_static_energy import ( + first_guess_moist_static_energy, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_start_level": {}, + "cloud_top_level": {}, + "mass_detrainment_updraft_forced": {}, + "mass_entrainment_updraft_forced": {}, + "local_normalized_massflux_updraft": {}, + "local_cloud_moist_static_energy_forced": {}, + "normalized_massflux_updraft_forced": {}, + "local_env_moist_static_energy_forced": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoyancy": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.start_level.data[:] = inputs["local_start_level"] - 1 + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + locals.normalized_massflux_updraft.data[:] = inputs["local_normalized_massflux_updraft"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=first_guess_moist_static_energy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + start_level=locals.start_level, + cloud_top_level=state.output.cloud_top_level, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + normalized_massflux_updraft=locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + vapor_excess=locals.vapor_excess, + t_excess=locals.t_excess, + add_buoyancy=locals.add_buoyancy, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_start_level": locals.start_level.field[:] + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_normalized_massflux_updraft": locals.normalized_massflux_updraft.field[:], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced, + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_add_buoyancy": locals.add_buoyancy.field[:], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_FirstGuessMoistStaticEnergy_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_1.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_1.py new file mode 100644 index 000000000..2ee9c95f1 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_1.py @@ -0,0 +1,224 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.moist_static_energy import parcel_moist_static_energy +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "ocean_fraction": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoy": {}, + "p_forced": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_moist_static_energy_origin_level": {}, + "local_moist_static_energy_origin_level_forced": {}, + "updraft_origin_level": {}, + "t_perturbation": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoy"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.moist_static_energy_origin_level.data[:] = inputs["local_moist_static_energy_origin_level"] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.t_perturbation.data[:] = inputs["t_perturbation"] + + code = self.stencil_factory.from_dims_halo( + func=parcel_moist_static_energy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + t_excess=locals.t_excess, + vapor_excess=locals.vapor_excess, + add_buoyancy=locals.add_buoyancy, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + p=state.input_output.p_forced, + environmenet_moist_static_energy=locals.environment_moist_static_energy_cloud_levels, + environmenet_moist_static_energy_forced=locals.environment_moist_static_energy_cloud_levels_forced, + t_perturbation=state.output.t_perturbation, + moist_static_energy_origin_level=locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "ocean_fraction": state.input.ocean_fraction.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_add_buoy": locals.add_buoyancy.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_moist_static_energy_origin_level": locals.moist_static_energy_origin_level.field[:], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.field[ + : + ], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "t_perturbation": state.output.t_perturbation.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_1_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_2.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_2.py new file mode 100644 index 000000000..e94f93b1b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_ParcelMoistStaticEnergy_2.py @@ -0,0 +1,224 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.moist_static_energy import parcel_moist_static_energy +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "ocean_fraction": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoy": {}, + "p_forced": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_moist_static_energy_origin_level": {}, + "local_moist_static_energy_origin_level_forced": {}, + "updraft_origin_level": {}, + "t_perturbation": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoy"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.moist_static_energy_origin_level.data[:] = inputs["local_moist_static_energy_origin_level"] + locals.moist_static_energy_origin_level_forced.data[:] = inputs[ + "local_moist_static_energy_origin_level_forced" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.t_perturbation.data[:] = inputs["t_perturbation"] + + code = self.stencil_factory.from_dims_halo( + func=parcel_moist_static_energy, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + t_excess=locals.t_excess, + vapor_excess=locals.vapor_excess, + add_buoyancy=locals.add_buoyancy, + ocean_fraction=state.input.ocean_fraction, + updraft_origin_level=state.output.updraft_origin_level, + p=state.input_output.p_forced, + environmenet_moist_static_energy=locals.environment_moist_static_energy_cloud_levels, + environmenet_moist_static_energy_forced=locals.environment_moist_static_energy_cloud_levels_forced, + t_perturbation=state.output.t_perturbation, + moist_static_energy_origin_level=locals.moist_static_energy_origin_level, + moist_static_energy_origin_level_forced=locals.moist_static_energy_origin_level_forced, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "ocean_fraction": state.input.ocean_fraction.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_add_buoy": locals.add_buoyancy.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_moist_static_energy_origin_level": locals.moist_static_energy_origin_level.field[:], + "local_moist_static_energy_origin_level_forced": locals.moist_static_energy_origin_level_forced.field[ + : + ], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "t_perturbation": state.output.t_perturbation.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ParcelMoistStaticEnergy_2_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_StaticControl.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_StaticControl.py new file mode 100644 index 000000000..76ea89113 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/moist_static_energy/translate_GF2020_CumulusParameterization_StaticControl.py @@ -0,0 +1,271 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.moist_static_energy import StaticControl +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_start_level": {}, + "lcl_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_cloud_moist_static_energy_modified": {}, + "local_moist_static_energy_origin_level_modified": {}, + "local_env_moist_static_energy_modified": {}, + "local_env_moist_static_energy_cloud_levels_modified": {}, + "local_env_saturation_moist_static_energy_cloud_levels_modified": {}, + "mass_detrainment_updraft_forced": {}, + "mass_entrainment_updraft_forced": {}, + "local_normalized_massflux_updraft_modified": {}, + "local_partition_liquid_ice": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoyancy": {}, + "cloud_liquid_after_rain_forced": {}, + "local_d_buoyancy_modified": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.start_level.data[:] = inputs["local_start_level"] - 1 + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.cloud_moist_static_energy_modified.data[:] = inputs["local_cloud_moist_static_energy_modified"] + locals.moist_static_energy_origin_level_modified.data[:] = inputs[ + "local_moist_static_energy_origin_level_modified" + ] + locals.environment_moist_static_energy_modified.data[:] = inputs[ + "local_env_moist_static_energy_modified" + ] + locals.environment_moist_static_energy_cloud_levels_modified.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_modified" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_modified.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_modified" + ] + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + locals.normalized_massflux_updraft_modified.data[:] = inputs[ + "local_normalized_massflux_updraft_modified" + ] + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + locals.d_buoyancy_modified.data[:] = inputs["local_d_buoyancy_modified"] + + # initialize test code + code = StaticControl( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + start_level=locals.start_level, + lcl_level=state.output.lcl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + cloud_moist_static_energy_modified=locals.cloud_moist_static_energy_modified, + moist_static_energy_origin_level_modified=locals.moist_static_energy_origin_level_modified, + environment_moist_static_energy_modified=locals.environment_moist_static_energy_modified, + environment_moist_static_energy_cloud_levels_modified=locals.environment_moist_static_energy_cloud_levels_modified, + environment_saturation_moist_static_energy_cloud_levels_modified=locals.environment_saturation_moist_static_energy_cloud_levels_modified, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + normalized_massflux_updraft_modified=locals.normalized_massflux_updraft_modified, + partition_liquid_ice=locals.partition_liquid_ice, + vapor_excess=locals.vapor_excess, + t_excess=locals.t_excess, + add_buoyancy=locals.add_buoyancy, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + d_buoyancy_modified=locals.d_buoyancy_modified, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "local_start_level": locals.start_level.data[:] + 1, + "lcl_level": state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_cloud_moist_static_energy_modified": locals.cloud_moist_static_energy_modified.data[:], + "local_moist_static_energy_origin_level_modified": locals.moist_static_energy_origin_level_modified.data[ + : + ], + "local_env_moist_static_energy_modified": locals.environment_moist_static_energy_modified.data[:], + "local_env_moist_static_energy_cloud_levels_modified": locals.environment_moist_static_energy_cloud_levels_modified.data[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_modified": locals.environment_saturation_moist_static_energy_cloud_levels_modified.data[ + : + ], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_normalized_massflux_updraft_modified": locals.normalized_massflux_updraft_modified.data[:], + "local_partition_liquid_ice": locals.partition_liquid_ice.data[:], + "local_vapor_excess": locals.vapor_excess.data[:], + "local_t_excess": locals.t_excess.data[:], + "local_add_buoyancy": locals.add_buoyancy.data[:], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_d_buoyancy_modified": locals.d_buoyancy_modified.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_StaticControl_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_StaticControl_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_StaticControl_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_CloudDissipation.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_CloudDissipation.py new file mode 100644 index 000000000..92107fa46 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_CloudDissipation.py @@ -0,0 +1,260 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.precip import cloud_dissapation +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_hydrostatic_air_density": {}, + "geopotential_height_forced": {}, + "local_t_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "cloud_base_mass_flux_modified": {}, + "local_vapor_cloud_levels_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_vertical_velocity_3d": {}, + "scale_dependence_factor": {}, + "dtdt": {}, + "dvapordt": {}, + "dcloudicedt": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.hydrostatic_air_density.data[:] = inputs["local_hydrostatic_air_density"] + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.cloud_base_mass_flux_modified.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_base_mass_flux_modified" + ] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.vertical_velocity_3d.data[:] = inputs["local_vertical_velocity_3d"] + state.output.scale_dependence_factor.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "scale_dependence_factor" + ] + state.output.dtdt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dtdt"] + state.output.dvapordt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dvapordt"] + state.output.dcloudicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dcloudicedt"] + + code = self.stencil_factory.from_dims_halo( + func=cloud_dissapation, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_CLOUD_DISSIPATION": cumulus_parameterization_config.USE_CLOUD_DISSIPATION, + "DTIME": cumulus_parameterization_config.DTIME, + "COUPLE_MICROPHYSICS": cumulus_parameterization_config.COUPLE_MICROPHYSICS, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + hydrostatic_air_density=locals.hydrostatic_air_density, + geopotential_height_forced=state.input_output.geopotential_height_forced, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + vertical_velocity_3d=locals.vertical_velocity_3d, + scale_dependence_factor=state.output.scale_dependence_factor, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dcloudicedt=state.output.dcloudicedt, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_hydrostatic_air_density": locals.hydrostatic_air_density.field[:], + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_vertical_velocity_3d": locals.vertical_velocity_3d.field[:], + "scale_dependence_factor": state.output.scale_dependence_factor.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "dtdt": state.output.dtdt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dvapordt": state.output.dvapordt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dcloudicedt": state.output.dcloudicedt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudDissipation_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudDissipation_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_CloudDissipation_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels.py new file mode 100644 index 000000000..10ac81e02 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels.py @@ -0,0 +1,241 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.get_levels import ( + find_detrainmet_start_level, + find_maximum_updraft_origin_level, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.precip import partition_liquid_ice +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_t_new": {}, + "topography_height_no_negative": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "p_cloud_levels_forced": {}, + "local_partition_liquid_ice": {}, + "local_melting_layer": {}, + "convection_fraction": {}, + "surface_type": {}, + "local_maximum_updraft_origin_level": {}, + "local_detrainment_start_level": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.t_new.data[:] = inputs["local_t_new"] + state.input_output.topography_height_no_negative.data[:] = inputs["topography_height_no_negative"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + locals.melting_layer.data[:] = inputs["local_melting_layer"] + state.input.convection_fraction.data[:] = inputs["convection_fraction"] + state.input.surface_type.data[:] = inputs["surface_type"] + locals.maximum_updraft_origin_level.data[:] = inputs["local_maximum_updraft_origin_level"] - 1 + locals.detrainment_start_level.data[:] = inputs["local_detrainment_start_level"] - 1 + + code_part_1 = self.stencil_factory.from_dims_halo( + func=partition_liquid_ice, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "MELT_GLAC": cumulus_parameterization_config.MELT_GLAC, + "FRAC_MODIS": cumulus_parameterization_config.FRAC_MODIS, + }, + ) + + code_part_2 = self.stencil_factory.from_dims_halo( + func=find_maximum_updraft_origin_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + code_part_3 = self.stencil_factory.from_dims_halo( + func=find_detrainmet_start_level, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code_part_1( + t=locals.t_new, + p=state.output.p_cloud_levels_forced, + geopotential_height=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + surface_type=state.input.surface_type, + convection_fraction=state.input.convection_fraction, + error_code=state.output.error_code, + melting_layer=locals.melting_layer, + part_liquid_ice=locals.partition_liquid_ice, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + code_part_2( + geopotential_height=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + MAX_UPDRAFT_ORIGIN_HEIGHT=plume_dependent_constants.MAX_UPDRAFT_ORIGIN_HEIGHT, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + code_part_3( + geopotential_height=locals.geopotential_height_cloud_levels_forced, + topography_height_no_negative=state.input_output.topography_height_no_negative, + error_code=state.output.error_code, + detrainment_start_level=locals.detrainment_start_level, + DETRAINMENT_CRITICAL_DEPTH=plume_dependent_constants.DETRAINMENT_CRITICAL_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_t_new": locals.t_new.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_partition_liquid_ice": locals.partition_liquid_ice.field[:], + "local_melting_layer": locals.melting_layer.field[:], + "convection_fraction": state.input.convection_fraction.field[:], + "surface_type": state.input.surface_type.field[:], + "local_maximum_updraft_origin_level": locals.maximum_updraft_origin_level.field[:] + 1, + "local_detrainment_start_level": locals.detrainment_start_level.field[:] + 1, + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PartitionLiquidIceAndGetLevels_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PrecipitationFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PrecipitationFlux.py new file mode 100644 index 000000000..60714c17b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_PrecipitationFlux.py @@ -0,0 +1,204 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.precip import get_precip_fluxes +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "cloud_base_mass_flux_modified": {}, + "epsilon_forced": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "local_precipitation_flux": {}, + "local_evaporation_flux": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.cloud_base_mass_flux_modified.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_base_mass_flux_modified" + ] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + locals.precipitation_flux.data[:] = inputs["local_precipitation_flux"] + locals.evaporation_flux.data[:] = inputs["local_evaporation_flux"] + + code = self.stencil_factory.from_dims_halo( + func=get_precip_fluxes, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + epsilon_forced=state.output.epsilon_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + precipitation_flux=locals.precipitation_flux, + evaporation_flux=locals.evaporation_flux, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "epsilon_forced": state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_precipitation_flux": locals.precipitation_flux.data[:], + "local_evaporation_flux": locals.evaporation_flux.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrecipitationFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrecipitationFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrecipitationFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_RainEvaporationBelowCloudBase.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_RainEvaporationBelowCloudBase.py new file mode 100644 index 000000000..be88dedbc --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/precip/translate_GF2020_CumulusParameterization_RainEvaporationBelowCloudBase.py @@ -0,0 +1,265 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.precip import rain_evaporation_below_cloud_base +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "ocean_fraction": {}, + "p_cloud_levels_forced": {}, + "p_surface": {}, + "local_t_cloud_levels": {}, + "local_vapor_cloud_levels_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels": {}, + "epsilon_forced": {}, + "cloud_base_mass_flux_modified": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "precip": {}, + "local_precipitation_flux": {}, + "local_evaporation_flux": {}, + "local_evaporation_below_cloud_base": {}, + "dtdt": {}, + "dvapordt": {}, + "dbuoyancydt": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.input_output.p_surface.data[:] = inputs["p_surface"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels" + ] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + state.output.cloud_base_mass_flux_modified.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_base_mass_flux_modified" + ] + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + state.output.precip.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["precip"] + locals.precipitation_flux.data[:] = inputs["local_precipitation_flux"] + locals.evaporation_flux.data[:] = inputs["local_evaporation_flux"] + locals.evaporation_below_cloud_base.data[:] = inputs["local_evaporation_below_cloud_base"] + state.output.dtdt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dtdt"] + state.output.dvapordt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dvapordt"] + state.output.dbuoyancydt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dbuoyancydt"] + + code = self.stencil_factory.from_dims_halo( + func=rain_evaporation_below_cloud_base, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + ocean_fraction=state.input.ocean_fraction, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + t_cloud_levels=locals.t_cloud_levels, + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + environment_saturation_mixing_ratio_cloud_levels=locals.environment_saturation_mixing_ratio_cloud_levels, + epsilon_forced=state.output.epsilon_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + precip=state.output.precip, + precipitation_flux=locals.precipitation_flux, + evaporation_flux=locals.evaporation_flux, + evaporation_below_cloud_base=locals.evaporation_below_cloud_base, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dbuoyancydt=state.output.dbuoyancydt, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "ocean_fraction": state.input.ocean_fraction.field[:], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "p_surface": state.input_output.p_surface.field[:], + "local_t_cloud_levels": locals.t_cloud_levels.field[:], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "local_env_saturation_mixing_ratio_cloud_levels": locals.environment_saturation_mixing_ratio_cloud_levels.field[ + : + ], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "precip": state.output.precip.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_precipitation_flux": locals.precipitation_flux.field[:], + "local_evaporation_flux": locals.evaporation_flux.field[:], + "local_evaporation_below_cloud_base": locals.evaporation_below_cloud_base.field[:], + "dtdt": state.output.dtdt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dvapordt": state.output.dvapordt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dbuoyancydt": state.output.dbuoyancydt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_RainEvaporationBelowCloudBase_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_DeepPrecipitationOutput.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_DeepPrecipitationOutput.py new file mode 100644 index 000000000..b88db7efc --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_DeepPrecipitationOutput.py @@ -0,0 +1,174 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import deep_precipitation_output +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "local_precipitation_flux": {}, + "convective_precip_flux": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.precipitation_flux.data[:] = inputs["local_precipitation_flux"] + state.output.convective_precip_flux.data[:] = inputs["convective_precip_flux"] + + code = self.stencil_factory.from_dims_halo( + func=deep_precipitation_output, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + precipitation_flux=locals.precipitation_flux, + convective_precip_flux=state.output.convective_precip_flux, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_precipitation_flux": locals.precipitation_flux.data[:], + "convective_precip_flux": state.output.convective_precip_flux.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_DeepPrecipitationOutput_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_EnsembleOutputAndFeedback.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_EnsembleOutputAndFeedback.py new file mode 100644 index 000000000..c89147fa6 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_EnsembleOutputAndFeedback.py @@ -0,0 +1,389 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import ensemble_output_and_feedback +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_error_code_2": {}, + "local_error_code_3": {}, + "cloud_top_level": {}, + "updraft_lfc_level": {}, + "p_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "precip": {}, + "local_effective_condensate_to_fall_forced": {}, + "cloud_base_mass_flux_modified": {}, + "scale_dependence_factor": {}, + "local_ocean_fraction": {}, + "local_f_dicycle_modified": {}, + "local_del_u_cloud_ensemble": {}, + "local_del_v_cloud_ensemble": {}, + "local_del_t_cloud_ensemble": {}, + "local_del_vapor_cloud_ensemble": {}, + "local_del_cloud_liquid_cloud_ensemble": {}, + "local_del_buoyancy_cloud_ensemble": {}, + "local_del_convective_ice_cloud_ensemble": {}, + "local_del_large_scale_ice_cloud_ensemble": {}, + "local_del_convective_liquid_cloud_ensemble": {}, + "local_del_large_scale_liquid_cloud_ensemble": {}, + "local_del_convective_cloud_fraction_cloud_ensemble": {}, + "local_del_large_scale_cloud_fraction_cloud_ensemble": {}, + "dtdt": {}, + "dvapordt": {}, + "dcloudicedt": {}, + "dudt": {}, + "dvdt": {}, + "dbuoyancydt": {}, + "dconvectiveicedt": {}, + "dlargescaleicedt": {}, + "dconvectiveliquiddt": {}, + "dlargescaleliquiddt": {}, + "dconvectivecloudfractiondt": {}, + "dlargescalecloudfractiondt": {}, + "local_mass_flux_ensemble": {}, + "local_precipitation_ensemble": {}, + "local_xff_mid": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.error_code_2.data[:] = inputs["local_error_code_2"] + locals.error_code_3.data[:] = inputs["local_error_code_3"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.precip.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["precip"] + locals.effective_condensate_to_fall_forced.data[:] = inputs[ + "local_effective_condensate_to_fall_forced" + ] + state.output.cloud_base_mass_flux_modified.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_base_mass_flux_modified" + ] + state.output.scale_dependence_factor.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "scale_dependence_factor" + ] + locals.ocean_fraction.data[:] = inputs["local_ocean_fraction"] + locals.f_dicycle_modified.data[:] = inputs["local_f_dicycle_modified"] + locals.del_u_cloud_ensemble.data[:] = inputs["local_del_u_cloud_ensemble"] + locals.del_v_cloud_ensemble.data[:] = inputs["local_del_v_cloud_ensemble"] + locals.del_t_cloud_ensemble.data[:] = inputs["local_del_t_cloud_ensemble"] + locals.del_vapor_cloud_ensemble.data[:] = inputs["local_del_vapor_cloud_ensemble"] + locals.del_cloud_liquid_cloud_ensemble.data[:] = inputs["local_del_cloud_liquid_cloud_ensemble"] + locals.del_buoyancy_cloud_ensemble.data[:] = inputs["local_del_buoyancy_cloud_ensemble"] + locals.del_convective_ice_cloud_ensemble.data[:] = inputs["local_del_convective_ice_cloud_ensemble"] + locals.del_large_scale_ice_cloud_ensemble.data[:] = inputs["local_del_large_scale_ice_cloud_ensemble"] + locals.del_convective_liquid_cloud_ensemble.data[:] = inputs[ + "local_del_convective_liquid_cloud_ensemble" + ] + locals.del_large_scale_liquid_cloud_ensemble.data[:] = inputs[ + "local_del_large_scale_liquid_cloud_ensemble" + ] + locals.del_convective_cloud_fraction_cloud_ensemble.data[:] = inputs[ + "local_del_convective_cloud_fraction_cloud_ensemble" + ] + locals.del_large_scale_cloud_fraction_cloud_ensemble.data[:] = inputs[ + "local_del_large_scale_cloud_fraction_cloud_ensemble" + ] + state.output.dtdt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dtdt"] + state.output.dvapordt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dvapordt"] + state.output.dcloudicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dcloudicedt"] + state.output.dudt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dudt"] + state.output.dvdt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dvdt"] + state.output.dbuoyancydt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dbuoyancydt"] + state.output.dconvectiveicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dconvectiveicedt" + ] + state.output.dlargescaleicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dlargescaleicedt" + ] + state.output.dconvectiveliquiddt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dconvectiveliquiddt" + ] + state.output.dlargescaleliquiddt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dlargescaleliquiddt" + ] + state.output.dconvectivecloudfractiondt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dconvectivecloudfractiondt" + ] + state.output.dlargescalecloudfractiondt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "dlargescalecloudfractiondt" + ] + locals.mass_flux_ensemble.data[:] = inputs["local_mass_flux_ensemble"][:, :, 0:16] + locals.precipitation_ensemble.data[:] = inputs["local_precipitation_ensemble"][:, :, 0:16] + locals.xff_mid.data[:] = inputs["local_xff_mid"][:, :, 0:16] + + code = self.stencil_factory.from_dims_halo( + func=ensemble_output_and_feedback, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "DTIME": cumulus_parameterization_config.DTIME, + "MAX_TEMP_VAPOR_TENDENCY": cumulus_parameterization_config.MAX_TEMP_VAPOR_TENDENCY, + "APPLY_SUBSIDENCE_MICROPHYSICS": config.APPLY_SUBSIDENCE_MICROPHYSICS, + "USE_SMOOTH_TENDENCIES": cumulus_parameterization_config.USE_SMOOTH_TENDENCIES, + }, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + error_code_2=locals.error_code_2, + error_code_3=locals.error_code_3, + cloud_top_level=state.output.cloud_top_level, + updraft_lfc_level=state.output.updraft_lfc_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + precip=state.output.precip, + effective_condensate_to_fall_forced=locals.effective_condensate_to_fall_forced, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + scale_dependence_factor=state.output.scale_dependence_factor, + ocean_fraction=locals.ocean_fraction, + f_dicycle_modified=locals.f_dicycle_modified, + del_u_cloud_ensemble=locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=locals.del_v_cloud_ensemble, + del_t_cloud_ensemble=locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=locals.del_cloud_liquid_cloud_ensemble, + del_buoyancy_cloud_ensemble=locals.del_buoyancy_cloud_ensemble, + del_convective_ice_cloud_ensemble=locals.del_convective_ice_cloud_ensemble, + del_large_scale_ice_cloud_ensemble=locals.del_large_scale_ice_cloud_ensemble, + del_convective_liquid_cloud_ensemble=locals.del_convective_liquid_cloud_ensemble, + del_large_scale_liquid_cloud_ensemble=locals.del_large_scale_liquid_cloud_ensemble, + del_convective_cloud_fraction_cloud_ensemble=locals.del_convective_cloud_fraction_cloud_ensemble, + del_large_scale_cloud_fraction_cloud_ensemble=locals.del_large_scale_cloud_fraction_cloud_ensemble, + dtdt=state.output.dtdt, + dvapordt=state.output.dvapordt, + dcloudicedt=state.output.dcloudicedt, + dudt=state.output.dudt, + dvdt=state.output.dvdt, + dbuoyancydt=state.output.dbuoyancydt, + dconvectiveicedt=state.output.dconvectiveicedt, + dlargescaleicedt=state.output.dlargescaleicedt, + dconvectiveliquiddt=state.output.dconvectiveliquiddt, + dlargescaleliquiddt=state.output.dlargescaleliquiddt, + dconvectivecloudfractiondt=state.output.dconvectivecloudfractiondt, + dlargescalecloudfractiondt=state.output.dlargescalecloudfractiondt, + mass_flux_ensemble=locals.mass_flux_ensemble, + precipitation_ensemble=locals.precipitation_ensemble, + xff_mid=locals.xff_mid, + CLOSURE_CHOICE=plume_dependent_constants.CLOSURE_CHOICE, + CLOUD_BASE_MASS_FLUX_FACTOR=plume_dependent_constants.CLOUD_BASE_MASS_FLUX_FACTOR, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_error_code_2": locals.error_code_2.field[:], + "local_error_code_3": locals.error_code_3.field[:], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "precip": state.output.precip.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_effective_condensate_to_fall_forced": locals.effective_condensate_to_fall_forced.field[:], + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "scale_dependence_factor": state.output.scale_dependence_factor.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_ocean_fraction": locals.ocean_fraction.field[:], + "local_f_dicycle_modified": locals.f_dicycle_modified.field[:], + "local_del_u_cloud_ensemble": locals.del_u_cloud_ensemble.field[:], + "local_del_v_cloud_ensemble": locals.del_v_cloud_ensemble.field[:], + "local_del_t_cloud_ensemble": locals.del_t_cloud_ensemble.field[:], + "local_del_vapor_cloud_ensemble": locals.del_vapor_cloud_ensemble.field[:], + "local_del_cloud_liquid_cloud_ensemble": locals.del_cloud_liquid_cloud_ensemble.field[:], + "local_del_buoyancy_cloud_ensemble": locals.del_buoyancy_cloud_ensemble.field[:], + "local_del_convective_ice_cloud_ensemble": locals.del_convective_ice_cloud_ensemble.field[:], + "local_del_large_scale_ice_cloud_ensemble": locals.del_large_scale_ice_cloud_ensemble.field[:], + "local_del_convective_liquid_cloud_ensemble": locals.del_convective_liquid_cloud_ensemble.field[ + : + ], + "local_del_large_scale_liquid_cloud_ensemble": locals.del_large_scale_liquid_cloud_ensemble.field[ + : + ], + "local_del_convective_cloud_fraction_cloud_ensemble": locals.del_convective_cloud_fraction_cloud_ensemble.field[ + : + ], + "local_del_large_scale_cloud_fraction_cloud_ensemble": locals.del_large_scale_cloud_fraction_cloud_ensemble.field[ + : + ], + "dtdt": state.output.dtdt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dvapordt": state.output.dvapordt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dcloudicedt": state.output.dcloudicedt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dudt": state.output.dudt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dvdt": state.output.dvdt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dbuoyancydt": state.output.dbuoyancydt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dconvectiveicedt": state.output.dconvectiveicedt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "dlargescaleicedt": state.output.dlargescaleicedt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "dconvectiveliquiddt": state.output.dconvectiveliquiddt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "dlargescaleliquiddt": state.output.dlargescaleliquiddt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "dconvectivecloudfractiondt": state.output.dconvectivecloudfractiondt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "dlargescalecloudfractiondt": state.output.dlargescalecloudfractiondt.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_mass_flux_ensemble": locals.mass_flux_ensemble.field[:], + "local_precipitation_ensemble": locals.precipitation_ensemble.field[:], + "local_xff_mid": locals.xff_mid.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_EnsembleOutputAndFeedback_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputUpdraftTemperature.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputUpdraftTemperature.py new file mode 100644 index 000000000..05f3f807b --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputUpdraftTemperature.py @@ -0,0 +1,171 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import output_updraft_temperature +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_t_cloud_levels": {}, + "local_updraft_column_temperature_forced": {}, + "t_updraft": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + state.output.t_updraft.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["t_updraft"] + + code = self.stencil_factory.from_dims_halo( + func=output_updraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_column_temperature_forced=locals.updraft_column_temperature_forced, + t_cloud_levels=locals.t_cloud_levels, + t_updraft=state.output.t_updraft, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_t_cloud_levels": locals.t_cloud_levels.data[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.data[:], + "t_updraft": state.output.t_updraft.data[:, :, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputUpdraftTemperature_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations.py new file mode 100644 index 000000000..c61831d88 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations.py @@ -0,0 +1,218 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import ( + OutputWorkfunctionsAndPrecipConcentrations, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "convection_fraction": {}, + "surface_type": {}, + "cloud_workfunction_0": {}, + "cloud_workfunction_1": {}, + "local_cloud_workfunction_0": {}, + "local_cloud_workfunction_1": {}, + "air_density": {}, + "local_updraft_column_temperature_forced": {}, + "dcloudicedt": {}, + "dnliquiddt": {}, + "dnicedt": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input.convection_fraction.data[:] = inputs["convection_fraction"] + state.input.surface_type.data[:] = inputs["surface_type"] + state.output.cloud_workfunction_0.data[:] = inputs["cloud_workfunction_0"] + state.output.cloud_workfunction_1.data[:] = inputs["cloud_workfunction_1"] + locals.cloud_workfunction_0.data[:] = inputs["local_cloud_workfunction_0"] + locals.cloud_workfunction_1.data[:] = inputs["local_cloud_workfunction_1"] + state.input_output.air_density.data[:] = inputs["air_density"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + state.output.dcloudicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dcloudicedt"] + state.output.dnliquiddt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dnliquiddt"] + state.output.dnicedt.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs["dnicedt"] + + code = OutputWorkfunctionsAndPrecipConcentrations( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + convection_fraction=state.input.convection_fraction, + surface_type=state.input.surface_type, + cloud_workfunction_0_output=state.output.cloud_workfunction_0, + cloud_workfunction_1_output=state.output.cloud_workfunction_1, + cloud_workfunction_0=locals.cloud_workfunction_0, + cloud_workfunction_1=locals.cloud_workfunction_1, + air_density=state.input_output.air_density, + updraft_column_temperature_forced=locals.updraft_column_temperature_forced, + dcloudicedt=state.output.dcloudicedt, + dnliquiddt=state.output.dnliquiddt, + dnicedt=state.output.dnicedt, + plume_dependent_constants=plume_dependent_constants, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "convection_fraction": state.input.convection_fraction.field[:], + "surface_type": state.input.surface_type.field[:], + "cloud_workfunction_0": state.output.cloud_workfunction_0.field[:], + "cloud_workfunction_1": state.output.cloud_workfunction_1.field[:], + "local_cloud_workfunction_0": locals.cloud_workfunction_0.field[:], + "local_cloud_workfunction_1": locals.cloud_workfunction_1.field[:], + "air_density": state.input_output.air_density.field[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.field[:], + "dcloudicedt": state.output.dcloudicedt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dnliquiddt": state.output.dnliquiddt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + "dnicedt": state.output.dnicedt.field[:, :, :, plume_dependent_constants.PLUME_INDEX], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_shallow( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_mid( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_OutputWorkfunctionsAndPrecipConcentrations_deep( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_PrepareOutput.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_PrepareOutput.py new file mode 100644 index 000000000..b654b2ace --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_PrepareOutput.py @@ -0,0 +1,275 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import prepare_output +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_base_mass_flux_modified": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_total_normalized_integrated_evaporate_forced": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "mass_entrainment_updraft_forced": {}, + "mass_detrainment_updraft_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_environment_massflux": {}, + "local_vapor_tendency_from_environmental_subsidence": {}, + "local_moist_static_energy_tendency_from_environmental_subsidence": {}, + "local_t_tendency_from_environmental_subsidence": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_base_mass_flux_modified.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_base_mass_flux_modified" + ] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.total_normalized_integrated_evaporate_forced.data[:] = inputs[ + "local_total_normalized_integrated_evaporate_forced" + ] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_entrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_entrainment_downdraft_forced"] + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.environment_massflux.data[:] = inputs["local_environment_massflux"] + locals.vapor_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_vapor_tendency_from_environmental_subsidence" + ] + locals.moist_static_energy_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_moist_static_energy_tendency_from_environmental_subsidence" + ] + locals.t_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_t_tendency_from_environmental_subsidence" + ] + + code = self.stencil_factory.from_dims_halo( + func=prepare_output, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_base_mass_flux_modified=state.output.cloud_base_mass_flux_modified, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + total_normalized_integrated_evaporate_forced=locals.total_normalized_integrated_evaporate_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_downdraft_forced=state.output.mass_entrainment_downdraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + environment_massflux=locals.environment_massflux, + vapor_tendency_from_environmental_subsidence=locals.vapor_tendency_from_environmental_subsidence, + moist_static_energy_tendency_from_environmental_subsidence=locals.moist_static_energy_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=locals.t_tendency_from_environmental_subsidence, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_total_normalized_integrated_evaporate_forced": locals.total_normalized_integrated_evaporate_forced.field[ + :, + ], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_environment_massflux": locals.environment_massflux.field[:], + "local_vapor_tendency_from_environmental_subsidence": locals.vapor_tendency_from_environmental_subsidence.field[ + : + ], + "local_moist_static_energy_tendency_from_environmental_subsidence": locals.moist_static_energy_tendency_from_environmental_subsidence.field[ + : + ], + "local_t_tendency_from_environmental_subsidence": locals.t_tendency_from_environmental_subsidence.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrepareOutput_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrepareOutput_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_PrepareOutput_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_TotalEvaporationFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_TotalEvaporationFlux.py new file mode 100644 index 000000000..181b0a9c5 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/prepare_output/translate_GF2020_CumulusParameterization_TotalEvaporationFlux.py @@ -0,0 +1,182 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.prepare_output import total_evaporation_flux +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "p_cloud_levels_forced": {}, + "local_evaporation_flux": {}, + "evaporation_sublimation_tendency": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.evaporation_flux.data[:] = inputs["local_evaporation_flux"] + state.output.evaporation_sublimation_tendency.data[:] = inputs["evaporation_sublimation_tendency"] + + code = self.stencil_factory.from_dims_halo( + func=total_evaporation_flux, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + evaporation_flux=locals.evaporation_flux, + evaporation_sublimation_tendency=state.output.evaporation_sublimation_tendency, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_evaporation_flux": locals.evaporation_flux.data[:], + "evaporation_sublimation_tendency": state.output.evaporation_sublimation_tendency.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_TotalEvaporationFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_C1DProfile.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_C1DProfile.py new file mode 100644 index 000000000..ee3270195 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_C1DProfile.py @@ -0,0 +1,328 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.profiles import C1DProfile +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_start_level": {}, + "lcl_level": {}, + "error_code": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "condensate_to_fall_forced": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_updraft_column_temperature_forced": {}, + "ocean_fraction": {}, + "convection_fraction": {}, + "surface_type": {}, + "p_forced": {}, + "local_p_cloud_levels": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_detrainment_function_updraft": {}, + "local_d_buoyancy_forced": {}, + "local_cloud_liquid_before_rain_forced": {}, + "local_t_cloud_levels": {}, + "local_vapor_forced": {}, + "local_gamma_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "updraft_origin_level": {}, + "local_vapor_cloud_levels_forced": {}, + "local_vapor_excess": {}, + "air_density": {}, + "local_mass_entrainment_updraft": {}, + "local_mass_detrainment_updraft": {}, + "local_psum": {}, + "local_psumh": {}, + "local_c1d": {}, + "local_add_buoyancy": {}, + "local_vertical_velocity_3d": {}, + "local_vertical_velocity_2d": {}, + "convective_scale_velocity": {}, + "entrainment_rate": {}, + "geopotential_height_forced": {}, + "local_t_cloud_levels_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.start_level.data[:] = inputs["local_start_level"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input.convection_fraction.data[:] = inputs["convection_fraction"] + state.input.surface_type.data[:] = inputs["surface_type"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.p_cloud_levels.data[:] = inputs["local_p_cloud_levels"] + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "updraft_lfc_level" + ] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "cloud_top_level" + ] + locals.detrainment_function_updraft.data[:] = inputs["local_detrainment_function_updraft"] + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.cloud_liquid_before_rain_forced.data[:] = inputs["local_cloud_liquid_before_rain_forced"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "updraft_origin_level" + ] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + state.input_output.air_density.data[:] = inputs["air_density"] + locals.mass_entrainment_updraft.data[:] = inputs["local_mass_entrainment_updraft"] + locals.mass_detrainment_updraft.data[:] = inputs["local_mass_detrainment_updraft"] + locals.psum.data[:] = inputs["local_psum"] + locals.psumh.data[:] = inputs["local_psumh"] + locals.c1d.data[:] = inputs["local_c1d"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + locals.vertical_velocity_3d.data[:] = inputs["local_vertical_velocity_3d"] + locals.vertical_velocity_2d.data[:] = inputs["local_vertical_velocity_2d"] + state.input_output.convective_scale_velocity.data[:] = inputs["convective_scale_velocity"] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + + # initialize test code + code = C1DProfile( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + state=state, + locals=locals, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "local_start_level": locals.start_level.field[:], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.field[ + : + ], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.field[:], + "ocean_fraction": state.input.ocean_fraction.field[:], + "convection_fraction": state.input.convection_fraction.field[:], + "surface_type": state.input.surface_type.field[:], + "p_forced": state.input_output.p_forced.field[:], + "local_p_cloud_levels": locals.p_cloud_levels.field[:], + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "cloud_top_level": state.output.cloud_top_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_detrainment_function_updraft": locals.detrainment_function_updraft.field[:], + "local_d_buoyancy_forced": locals.d_buoyancy_forced.field[:], + "local_cloud_liquid_before_rain_forced": locals.cloud_liquid_before_rain_forced.field[:], + "local_t_cloud_levels": locals.t_cloud_levels.field[:], + "local_vapor_forced": locals.vapor_forced.field[:], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.field[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.field[ + : + ], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "air_density": state.input_output.air_density.field[:], + "local_mass_entrainment_updraft": locals.mass_entrainment_updraft.field[:], + "local_mass_detrainment_updraft": locals.mass_detrainment_updraft.field[:], + "local_psum": locals.psum.field[:], + "local_psumh": locals.psumh.field[:], + "local_c1d": locals.c1d.field[:], + "local_add_buoyancy": locals.add_buoyancy.field[:], + "local_vertical_velocity_3d": locals.vertical_velocity_3d.field[:], + "local_vertical_velocity_2d": locals.vertical_velocity_2d.field[:], + "convective_scale_velocity": state.input_output.convective_scale_velocity.field[:], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_C1DProfile_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_C1DProfile_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_C1DProfile_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_MeltingProfile.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_MeltingProfile.py new file mode 100644 index 000000000..f1932db40 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/profiles/translate_GF2020_CumulusParameterization_MeltingProfile.py @@ -0,0 +1,198 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.profiles import melting_profile +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_melting_layer": {}, + "local_partition_liquid_ice": {}, + "p_cloud_levels_forced": {}, + "condensate_to_fall_forced": {}, + "local_melting": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + + locals.melting_layer.data[:] = inputs["local_melting_layer"] + + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + + locals.melting.data[:] = inputs["local_melting"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=melting_profile, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "MELT_GLAC": cumulus_parameterization_config.MELT_GLAC, + }, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + melting_layer=locals.melting_layer, + partition_liquid_ice=locals.partition_liquid_ice, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + melting=locals.melting, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_melting_layer": locals.melting_layer.field[:], + "local_melting": locals.melting.field[:], + "local_partition_liquid_ice": locals.partition_liquid_ice.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_MeltingProfile_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_MeltingProfile_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_MeltingProfile_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/setup/translate_GF2020_CumulusParameterization_Setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/setup/translate_GF2020_CumulusParameterization_Setup.py new file mode 100644 index 000000000..10b9e9f30 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/setup/translate_GF2020_CumulusParameterization_Setup.py @@ -0,0 +1,433 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.setup import Setup +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "t_excess": {}, + "vapor_excess": {}, + "ocean_fraction": {}, + "t_old": {}, + "vapor_old": {}, + "grid_scale_forcing_t": {}, + "grid_scale_forcing_vapor": {}, + "subgrid_scale_forcing_t": {}, + "subgrid_scale_forcing_vapor": {}, + "geopotential_height_forced": {}, + "epsilon_forced": {}, + "precip": {}, + "scale_dependence_factor": {}, + "lightning_density": {}, + "seed_convection": {}, + "error_code": {}, + "grid_length": {}, + "lateral_entrainment_rate": {}, + "entrainment_rate": {}, + } + + out_vars.update(in_vars["data_vars"]) + del ( + out_vars["t_excess"], + out_vars["vapor_excess"], + out_vars["ocean_fraction"], + out_vars["t_old"], + out_vars["vapor_old"], + out_vars["grid_scale_forcing_t"], + out_vars["grid_scale_forcing_vapor"], + out_vars["subgrid_scale_forcing_t"], + out_vars["subgrid_scale_forcing_vapor"], + out_vars["seed_convection"], + ) + out_vars.update( + { + "kstabm": {}, + "local_t_excess": {}, + "local_vapor_excess": {}, + "local_t_new": {}, + "local_vapor_forced": {}, + "local_t_new_pbl": {}, + "local_vapor_forced_pbl": {}, + "local_dmoist_static_energydt": {}, + "local_maximum_updraft_origin_level": {}, + "local_ocean_fraction": {}, + "local_error_code_2": {}, + "local_error_code_3": {}, + "local_cap_max": {}, + "local_cap_max_increment": {}, + "local_geopotential_height": {}, + "local_geopotential_height_modified": {}, + "local_cloud_workfunction_0": {}, + "local_cloud_workfunction_0_pbl": {}, + "local_cloud_workfunction_1": {}, + "local_cloud_workfunction_1_pbl": {}, + "local_cloud_workfunction_1_fa": {}, + "local_cloud_workfunction_2": {}, + "local_cloud_workfunction_3": {}, + "local_cin_1": {}, + "local_k_x_modified": {}, + "local_epsilon": {}, + "local_epsilon_min": {}, + "local_epsilon_max": {}, + "local_pbl_time_scale": {}, + "local_t_wetbulb": {}, + "local_vapor_wetbulb": {}, + "local_cape_removal_time_scale": {}, + "local_f_dicycle_modified": {}, + "local_add_buoyancy": {}, + "local_cloud_moist_static_energy_downdraft_forced": {}, + "local_downdraft_saturation_vapor_forced": {}, + "local_cloud_moist_static_energy_forced_transported": {}, + "local_c1d": {}, + "local_evaporation_below_cloud_base": {}, + "local_mass_flux_ensemble": {}, + "local_precipitation_ensemble": {}, + "local_scale_dependence_factor_downdraft": {}, + "local_random_number": {}, + "local_detrainment_function_updraft": {}, + "local_arbitrary_numerical_parameter": {}, + } + ) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.input.t_excess.data[:] = inputs["t_excess"] + state.input.vapor_excess.data[:] = inputs["vapor_excess"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input_output.t_old.data[:] = inputs["t_old"] + state.input_output.vapor_old.data[:] = inputs["vapor_old"] + state.input.grid_scale_forcing_t.data[:] = inputs["grid_scale_forcing_t"] + state.input.grid_scale_forcing_vapor.data[:] = inputs["grid_scale_forcing_vapor"] + state.input.subgrid_scale_forcing_t.data[:] = inputs["subgrid_scale_forcing_t"] + state.input.subgrid_scale_forcing_vapor.data[:] = inputs["subgrid_scale_forcing_vapor"] + state.input_output.geopotential_height_forced.data[:] = inputs["geopotential_height_forced"] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + state.output.precip.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["precip"] + state.output.scale_dependence_factor.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "scale_dependence_factor" + ] + state.output.lightning_density.data[:] = inputs["lightning_density"] + state.input.seed_convection.data[:] = inputs["seed_convection"] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input_output.grid_length.data[:] = inputs["grid_length"] + state.input.lateral_entrainment_rate.data[:] = inputs["lateral_entrainment_rate"] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + + # prefill locals with nans to replicate fortran initalization + locals.t_excess.field[:] = np.nan + locals.vapor_excess.field[:] = np.nan + locals.t_new.field[:] = np.nan + locals.vapor_forced.field[:] = np.nan + locals.t_new_pbl.field[:] = np.nan + locals.vapor_forced_pbl.field[:] = np.nan + locals.dmoist_static_energydt.field[:] = np.nan + # locals.maximum_updraft_origin_level.field[:] = np.nan + locals.ocean_fraction.field[:] = np.nan + # locals.error_code_2.field[:] = np.nan + # locals.error_code_3.field[:] = np.nan + locals.cap_max.field[:] = np.nan + locals.cap_max_increment.field[:] = np.nan + locals.geopotential_height.field[:] = np.nan + locals.geopotential_height_modified.field[:] = np.nan + locals.cloud_workfunction_0.field[:] = np.nan + locals.cloud_workfunction_0_pbl.field[:] = np.nan + locals.cloud_workfunction_1.field[:] = np.nan + locals.cloud_workfunction_1_pbl.field[:] = np.nan + locals.cloud_workfunction_1_fa.field[:] = np.nan + locals.cloud_workfunction_2.field[:] = np.nan + locals.cloud_workfunction_3.field[:] = np.nan + locals.cin_1.field[:] = np.nan + locals.k_x_modified.field[:] = np.nan + locals.epsilon.field[:] = np.nan + locals.epsilon_min.field[:] = np.nan + locals.epsilon_max.field[:] = np.nan + locals.pbl_time_scale.field[:] = np.nan + locals.t_wetbulb.field[:] = np.nan + locals.vapor_wetbulb.field[:] = np.nan + locals.cape_removal_time_scale.field[:] = np.nan + locals.f_dicycle_modified.field[:] = np.nan + locals.add_buoyancy.field[:] = np.nan + locals.cloud_moist_static_energy_downdraft_forced.field[:] = np.nan + locals.downdraft_saturation_vapor_forced.field[:] = np.nan + locals.cloud_moist_static_energy_forced_transported.field[:] = np.nan + locals.c1d.field[:] = np.nan + locals.evaporation_below_cloud_base.field[:] = np.nan + locals.mass_flux_ensemble.field[:] = np.nan + locals.precipitation_ensemble.field[:] = np.nan + locals.scale_dependence_factor_downdraft.field[:] = np.nan + locals.random_number.field[:] = np.nan + locals.detrainment_function_updraft.field[:] = np.nan + locals.arbitrary_numerical_parameter.field[:] = np.nan + + # initialize test code + code = Setup( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + code( + error_code=state.output.error_code, + error_code_2=locals.error_code_2, + error_code_3=locals.error_code_3, + maximum_updraft_origin_level=locals.maximum_updraft_origin_level, + kstabm=state.output.kstabm, + t_excess=state.input.t_excess, + t_excess_local=locals.t_excess, + vapor_excess=state.input.vapor_excess, + vapor_excess_local=locals.vapor_excess, + ocean_fraction=state.input.ocean_fraction, + ocean_fraction_local=locals.ocean_fraction, + t_old=state.input_output.t_old, + t_new=locals.t_new, + t_new_pbl=locals.t_new_pbl, + vapor_old=state.input_output.vapor_old, + vapor_forced=locals.vapor_forced, + vapor_forced_pbl=locals.vapor_forced_pbl, + downdraft_saturation_vapor_forced=locals.downdraft_saturation_vapor_forced, + grid_scale_forcing_t=state.input.grid_scale_forcing_t, + grid_scale_forcing_vapor=state.input.grid_scale_forcing_vapor, + subgrid_scale_forcing_t=state.input.subgrid_scale_forcing_t, + subgrid_scale_forcing_vapor=state.input.subgrid_scale_forcing_vapor, + dmoist_static_energydt=locals.dmoist_static_energydt, + cloud_moist_static_energy_downdraft_forced=locals.cloud_moist_static_energy_downdraft_forced, + cloud_moist_static_energy_forced_transported=locals.cloud_moist_static_energy_forced_transported, + cap_max=locals.cap_max, + cap_max_increment=locals.cap_max_increment, + geopotential_height=state.input_output.geopotential_height_forced, + geopotential_height_local=locals.geopotential_height, + geopotential_height_modified_local=locals.geopotential_height_modified, + cloud_workfunction_0=locals.cloud_workfunction_0, + cloud_workfunction_1=locals.cloud_workfunction_1, + cloud_workfunction_2=locals.cloud_workfunction_2, + cloud_workfunction_3=locals.cloud_workfunction_3, + cloud_workfunction_0_pbl=locals.cloud_workfunction_0_pbl, + cloud_workfunction_1_pbl=locals.cloud_workfunction_1_pbl, + cloud_workfunction_1_fa=locals.cloud_workfunction_1_fa, + cin_1=locals.cin_1, + k_x_modified=locals.k_x_modified, + epsilon_forced=state.output.epsilon_forced, + epsilon_local=locals.epsilon, + epsilon_min=locals.epsilon_min, + epsilon_max=locals.epsilon_max, + pbl_time_scale=locals.pbl_time_scale, + t_wetbulb=locals.t_wetbulb, + vapor_wetbulb=locals.vapor_wetbulb, + cape_removal_time_scale=locals.cape_removal_time_scale, + f_dicycle_modified=locals.f_dicycle_modified, + add_buoyancy=locals.add_buoyancy, + scale_dependence_factor=state.output.scale_dependence_factor, + scale_dependence_factor_downdraft=locals.scale_dependence_factor_downdraft, + c1d=locals.c1d, + evaporation_below_cloud_base=locals.evaporation_below_cloud_base, + mass_flux_ensemble=locals.mass_flux_ensemble, + precipitation_ensemble=locals.precipitation_ensemble, + precip=state.output.precip, + lightning_density=state.output.lightning_density, + seed_convection=state.input.seed_convection, + grid_length=state.input_output.grid_length, + random_number=locals.random_number, + lateral_entrainment_rate=state.input.lateral_entrainment_rate, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=locals.detrainment_function_updraft, + arbitrary_numerical_parameter=locals.arbitrary_numerical_parameter, + plume_dependent_constants=plume_dependent_constants, + plume=plume, + ) + + # write output + outputs = { + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "precip": state.output.precip.field[:, :, plume_dependent_constants.PLUME_INDEX], + "scale_dependence_factor": state.output.scale_dependence_factor.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "lightning_density": state.output.lightning_density.field[:], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "grid_length": state.input_output.grid_length.field[:], + "lateral_entrainment_rate": state.input.lateral_entrainment_rate.field[:], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "kstabm": state.output.kstabm.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "local_t_excess": locals.t_excess.field[:], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_new": locals.t_new.field[:], + "local_vapor_forced": locals.vapor_forced.field[:], + "local_t_new_pbl": locals.t_new_pbl.field[:], + "local_vapor_forced_pbl": locals.vapor_forced_pbl.field[:], + "local_dmoist_static_energydt": locals.dmoist_static_energydt.field[:], + "local_maximum_updraft_origin_level": locals.maximum_updraft_origin_level.field[:] + 1, + "local_ocean_fraction": locals.ocean_fraction.field[:], + "local_error_code_2": locals.error_code_2.field[:], + "local_error_code_3": locals.error_code_3.field[:], + "local_cap_max": locals.cap_max.field[:], + "local_cap_max_increment": locals.cap_max_increment.field[:], + "local_geopotential_height": locals.geopotential_height.field[:], + "local_geopotential_height_modified": locals.geopotential_height_modified.field[:], + "local_cloud_workfunction_0": locals.cloud_workfunction_0.field[:], + "local_cloud_workfunction_0_pbl": locals.cloud_workfunction_0_pbl.field[:], + "local_cloud_workfunction_1": locals.cloud_workfunction_1.field[:], + "local_cloud_workfunction_1_pbl": locals.cloud_workfunction_1_pbl.field[:], + "local_cloud_workfunction_1_fa": locals.cloud_workfunction_1_fa.field[:], + "local_cloud_workfunction_2": locals.cloud_workfunction_2.field[:], + "local_cloud_workfunction_3": locals.cloud_workfunction_3.field[:], + "local_cin_1": locals.cin_1.field[:], + "local_k_x_modified": locals.k_x_modified.field[:], + "local_epsilon": locals.epsilon.field[:], + "local_epsilon_min": locals.epsilon_min.field[:], + "local_epsilon_max": locals.epsilon_max.field[:], + "local_pbl_time_scale": locals.pbl_time_scale.field[:], + "local_t_wetbulb": locals.t_wetbulb.field[:], + "local_vapor_wetbulb": locals.vapor_wetbulb.field[:], + "local_cape_removal_time_scale": locals.cape_removal_time_scale.field[:], + "local_f_dicycle_modified": locals.f_dicycle_modified.field[:], + "local_add_buoyancy": locals.add_buoyancy.field[:], + "local_cloud_moist_static_energy_downdraft_forced": locals.cloud_moist_static_energy_downdraft_forced.field[ + : + ], + "local_downdraft_saturation_vapor_forced": locals.downdraft_saturation_vapor_forced.field[:], + "local_cloud_moist_static_energy_forced_transported": locals.cloud_moist_static_energy_forced_transported.field[ + : + ], + "local_c1d": locals.c1d.field[:], + "local_evaporation_below_cloud_base": locals.evaporation_below_cloud_base.field[:], + "local_mass_flux_ensemble": locals.mass_flux_ensemble.field[:], + "local_precipitation_ensemble": locals.precipitation_ensemble.field[:], + "local_scale_dependence_factor_downdraft": locals.scale_dependence_factor_downdraft.field[:], + "local_random_number": locals.random_number.field[:], + "local_detrainment_function_updraft": locals.detrainment_function_updraft.field[:], + "local_arbitrary_numerical_parameter": locals.arbitrary_numerical_parameter.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_Setup_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_Setup_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_Setup_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/smoothing/translate_GF2020_CumulusParameterization_SmoothTendencies.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/smoothing/translate_GF2020_CumulusParameterization_SmoothTendencies.py new file mode 100644 index 000000000..debde8178 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/smoothing/translate_GF2020_CumulusParameterization_SmoothTendencies.py @@ -0,0 +1,202 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.smoothing import smooth_tendencies +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "p_cloud_levels_forced": {}, + "local_del_moist_static_energy_cloud_ensemble": {}, + "local_del_vapor_cloud_ensemble": {}, + "local_del_cloud_liquid_cloud_ensemble": {}, + "local_del_u_cloud_ensemble": {}, + "local_del_v_cloud_ensemble": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.del_moist_static_energy_cloud_ensemble.data[:] = inputs[ + "local_del_moist_static_energy_cloud_ensemble" + ] + locals.del_vapor_cloud_ensemble.data[:] = inputs["local_del_vapor_cloud_ensemble"] + locals.del_cloud_liquid_cloud_ensemble.data[:] = inputs["local_del_cloud_liquid_cloud_ensemble"] + locals.del_u_cloud_ensemble.data[:] = inputs["local_del_u_cloud_ensemble"] + locals.del_v_cloud_ensemble.data[:] = inputs["local_del_v_cloud_ensemble"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=smooth_tendencies, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"USE_SMOOTH_TENDENCIES": cumulus_parameterization_config.USE_SMOOTH_TENDENCIES}, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + del_moist_static_energy_cloud_ensemble=locals.del_moist_static_energy_cloud_ensemble, + del_vapor_cloud_ensemble=locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=locals.del_cloud_liquid_cloud_ensemble, + del_u_cloud_ensemble=locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=locals.del_v_cloud_ensemble, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_del_moist_static_energy_cloud_ensemble": locals.del_moist_static_energy_cloud_ensemble.data[ + : + ], + "local_del_vapor_cloud_ensemble": locals.del_vapor_cloud_ensemble.data[:], + "local_del_cloud_liquid_cloud_ensemble": locals.del_cloud_liquid_cloud_ensemble.data[:], + "local_del_u_cloud_ensemble": locals.del_u_cloud_ensemble.data[:], + "local_del_v_cloud_ensemble": locals.del_v_cloud_ensemble.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_SmoothTendencies_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_SmoothTendencies_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_SmoothTendencies_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/translate_GF2020_CumulusParameterization.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/translate_GF2020_CumulusParameterization.py new file mode 100644 index 000000000..92ed4cf1e --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/translate_GF2020_CumulusParameterization.py @@ -0,0 +1,530 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import NUMBER_OF_PLUMES +from pyMoist.convection.GF_2020.cumulus_parameterization.cumulus_parameterization import ( + GF2020CumulusParameterization, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class TranslateGF2020_CumulusParameterization(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = {} + + self.out_vars = self.in_vars["data_vars"].copy() + self.out_vars.update( + { + "t_excess": {}, + "vapor_excess": {}, + "grid_scale_forcing_t": {}, + "grid_scale_forcing_vapor": {}, + "subgrid_scale_forcing_t": {}, + "subgrid_scale_forcing_vapor": {}, + "seed_convection": {}, + "saturation_water_vapor": {}, + "ocean_fraction": {}, + "convection_fraction": {}, + "surface_type": {}, + "lateral_entrainment_rate": {}, + "last_error_code": {}, + "dtdt": {}, + "dvapordt": {}, + "dcloudicedt": {}, + "dudt": {}, + "dvdt": {}, + "dnliquiddt": {}, + "dnicedt": {}, + "dbuoyancydt": {}, + "dconvectiveicedt": {}, + "dlargescaleicedt": {}, + "dconvectiveliquiddt": {}, + "dlargescaleliquiddt": {}, + "dconvectivecloudfractiondt": {}, + "dlargescalecloudfractiondt": {}, + "error_code": {}, + "downdraft_origin_level": {}, + "lcl_level": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "kstabi": {}, + "kstabm": {}, + "precip": {}, + "cloud_base_mass_flux_modified": {}, + "epsilon_forced": {}, + "total_normalized_integrated_condensate_forced": {}, + "scale_dependence_factor": {}, + "p_cloud_levels_forced": {}, + "entrainment_rate": {}, + "mass_entrainment_updraft_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_updraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "t_updraft": {}, + "convective_cloud_fraction_output": {}, + "cloud_workfunction_0": {}, + "cloud_workfunction_1": {}, + "cloud_workfunction_2": {}, + "cloud_workfunction_3": {}, + "cloud_workfunction_1_pbl": {}, + "cloud_workfunction_1_cin": {}, + "cape_removal_time_scale": {}, + "pbl_time_scale": {}, + "lightning_density": {}, + "evaporation_sublimation_tendency": {}, + "convective_precip_flux": {}, + "t_perturbation": {}, + "grid_length": {}, + "pbl_level": {}, + "ccn": {}, + "air_density": {}, + "omega": {}, + "topography_height_no_negative": {}, + "sensible_heat_flux": {}, + "latent_heat_flux": {}, + "longitude_degrees": {}, + "latitude_degrees": {}, + "t_old": {}, + "vapor_old": {}, + "t_modified_by_advection": {}, + "vapor_modified_by_advection": {}, + "geopotential_height_forced": {}, + "p_forced": {}, + "p_surface": {}, + "t_surface": {}, + "u": {}, + "v": {}, + "w": {}, + "mass": {}, + "convective_scale_velocity": {}, + "buoyancy_excess": {}, + "large_scale_ice": {}, + "convective_ice": {}, + "large_scale_liquid": {}, + "convective_liquid": {}, + "large_scale_cloud_fraction": {}, + "convective_cloud_fraction": {}, + # "chemistry_tracers": {}, + # "chemistry_tracers_output": {}, + } + ) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers_input = data_loader.load("GF2020_ConvectionTracers") + + # workaround because translate test cannot read in 4d fields + self.manual_inputs = data_loader.load("GF2020_CumulusParameterization-In") + + def compute_func(self, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **self.constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**self.cu_param_constants) + + # initialize convection tracers + convection_tracers = ConvectionTracers.ones( + self.quantity_factory, + data_dimensions={ + "convection_tracers": config.NUMBER_OF_TRACERS, + "size_three_dimension": 3, + "size_four_dimension": 4, + }, + ) + + convection_tracers.tracers.field[:] = np.moveaxis(self.convection_tracers_input["tracers"], 0, 3) + convection_tracers.vect_hcts.field[:] = self.convection_tracers_input["vect_hcts"] + convection_tracers.kc_scal.field[:] = self.convection_tracers_input["kc_scal"] + convection_tracers.fscav.field[:] = self.convection_tracers_input["fscav"] + convection_tracers.convfaci2g.field[:] = self.convection_tracers_input["convfaci2g"] + convection_tracers.retfactor.field[:] = self.convection_tracers_input["retfactor"] + convection_tracers.liq_and_gas.field[:] = self.convection_tracers_input["liq_and_gas"] + convection_tracers.online_cldliq.field[:] = self.convection_tracers_input["online_cldliq"] + convection_tracers.online_vud.field[:] = self.convection_tracers_input["online_vud"] + convection_tracers.ftemp_threshold.field[:] = self.convection_tracers_input["ftemp_threshold"] + convection_tracers.use_gcc_washout.field[:] = self.convection_tracers_input["use_gcc_washout"] + convection_tracers.use_gocart.field[:] = self.convection_tracers_input["use_gocart"] + convection_tracers.is_wetdep.field[:] = self.convection_tracers_input["is_wetdep"] + + # initialize pyMoist saturation tables + saturation_tables = SaturationVaporPressureTable(self.stencil_factory.backend) + + # initialize state + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill state with input data + if cumulus_parameterization_config.PLUME_ORDER == 0: + # fortran data is saved [deep, shallow, mid] in plume dimension + # must be reordered to [shallow, mid, deep] for python + # input fields + state.input.t_excess.field[:] = self.manual_inputs["t_excess"] + state.input.vapor_excess.field[:] = self.manual_inputs["vapor_excess"] + state.input.grid_scale_forcing_t.field[:] = self.manual_inputs["grid_scale_forcing_t"] + state.input.grid_scale_forcing_vapor.field[:] = self.manual_inputs["grid_scale_forcing_vapor"] + state.input.subgrid_scale_forcing_t.field[:] = self.manual_inputs["subgrid_scale_forcing_t"] + state.input.subgrid_scale_forcing_vapor.field[:] = self.manual_inputs[ + "subgrid_scale_forcing_vapor" + ] + state.input.seed_convection.field[:] = self.manual_inputs["seed_convection"] + state.input.saturation_water_vapor.field[:] = self.manual_inputs["saturation_water_vapor"] + state.input.ocean_fraction.field[:] = self.manual_inputs["ocean_fraction"] + state.input.convection_fraction.field[:] = self.manual_inputs["convection_fraction"] + state.input.surface_type.field[:] = self.manual_inputs["surface_type"] + state.input.lateral_entrainment_rate.field[:] = self.manual_inputs["lateral_entrainment_rate"] + state.input.last_error_code.field[:] = self.manual_inputs["last_error_code"] + # output fields + state.output.dtdt.field[:] = self.manual_inputs["dtdt"][:, :, :, [1, 2, 0]] + state.output.dvapordt.field[:] = self.manual_inputs["dvapordt"][:, :, :, [1, 2, 0]] + state.output.dcloudicedt.field[:] = self.manual_inputs["dcloudicedt"][:, :, :, [1, 2, 0]] + state.output.dudt.field[:] = self.manual_inputs["dudt"][:, :, :, [1, 2, 0]] + state.output.dvdt.field[:] = self.manual_inputs["dvdt"][:, :, :, [1, 2, 0]] + state.output.dnliquiddt.field[:] = self.manual_inputs["dnliquiddt"][:, :, :, [1, 2, 0]] + state.output.dnicedt.field[:] = self.manual_inputs["dnicedt"][:, :, :, [1, 2, 0]] + state.output.dbuoyancydt.field[:] = self.manual_inputs["dbuoyancydt"][:, :, :, [1, 2, 0]] + state.output.dconvectiveicedt.field[:] = self.manual_inputs["dconvectiveicedt"][ + :, :, :, [1, 2, 0] + ] + state.output.dlargescaleicedt.field[:] = self.manual_inputs["dlargescaleicedt"][ + :, :, :, [1, 2, 0] + ] + state.output.dconvectiveliquiddt.field[:] = self.manual_inputs["dconvectiveliquiddt"][ + :, :, :, [1, 2, 0] + ] + state.output.dlargescaleliquiddt.field[:] = self.manual_inputs["dlargescaleliquiddt"][ + :, :, :, [1, 2, 0] + ] + state.output.dconvectivecloudfractiondt.field[:] = self.manual_inputs[ + "dconvectivecloudfractiondt" + ][:, :, :, [1, 2, 0]] + state.output.dlargescalecloudfractiondt.field[:] = self.manual_inputs[ + "dlargescalecloudfractiondt" + ][:, :, :, [1, 2, 0]] + state.output.error_code.field[:] = self.manual_inputs["error_code"][:, :, [1, 2, 0]] + state.output.downdraft_origin_level.field[:] = ( + self.manual_inputs["downdraft_origin_level"][:, :, [1, 2, 0]] - 1 + ) + state.output.lcl_level.field[:] = self.manual_inputs["lcl_level"][:, :, [1, 2, 0]] - 1 + state.output.updraft_origin_level.field[:] = ( + self.manual_inputs["updraft_origin_level"][:, :, [1, 2, 0]] - 1 + ) + state.output.updraft_lfc_level.field[:] = ( + self.manual_inputs["updraft_lfc_level"][:, :, [1, 2, 0]] - 1 + ) + state.output.cloud_top_level.field[:] = self.manual_inputs["cloud_top_level"][:, :, [1, 2, 0]] - 1 + state.output.kstabi.field[:] = self.manual_inputs["kstabi"][:, :, [1, 2, 0]] - 1 + state.output.kstabm.field[:] = self.manual_inputs["kstabm"][:, :, [1, 2, 0]] - 1 + state.output.precip.field[:] = self.manual_inputs["precip"][:, :, [1, 2, 0]] + state.output.cloud_base_mass_flux_modified.field[:] = self.manual_inputs[ + "cloud_base_mass_flux_modified" + ][:, :, [1, 2, 0]] + state.output.epsilon_forced.field[:] = self.manual_inputs["epsilon_forced"][:, :, [1, 2, 0]] + state.output.total_normalized_integrated_condensate_forced.field[:] = self.manual_inputs[ + "total_normalized_integrated_condensate_forced" + ][:, :, [1, 2, 0]] + state.output.scale_dependence_factor.field[:] = self.manual_inputs["scale_dependence_factor"][ + :, :, [1, 2, 0] + ] + state.output.p_cloud_levels_forced.field[:] = self.manual_inputs["p_cloud_levels_forced"][ + :, :, :, [1, 2, 0] + ] + state.output.entrainment_rate.field[:] = self.manual_inputs["entrainment_rate"] + state.output.mass_entrainment_updraft_forced.field[:] = self.manual_inputs[ + "mass_entrainment_updraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.mass_entrainment_downdraft_forced.field[:] = self.manual_inputs[ + "mass_entrainment_downdraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.mass_detrainment_updraft_forced.field[:] = self.manual_inputs[ + "mass_detrainment_updraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.mass_detrainment_downdraft_forced.field[:] = self.manual_inputs[ + "mass_detrainment_downdraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.normalized_massflux_updraft_forced.field[:] = self.manual_inputs[ + "normalized_massflux_updraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.normalized_massflux_downdraft_forced.field[:] = self.manual_inputs[ + "normalized_massflux_downdraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.condensate_to_fall_forced.field[:] = self.manual_inputs["condensate_to_fall_forced"][ + :, :, :, [1, 2, 0] + ] + state.output.evaporate_in_downdraft_forced.field[:] = self.manual_inputs[ + "evaporate_in_downdraft_forced" + ][:, :, :, [1, 2, 0]] + state.output.cloud_liquid_after_rain_forced.field[:] = self.manual_inputs[ + "cloud_liquid_after_rain_forced" + ][:, :, :, [1, 2, 0]] + state.output.t_updraft.field[:] = self.manual_inputs["t_updraft"][:, :, :, [1, 2, 0]] + state.output.convective_cloud_fraction.field[:] = self.manual_inputs[ + "convective_cloud_fraction_output" + ][:, :, :, [1, 2, 0]] + state.output.cloud_workfunction_0.field[:] = self.manual_inputs["cloud_workfunction_0"] + state.output.cloud_workfunction_1.field[:] = self.manual_inputs["cloud_workfunction_1"] + state.output.cloud_workfunction_2.field[:] = self.manual_inputs["cloud_workfunction_2"] + state.output.cloud_workfunction_3.field[:] = self.manual_inputs["cloud_workfunction_3"] + state.output.cloud_workfunction_1_pbl.field[:] = self.manual_inputs["cloud_workfunction_1_pbl"] + state.output.cloud_workfunction_1_cin.field[:] = self.manual_inputs["cloud_workfunction_1_cin"] + state.output.cape_removal_time_scale.field[:] = self.manual_inputs["cape_removal_time_scale"] + state.output.pbl_time_scale.field[:] = self.manual_inputs["pbl_time_scale"] + state.output.lightning_density.field[:] = self.manual_inputs["lightning_density"] + state.output.evaporation_sublimation_tendency.field[:] = self.manual_inputs[ + "evaporation_sublimation_tendency" + ] + state.output.convective_precip_flux.field[:] = self.manual_inputs["convective_precip_flux"] + state.output.t_perturbation.field[:] = self.manual_inputs["t_perturbation"] + # input/output fields + state.input_output.grid_length.field[:] = self.manual_inputs["grid_length"] + state.input_output.pbl_level.field[:] = self.manual_inputs["pbl_level"] - 1 + state.input_output.ccn.field[:] = self.manual_inputs["ccn"] + state.input_output.air_density.field[:] = self.manual_inputs["air_density"] + state.input_output.omega.field[:] = self.manual_inputs["omega"] + state.input_output.topography_height_no_negative.field[:] = self.manual_inputs[ + "topography_height_no_negative" + ] + state.input_output.sensible_heat_flux.field[:] = self.manual_inputs["sensible_heat_flux"] + state.input_output.latent_heat_flux.field[:] = self.manual_inputs["latent_heat_flux"] + state.input_output.longitude_degrees.field[:] = self.manual_inputs["longitude_degrees"] + state.input_output.latitude_degrees.field[:] = self.manual_inputs["latitude_degrees"] + state.input_output.t_old.field[:] = self.manual_inputs["t_old"] + state.input_output.vapor_old.field[:] = self.manual_inputs["vapor_old"] + state.input_output.t_modified_by_advection.field[:] = self.manual_inputs[ + "t_modified_by_advection" + ] + state.input_output.vapor_modified_by_advection.field[:] = self.manual_inputs[ + "vapor_modified_by_advection" + ] + state.input_output.geopotential_height_forced.field[:] = self.manual_inputs[ + "geopotential_height_forced" + ] + state.input_output.p_forced.field[:] = self.manual_inputs["p_forced"] + state.input_output.p_surface.field[:] = self.manual_inputs["p_surface"] + state.input_output.t_surface.field[:] = self.manual_inputs["t_surface"] + state.input_output.u.field[:] = self.manual_inputs["u"] + state.input_output.v.field[:] = self.manual_inputs["v"] + state.input_output.w.field[:] = self.manual_inputs["w"] + state.input_output.mass.field[:] = self.manual_inputs["mass"] + state.input_output.convective_scale_velocity.field[:] = self.manual_inputs[ + "convective_scale_velocity" + ] + state.input_output.buoyancy_excess.field[:] = self.manual_inputs["buoyancy_excess"] + state.input_output.large_scale_ice.field[:] = self.manual_inputs["large_scale_ice"] + state.input_output.convective_ice.field[:] = self.manual_inputs["convective_ice"] + state.input_output.large_scale_liquid.field[:] = self.manual_inputs["large_scale_liquid"] + state.input_output.convective_liquid.field[:] = self.manual_inputs["convective_liquid"] + state.input_output.large_scale_cloud_fraction.field[:] = self.manual_inputs[ + "large_scale_cloud_fraction" + ] + state.input_output.convective_cloud_fraction.field[:] = self.manual_inputs[ + "convective_cloud_fraction" + ] + state.input_output.chemistry_tracers.field[:] = self.manual_inputs["chemistry_tracers"] + chemistry_tracers_input_5d = np.full( + state.input_output.chemistry_tracers_output.field[:].shape, np.nan + ) + for plume in range(NUMBER_OF_PLUMES): + chemistry_tracers_input_5d[:, :, :, plume, :] = self.manual_inputs[ + "chemistry_tracers_output" + ][ + :, + :, + :, + plume * config.NUMBER_OF_TRACERS : plume * config.NUMBER_OF_TRACERS + + config.NUMBER_OF_TRACERS, + ] + state.input_output.chemistry_tracers_output.field[:] = chemistry_tracers_input_5d[ + :, :, :, [1, 2, 0], : + ] + + # initialize the test subject + code = GF2020CumulusParameterization( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + saturation_tables=saturation_tables, + ) + + # call the test subject + code(state, convection_tracers) + + # collapse plume dim for chemistry_tracers_output + # NOTE ideally this has no numpy dependency + chemistry_tracers_output_5d_reordered = state.input_output.chemistry_tracers_output.field[ + :, :, :, [2, 0, 1], : + ] + chemistry_tracers_output_4d = np.full( + self.manual_inputs["chemistry_tracers_output"].shape, np.nan + ) + for plume in range(NUMBER_OF_PLUMES): + chemistry_tracers_output_4d[ + :, + :, + :, + plume * config.NUMBER_OF_TRACERS : plume * config.NUMBER_OF_TRACERS + + config.NUMBER_OF_TRACERS, + ] = chemistry_tracers_output_5d_reordered[:, :, :, plume, :] + + outputs = { + # input fields + "t_excess": state.input.t_excess.field[:], + "vapor_excess": state.input.vapor_excess.field[:], + "grid_scale_forcing_t": state.input.grid_scale_forcing_t.field[:], + "grid_scale_forcing_vapor": state.input.grid_scale_forcing_vapor.field[:], + "subgrid_scale_forcing_t": state.input.subgrid_scale_forcing_t.field[:], + "subgrid_scale_forcing_vapor": state.input.subgrid_scale_forcing_vapor.field[:], + "seed_convection": state.input.seed_convection.field[:], + "saturation_water_vapor": state.input.saturation_water_vapor.field[:], + "ocean_fraction": state.input.ocean_fraction.field[:], + "convection_fraction": state.input.convection_fraction.field[:], + "surface_type": state.input.surface_type.field[:], + "lateral_entrainment_rate": state.input.lateral_entrainment_rate.field[:], + "last_error_code": state.input.last_error_code.field[:], + # output fields + "dtdt": state.output.dtdt.field[:, :, :, [2, 0, 1]], + "dvapordt": state.output.dvapordt.field[:, :, :, [2, 0, 1]], + "dcloudicedt": state.output.dcloudicedt.field[:, :, :, [2, 0, 1]], + "dudt": state.output.dudt.field[:, :, :, [2, 0, 1]], + "dvdt": state.output.dvdt.field[:, :, :, [2, 0, 1]], + "dnliquiddt": state.output.dnliquiddt.field[:, :, :, [2, 0, 1]], + "dnicedt": state.output.dnicedt.field[:, :, :, [2, 0, 1]], + "dbuoyancydt": state.output.dbuoyancydt.field[:, :, :, [2, 0, 1]], + "dconvectiveicedt": state.output.dconvectiveicedt.field[:, :, :, [2, 0, 1]], + "dlargescaleicedt": state.output.dlargescaleicedt.field[:, :, :, [2, 0, 1]], + "dconvectiveliquiddt": state.output.dconvectiveliquiddt.field[:, :, :, [2, 0, 1]], + "dlargescaleliquiddt": state.output.dlargescaleliquiddt.field[:, :, :, [2, 0, 1]], + "dconvectivecloudfractiondt": state.output.dconvectivecloudfractiondt.field[ + :, :, :, [2, 0, 1] + ], + "dlargescalecloudfractiondt": state.output.dlargescalecloudfractiondt.field[ + :, :, :, [2, 0, 1] + ], + "error_code": state.output.error_code.field[:, :, [2, 0, 1]], + "downdraft_origin_level": state.output.downdraft_origin_level.field[:, :, [2, 0, 1]] + 1, + "lcl_level": state.output.lcl_level.field[:, :, [2, 0, 1]] + 1, + "updraft_origin_level": state.output.updraft_origin_level.field[:, :, [2, 0, 1]] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[:, :, [2, 0, 1]] + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, [2, 0, 1]] + 1, + "kstabi": state.output.kstabi.field[:, :, [2, 0, 1]] + 1, + "kstabm": state.output.kstabm.field[:, :, [2, 0, 1]] + 1, + "precip": state.output.precip.field[:, :, [2, 0, 1]], + "cloud_base_mass_flux_modified": state.output.cloud_base_mass_flux_modified.field[ + :, :, [2, 0, 1] + ], + "epsilon_forced": state.output.epsilon_forced.field[:, :, [2, 0, 1]], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.field[ + :, :, [2, 0, 1] + ], + "scale_dependence_factor": state.output.scale_dependence_factor.field[:, :, [2, 0, 1]], + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[:, :, :, [2, 0, 1]], + "entrainment_rate": state.output.entrainment_rate.field[:, :, :, [2, 0, 1]], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_entrainment_downdraft_forced": state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[:, :, :, [2, 0, 1]], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, [2, 0, 1] + ], + "t_updraft": state.output.t_updraft.field[:, :, :, [2, 0, 1]], + "convective_cloud_fraction_output": state.output.convective_cloud_fraction.field[ + :, :, :, [2, 0, 1] + ], + "cloud_workfunction_0": state.output.cloud_workfunction_0.field[:], + "cloud_workfunction_1": state.output.cloud_workfunction_1.field[:], + "cloud_workfunction_2": state.output.cloud_workfunction_2.field[:], + "cloud_workfunction_3": state.output.cloud_workfunction_3.field[:], + "cloud_workfunction_1_pbl": state.output.cloud_workfunction_1_pbl.field[:], + "cloud_workfunction_1_cin": state.output.cloud_workfunction_1_cin.field[:], + "cape_removal_time_scale": state.output.cape_removal_time_scale.field[:], + "pbl_time_scale": state.output.pbl_time_scale.field[:], + "lightning_density": state.output.lightning_density.field[:], + "evaporation_sublimation_tendency": state.output.evaporation_sublimation_tendency.field[:], + "convective_precip_flux": state.output.convective_precip_flux.field[:], + "t_perturbation": state.output.t_perturbation.field[:], + # input/output fields + "grid_length": state.input_output.grid_length.field[:], + "pbl_level": state.input_output.pbl_level.field[:] + 1, + "ccn": state.input_output.ccn.field[:], + "air_density": state.input_output.air_density.field[:], + "omega": state.input_output.omega.field[:], + "topography_height_no_negative": state.input_output.topography_height_no_negative.field[:], + "sensible_heat_flux": state.input_output.sensible_heat_flux.field[:], + "latent_heat_flux": state.input_output.latent_heat_flux.field[:], + "longitude_degrees": state.input_output.longitude_degrees.field[:], + "latitude_degrees": state.input_output.latitude_degrees.field[:], + "t_old": state.input_output.t_old.field[:], + "vapor_old": state.input_output.vapor_old.field[:], + "t_modified_by_advection": state.input_output.t_modified_by_advection.field[:], + "vapor_modified_by_advection": state.input_output.vapor_modified_by_advection.field[:], + "geopotential_height_forced": state.input_output.geopotential_height_forced.field[:], + "p_forced": state.input_output.p_forced.field[:], + "p_surface": state.input_output.p_surface.field[:], + "t_surface": state.input_output.t_surface.field[:], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "w": state.input_output.w.field[:], + "mass": state.input_output.mass.field[:], + "convective_scale_velocity": state.input_output.convective_scale_velocity.field[:], + "buoyancy_excess": state.input_output.buoyancy_excess.field[:], + "large_scale_ice": state.input_output.large_scale_ice.field[:], + "convective_ice": state.input_output.convective_ice.field[:], + "large_scale_liquid": state.input_output.large_scale_liquid.field[:], + "convective_liquid": state.input_output.convective_liquid.field[:], + "large_scale_cloud_fraction": state.input_output.large_scale_cloud_fraction.field[:], + "convective_cloud_fraction": state.input_output.convective_cloud_fraction.field[:], + "chemistry_tracers": state.input_output.chemistry_tracers.field[:], + "chemistry_tracers_output": chemistry_tracers_output_4d, + } + else: + raise NotImplementedError( + "Plume order unsupported. Please implement a way to transfer the fortran shape into python shape." + ) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/triggers/translate_GF2020_CumulusParameterization_ConvectionTrigger.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/triggers/translate_GF2020_CumulusParameterization_ConvectionTrigger.py new file mode 100644 index 000000000..e3b4668d1 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/triggers/translate_GF2020_CumulusParameterization_ConvectionTrigger.py @@ -0,0 +1,171 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.triggers import convection_trigger + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "convective_scale_velocity": {}, + "local_cin_0": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.input_output.convective_scale_velocity.data[:] = inputs["convective_scale_velocity"] + locals.cin_0.data[:] = inputs["local_cin_0"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=convection_trigger, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"DICYCLE": cumulus_parameterization_config.DIURNAL_CYCLE}, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + convective_scale_velosity=state.input_output.convective_scale_velocity, + cin_0=locals.cin_0, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "convective_scale_velocity": state.input_output.convective_scale_velocity.field[:], + "local_cin_0": locals.cin_0.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectionTrigger_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectionTrigger_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_ConvectionTrigger_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble.py new file mode 100644 index 000000000..c35e14683 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble.py @@ -0,0 +1,260 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import ( + UpdateWorkfunctionAndPrecipitationEnsemble, +) + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_geopotential_height_cloud_levels_modified": {}, + "local_normalized_massflux_updraft_modified": {}, + "local_d_buoyancy_modified": {}, + "local_gamma_cloud_levels": {}, + "local_t_cloud_levels_modified": {}, + "local_cloud_workfunction_0_modified": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "epsilon_forced": {}, + "local_precipitation_ensemble": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.geopotential_height_cloud_levels_modified.data[:] = inputs[ + "local_geopotential_height_cloud_levels_modified" + ] + locals.normalized_massflux_updraft_modified.data[:] = inputs[ + "local_normalized_massflux_updraft_modified" + ] + locals.d_buoyancy_modified.data[:] = inputs["local_d_buoyancy_modified"] + locals.gamma_cloud_levels.data[:] = inputs["local_gamma_cloud_levels"] + locals.t_cloud_levels_modified.data[:] = inputs["local_t_cloud_levels_modified"] + locals.cloud_workfunction_0_modified.data[:] = inputs["local_cloud_workfunction_0_modified"] + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + if MAXENS1 * MAXENS2 * MAXENS3 != 16: + raise NotImplementedError( + "Due to limitations in the translate test, the size of the ensemble" + "dimension must be set manually in the translate test. Please modify, then disable this error" + "manually to continue." + ) + locals.precipitation_ensemble.data[:] = inputs["local_precipitation_ensemble"][:, :, 0:16] + + # initialize test code + code = UpdateWorkfunctionAndPrecipitationEnsemble( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels_modified=locals.geopotential_height_cloud_levels_modified, + normalized_massflux_updraft_modified=locals.normalized_massflux_updraft_modified, + d_buoyancy_modified=locals.d_buoyancy_modified, + gamma_cloud_levels=locals.gamma_cloud_levels, + t_cloud_levels_modified=locals.t_cloud_levels_modified, + cloud_workfunction_0_modified=locals.cloud_workfunction_0_modified, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + epsilon_forced=state.output.epsilon_forced, + precipitation_ensemble=locals.precipitation_ensemble, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_origin_level": state.output.updraft_origin_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_geopotential_height_cloud_levels_modified": locals.geopotential_height_cloud_levels_modified.field[ + : + ], + "local_normalized_massflux_updraft_modified": locals.normalized_massflux_updraft_modified.field[ + : + ], + "local_d_buoyancy_modified": locals.d_buoyancy_modified.field[:], + "local_gamma_cloud_levels": locals.gamma_cloud_levels.field[:], + "local_t_cloud_levels_modified": locals.t_cloud_levels_modified.field[:], + "local_cloud_workfunction_0_modified": locals.cloud_workfunction_0_modified.field[:], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_precipitation_ensemble": locals.precipitation_ensemble.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_shallow( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_mid( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdateWorkfunctionAndPrecipitationEnsemble_deep( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftCIN.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftCIN.py new file mode 100644 index 000000000..a84a8bd39 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftCIN.py @@ -0,0 +1,244 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import UpdraftCIN + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_geopotential_height_cloud_levels": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_normalized_massflux_updraft": {}, + "normalized_massflux_updraft_forced": {}, + "local_d_buoyancy": {}, + "local_d_buoyancy_forced": {}, + "local_gamma_cloud_levels": {}, + "local_gamma_cloud_levels_forced": {}, + "local_t_cloud_levels": {}, + "local_t_cloud_levels_forced": {}, + "local_cin_0": {}, + "local_cin_1": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.normalized_massflux_updraft.data[:] = inputs["local_normalized_massflux_updraft"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.d_buoyancy.data[:] = inputs["local_d_buoyancy"] + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.gamma_cloud_levels.data[:] = inputs["local_gamma_cloud_levels"] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + locals.cin_0.data[:] = inputs["local_cin_0"] + locals.cin_1.data[:] = inputs["local_cin_1"] + + # initialize test code + code = UpdraftCIN( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + d_buoyancy=locals.d_buoyancy, + d_buoyancy_forced=locals.d_buoyancy_forced, + gamma_cloud_levels=locals.gamma_cloud_levels, + gamma_cloud_levels_forced=locals.gamma_cloud_levels_forced, + t_cloud_levels=locals.t_cloud_levels, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + cin_0=locals.cin_0, + cin_1=locals.cin_1, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.data[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.data[ + : + ], + "local_normalized_massflux_updraft": locals.normalized_massflux_updraft.data[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_d_buoyancy": locals.d_buoyancy.data[:], + "local_d_buoyancy_forced": locals.d_buoyancy_forced.data[:], + "local_gamma_cloud_levels": locals.gamma_cloud_levels.data[:], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.data[:], + "local_t_cloud_levels": locals.t_cloud_levels.data[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.data[:], + "local_cin_0": locals.cin_0.data[:], + "local_cin_1": locals.cin_1.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftCIN_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftCIN_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftCIN_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftInitialWorkfunctions.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftInitialWorkfunctions.py new file mode 100644 index 000000000..07bf481ac --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftInitialWorkfunctions.py @@ -0,0 +1,244 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import UpdraftInitialWorkfunctions + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "local_geopotential_height_cloud_levels": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_normalized_massflux_updraft": {}, + "normalized_massflux_updraft_forced": {}, + "local_d_buoyancy": {}, + "local_d_buoyancy_forced": {}, + "local_gamma_cloud_levels": {}, + "local_gamma_cloud_levels_forced": {}, + "local_t_cloud_levels": {}, + "local_t_cloud_levels_forced": {}, + "local_cloud_workfunction_0": {}, + "local_cloud_workfunction_1": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.geopotential_height_cloud_levels.data[:] = inputs["local_geopotential_height_cloud_levels"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.normalized_massflux_updraft.data[:] = inputs["local_normalized_massflux_updraft"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.d_buoyancy.data[:] = inputs["local_d_buoyancy"] + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.gamma_cloud_levels.data[:] = inputs["local_gamma_cloud_levels"] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + locals.cloud_workfunction_0.data[:] = inputs["local_cloud_workfunction_0"] + locals.cloud_workfunction_1.data[:] = inputs["local_cloud_workfunction_1"] + + # initialize test code + code = UpdraftInitialWorkfunctions( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + geopotential_height_cloud_levels=locals.geopotential_height_cloud_levels, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft=locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + d_buoyancy=locals.d_buoyancy, + d_buoyancy_forced=locals.d_buoyancy_forced, + gamma_cloud_levels=locals.gamma_cloud_levels, + gamma_cloud_levels_forced=locals.gamma_cloud_levels_forced, + t_cloud_levels=locals.t_cloud_levels, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + cloud_workfunction_0=locals.cloud_workfunction_0, + cloud_workfunction_1=locals.cloud_workfunction_1, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_geopotential_height_cloud_levels": locals.geopotential_height_cloud_levels.data[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.data[ + : + ], + "local_normalized_massflux_updraft": locals.normalized_massflux_updraft.data[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_d_buoyancy": locals.d_buoyancy.data[:], + "local_d_buoyancy_forced": locals.d_buoyancy_forced.data[:], + "local_gamma_cloud_levels": locals.gamma_cloud_levels.data[:], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.data[:], + "local_t_cloud_levels": locals.t_cloud_levels.data[:], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.data[:], + "local_cloud_workfunction_0": locals.cloud_workfunction_0.data[:], + "local_cloud_workfunction_1": locals.cloud_workfunction_1.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftInitialWorkfunctions_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMassFlux.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMassFlux.py new file mode 100644 index 000000000..6ea5f2a8c --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMassFlux.py @@ -0,0 +1,235 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import UpdraftMassFlux + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "updraft_origin_level": {}, + "cloud_top_level": {}, + "pbl_level": {}, + "updraft_lfc_level": {}, + "lcl_level": {}, + "p_cloud_levels_forced": {}, + "p_surface": {}, + "ocean_fraction": {}, + "local_normalized_massflux_updraft": {}, + "normalized_massflux_updraft_forced": {}, + "local_normalized_massflux_updraft_modified": {}, + "local_random_number": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input_output.pbl_level.data[:] = inputs["pbl_level"] - 1 + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + state.input_output.p_surface.data[:] = inputs["p_surface"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + locals.normalized_massflux_updraft.data[:] = inputs["local_normalized_massflux_updraft"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.normalized_massflux_updraft_modified.data[:] = inputs[ + "local_normalized_massflux_updraft_modified" + ] + locals.random_number.data[:] = inputs["local_random_number"] + + # initialize test code + code = UpdraftMassFlux( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + updraft_origin_level=state.output.updraft_origin_level, + cloud_top_level=state.output.cloud_top_level, + pbl_level=state.input_output.pbl_level, + updraft_lfc_level=state.output.updraft_lfc_level, + lcl_level=state.output.lcl_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + p_surface=state.input_output.p_surface, + ocean_fraction=state.input.ocean_fraction, + normalized_massflux_updraft=locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_updraft_modified=locals.normalized_massflux_updraft_modified, + random_number=locals.random_number, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "pbl_level": state.input_output.pbl_level.data[:] + 1, + "updraft_lfc_level": state.output.updraft_lfc_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "lcl_level": state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "p_surface": state.input_output.p_surface.data[:], + "ocean_fraction": state.input.ocean_fraction.data[:], + "local_normalized_massflux_updraft": locals.normalized_massflux_updraft.data[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_normalized_massflux_updraft_modified": locals.normalized_massflux_updraft_modified.data[:], + "local_random_number": locals.random_number.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMassFlux_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMassFlux_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMassFlux_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget.py new file mode 100644 index 000000000..87eca61c9 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget.py @@ -0,0 +1,330 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import ( + updraft_moist_static_energy_and_momentum_budget, +) + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_start_level": {}, + "cloud_top_level": {}, + "p_forced": {}, + "local_env_moist_static_energy": {}, + "local_env_moist_static_energy_forced": {}, + "local_env_moist_static_energy_cloud_levels": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_env_saturation_moist_static_energy_cloud_levels": {}, + "local_env_saturation_moist_static_energy_cloud_levels_forced": {}, + "local_cloud_moist_static_energy": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_normalized_massflux_updraft": {}, + "normalized_massflux_updraft_forced": {}, + "local_mass_entrainment_updraft": {}, + "local_mass_detrainment_updraft": {}, + "local_mass_entrainment_u_updraft": {}, + "local_mass_detrainment_u_updraft": {}, + "mass_detrainment_updraft_forced": {}, + "mass_entrainment_updraft_forced": {}, + "u": {}, + "v": {}, + "local_u_c": {}, + "local_v_c": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_partition_liquid_ice": {}, + "cloud_liquid_after_rain_forced": {}, + "local_vapor_excess": {}, + "local_t_excess": {}, + "local_add_buoyancy": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.start_level.data[:] = inputs["local_start_level"] - 1 + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.input_output.p_forced.data[:] = inputs["p_forced"] + locals.environment_moist_static_energy.data[:] = inputs["local_env_moist_static_energy"] + locals.environment_moist_static_energy_forced.data[:] = inputs["local_env_moist_static_energy_forced"] + locals.environment_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels" + ] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.environment_saturation_moist_static_energy_cloud_levels.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels" + ] + locals.environment_saturation_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_moist_static_energy_cloud_levels_forced" + ] + locals.cloud_moist_static_energy.data[:] = inputs["local_cloud_moist_static_energy"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.normalized_massflux_updraft.data[:] = inputs["local_normalized_massflux_updraft"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.mass_entrainment_updraft.data[:] = inputs["local_mass_entrainment_updraft"] + locals.mass_detrainment_updraft.data[:] = inputs["local_mass_detrainment_updraft"] + locals.mass_entrainment_u_updraft.data[:] = inputs["local_mass_entrainment_u_updraft"] + locals.mass_detrainment_u_updraft.data[:] = inputs["local_mass_detrainment_u_updraft"] + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_entrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_entrainment_updraft_forced"] + ) + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.u_c.data[:] = inputs["local_u_c"] + locals.v_c.data[:] = inputs["local_v_c"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + locals.t_excess.data[:] = inputs["local_t_excess"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=updraft_moist_static_energy_and_momentum_budget, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES": cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + "PRESSURE_GRADIENT_CONSTANT": cumulus_parameterization_config.PRESSURE_GRADIENT_CONSTANT, + }, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + start_level=locals.start_level, + cloud_top_level=state.output.cloud_top_level, + p_forced=state.input_output.p_forced, + environment_moist_static_energy=locals.environment_moist_static_energy, + environment_moist_static_energy_forced=locals.environment_moist_static_energy_forced, + environment_moist_static_energy_cloud_levels=locals.environment_moist_static_energy_cloud_levels, + environment_moist_static_energy_cloud_levels_forced=locals.environment_moist_static_energy_cloud_levels_forced, + environment_saturation_moist_static_energy_cloud_levels=locals.environment_saturation_moist_static_energy_cloud_levels, + environment_saturation_moist_static_energy_cloud_levels_forced=locals.environment_saturation_moist_static_energy_cloud_levels_forced, + cloud_moist_static_energy=locals.cloud_moist_static_energy, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + normalized_massflux_updraft=locals.normalized_massflux_updraft, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + mass_entrainment_updraft=locals.mass_entrainment_updraft, + mass_detrainment_updraft=locals.mass_detrainment_updraft, + mass_entrainment_u_updraft=locals.mass_entrainment_u_updraft, + mass_detrainment_u_updraft=locals.mass_detrainment_u_updraft, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_entrainment_updraft_forced=state.output.mass_entrainment_updraft_forced, + u=state.input_output.u, + v=state.input_output.v, + u_c=locals.u_c, + v_c=locals.v_c, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + partition_liquid_ice=locals.partition_liquid_ice, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + vapor_excess=locals.vapor_excess, + t_excess=locals.t_excess, + add_buoyancy=locals.add_buoyancy, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_start_level": locals.start_level.field[:] + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "p_forced": state.input_output.p_forced.field[:], + "local_env_moist_static_energy": locals.environment_moist_static_energy.field[:], + "local_env_moist_static_energy_forced": locals.environment_moist_static_energy_forced.field[:], + "local_env_moist_static_energy_cloud_levels": locals.environment_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels": locals.environment_saturation_moist_static_energy_cloud_levels.field[ + : + ], + "local_env_saturation_moist_static_energy_cloud_levels_forced": locals.environment_saturation_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_cloud_moist_static_energy": locals.cloud_moist_static_energy.field[:], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_normalized_massflux_updraft": locals.normalized_massflux_updraft.field[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_mass_entrainment_updraft": locals.mass_entrainment_updraft.field[:], + "local_mass_detrainment_updraft": locals.mass_detrainment_updraft.field[:], + "local_mass_entrainment_u_updraft": locals.mass_entrainment_u_updraft.field[:], + "local_mass_detrainment_u_updraft": locals.mass_detrainment_u_updraft.field[:], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_entrainment_updraft_forced": state.output.mass_entrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_u_c": locals.u_c.field[:], + "local_v_c": locals.v_c.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_partition_liquid_ice": locals.partition_liquid_ice.field[:], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vapor_excess": locals.vapor_excess.field[:], + "local_t_excess": locals.t_excess.field[:], + "local_add_buoyancy": locals.add_buoyancy.field[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_shallow( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_mid( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoistStaticEnergyAndMomentumBudget_deep( + TranslateFortranData2Py +): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoisture.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoisture.py new file mode 100644 index 000000000..e2ecaa1ba --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftMoisture.py @@ -0,0 +1,334 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import updraft_moisture + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_start_level": {}, + "error_code": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "condensate_to_fall_forced": {}, + "total_normalized_integrated_condensate_forced": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_updraft_column_temperature_forced": {}, + "ocean_fraction": {}, + "convection_fraction": {}, + "surface_type": {}, + "p_forced": {}, + "cloud_top_level": {}, + "local_d_buoyancy_forced": {}, + "local_cloud_liquid_before_rain_forced": {}, + "local_t_cloud_levels": {}, + "local_vapor_forced": {}, + "local_gamma_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "local_env_saturation_mixing_ratio_cloud_levels_forced": {}, + "updraft_origin_level": {}, + "local_vapor_cloud_levels_forced": {}, + "local_vapor_excess": {}, + "ccn": {}, + "local_mass_entrainment_updraft": {}, + "local_mass_detrainment_updraft": {}, + "local_psum": {}, + "local_psumh": {}, + "local_c1d": {}, + "local_add_buoyancy": {}, + "local_vertical_velocity_3d": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.start_level.data[:] = inputs["local_start_level"] - 1 + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["total_normalized_integrated_condensate_forced"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + state.input.ocean_fraction.data[:] = inputs["ocean_fraction"] + state.input.convection_fraction.data[:] = inputs["convection_fraction"] + state.input.surface_type.data[:] = inputs["surface_type"] + state.input_output.p_forced.data[:] = inputs["p_forced"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + locals.d_buoyancy_forced.data[:] = inputs["local_d_buoyancy_forced"] + locals.cloud_liquid_before_rain_forced.data[:] = inputs["local_cloud_liquid_before_rain_forced"] + locals.t_cloud_levels.data[:] = inputs["local_t_cloud_levels"] + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + locals.gamma_cloud_levels_forced.data[:] = inputs["local_gamma_cloud_levels_forced"] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[:] = inputs[ + "local_env_saturation_mixing_ratio_cloud_levels_forced" + ] + state.output.updraft_origin_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_origin_level"] - 1 + ) + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.vapor_excess.data[:] = inputs["local_vapor_excess"] + state.input_output.ccn.data[:] = inputs["ccn"] + locals.mass_entrainment_updraft.data[:] = inputs["local_mass_entrainment_updraft"] + locals.mass_detrainment_updraft.data[:] = inputs["local_mass_detrainment_updraft"] + locals.psum.data[:] = inputs["local_psum"] + locals.psumh.data[:] = inputs["local_psumh"] + locals.c1d.data[:] = inputs["local_c1d"] + locals.add_buoyancy.data[:] = inputs["local_add_buoyancy"] + locals.vertical_velocity_3d.data[:] = inputs["local_vertical_velocity_3d"] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=updraft_moisture, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "BOUNDARY_CONDITION_METHOD": cumulus_parameterization_config.BOUNDARY_CONDITION_METHOD, + "USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES": cumulus_parameterization_config.USE_LINEAR_SUBCLOUD_MOISTURE_FLUXES, + "AUTOCONV": config.AUTOCONV, + "CRITICAL_MIXING_RATIO_OVER_OCEAN": cumulus_parameterization_config.CRITICAL_MIXING_RATIO_OVER_OCEAN, + "CRITICAL_MIXING_RATIO_OVER_LAND": cumulus_parameterization_config.CRITICAL_MIXING_RATIO_OVER_LAND, + "FRAC_MODIS": cumulus_parameterization_config.FRAC_MODIS, + "ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF, + }, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + start_level=locals.start_level, + error_code=state.output.error_code, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=locals.cloud_total_water_after_entrainment_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + total_normalized_integrated_condensate_forced=state.output.total_normalized_integrated_condensate_forced, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + updraft_column_temperature_forced=locals.updraft_column_temperature_forced, + ocean_fraction=state.input.ocean_fraction, + convection_fraction=state.input.convection_fraction, + surface_type=state.input.surface_type, + p_forced=state.input_output.p_forced, + cloud_top_level=state.output.cloud_top_level, + d_buoyancy_forced=locals.d_buoyancy_forced, + cloud_liquid_before_rain_forced=locals.cloud_liquid_before_rain_forced, + t_cloud_levels=locals.t_cloud_levels, + vapor_forced=locals.vapor_forced, + gamma_cloud_levels_forced=locals.gamma_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + environment_saturation_mixing_ratio_cloud_levels_forced=locals.environment_saturation_mixing_ratio_cloud_levels_forced, + updraft_origin_level=state.output.updraft_origin_level, + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + vapor_excess=locals.vapor_excess, + ccn=state.input_output.ccn, + mass_entrainment_updraft=locals.mass_entrainment_updraft, + mass_detrainment_updraft=locals.mass_detrainment_updraft, + psum=locals.psum, + psumh=locals.psumh, + c1d=locals.c1d, + add_buoyancy=locals.add_buoyancy, + vertical_velocity_3d=locals.vertical_velocity_3d, + C0=plume_dependent_constants.C0, + AVERAGE_LAYER_DEPTH=plume_dependent_constants.AVERAGE_LAYER_DEPTH, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "local_start_level": locals.start_level.data[:] + 1, + "error_code": state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.data[ + : + ], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.data[ + : + ], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "total_normalized_integrated_condensate_forced": state.output.total_normalized_integrated_condensate_forced.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.data[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.data[:], + "ocean_fraction": state.input.ocean_fraction.data[:], + "convection_fraction": state.input.convection_fraction.data[:], + "surface_type": state.input.surface_type.data[:], + "p_forced": state.input_output.p_forced.data[:], + "cloud_top_level": state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "local_d_buoyancy_forced": locals.d_buoyancy_forced.data[:], + "local_cloud_liquid_before_rain_forced": locals.cloud_liquid_before_rain_forced.data[:], + "local_t_cloud_levels": locals.t_cloud_levels.data[:], + "local_vapor_forced": locals.vapor_forced.data[:], + "local_gamma_cloud_levels_forced": locals.gamma_cloud_levels_forced.data[:], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_env_saturation_mixing_ratio_cloud_levels_forced": locals.environment_saturation_mixing_ratio_cloud_levels_forced.data[ + : + ], + "updraft_origin_level": state.output.updraft_origin_level.data[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.data[:], + "local_vapor_excess": locals.vapor_excess.data[:], + "ccn": state.input_output.ccn.data[:], + "local_mass_entrainment_updraft": locals.mass_entrainment_updraft.data[:], + "local_mass_detrainment_updraft": locals.mass_detrainment_updraft.data[:], + "local_psum": locals.psum.data[:], + "local_psumh": locals.psumh.data[:], + "local_c1d": locals.c1d.data[:], + "local_add_buoyancy": locals.add_buoyancy.data[:], + "local_vertical_velocity_3d": locals.vertical_velocity_3d.data[:], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoisture_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoisture_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftMoisture_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftTemperature.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftTemperature.py new file mode 100644 index 000000000..7e30ea0ad --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftTemperature.py @@ -0,0 +1,191 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.updraft import updraft_temperature + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_t_cloud_levels_forced": {}, + "local_updraft_column_temperature_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=updraft_temperature, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + error_code=state.output.error_code, + updraft_column_temperature_forced=locals.updraft_column_temperature_forced, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=locals.cloud_total_water_after_entrainment_forced, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.field[:], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftTemperature_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftTemperature_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftTemperature_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftVerticalVelocity.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftVerticalVelocity.py new file mode 100644 index 000000000..e4d4caa7a --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/updraft/translate_GF2020_CumulusParameterization_UpdraftVerticalVelocity.py @@ -0,0 +1,243 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.shared_stencils import updraft_vertical_velocity +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "local_vertical_velocity_3d": {}, + "local_vertical_velocity_2d": {}, + "convective_scale_velocity": {}, + "entrainment_rate": {}, + "error_code": {}, + "local_detrainment_function_updraft": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "local_t_cloud_levels_forced": {}, + "local_updraft_column_temperature_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "local_vapor_forced": {}, + "lcl_level": {}, + "cloud_top_level": {}, + "updraft_lfc_level": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + locals.vertical_velocity_3d.data[:] = inputs["local_vertical_velocity_3d"] + locals.vertical_velocity_2d.data[:] = inputs["local_vertical_velocity_2d"] + state.input_output.convective_scale_velocity.data[:] = inputs["convective_scale_velocity"] + state.output.entrainment_rate.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "entrainment_rate" + ] + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + locals.detrainment_function_updraft.data[:] = inputs["local_detrainment_function_updraft"] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + locals.t_cloud_levels_forced.data[:] = inputs["local_t_cloud_levels_forced"] + locals.updraft_column_temperature_forced.data[:] = inputs["local_updraft_column_temperature_forced"] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + locals.vapor_forced.data[:] = inputs["local_vapor_forced"] + state.output.lcl_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["lcl_level"] - 1 + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.updraft_lfc_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["updraft_lfc_level"] - 1 + ) + + # initialize test code + code = self.stencil_factory.from_dims_halo( + func=updraft_vertical_velocity, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={"ZERO_DIFF": cumulus_parameterization_config.ZERO_DIFF}, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + if cumulus_parameterization_config.FIRST_GUESS_W == 0: + code( + vertical_velocity_3d=locals.vertical_velocity_3d, + vertical_velocity_2d=locals.vertical_velocity_2d, + convective_scale_velocity=state.input_output.convective_scale_velocity, + entrainment_rate=state.output.entrainment_rate, + detrainment_function_updraft=locals.detrainment_function_updraft, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + t_cloud_levels_forced=locals.t_cloud_levels_forced, + updraft_column_temperature_forced=locals.updraft_column_temperature_forced, + cloud_total_water_after_entrainment_forced=locals.cloud_total_water_after_entrainment_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + vapor_forced=locals.vapor_forced, + updraft_lfc_level=state.output.updraft_lfc_level, + cloud_top_level=state.output.cloud_top_level, + error_code=state.output.error_code, + plume=plume_dependent_constants.PLUME_INDEX, + ) + + # write output + outputs = { + "local_vertical_velocity_3d": locals.vertical_velocity_3d.field[:], + "local_vertical_velocity_2d": locals.vertical_velocity_2d.field[:], + "convective_scale_velocity": state.input_output.convective_scale_velocity.field[:], + "entrainment_rate": state.output.entrainment_rate.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_detrainment_function_updraft": locals.detrainment_function_updraft.field[:], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "local_t_cloud_levels_forced": locals.t_cloud_levels_forced.field[:], + "local_updraft_column_temperature_forced": locals.updraft_column_temperature_forced.field[:], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.field[ + : + ], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_vapor_forced": locals.vapor_forced.field[:], + "lcl_level": state.output.lcl_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + 1, + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "updraft_lcf_level": state.output.updraft_lfc_level.field[ + :, :, plume_dependent_constants.PLUME_INDEX + ] + + 1, + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_UpdraftVerticalVelocity_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/vertical_discretization/translate_GF2020_CumulusParameterization_VerticalDiscretization.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/vertical_discretization/translate_GF2020_CumulusParameterization_VerticalDiscretization.py new file mode 100644 index 000000000..002a4e3c3 --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/cumulus_parameterization/vertical_discretization/translate_GF2020_CumulusParameterization_VerticalDiscretization.py @@ -0,0 +1,396 @@ +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import ( + MAXENS1, + MAXENS2, + MAXENS3, + NUMBER_OF_PLUMES, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.locals import GF2020CumulusParameterizationLocals +from pyMoist.convection.GF_2020.cumulus_parameterization.plume_dependent_constants import ( + GF2020PlumeDependentConstants, +) +from pyMoist.convection.GF_2020.cumulus_parameterization.setup.set_constants import set_constants +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.cumulus_parameterization.vertical_discretization import VerticalDiscretization + + +class TestCore: + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + in_vars: dict, + out_vars: dict, + ): + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + in_vars["data_vars"] = { + "error_code": {}, + "cloud_top_level": {}, + "p_cloud_levels_forced": {}, + "local_geopotential_height_cloud_levels_forced": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "local_environment_massflux": {}, + "mass_detrainment_updraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "local_c1d": {}, + "u": {}, + "v": {}, + "local_u_cloud_levels": {}, + "local_v_cloud_levels": {}, + "local_u_c": {}, + "local_v_c": {}, + "local_u_c_downdraft": {}, + "local_v_c_downdraft": {}, + "local_cloud_moist_static_energy_forced": {}, + "local_cloud_moist_static_energy_downdraft_forced": {}, + "local_env_moist_static_energy_cloud_levels_forced": {}, + "local_vapor_cloud_levels_forced": {}, + "local_cloud_total_water_after_entrainment_forced": {}, + "local_cloud_total_water_after_entrainment_downdraft_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "local_melting": {}, + "local_partition_liquid_ice": {}, + "epsilon_forced": {}, + "local_d_buoyancy_downdraft_forced": {}, + "local_del_u_cloud_ensemble": {}, + "local_del_v_cloud_ensemble": {}, + "local_del_moist_static_energy_cloud_ensemble": {}, + "local_del_t_cloud_ensemble": {}, + "local_del_vapor_cloud_ensemble": {}, + "local_del_cloud_liquid_cloud_ensemble": {}, + "local_del_buoyancy_cloud_ensemble": {}, + "local_moist_static_energy_tendency_from_environmental_subsidence": {}, + "local_vapor_tendency_from_environmental_subsidence": {}, + "local_t_tendency_from_environmental_subsidence": {}, + } + + out_vars.update(in_vars["data_vars"]) + + def __call__(self, constants: dict, cu_param_constants: dict, plume: str, **inputs): + # initialize constants + config = GF2020Config(SINGLE_COLUMN_MODE=False, **constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**cu_param_constants) + plume_dependent_constants = GF2020PlumeDependentConstants() + plume_dependent_constants = set_constants( + cumulus_parameterization_config, plume_dependent_constants, plume + ) + + # initialize dataclasses + state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + locals = GF2020CumulusParameterizationLocals.zeros( + self.quantity_factory, + data_dimensions={ + "ensemble_1": MAXENS1, + "ensemble_2": MAXENS2, + "ensemble_3": MAXENS3, + "ensemble_members": MAXENS1 * MAXENS2 * MAXENS3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill relevant parts of dataclasses + state.output.error_code.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs["error_code"] + state.output.cloud_top_level.data[:, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_top_level"] - 1 + ) + state.output.p_cloud_levels_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "p_cloud_levels_forced" + ] + locals.geopotential_height_cloud_levels_forced.data[:] = inputs[ + "local_geopotential_height_cloud_levels_forced" + ] + state.output.normalized_massflux_updraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_updraft_forced"] + state.output.normalized_massflux_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["normalized_massflux_downdraft_forced"] + locals.environment_massflux.data[:] = inputs["local_environment_massflux"] + state.output.mass_detrainment_updraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["mass_detrainment_updraft_forced"] + ) + state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ] = inputs["mass_detrainment_downdraft_forced"] + locals.c1d.data[:] = inputs["local_c1d"] + state.input_output.u.data[:] = inputs["u"] + state.input_output.v.data[:] = inputs["v"] + locals.u_cloud_levels.data[:] = inputs["local_u_cloud_levels"] + locals.v_cloud_levels.data[:] = inputs["local_v_cloud_levels"] + locals.u_c.data[:] = inputs["local_u_c"] + locals.v_c.data[:] = inputs["local_v_c"] + locals.u_c_downdraft.data[:] = inputs["local_u_c_downdraft"] + locals.v_c_downdraft.data[:] = inputs["local_v_c_downdraft"] + locals.cloud_moist_static_energy_forced.data[:] = inputs["local_cloud_moist_static_energy_forced"] + locals.cloud_moist_static_energy_downdraft_forced.data[:] = inputs[ + "local_cloud_moist_static_energy_downdraft_forced" + ] + locals.environment_moist_static_energy_cloud_levels_forced.data[:] = inputs[ + "local_env_moist_static_energy_cloud_levels_forced" + ] + locals.vapor_cloud_levels_forced.data[:] = inputs["local_vapor_cloud_levels_forced"] + locals.cloud_total_water_after_entrainment_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_forced" + ] + locals.cloud_total_water_after_entrainment_downdraft_forced.data[:] = inputs[ + "local_cloud_total_water_after_entrainment_downdraft_forced" + ] + state.output.cloud_liquid_after_rain_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["cloud_liquid_after_rain_forced"] + ) + state.output.condensate_to_fall_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "condensate_to_fall_forced" + ] + state.output.evaporate_in_downdraft_forced.data[:, :, :, plume_dependent_constants.PLUME_INDEX] = ( + inputs["evaporate_in_downdraft_forced"] + ) + locals.melting.data[:] = inputs["local_melting"] + locals.partition_liquid_ice.data[:] = inputs["local_partition_liquid_ice"] + state.output.epsilon_forced.data[:, :, plume_dependent_constants.PLUME_INDEX] = inputs[ + "epsilon_forced" + ] + locals.d_buoyancy_downdraft_forced.data[:] = inputs["local_d_buoyancy_downdraft_forced"] + locals.del_u_cloud_ensemble.data[:] = inputs["local_del_u_cloud_ensemble"] + locals.del_v_cloud_ensemble.data[:] = inputs["local_del_v_cloud_ensemble"] + locals.del_moist_static_energy_cloud_ensemble.data[:] = inputs[ + "local_del_moist_static_energy_cloud_ensemble" + ] + locals.del_t_cloud_ensemble.data[:] = inputs["local_del_t_cloud_ensemble"] + locals.del_vapor_cloud_ensemble.data[:] = inputs["local_del_vapor_cloud_ensemble"] + locals.del_cloud_liquid_cloud_ensemble.data[:] = inputs["local_del_cloud_liquid_cloud_ensemble"] + locals.del_buoyancy_cloud_ensemble.data[:] = inputs["local_del_buoyancy_cloud_ensemble"] + locals.moist_static_energy_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_moist_static_energy_tendency_from_environmental_subsidence" + ] + locals.vapor_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_vapor_tendency_from_environmental_subsidence" + ] + locals.t_tendency_from_environmental_subsidence.data[:] = inputs[ + "local_t_tendency_from_environmental_subsidence" + ] + + # initialize test code + code = VerticalDiscretization( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + ) + + # call test code + if plume_dependent_constants.ENABLE_PLUME == 1: + code( + error_code=state.output.error_code, + cloud_top_level=state.output.cloud_top_level, + p_cloud_levels_forced=state.output.p_cloud_levels_forced, + geopotential_height_cloud_levels_forced=locals.geopotential_height_cloud_levels_forced, + normalized_massflux_updraft_forced=state.output.normalized_massflux_updraft_forced, + normalized_massflux_downdraft_forced=state.output.normalized_massflux_downdraft_forced, + environment_massflux=locals.environment_massflux, + mass_detrainment_updraft_forced=state.output.mass_detrainment_updraft_forced, + mass_detrainment_downdraft_forced=state.output.mass_detrainment_downdraft_forced, + c1d=locals.c1d, + u=state.input_output.u, + v=state.input_output.v, + u_cloud_levels=locals.u_cloud_levels, + v_cloud_levels=locals.v_cloud_levels, + u_c=locals.u_c, + v_c=locals.v_c, + u_c_downdraft=locals.u_c_downdraft, + v_c_downdraft=locals.v_c_downdraft, + cloud_moist_static_energy_forced=locals.cloud_moist_static_energy_forced, + cloud_moist_static_energy_downdraft_forced=locals.cloud_moist_static_energy_downdraft_forced, + environment_moist_static_energy_cloud_levels_forced=locals.environment_moist_static_energy_cloud_levels_forced, + vapor_cloud_levels_forced=locals.vapor_cloud_levels_forced, + cloud_total_water_after_entrainment_forced=locals.cloud_total_water_after_entrainment_forced, + cloud_total_water_after_entrainment_downdraft_forced=locals.cloud_total_water_after_entrainment_downdraft_forced, + cloud_liquid_after_rain_forced=state.output.cloud_liquid_after_rain_forced, + condensate_to_fall_forced=state.output.condensate_to_fall_forced, + evaporate_in_downdraft_forced=state.output.evaporate_in_downdraft_forced, + melting=locals.melting, + partition_liquid_ice=locals.partition_liquid_ice, + epsilon_forced=state.output.epsilon_forced, + d_buoyancy_downdraft_forced=locals.d_buoyancy_downdraft_forced, + del_u_cloud_ensemble=locals.del_u_cloud_ensemble, + del_v_cloud_ensemble=locals.del_v_cloud_ensemble, + del_moist_static_energy_cloud_ensemble=locals.del_moist_static_energy_cloud_ensemble, + del_t_cloud_ensemble=locals.del_t_cloud_ensemble, + del_vapor_cloud_ensemble=locals.del_vapor_cloud_ensemble, + del_cloud_liquid_cloud_ensemble=locals.del_cloud_liquid_cloud_ensemble, + del_buoyancy_cloud_ensemble=locals.del_buoyancy_cloud_ensemble, + moist_static_energy_tendency_from_environmental_subsidence=locals.moist_static_energy_tendency_from_environmental_subsidence, + vapor_tendency_from_environmental_subsidence=locals.vapor_tendency_from_environmental_subsidence, + t_tendency_from_environmental_subsidence=locals.t_tendency_from_environmental_subsidence, + plume_dependent_constants=plume_dependent_constants, + ) + + # write output + outputs = { + "error_code": state.output.error_code.field[:, :, plume_dependent_constants.PLUME_INDEX], + "cloud_top_level": state.output.cloud_top_level.field[:, :, plume_dependent_constants.PLUME_INDEX] + + 1, + "p_cloud_levels_forced": state.output.p_cloud_levels_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_geopotential_height_cloud_levels_forced": locals.geopotential_height_cloud_levels_forced.field[ + : + ], + "normalized_massflux_updraft_forced": state.output.normalized_massflux_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "normalized_massflux_downdraft_forced": state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_environment_massflux": locals.environment_massflux.field[:], + "mass_detrainment_updraft_forced": state.output.mass_detrainment_updraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "mass_detrainment_downdraft_forced": state.output.mass_detrainment_downdraft_forced.data[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_c1d": locals.c1d.field[:], + "u": state.input_output.u.field[:], + "v": state.input_output.v.field[:], + "local_u_cloud_levels": locals.u_cloud_levels.field[:], + "local_v_cloud_levels": locals.v_cloud_levels.field[:], + "local_u_c": locals.u_c.field[:], + "local_v_c": locals.v_c.field[:], + "local_u_c_downdraft": locals.u_c_downdraft.field[:], + "local_v_c_downdraft": locals.v_c_downdraft.field[:], + "local_cloud_moist_static_energy_forced": locals.cloud_moist_static_energy_forced.field[:], + "local_cloud_moist_static_energy_downdraft_forced": locals.cloud_moist_static_energy_downdraft_forced.field[ + : + ], + "local_env_moist_static_energy_cloud_levels_forced": locals.environment_moist_static_energy_cloud_levels_forced.field[ + : + ], + "local_vapor_cloud_levels_forced": locals.vapor_cloud_levels_forced.data[:], + "local_cloud_total_water_after_entrainment_forced": locals.cloud_total_water_after_entrainment_forced.field[ + : + ], + "local_cloud_total_water_after_entrainment_downdraft_forced": locals.cloud_total_water_after_entrainment_downdraft_forced.field[ + : + ], + "cloud_liquid_after_rain_forced": state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "condensate_to_fall_forced": state.output.condensate_to_fall_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "evaporate_in_downdraft_forced": state.output.evaporate_in_downdraft_forced.field[ + :, :, :, plume_dependent_constants.PLUME_INDEX + ], + "local_melting": locals.melting.field[:], + "local_partition_liquid_ice": locals.partition_liquid_ice.field[:], + "epsilon_forced": state.output.epsilon_forced.field[:, :, plume_dependent_constants.PLUME_INDEX], + "local_d_buoyancy_downdraft_forced": locals.d_buoyancy_downdraft_forced.data[:], + "local_del_u_cloud_ensemble": locals.del_u_cloud_ensemble.field[:], + "local_del_v_cloud_ensemble": locals.del_v_cloud_ensemble.field[:], + "local_del_moist_static_energy_cloud_ensemble": locals.del_moist_static_energy_cloud_ensemble.field[ + : + ], + "local_del_t_cloud_ensemble": locals.del_t_cloud_ensemble.field[:], + "local_del_vapor_cloud_ensemble": locals.del_vapor_cloud_ensemble.field[:], + "local_del_cloud_liquid_cloud_ensemble": locals.del_cloud_liquid_cloud_ensemble.field[:], + "local_del_buoyancy_cloud_ensemble": locals.del_buoyancy_cloud_ensemble.field[:], + "local_moist_static_energy_tendency_from_environmental_subsidence": locals.moist_static_energy_tendency_from_environmental_subsidence.field[ + : + ], + "local_vapor_tendency_from_environmental_subsidence": locals.vapor_tendency_from_environmental_subsidence.field[ + : + ], + "local_t_tendency_from_environmental_subsidence": locals.t_tendency_from_environmental_subsidence.field[ + : + ], + } + + return outputs + + +class TranslateGF2020_CumulusParameterization_VerticalDiscretization_shallow(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "shallow", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_VerticalDiscretization_mid(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "mid", **inputs) + + return outputs + + +class TranslateGF2020_CumulusParameterization_VerticalDiscretization_deep(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.test_core = TestCore(grid, namelist, stencil_factory, self.in_vars, self.out_vars) + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + + def compute_func(self, **inputs): + outputs = self.test_core(self.constants, self.cu_param_constants, "deep", **inputs) + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020.py new file mode 100644 index 000000000..d9846acae --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020.py @@ -0,0 +1,405 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.GF_2020 import GF2020 +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class TranslateGF2020(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "latitude_bugworkaroundname": {}, + "longitude_bugworkaroundname": {}, + "p_interface_bugworkaroundname": {}, + "t_bugworkaroundname": {}, + "u_bugworkaroundname": {}, + "v_bugworkaroundname": {}, + "w_bugworkaroundname": {}, + "omega_bugworkaroundname": {}, + "t_2m_bugworkaroundname": {}, + "specific_humidity_2m_bugworkaroundname": {}, + "t_surface_bugworkaroundname": {}, + "specific_humidity_surface_bugworkaroundname": {}, + "vapor_bugworkaroundname": {}, + "convective_liquid_bugworkaroundname": {}, + "convective_ice_bugworkaroundname": {}, + "large_scale_liquid_bugworkaroundname": {}, + "large_scale_ice_bugworkaroundname": {}, + "convective_cloud_fraction_bugworkaroundname": {}, + "large_scale_cloud_fraction_bugworkaroundname": {}, + "p_interface_timestep_start_bugworkaroundname": {}, + "t_timestep_start_bugworkaroundname": {}, + "u_timestep_start_bugworkaroundname": {}, + "v_timestep_start_bugworkaroundname": {}, + "vapor_timestep_start_bugworkaroundname": {}, + "geopotential_height_interface_bugworkaroundname": {}, + "geopotential_height_surface_bugworkaroundname": {}, + "area_bugworkaroundname": {}, + "pbl_level_bugworkaroundname": {}, + "convection_fraction_bugworkaroundname": {}, + "surface_type_bugworkaroundname": {}, + "seed_convection_bugworkaroundname": {}, + "land_fraction_bugworkaroundname": {}, + "scalar_diffusivity_bugworkaroundname": {}, + "buoyancy_bugworkaroundname": {}, + "convective_precipitation_GF_bugworkaroundname": {}, + "convective_precipitation_RAS_bugworkaroundname": {}, + "sensible_heat_flux_bugworkaroundname": {}, + "total_water_flux_deep_convection_interface_bugworkaroundname": {}, + "evaporation_bugworkaroundname": {}, + "convective_condensate_source_bugworkaroundname": {}, + "convective_condensate_grid_mean_bugworkaroundname": {}, + "entrainment_parameter_bugworkaroundname": {}, + "lateral_entrainment_rate_bugworkaroundname": {}, + "lateral_entrainment_rate_shallow_bugworkaroundname": {}, + "lateral_entrainment_rate_mid_bugworkaroundname": {}, + "lateral_entrainment_rate_deep_bugworkaroundname": {}, + "updraft_areal_fraction_bugworkaroundname": {}, + "updraft_vertical_velocity_bugworkaroundname": {}, + "dtdt_shortwave_bugworkaroundname": {}, + "dtdt_longwave_bugworkaroundname": {}, + "dspecific_humiditydt_pbl_bugworkaroundname": {}, + "dtdt_pbl_bugworkaroundname": {}, + "dtdt_from_dynamics_bugworkaroundname": {}, + "dvapordt_from_dynamics_bugworkaroundname": {}, + "sigma_mid_bugworkaroundname": {}, + "sigma_deep_bugworkaroundname": {}, + "total_precipitable_water_initial_bugworkaroundname": {}, + "saturation_total_precipitable_water_initial_bugworkaroundname": {}, + "dvapordt_deep_convection_bugworkaroundname": {}, + "dtdt_deep_convection_bugworkaroundname": {}, + "dudt_deep_convection_bugworkaroundname": {}, + "dvdt_deep_convection_bugworkaroundname": {}, + "pressure_shallow_convective_cloud_top_bugworkaroundname": {}, + "pressure_mid_convective_cloud_top_bugworkaroundname": {}, + "pressure_deep_convective_cloud_top_bugworkaroundname": {}, + "mass_flux_shallow_bugworkaroundname": {}, + "mass_flux_mid_bugworkaroundname": {}, + "mass_flux_deep_updraft_bugworkaroundname": {}, + "mass_flux_deep_updraft_interface_bugworkaroundname": {}, + "mass_flux_deep_updraft_detrained_bugworkaroundname": {}, + "mass_flux_deep_downdraft_bugworkaroundname": {}, + "mass_flux_cloud_base_bugworkaroundname": {}, + "mass_flux_cloud_base_shallow_bugworkaroundname": {}, + "mass_flux_cloud_base_mid_bugworkaroundname": {}, + "mass_flux_cloud_base_deep_bugworkaroundname": {}, + "convection_code_shallow_bugworkaroundname": {}, + "convection_code_mid_bugworkaroundname": {}, + "convection_code_deep_bugworkaroundname": {}, + "cloud_workfunction_0_bugworkaroundname": {}, + "cloud_workfunction_1_bugworkaroundname": {}, + "cloud_workfunction_2_bugworkaroundname": {}, + "cloud_workfunction_3_bugworkaroundname": {}, + "cloud_workfunction_1_pbl_bugworkaroundname": {}, + "cloud_workfunction_1_cin_bugworkaroundname": {}, + "pbl_time_scale_bugworkaroundname": {}, + "cape_removal_time_scale_bugworkaroundname": {}, + "lighting_density_bugworkaroundname": {}, + "convection_tracer_bugworkaroundname": {}, + } + + self.out_vars = self.in_vars["data_vars"].copy() + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers_input = data_loader.load("GF2020_ConvectionTracers") + + # workaround because translate test cannot read in 4d fields + self.manual_inputs = data_loader.load("GF2020_CumulusParameterization-Out") + + # load data from GF2020-In to fill in unmodified fields from the state + self.unmodified_state = data_loader.load("GF2020-In") + + def compute(self, inputs): + config = GF2020Config(SINGLE_COLUMN_MODE=False, **self.constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**self.cu_param_constants) + + # initialize saturation tables + saturation_tables = SaturationVaporPressureTable(self.stencil_factory.backend) + + # initialize convection tracers + convection_tracers = ConvectionTracers.ones( + self.quantity_factory, + data_dimensions={ + "convection_tracers": config.NUMBER_OF_TRACERS, + "size_three_dimension": 3, + "size_four_dimension": 4, + }, + ) + + convection_tracers.tracers.field[:] = np.moveaxis(self.convection_tracers_input["tracers"], 0, 3) + convection_tracers.vect_hcts.field[:] = self.convection_tracers_input["vect_hcts"] + convection_tracers.kc_scal.field[:] = self.convection_tracers_input["kc_scal"] + convection_tracers.fscav.field[:] = self.convection_tracers_input["fscav"] + convection_tracers.convfaci2g.field[:] = self.convection_tracers_input["convfaci2g"] + convection_tracers.retfactor.field[:] = self.convection_tracers_input["retfactor"] + convection_tracers.liq_and_gas.field[:] = self.convection_tracers_input["liq_and_gas"] + convection_tracers.online_cldliq.field[:] = self.convection_tracers_input["online_cldliq"] + convection_tracers.online_vud.field[:] = self.convection_tracers_input["online_vud"] + convection_tracers.ftemp_threshold.field[:] = self.convection_tracers_input["ftemp_threshold"] + convection_tracers.use_gcc_washout.field[:] = self.convection_tracers_input["use_gcc_washout"] + convection_tracers.use_gocart.field[:] = self.convection_tracers_input["use_gocart"] + convection_tracers.is_wetdep.field[:] = self.convection_tracers_input["is_wetdep"] + + # initialize GF2020 state + state = GF2020State.zeros(self.quantity_factory) + + state.latitude.field[:] = inputs["latitude_bugworkaroundname"] + state.longitude.field[:] = inputs["longitude_bugworkaroundname"] + state.p_interface.field[:] = inputs["p_interface_bugworkaroundname"] + state.t.field[:] = inputs["t_bugworkaroundname"] + state.u.field[:] = inputs["u_bugworkaroundname"] + state.v.field[:] = inputs["v_bugworkaroundname"] + state.w.field[:] = inputs["w_bugworkaroundname"] + state.omega.field[:] = inputs["omega_bugworkaroundname"] + state.t_2m.field[:] = inputs["t_2m_bugworkaroundname"] + state.specific_humidity_2m.field[:] = inputs["specific_humidity_2m_bugworkaroundname"] + state.t_surface.field[:] = inputs["t_surface_bugworkaroundname"] + state.specific_humidity_surface.field[:] = inputs["specific_humidity_surface_bugworkaroundname"] + state.vapor.field[:] = inputs["vapor_bugworkaroundname"] + state.convective_liquid.field[:] = inputs["convective_liquid_bugworkaroundname"] + state.convective_ice.field[:] = inputs["convective_ice_bugworkaroundname"] + state.large_scale_liquid.field[:] = inputs["large_scale_liquid_bugworkaroundname"] + state.large_scale_ice.field[:] = inputs["large_scale_ice_bugworkaroundname"] + state.convective_cloud_fraction.field[:] = inputs["convective_cloud_fraction_bugworkaroundname"] + state.large_scale_cloud_fraction.field[:] = inputs["large_scale_cloud_fraction_bugworkaroundname"] + state.p_interface_timestep_start.field[:] = inputs["p_interface_timestep_start_bugworkaroundname"] + state.t_timestep_start.field[:] = inputs["t_timestep_start_bugworkaroundname"] + state.u_timestep_start.field[:] = inputs["u_timestep_start_bugworkaroundname"] + state.v_timestep_start.field[:] = inputs["v_timestep_start_bugworkaroundname"] + state.vapor_timestep_start.field[:] = inputs["vapor_timestep_start_bugworkaroundname"] + state.geopotential_height_interface.field[:] = inputs[ + "geopotential_height_interface_bugworkaroundname" + ] + state.geopotential_height_surface.field[:] = inputs["geopotential_height_surface_bugworkaroundname"] + state.area.field[:] = inputs["area_bugworkaroundname"] + state.pbl_level.field[:] = inputs["pbl_level_bugworkaroundname"] - 1 + state.convection_fraction.field[:] = inputs["convection_fraction_bugworkaroundname"] + state.surface_type.field[:] = inputs["surface_type_bugworkaroundname"] + state.seed_convection.field[:] = inputs["seed_convection_bugworkaroundname"] + state.land_fraction.field[:] = inputs["land_fraction_bugworkaroundname"] + state.scalar_diffusivity.field[:] = inputs["scalar_diffusivity_bugworkaroundname"] + state.buoyancy.field[:] = inputs["buoyancy_bugworkaroundname"] + state.convective_precipitation_GF.field[:] = inputs["convective_precipitation_GF_bugworkaroundname"] + state.convective_precipitation_RAS.field[:] = inputs["convective_precipitation_RAS_bugworkaroundname"] + state.sensible_heat_flux.field[:] = inputs["sensible_heat_flux_bugworkaroundname"] + state.total_water_flux_deep_convection_interface.field[:] = inputs[ + "total_water_flux_deep_convection_interface_bugworkaroundname" + ] + state.evaporation.field[:] = inputs["evaporation_bugworkaroundname"] + state.convective_condensate_source.field[:] = inputs["convective_condensate_source_bugworkaroundname"] + state.convective_condensate_grid_mean.field[:] = inputs[ + "convective_condensate_grid_mean_bugworkaroundname" + ] + state.entrainment_parameter.field[:] = inputs["entrainment_parameter_bugworkaroundname"] + state.lateral_entrainment_rate.field[:] = inputs["lateral_entrainment_rate_bugworkaroundname"] + state.lateral_entrainment_rate_shallow.field[:] = inputs[ + "lateral_entrainment_rate_shallow_bugworkaroundname" + ] + state.lateral_entrainment_rate_mid.field[:] = inputs["lateral_entrainment_rate_mid_bugworkaroundname"] + state.lateral_entrainment_rate_deep.field[:] = inputs[ + "lateral_entrainment_rate_deep_bugworkaroundname" + ] + state.updraft_areal_fraction.field[:] = inputs["updraft_areal_fraction_bugworkaroundname"] + state.updraft_vertical_velocity.field[:] = inputs["updraft_vertical_velocity_bugworkaroundname"] + state.dtdt_shortwave.field[:] = inputs["dtdt_shortwave_bugworkaroundname"] + state.dtdt_longwave.field[:] = inputs["dtdt_longwave_bugworkaroundname"] + state.dspecific_humiditydt_pbl.field[:] = inputs["dspecific_humiditydt_pbl_bugworkaroundname"] + state.dtdt_pbl.field[:] = inputs["dtdt_pbl_bugworkaroundname"] + state.dtdt_from_dynamics.field[:] = inputs["dtdt_from_dynamics_bugworkaroundname"] + state.dvapordt_from_dynamics.field[:] = inputs["dvapordt_from_dynamics_bugworkaroundname"] + state.sigma_mid.field[:] = inputs["sigma_mid_bugworkaroundname"] + state.sigma_deep.field[:] = inputs["sigma_deep_bugworkaroundname"] + state.total_precipitable_water_initial.field[:] = inputs[ + "total_precipitable_water_initial_bugworkaroundname" + ] + state.saturation_total_precipitable_water_initial.field[:] = inputs[ + "saturation_total_precipitable_water_initial_bugworkaroundname" + ] + state.dvapordt_deep_convection.field[:] = inputs["dvapordt_deep_convection_bugworkaroundname"] + state.dtdt_deep_convection.field[:] = inputs["dtdt_deep_convection_bugworkaroundname"] + state.dudt_deep_convection.field[:] = inputs["dudt_deep_convection_bugworkaroundname"] + state.dvdt_deep_convection.field[:] = inputs["dvdt_deep_convection_bugworkaroundname"] + state.pressure_shallow_convective_cloud_top.field[:] = inputs[ + "pressure_shallow_convective_cloud_top_bugworkaroundname" + ] + state.pressure_mid_convective_cloud_top.field[:] = inputs[ + "pressure_mid_convective_cloud_top_bugworkaroundname" + ] + state.pressure_deep_convective_cloud_top.field[:] = inputs[ + "pressure_deep_convective_cloud_top_bugworkaroundname" + ] + state.mass_flux_shallow.field[:] = inputs["mass_flux_shallow_bugworkaroundname"] + state.mass_flux_mid.field[:] = inputs["mass_flux_mid_bugworkaroundname"] + state.mass_flux_deep_updraft.field[:] = inputs["mass_flux_deep_updraft_bugworkaroundname"] + state.mass_flux_deep_updraft_interface.field[:] = inputs[ + "mass_flux_deep_updraft_interface_bugworkaroundname" + ] + state.mass_flux_deep_updraft_detrained.field[:] = inputs[ + "mass_flux_deep_updraft_detrained_bugworkaroundname" + ] + state.mass_flux_deep_downdraft.field[:] = inputs["mass_flux_deep_downdraft_bugworkaroundname"] + state.mass_flux_cloud_base.field[:] = inputs["mass_flux_cloud_base_bugworkaroundname"] + state.mass_flux_cloud_base_shallow.field[:] = inputs["mass_flux_cloud_base_shallow_bugworkaroundname"] + state.mass_flux_cloud_base_mid.field[:] = inputs["mass_flux_cloud_base_mid_bugworkaroundname"] + state.mass_flux_cloud_base_deep.field[:] = inputs["mass_flux_cloud_base_deep_bugworkaroundname"] + state.convection_code_shallow.field[:] = inputs["convection_code_shallow_bugworkaroundname"] + state.convection_code_mid.field[:] = inputs["convection_code_mid_bugworkaroundname"] + state.convection_code_deep.field[:] = inputs["convection_code_deep_bugworkaroundname"] + state.cloud_workfunction_0.field[:] = inputs["cloud_workfunction_0_bugworkaroundname"] + state.cloud_workfunction_1.field[:] = inputs["cloud_workfunction_1_bugworkaroundname"] + state.cloud_workfunction_2.field[:] = inputs["cloud_workfunction_2_bugworkaroundname"] + state.cloud_workfunction_3.field[:] = inputs["cloud_workfunction_3_bugworkaroundname"] + state.cloud_workfunction_1_pbl.field[:] = inputs["cloud_workfunction_1_pbl_bugworkaroundname"] + state.cloud_workfunction_1_cin.field[:] = inputs["cloud_workfunction_1_cin_bugworkaroundname"] + state.pbl_time_scale.field[:] = inputs["pbl_time_scale_bugworkaroundname"] + state.cape_removal_time_scale.field[:] = inputs["cape_removal_time_scale_bugworkaroundname"] + state.lightning_density.field[:] = inputs["lighting_density_bugworkaroundname"] + state.convection_tracer.field[:] = inputs["convection_tracer_bugworkaroundname"] + + # initialize test code + code = GF2020( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + saturation_tables=saturation_tables, + ) + + # run test code + code( + state=state, + convection_tracers=convection_tracers, + ) + + # fill output dictionary for testing + outputs = { + "latitude_bugworkaroundname": state.latitude.field[:], + "longitude_bugworkaroundname": state.longitude.field[:], + "p_interface_bugworkaroundname": state.p_interface.field[:], + "t_bugworkaroundname": state.t.field[:], + "u_bugworkaroundname": state.u.field[:], + "v_bugworkaroundname": state.v.field[:], + "w_bugworkaroundname": state.w.field[:], + "omega_bugworkaroundname": state.omega.field[:], + "t_2m_bugworkaroundname": state.t_2m.field[:], + "specific_humidity_2m_bugworkaroundname": state.specific_humidity_2m.field[:], + "t_surface_bugworkaroundname": state.t_surface.field[:], + "specific_humidity_surface_bugworkaroundname": state.specific_humidity_surface.field[:], + "vapor_bugworkaroundname": state.vapor.field[:], + "convective_liquid_bugworkaroundname": state.convective_liquid.field[:], + "convective_ice_bugworkaroundname": state.convective_ice.field[:], + "large_scale_liquid_bugworkaroundname": state.large_scale_liquid.field[:], + "large_scale_ice_bugworkaroundname": state.large_scale_ice.field[:], + "convective_cloud_fraction_bugworkaroundname": state.convective_cloud_fraction.field[:], + "large_scale_cloud_fraction_bugworkaroundname": state.large_scale_cloud_fraction.field[:], + "p_interface_timestep_start_bugworkaroundname": state.p_interface_timestep_start.field[:], + "t_timestep_start_bugworkaroundname": state.t_timestep_start.field[:], + "u_timestep_start_bugworkaroundname": state.u_timestep_start.field[:], + "v_timestep_start_bugworkaroundname": state.v_timestep_start.field[:], + "vapor_timestep_start_bugworkaroundname": state.vapor_timestep_start.field[:], + "geopotential_height_interface_bugworkaroundname": state.geopotential_height_interface.field[:], + "geopotential_height_surface_bugworkaroundname": state.geopotential_height_surface.field[:], + "area_bugworkaroundname": state.area.field[:], + "pbl_level_bugworkaroundname": state.pbl_level.field[:] + 1, + "convection_fraction_bugworkaroundname": state.convection_fraction.field[:], + "surface_type_bugworkaroundname": state.surface_type.field[:], + "seed_convection_bugworkaroundname": state.seed_convection.field[:], + "land_fraction_bugworkaroundname": state.land_fraction.field[:], + "scalar_diffusivity_bugworkaroundname": state.scalar_diffusivity.field[:], + "buoyancy_bugworkaroundname": state.buoyancy.field[:], + "convective_precipitation_GF_bugworkaroundname": state.convective_precipitation_GF.field[:], + "convective_precipitation_RAS_bugworkaroundname": state.convective_precipitation_RAS.field[:], + "sensible_heat_flux_bugworkaroundname": state.sensible_heat_flux.field[:], + "total_water_flux_deep_convection_interface_bugworkaroundname": state.total_water_flux_deep_convection_interface.field[ + : + ], + "evaporation_bugworkaroundname": state.evaporation.field[:], + "convective_condensate_source_bugworkaroundname": state.convective_condensate_source.field[:], + "convective_condensate_grid_mean_bugworkaroundname": state.convective_condensate_grid_mean.field[ + : + ], + "entrainment_parameter_bugworkaroundname": state.entrainment_parameter.field[:], + "lateral_entrainment_rate_bugworkaroundname": state.lateral_entrainment_rate.field[:], + "lateral_entrainment_rate_shallow_bugworkaroundname": state.lateral_entrainment_rate_shallow.field[ + : + ], + "lateral_entrainment_rate_mid_bugworkaroundname": state.lateral_entrainment_rate_mid.field[:], + "lateral_entrainment_rate_deep_bugworkaroundname": state.lateral_entrainment_rate_deep.field[:], + "updraft_areal_fraction_bugworkaroundname": state.updraft_areal_fraction.field[:], + "updraft_vertical_velocity_bugworkaroundname": state.updraft_vertical_velocity.field[:], + "dtdt_shortwave_bugworkaroundname": state.dtdt_shortwave.field[:], + "dtdt_longwave_bugworkaroundname": state.dtdt_longwave.field[:], + "dspecific_humiditydt_pbl_bugworkaroundname": state.dspecific_humiditydt_pbl.field[:], + "dtdt_pbl_bugworkaroundname": state.dtdt_pbl.field[:], + "dtdt_from_dynamics_bugworkaroundname": state.dtdt_from_dynamics.field[:], + "dvapordt_from_dynamics_bugworkaroundname": state.dvapordt_from_dynamics.field[:], + "sigma_mid_bugworkaroundname": state.sigma_mid.field[:], + "sigma_deep_bugworkaroundname": state.sigma_deep.field[:], + "total_precipitable_water_initial_bugworkaroundname": state.total_precipitable_water_initial.field[ + : + ], + "saturation_total_precipitable_water_initial_bugworkaroundname": state.saturation_total_precipitable_water_initial.field[ + : + ], + "dvapordt_deep_convection_bugworkaroundname": state.dvapordt_deep_convection.field[:], + "dtdt_deep_convection_bugworkaroundname": state.dtdt_deep_convection.field[:], + "dudt_deep_convection_bugworkaroundname": state.dudt_deep_convection.field[:], + "dvdt_deep_convection_bugworkaroundname": state.dvdt_deep_convection.field[:], + "pressure_shallow_convective_cloud_top_bugworkaroundname": state.pressure_shallow_convective_cloud_top.field[ + : + ], + "pressure_mid_convective_cloud_top_bugworkaroundname": state.pressure_mid_convective_cloud_top.field[ + : + ], + "pressure_deep_convective_cloud_top_bugworkaroundname": state.pressure_deep_convective_cloud_top.field[ + : + ], + "mass_flux_shallow_bugworkaroundname": state.mass_flux_shallow.field[:], + "mass_flux_mid_bugworkaroundname": state.mass_flux_mid.field[:], + "mass_flux_deep_updraft_bugworkaroundname": state.mass_flux_deep_updraft.field[:], + "mass_flux_deep_updraft_interface_bugworkaroundname": state.mass_flux_deep_updraft_interface.field[ + : + ], + "mass_flux_deep_updraft_detrained_bugworkaroundname": state.mass_flux_deep_updraft_detrained.field[ + : + ], + "mass_flux_deep_downdraft_bugworkaroundname": state.mass_flux_deep_downdraft.field[:], + "mass_flux_cloud_base_bugworkaroundname": state.mass_flux_cloud_base.field[:], + "mass_flux_cloud_base_shallow_bugworkaroundname": state.mass_flux_cloud_base_shallow.field[:], + "mass_flux_cloud_base_mid_bugworkaroundname": state.mass_flux_cloud_base_mid.field[:], + "mass_flux_cloud_base_deep_bugworkaroundname": state.mass_flux_cloud_base_deep.field[:], + "convection_code_shallow_bugworkaroundname": state.convection_code_shallow.field[:], + "convection_code_mid_bugworkaroundname": state.convection_code_mid.field[:], + "convection_code_deep_bugworkaroundname": state.convection_code_deep.field[:], + "cloud_workfunction_0_bugworkaroundname": state.cloud_workfunction_0.field[:], + "cloud_workfunction_1_bugworkaroundname": state.cloud_workfunction_1.field[:], + "cloud_workfunction_2_bugworkaroundname": state.cloud_workfunction_2.field[:], + "cloud_workfunction_3_bugworkaroundname": state.cloud_workfunction_3.field[:], + "cloud_workfunction_1_pbl_bugworkaroundname": state.cloud_workfunction_1_pbl.field[:], + "cloud_workfunction_1_cin_bugworkaroundname": state.cloud_workfunction_1_cin.field[:], + "pbl_time_scale_bugworkaroundname": state.pbl_time_scale.field[:], + "cape_removal_time_scale_bugworkaroundname": state.cape_removal_time_scale.field[:], + "lighting_density_bugworkaroundname": state.lightning_density.field[:], + "convection_tracer_bugworkaroundname": state.convection_tracer.field[:], + } + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Finalize.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Finalize.py new file mode 100644 index 000000000..b8b08aaef --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Finalize.py @@ -0,0 +1,846 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.config import GF2020CumulusParameterizationConfig +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import NUMBER_OF_PLUMES +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.finalize import GF2020Finalize +from pyMoist.convection.GF_2020.locals import GF2020Locals +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class TranslateGF2020_Finalize(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + # fields saved midway through GF2020 fortran + "internal_lateral_entrainment_rate": {}, + "local_edge_height_above_surface": {}, + "local_layer_height_above_surface": {}, + "local_p": {}, + "local_p_kappa": {}, + "local_th": {}, + "local_mass": {}, + "local_modified_area": {}, + "local_vertical_velocity": {}, + "local_dz": {}, + "local_air_density": {}, + "local_scalar_diffusivity": {}, + "local_vapor_current": {}, + "local_p_flipped": {}, + "local_vapor_flipped": {}, + # NOTE these are disabled and loaded manually because the translate test cannot handle 4D data + # # CumulusParameterization state - input + # "t_excess": {}, + # "vapor_excess": {}, + # "grid_scale_forcing_t": {}, + # "grid_scale_forcing_vapor": {}, + # "subgrid_scale_forcing_t": {}, + # "subgrid_scale_forcing_vapor": {}, + # "seed_convection": {}, + # "saturation_water_vapor": {}, + # "ocean_fraction": {}, + # "convection_fraction": {}, + # "surface_type": {}, + # "lateral_entrainment_rate": {}, + # "last_error_code": {}, + # # CumulusParameterization state - output + # "dtdt": {}, + # "dvapordt": {}, + # "dcloudicedt": {}, + # "dudt": {}, + # "dvdt": {}, + # "dnliquiddt": {}, + # "dnicedt": {}, + # "dbuoyancydt": {}, + # "dconvectiveicedt": {}, + # "dlargescaleicedt": {}, + # "dconvectiveliquiddt": {}, + # "dlargescaleliquiddt": {}, + # "dconvectivecloudfractiondt": {}, + # "dlargescalecloudfractiondt": {}, + # "error_code": {}, + # "downdraft_origin_level": {}, + # "lcl_level": {}, + # "updraft_origin_level": {}, + # "updraft_lfc_level": {}, + # "cloud_top_level": {}, + # "kstabi": {}, + # "kstabm": {}, + # "precip": {}, + # "cloud_base_mass_flux_modified": {}, + # "epsilon_forced": {}, + # "total_normalized_integrated_condensate_forced": {}, + # "scale_dependence_factor": {}, + # "p_cloud_levels_forced": {}, + # "entrainment_rate": {}, + # "mass_entrainment_updraft_forced": {}, + # "mass_entrainment_downdraft_forced": {}, + # "mass_detrainment_updraft_forced": {}, + # "mass_detrainment_downdraft_forced": {}, + # "normalized_massflux_updraft_forced": {}, + # "normalized_massflux_downdraft_forced": {}, + # "condensate_to_fall_forced": {}, + # "evaporate_in_downdraft_forced": {}, + # "cloud_liquid_after_rain_forced": {}, + # "t_updraft": {}, + # "convective_cloud_fraction_output": {}, + # "cloud_workfunction_0": {}, + # "cloud_workfunction_1": {}, + # "cloud_workfunction_2": {}, + # "cloud_workfunction_3": {}, + # "cloud_workfunction_1_pbl": {}, + # "cloud_workfunction_1_cin": {}, + # "cape_removal_time_scale": {}, + # "pbl_time_scale": {}, + # "lightning_density": {}, + # "evaporation_sublimation_tendency": {}, + # "convective_precip_flux": {}, + # "t_perturbation": {}, + # # CumulusParameterization state - input-output + # "grid_length": {}, + # "pbl_level": {}, + # "ccn": {}, + # "air_density": {}, + # "omega": {}, + # "topography_height_no_negative": {}, + # "sensible_heat_flux": {}, + # "latent_heat_flux": {}, + # "longitude_degrees": {}, + # "latitude_degrees": {}, + # "t_old": {}, + # "vapor_old": {}, + # "t_modified_by_advection": {}, + # "vapor_modified_by_advection": {}, + # "geopotential_height_forced": {}, + # "p_forced": {}, + # "p_surface": {}, + # "t_surface": {}, + # "u": {}, + # "v": {}, + # "w": {}, + # "mass": {}, + # "convective_scale_velocity": {}, + # "buoyancy_excess": {}, + # "large_scale_ice": {}, + # "convective_ice": {}, + # "large_scale_liquid": {}, + # "convective_liquid": {}, + # "large_scale_cloud_fraction": {}, + # "convective_cloud_fraction": {}, + # "chemistry_tracers": {}, + # "chemistry_tracers_output": {}, + } + + # NOTE disabled fields are nan in fortran - zero in python, disabled so the test passes + self.out_vars: dict = { + "latitude_bugworkaroundname": {}, + "longitude_bugworkaroundname": {}, + "p_interface_bugworkaroundname": {}, + "t_bugworkaroundname": {}, + "u_bugworkaroundname": {}, + "v_bugworkaroundname": {}, + "w_bugworkaroundname": {}, + "omega_bugworkaroundname": {}, + "t_2m_bugworkaroundname": {}, + "specific_humidity_2m_bugworkaroundname": {}, + "t_surface_bugworkaroundname": {}, + "specific_humidity_surface_bugworkaroundname": {}, + "vapor_bugworkaroundname": {}, + "convective_liquid_bugworkaroundname": {}, + "convective_ice_bugworkaroundname": {}, + "large_scale_liquid_bugworkaroundname": {}, + "large_scale_ice_bugworkaroundname": {}, + "convective_cloud_fraction_bugworkaroundname": {}, + "large_scale_cloud_fraction_bugworkaroundname": {}, + "p_interface_timestep_start_bugworkaroundname": {}, + "t_timestep_start_bugworkaroundname": {}, + "u_timestep_start_bugworkaroundname": {}, + "v_timestep_start_bugworkaroundname": {}, + "vapor_timestep_start_bugworkaroundname": {}, + "geopotential_height_interface_bugworkaroundname": {}, + "geopotential_height_surface_bugworkaroundname": {}, + "area_bugworkaroundname": {}, + "pbl_level_bugworkaroundname": {}, + "convection_fraction_bugworkaroundname": {}, + "surface_type_bugworkaroundname": {}, + "seed_convection_bugworkaroundname": {}, + "land_fraction_bugworkaroundname": {}, + "scalar_diffusivity_bugworkaroundname": {}, + "buoyancy_bugworkaroundname": {}, + "convective_precipitation_GF_bugworkaroundname": {}, + "convective_precipitation_RAS_bugworkaroundname": {}, + "sensible_heat_flux_bugworkaroundname": {}, + "total_water_flux_deep_convection_interface_bugworkaroundname": {}, + "evaporation_bugworkaroundname": {}, + "sublimation_of_convective_precipitation_bugworkaroundname": {}, + "evaporation_of_convective_precipitation_bugworkaroundname": {}, + "ice_precip_flux_interface_bugworkaroundname": {}, + "liquid_precip_flux_interface_bugworkaroundname": {}, + "convective_condensate_source_bugworkaroundname": {}, + "convective_condensate_grid_mean_bugworkaroundname": {}, + "entrainment_parameter_bugworkaroundname": {}, + "lateral_entrainment_rate_bugworkaroundname": {}, + "lateral_entrainment_rate_shallow_bugworkaroundname": {}, + "lateral_entrainment_rate_mid_bugworkaroundname": {}, + "lateral_entrainment_rate_deep_bugworkaroundname": {}, + "updraft_areal_fraction_bugworkaroundname": {}, + "updraft_vertical_velocity_bugworkaroundname": {}, + "dtdt_shortwave_bugworkaroundname": {}, + "dtdt_longwave_bugworkaroundname": {}, + "dspecific_humiditydt_pbl_bugworkaroundname": {}, + "dtdt_pbl_bugworkaroundname": {}, + "dtdt_from_dynamics_bugworkaroundname": {}, + "dvapordt_from_dynamics_bugworkaroundname": {}, + "sigma_mid_bugworkaroundname": {}, + "sigma_deep_bugworkaroundname": {}, + "total_precipitable_water_initial_bugworkaroundname": {}, + "saturation_total_precipitable_water_initial_bugworkaroundname": {}, + "dvapordt_deep_convection_bugworkaroundname": {}, + "dtdt_deep_convection_bugworkaroundname": {}, + "dudt_deep_convection_bugworkaroundname": {}, + "dvdt_deep_convection_bugworkaroundname": {}, + "dliquiddt_deep_convection_bugworkaroundname": {}, + "dicedt_deep_convection_bugworkaroundname": {}, + "dcloudfractiondt_deep_convection_bugworkaroundname": {}, + "pressure_shallow_convective_cloud_top_bugworkaroundname": {}, + "pressure_mid_convective_cloud_top_bugworkaroundname": {}, + "pressure_deep_convective_cloud_top_bugworkaroundname": {}, + "mass_flux_shallow_bugworkaroundname": {}, + "mass_flux_mid_bugworkaroundname": {}, + "mass_flux_deep_updraft_bugworkaroundname": {}, + "mass_flux_deep_updraft_interface_bugworkaroundname": {}, + "mass_flux_deep_updraft_detrained_bugworkaroundname": {}, + "mass_flux_deep_downdraft_bugworkaroundname": {}, + "mass_flux_cloud_base_bugworkaroundname": {}, + "mass_flux_cloud_base_shallow_bugworkaroundname": {}, + "mass_flux_cloud_base_mid_bugworkaroundname": {}, + "mass_flux_cloud_base_deep_bugworkaroundname": {}, + "total_cumulative_mass_flux_interface_bugworkaroundname": {}, + "total_detraining_mass_flux_bugworkaroundname": {}, + "convection_code_shallow_bugworkaroundname": {}, + "convection_code_mid_bugworkaroundname": {}, + "convection_code_deep_bugworkaroundname": {}, + "cloud_workfunction_0_bugworkaroundname": {}, + "cloud_workfunction_1_bugworkaroundname": {}, + "cloud_workfunction_2_bugworkaroundname": {}, + "cloud_workfunction_3_bugworkaroundname": {}, + "cloud_workfunction_1_pbl_bugworkaroundname": {}, + "cloud_workfunction_1_cin_bugworkaroundname": {}, + "pbl_time_scale_bugworkaroundname": {}, + "cape_removal_time_scale_bugworkaroundname": {}, + "lighting_density_bugworkaroundname": {}, + "convection_tracer_bugworkaroundname": {}, + } + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.cu_param_constants = data_loader.load("GF2020_CumulusParameterization-constants") + self.convection_tracers_input = data_loader.load("GF2020_ConvectionTracers") + + # workaround because translate test cannot read in 4d fields + self.manual_inputs = data_loader.load("GF2020_CumulusParameterization-Out") + + # load data from GF2020-In to fill in unmodified fields from the state + self.unmodified_state = data_loader.load("GF2020-In") + + # load data from GF2020-CumulusParameterization-In as a roundabout way to fill unsaved locals + self.cumulus_state_before_call = data_loader.load("GF2020_CumulusParameterization-In") + + def compute(self, inputs): + config = GF2020Config(SINGLE_COLUMN_MODE=False, **self.constants) + cumulus_parameterization_config = GF2020CumulusParameterizationConfig(**self.cu_param_constants) + + # initialize GF2020 state + state = GF2020State.zeros(self.quantity_factory) + + # initialize GF2020 CumulusParameterization state + cumulus_parameterization_state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # initialize GF2020 locals + locals = GF2020Locals.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": 3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill state with data not passed in/unmodified inside the cumulus parameterization core + state.latitude.field[:] = self.unmodified_state["latitude_bugworkaroundname"] + state.longitude.field[:] = self.unmodified_state["longitude_bugworkaroundname"] + state.p_interface.field[:] = self.unmodified_state["p_interface_bugworkaroundname"] + state.t.field[:] = self.unmodified_state["t_bugworkaroundname"] + state.u.field[:] = self.unmodified_state["u_bugworkaroundname"] + state.v.field[:] = self.unmodified_state["v_bugworkaroundname"] + state.w.field[:] = self.unmodified_state["w_bugworkaroundname"] + state.omega.field[:] = self.unmodified_state["omega_bugworkaroundname"] + state.t_2m.field[:] = self.unmodified_state["t_2m_bugworkaroundname"] + state.specific_humidity_2m.field[:] = self.unmodified_state["specific_humidity_2m_bugworkaroundname"] + state.t_surface.field[:] = self.unmodified_state["t_surface_bugworkaroundname"] + state.specific_humidity_surface.field[:] = self.unmodified_state[ + "specific_humidity_surface_bugworkaroundname" + ] + state.vapor.field[:] = self.unmodified_state["vapor_bugworkaroundname"] + state.convective_liquid.field[:] = self.unmodified_state["convective_liquid_bugworkaroundname"] + state.convective_ice.field[:] = self.unmodified_state["convective_ice_bugworkaroundname"] + state.convective_cloud_fraction.field[:] = self.unmodified_state[ + "convective_cloud_fraction_bugworkaroundname" + ] + state.large_scale_liquid.field[:] = self.unmodified_state["large_scale_liquid_bugworkaroundname"] + state.large_scale_ice.field[:] = self.unmodified_state["large_scale_ice_bugworkaroundname"] + state.large_scale_cloud_fraction.field[:] = self.unmodified_state[ + "large_scale_cloud_fraction_bugworkaroundname" + ] + state.p_interface_timestep_start.field[:] = self.unmodified_state[ + "p_interface_timestep_start_bugworkaroundname" + ] + state.t_timestep_start.field[:] = self.unmodified_state["t_timestep_start_bugworkaroundname"] + state.u_timestep_start.field[:] = self.unmodified_state["u_timestep_start_bugworkaroundname"] + state.v_timestep_start.field[:] = self.unmodified_state["v_timestep_start_bugworkaroundname"] + state.vapor_timestep_start.field[:] = self.unmodified_state["vapor_timestep_start_bugworkaroundname"] + state.geopotential_height_interface.field[:] = self.unmodified_state[ + "geopotential_height_interface_bugworkaroundname" + ] + state.geopotential_height_surface.field[:] = self.unmodified_state[ + "geopotential_height_surface_bugworkaroundname" + ] + state.area.field[:] = self.unmodified_state["area_bugworkaroundname"] + state.pbl_level.field[:] = self.unmodified_state["pbl_level_bugworkaroundname"] + state.convection_fraction.field[:] = self.unmodified_state["convection_fraction_bugworkaroundname"] + state.surface_type.field[:] = self.unmodified_state["surface_type_bugworkaroundname"] + state.seed_convection.field[:] = self.unmodified_state["seed_convection_bugworkaroundname"] + state.land_fraction.field[:] = self.unmodified_state["land_fraction_bugworkaroundname"] + state.scalar_diffusivity.field[:] = self.unmodified_state["scalar_diffusivity_bugworkaroundname"] + state.buoyancy.field[:] = self.unmodified_state["buoyancy_bugworkaroundname"] + # state.convective_precipitation_GF.field[:] = self.unmodified_state["convective_precipitation_GF_bugworkaroundname"] + # state.convective_precipitation_RAS.field[:] = self.unmodified_state["convective_precipitation_RAS_bugworkaroundname"] + state.sensible_heat_flux.field[:] = self.unmodified_state["sensible_heat_flux_bugworkaroundname"] + # state.total_water_flux_deep_convection.field[:] = self.unmodified_state["total_water_flux_deep_convection_interface_bugworkaroundname"] + # state.sublimation_of_convective_precipitation.field[:] = self.unmodified_state["sublimation_of_convective_precipitation_bugworkaroundname"] + # state.evaporation_of_convective_precipitation.field[:] = self.unmodified_state["evaporation_of_convective_precipitation_bugworkaroundname"] + # state.ice_precip_flux_interface.field[:] = self.unmodified_state["ice_precip_flux_interface_bugworkaroundname"] + # state.liquid_precip_flux_interface.field[:] = self.unmodified_state["liquid_precip_flux_interface_bugworkaroundname"] + state.evaporation.field[:] = self.unmodified_state["evaporation_bugworkaroundname"] + # state.convective_condensate_source.field[:] = self.unmodified_state["convective_condensate_source_bugworkaroundname"] + # state.convective_condensate_grid_mean.field[:] = self.unmodified_state["convective_condensate_grid_mean_bugworkaroundname"] + # state.entrainment_parameter.field[:] = self.unmodified_state["entrainment_parameter_bugworkaroundname"] + state.lateral_entrainment_rate.field[:] = inputs["internal_lateral_entrainment_rate"] + # state.lateral_entrainment_rate_shallow.field[:] = self.unmodified_state["lateral_entrainment_rate_shallow_bugworkaroundname"] + # state.lateral_entrainment_rate_mid.field[:] = self.unmodified_state["lateral_entrainment_rate_mid_bugworkaroundname"] + # state.lateral_entrainment_rate_deep.field[:] = self.unmodified_state["lateral_entrainment_rate_deep_bugworkaroundname"] + # state.updraft_areal_fraction.field[:] = self.unmodified_state["updraft_areal_fraction_bugworkaroundname"] + # state.updraft_vertical_velocity.field[:] = self.unmodified_state["updraft_vertical_velocity_bugworkaroundname"] + state.dtdt_shortwave.field[:] = self.unmodified_state["dtdt_shortwave_bugworkaroundname"] + state.dtdt_longwave.field[:] = self.unmodified_state["dtdt_longwave_bugworkaroundname"] + state.dspecific_humiditydt_pbl.field[:] = self.unmodified_state[ + "dspecific_humiditydt_pbl_bugworkaroundname" + ] + state.dtdt_pbl.field[:] = self.unmodified_state["dtdt_pbl_bugworkaroundname"] + state.dtdt_from_dynamics.field[:] = self.unmodified_state["dtdt_from_dynamics_bugworkaroundname"] + state.dvapordt_from_dynamics.field[:] = self.unmodified_state[ + "dvapordt_from_dynamics_bugworkaroundname" + ] + # state.sigma_mid.field[:] = self.unmodified_state["sigma_mid_bugworkaroundname"] + # state.sigma_deep.field[:] = self.unmodified_state["sigma_deep_bugworkaroundname"] + state.total_precipitable_water_initial.field[:] = self.unmodified_state[ + "total_precipitable_water_initial_bugworkaroundname" + ] + state.saturation_total_precipitable_water_initial.field[:] = self.unmodified_state[ + "saturation_total_precipitable_water_initial_bugworkaroundname" + ] + # state.dvapordt_deep_convection.field[:] = self.unmodified_state["dvapordt_deep_convection_bugworkaroundname"] + # state.dtdt_deep_convection.field[:] = self.unmodified_state["dtdt_deep_convection_bugworkaroundname"] + # state.dudt_deep_convection.field[:] = self.unmodified_state["dudt_deep_convection_bugworkaroundname"] + # state.dvdt_deep_convection.field[:] = self.unmodified_state["dvdt_deep_convection_bugworkaroundname"] + # state.dliquiddt_deep_convection.field[:] = self.unmodified_state["dliquiddt_deep_convection_bugworkaroundname"] + # state.dicedt_deep_convection.field[:] = self.unmodified_state["dicedt_deep_convection_bugworkaroundname"] + # state.dcloudfractiondt_deep_convection.field[:] = self.unmodified_state["dcloudfractiondt_deep_convection_bugworkaroundname"] + # state.pressure_shallow_convective_cloud_top.field[:] = self.unmodified_state["pressure_shallow_convective_cloud_top_bugworkaroundname"] + # state.pressure_mid_convective_cloud_top.field[:] = self.unmodified_state["pressure_mid_convective_cloud_top_bugworkaroundname"] + # state.pressure_deep_convective_cloud_top.field[:] = self.unmodified_state["pressure_deep_convective_cloud_top_bugworkaroundname"] + # state.mass_flux_shallow.field[:] = self.unmodified_state["mass_flux_shallow_bugworkaroundname"] + # state.mass_flux_mid.field[:] = self.unmodified_state["mass_flux_mid_bugworkaroundname"] + # state.mass_flux_deep_updraft.field[:] = self.unmodified_state["mass_flux_deep_updraft_bugworkaroundname"] + # state.mass_flux_deep_updraft_interface.field[:] = self.unmodified_state["mass_flux_deep_updraft_interface_bugworkaroundname"] + # state.mass_flux_deep_updraft_detrained.field[:] = self.unmodified_state["mass_flux_deep_updraft_detrained_bugworkaroundname"] + # state.mass_flux_deep_downdraft.field[:] = self.unmodified_state["mass_flux_deep_downdraft_bugworkaroundname"] + # state.mass_flux_cloud_base.field[:] = self.unmodified_state["mass_flux_cloud_base_bugworkaroundname"] + # state.mass_flux_cloud_base_shallow.field[:] = self.unmodified_state["mass_flux_cloud_base_shallow_bugworkaroundname"] + # state.mass_flux_cloud_base_mid.field[:] = self.unmodified_state["mass_flux_cloud_base_mid_bugworkaroundname"] + # state.mass_flux_cloud_base_deep.field[:] = self.unmodified_state["mass_flux_cloud_base_deep_bugworkaroundname"] + # state.convection_code_shallow.field[:] = self.unmodified_state["convection_code_shallow_bugworkaroundname"] + # state.convection_code_mid.field[:] = self.unmodified_state["convection_code_mid_bugworkaroundname"] + # state.convection_code_deep.field[:] = self.unmodified_state["convection_code_deep_bugworkaroundname"] + # state.cloud_workfunction_0.field[:] = self.unmodified_state["cloud_workfunction_0_bugworkaroundname"] + # state.cloud_workfunction_1.field[:] = self.unmodified_state["cloud_workfunction_1_bugworkaroundname"] + # state.cloud_workfunction_2.field[:] = self.unmodified_state["cloud_workfunction_2_bugworkaroundname"] + # state.cloud_workfunction_3.field[:] = self.unmodified_state["cloud_workfunction_3_bugworkaroundname"] + # state.cloud_workfunction_1_pbl.field[:] = self.unmodified_state["cloud_workfunction_1_pbl_bugworkaroundname"] + # state.cloud_workfunction_1_cin.field[:] = self.unmodified_state["cloud_workfunction_1_cin_bugworkaroundname"] + # state.pbl_time_scale.field[:] = self.unmodified_state["pbl_time_scale_bugworkaroundname"] + # state.cape_removal_time_scale.field[:] = self.unmodified_state["cape_removal_time_scale_bugworkaroundname"] + # state.lightning_density.field[:] = self.unmodified_state["lightning_density_bugworkaroundname"] + state.convection_tracer.field[:] = self.unmodified_state["convection_tracer_bugworkaroundname"] + + # fill locals with input data + locals.derived_state.edge_height_above_surface.field[:] = inputs["local_edge_height_above_surface"] + locals.derived_state.layer_height_above_surface.field[:] = inputs["local_layer_height_above_surface"] + locals.derived_state.p.field[:] = inputs["local_p"] + locals.derived_state.p_kappa.field[:] = inputs["local_p_kappa"] + locals.derived_state.th.field[:] = inputs["local_th"] + locals.derived_state.mass.field[:] = inputs["local_mass"] + locals.derived_state.modified_area.field[:] = inputs["local_modified_area"] + locals.derived_state.vertical_velocity.field[:] = inputs["local_vertical_velocity"] + locals.derived_state.dz.field[:] = inputs["local_dz"] + locals.derived_state.air_density.field[:] = inputs["local_air_density"] + locals.derived_state.scalar_diffusivity.field[:] = np.moveaxis( + inputs["local_scalar_diffusivity"], 0, 2 + ) + locals.flipped_copy.vapor_current.field[:] = np.moveaxis(inputs["local_vapor_current"], 0, 2) + locals.flipped_copy.p.field[:] = np.moveaxis(inputs["local_p_flipped"], 0, 2) + locals.flipped_copy.vapor.field[:] = np.moveaxis(inputs["local_vapor_flipped"], 0, 2) + + # fill CumlulusParameterizaiton State with input data + # input fields + cumulus_parameterization_state.input.t_excess.field[:] = self.manual_inputs["t_excess"] + cumulus_parameterization_state.input.vapor_excess.field[:] = self.manual_inputs["vapor_excess"] + cumulus_parameterization_state.input.grid_scale_forcing_t.field[:] = self.manual_inputs[ + "grid_scale_forcing_t" + ] + cumulus_parameterization_state.input.grid_scale_forcing_vapor.field[:] = self.manual_inputs[ + "grid_scale_forcing_vapor" + ] + cumulus_parameterization_state.input.subgrid_scale_forcing_t.field[:] = self.manual_inputs[ + "subgrid_scale_forcing_t" + ] + cumulus_parameterization_state.input.subgrid_scale_forcing_vapor.field[:] = self.manual_inputs[ + "subgrid_scale_forcing_vapor" + ] + cumulus_parameterization_state.input.seed_convection.field[:] = self.manual_inputs["seed_convection"] + cumulus_parameterization_state.input.saturation_water_vapor.field[:] = self.manual_inputs[ + "saturation_water_vapor" + ] + cumulus_parameterization_state.input.ocean_fraction.field[:] = self.manual_inputs["ocean_fraction"] + cumulus_parameterization_state.input.convection_fraction.field[:] = self.manual_inputs[ + "convection_fraction" + ] + cumulus_parameterization_state.input.surface_type.field[:] = self.manual_inputs["surface_type"] + cumulus_parameterization_state.input.lateral_entrainment_rate.field[:] = self.manual_inputs[ + "lateral_entrainment_rate" + ] + cumulus_parameterization_state.input.last_error_code.field[:] = self.manual_inputs["last_error_code"] + # output fields + cumulus_parameterization_state.output.dtdt.field[:] = self.manual_inputs["dtdt"][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dvapordt.field[:] = self.manual_inputs["dvapordt"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.dcloudicedt.field[:] = self.manual_inputs["dcloudicedt"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.dudt.field[:] = self.manual_inputs["dudt"][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dvdt.field[:] = self.manual_inputs["dvdt"][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dnliquiddt.field[:] = self.manual_inputs["dnliquiddt"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.dnicedt.field[:] = self.manual_inputs["dnicedt"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.dbuoyancydt.field[:] = self.manual_inputs["dbuoyancydt"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.dconvectiveicedt.field[:] = self.manual_inputs[ + "dconvectiveicedt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dlargescaleicedt.field[:] = self.manual_inputs[ + "dlargescaleicedt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dconvectiveliquiddt.field[:] = self.manual_inputs[ + "dconvectiveliquiddt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dlargescaleliquiddt.field[:] = self.manual_inputs[ + "dlargescaleliquiddt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dconvectivecloudfractiondt.field[:] = self.manual_inputs[ + "dconvectivecloudfractiondt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.dlargescalecloudfractiondt.field[:] = self.manual_inputs[ + "dlargescalecloudfractiondt" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.error_code.field[:] = self.manual_inputs["error_code"][ + :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.downdraft_origin_level.field[:] = ( + self.manual_inputs["downdraft_origin_level"][:, :, [1, 2, 0]] - 1 + ) + cumulus_parameterization_state.output.lcl_level.field[:] = ( + self.manual_inputs["lcl_level"][:, :, [1, 2, 0]] - 1 + ) + cumulus_parameterization_state.output.updraft_origin_level.field[:] = ( + self.manual_inputs["updraft_origin_level"][:, :, [1, 2, 0]] - 1 + ) + cumulus_parameterization_state.output.updraft_lfc_level.field[:] = ( + self.manual_inputs["updraft_lfc_level"][:, :, [1, 2, 0]] - 1 + ) + cumulus_parameterization_state.output.cloud_top_level.field[:] = ( + self.manual_inputs["cloud_top_level"][:, :, [1, 2, 0]] - 1 + ) + cumulus_parameterization_state.output.kstabi.field[:] = self.manual_inputs["kstabi"][:, :, [1, 2, 0]] + cumulus_parameterization_state.output.kstabm.field[:] = self.manual_inputs["kstabm"][:, :, [1, 2, 0]] + cumulus_parameterization_state.output.precip.field[:] = self.manual_inputs["precip"][:, :, [1, 2, 0]] + cumulus_parameterization_state.output.cloud_base_mass_flux_modified.field[:] = self.manual_inputs[ + "cloud_base_mass_flux_modified" + ][:, :, [1, 2, 0]] + cumulus_parameterization_state.output.epsilon_forced.field[:] = self.manual_inputs["epsilon_forced"][ + :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.total_normalized_integrated_condensate_forced.field[:] = ( + self.manual_inputs["total_normalized_integrated_condensate_forced"][:, :, [1, 2, 0]] + ) + cumulus_parameterization_state.output.scale_dependence_factor.field[:] = self.manual_inputs[ + "scale_dependence_factor" + ][:, :, [1, 2, 0]] + cumulus_parameterization_state.output.p_cloud_levels_forced.field[:] = self.manual_inputs[ + "p_cloud_levels_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.entrainment_rate.field[:] = self.manual_inputs[ + "entrainment_rate" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.mass_entrainment_updraft_forced.field[:] = self.manual_inputs[ + "mass_entrainment_updraft_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.mass_entrainment_downdraft_forced.field[:] = self.manual_inputs[ + "mass_entrainment_downdraft_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.mass_detrainment_updraft_forced.field[:] = self.manual_inputs[ + "mass_detrainment_updraft_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.mass_detrainment_downdraft_forced.field[:] = self.manual_inputs[ + "mass_detrainment_downdraft_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.normalized_massflux_updraft_forced.field[:] = ( + self.manual_inputs["normalized_massflux_updraft_forced"][:, :, :, [1, 2, 0]] + ) + cumulus_parameterization_state.output.normalized_massflux_downdraft_forced.field[:] = ( + self.manual_inputs["normalized_massflux_downdraft_forced"][:, :, :, [1, 2, 0]] + ) + cumulus_parameterization_state.output.condensate_to_fall_forced.field[:] = self.manual_inputs[ + "condensate_to_fall_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.evaporate_in_downdraft_forced.field[:] = self.manual_inputs[ + "evaporate_in_downdraft_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.cloud_liquid_after_rain_forced.field[:] = self.manual_inputs[ + "cloud_liquid_after_rain_forced" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.t_updraft.field[:] = self.manual_inputs["t_updraft"][ + :, :, :, [1, 2, 0] + ] + cumulus_parameterization_state.output.convective_cloud_fraction.field[:] = self.manual_inputs[ + "convective_cloud_fraction_output" + ][:, :, :, [1, 2, 0]] + cumulus_parameterization_state.output.cloud_workfunction_0.field[:] = self.manual_inputs[ + "cloud_workfunction_0" + ] + cumulus_parameterization_state.output.cloud_workfunction_1.field[:] = self.manual_inputs[ + "cloud_workfunction_1" + ] + cumulus_parameterization_state.output.cloud_workfunction_2.field[:] = self.manual_inputs[ + "cloud_workfunction_2" + ] + cumulus_parameterization_state.output.cloud_workfunction_3.field[:] = self.manual_inputs[ + "cloud_workfunction_3" + ] + cumulus_parameterization_state.output.cloud_workfunction_1_pbl.field[:] = self.manual_inputs[ + "cloud_workfunction_1_pbl" + ] + cumulus_parameterization_state.output.cloud_workfunction_1_cin.field[:] = self.manual_inputs[ + "cloud_workfunction_1_cin" + ] + cumulus_parameterization_state.output.cape_removal_time_scale.field[:] = self.manual_inputs[ + "cape_removal_time_scale" + ] + cumulus_parameterization_state.output.pbl_time_scale.field[:] = self.manual_inputs["pbl_time_scale"] + cumulus_parameterization_state.output.lightning_density.field[:] = self.manual_inputs[ + "lightning_density" + ] + cumulus_parameterization_state.output.evaporation_sublimation_tendency.field[:] = self.manual_inputs[ + "evaporation_sublimation_tendency" + ] + cumulus_parameterization_state.output.convective_precip_flux.field[:] = self.manual_inputs[ + "convective_precip_flux" + ] + cumulus_parameterization_state.output.t_perturbation.field[:] = self.manual_inputs["t_perturbation"] + # input/output fields + cumulus_parameterization_state.input_output.grid_length.field[:] = self.manual_inputs["grid_length"] + cumulus_parameterization_state.input_output.pbl_level.field[:] = self.manual_inputs["pbl_level"] - 1 + cumulus_parameterization_state.input_output.ccn.field[:] = self.manual_inputs["ccn"] + cumulus_parameterization_state.input_output.air_density.field[:] = self.manual_inputs["air_density"] + cumulus_parameterization_state.input_output.omega.field[:] = self.manual_inputs["omega"] + cumulus_parameterization_state.input_output.topography_height_no_negative.field[:] = ( + self.manual_inputs["topography_height_no_negative"] + ) + cumulus_parameterization_state.input_output.sensible_heat_flux.field[:] = self.manual_inputs[ + "sensible_heat_flux" + ] + cumulus_parameterization_state.input_output.latent_heat_flux.field[:] = self.manual_inputs[ + "latent_heat_flux" + ] + cumulus_parameterization_state.input_output.longitude_degrees.field[:] = self.manual_inputs[ + "longitude_degrees" + ] + cumulus_parameterization_state.input_output.latitude_degrees.field[:] = self.manual_inputs[ + "latitude_degrees" + ] + cumulus_parameterization_state.input_output.t_old.field[:] = self.manual_inputs["t_old"] + cumulus_parameterization_state.input_output.vapor_old.field[:] = self.manual_inputs["vapor_old"] + cumulus_parameterization_state.input_output.t_modified_by_advection.field[:] = self.manual_inputs[ + "t_modified_by_advection" + ] + cumulus_parameterization_state.input_output.vapor_modified_by_advection.field[:] = self.manual_inputs[ + "vapor_modified_by_advection" + ] + cumulus_parameterization_state.input_output.geopotential_height_forced.field[:] = self.manual_inputs[ + "geopotential_height_forced" + ] + cumulus_parameterization_state.input_output.p_forced.field[:] = self.manual_inputs["p_forced"] + cumulus_parameterization_state.input_output.p_surface.field[:] = self.manual_inputs["p_surface"] + cumulus_parameterization_state.input_output.t_surface.field[:] = self.manual_inputs["t_surface"] + cumulus_parameterization_state.input_output.u.field[:] = self.manual_inputs["u"] + cumulus_parameterization_state.input_output.v.field[:] = self.manual_inputs["v"] + cumulus_parameterization_state.input_output.w.field[:] = self.manual_inputs["w"] + cumulus_parameterization_state.input_output.mass.field[:] = self.manual_inputs["mass"] + cumulus_parameterization_state.input_output.convective_scale_velocity.field[:] = self.manual_inputs[ + "convective_scale_velocity" + ] + cumulus_parameterization_state.input_output.buoyancy_excess.field[:] = self.manual_inputs[ + "buoyancy_excess" + ] + cumulus_parameterization_state.input_output.large_scale_ice.field[:] = self.manual_inputs[ + "large_scale_ice" + ] + cumulus_parameterization_state.input_output.convective_ice.field[:] = self.manual_inputs[ + "convective_ice" + ] + cumulus_parameterization_state.input_output.large_scale_liquid.field[:] = self.manual_inputs[ + "large_scale_liquid" + ] + cumulus_parameterization_state.input_output.convective_liquid.field[:] = self.manual_inputs[ + "convective_liquid" + ] + cumulus_parameterization_state.input_output.large_scale_cloud_fraction.field[:] = self.manual_inputs[ + "large_scale_cloud_fraction" + ] + cumulus_parameterization_state.input_output.convective_cloud_fraction.field[:] = self.manual_inputs[ + "convective_cloud_fraction" + ] + cumulus_parameterization_state.input_output.chemistry_tracers.field[:] = self.manual_inputs[ + "chemistry_tracers" + ] + chemistry_tracers_input_5d = np.full( + cumulus_parameterization_state.input_output.chemistry_tracers_output.field[:].shape, np.nan + ) + for plume in range(NUMBER_OF_PLUMES): + chemistry_tracers_input_5d[:, :, :, plume, :] = self.manual_inputs["chemistry_tracers_output"][ + :, + :, + :, + plume * config.NUMBER_OF_TRACERS : plume * config.NUMBER_OF_TRACERS + + config.NUMBER_OF_TRACERS, + ] + cumulus_parameterization_state.input_output.chemistry_tracers_output.field[:] = ( + chemistry_tracers_input_5d[:, :, :, [1, 2, 0], :] + ) + + # initialize convection tracers + convection_tracers = ConvectionTracers.ones( + self.quantity_factory, + data_dimensions={ + "convection_tracers": config.NUMBER_OF_TRACERS, + "size_three_dimension": 3, + "size_four_dimension": 4, + }, + ) + + convection_tracers.tracers.field[:] = np.moveaxis(self.convection_tracers_input["tracers"], 0, 3) + convection_tracers.vect_hcts.field[:] = self.convection_tracers_input["vect_hcts"] + convection_tracers.kc_scal.field[:] = self.convection_tracers_input["kc_scal"] + convection_tracers.fscav.field[:] = self.convection_tracers_input["fscav"] + convection_tracers.convfaci2g.field[:] = self.convection_tracers_input["convfaci2g"] + convection_tracers.retfactor.field[:] = self.convection_tracers_input["retfactor"] + convection_tracers.liq_and_gas.field[:] = self.convection_tracers_input["liq_and_gas"] + convection_tracers.online_cldliq.field[:] = self.convection_tracers_input["online_cldliq"] + convection_tracers.online_vud.field[:] = self.convection_tracers_input["online_vud"] + convection_tracers.ftemp_threshold.field[:] = self.convection_tracers_input["ftemp_threshold"] + convection_tracers.use_gcc_washout.field[:] = self.convection_tracers_input["use_gcc_washout"] + convection_tracers.use_gocart.field[:] = self.convection_tracers_input["use_gocart"] + convection_tracers.is_wetdep.field[:] = self.convection_tracers_input["is_wetdep"] + + saturation_tables = SaturationVaporPressureTable(self.stencil_factory.backend) + + code = GF2020Finalize( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + cumulus_parameterization_config=cumulus_parameterization_config, + saturation_tables=saturation_tables, + ) + + code( + state=state, + locals=locals, + cumulus_parameterization_state=cumulus_parameterization_state, + convection_tracers=convection_tracers, + ) + + # fill output dictionary for testing + outputs = { + "latitude_bugworkaroundname": state.latitude.field[:], + "longitude_bugworkaroundname": state.longitude.field[:], + "p_interface_bugworkaroundname": state.p_interface.field[:], + "t_bugworkaroundname": state.t.field[:], + "u_bugworkaroundname": state.u.field[:], + "v_bugworkaroundname": state.v.field[:], + "w_bugworkaroundname": state.w.field[:], + "omega_bugworkaroundname": state.omega.field[:], + "t_2m_bugworkaroundname": state.t_2m.field[:], + "specific_humidity_2m_bugworkaroundname": state.specific_humidity_2m.field[:], + "t_surface_bugworkaroundname": state.t_surface.field[:], + "specific_humidity_surface_bugworkaroundname": state.specific_humidity_surface.field[:], + "vapor_bugworkaroundname": state.vapor.field[:], + "convective_liquid_bugworkaroundname": state.convective_liquid.field[:], + "convective_ice_bugworkaroundname": state.convective_ice.field[:], + "large_scale_liquid_bugworkaroundname": state.large_scale_liquid.field[:], + "large_scale_ice_bugworkaroundname": state.large_scale_ice.field[:], + "convective_cloud_fraction_bugworkaroundname": state.convective_cloud_fraction.field[:], + "large_scale_cloud_fraction_bugworkaroundname": state.large_scale_cloud_fraction.field[:], + "p_interface_timestep_start_bugworkaroundname": state.p_interface_timestep_start.field[:], + "t_timestep_start_bugworkaroundname": state.t_timestep_start.field[:], + "u_timestep_start_bugworkaroundname": state.u_timestep_start.field[:], + "v_timestep_start_bugworkaroundname": state.v_timestep_start.field[:], + "vapor_timestep_start_bugworkaroundname": state.vapor_timestep_start.field[:], + "geopotential_height_interface_bugworkaroundname": state.geopotential_height_interface.field[:], + "geopotential_height_surface_bugworkaroundname": state.geopotential_height_surface.field[:], + "area_bugworkaroundname": state.area.field[:], + "pbl_level_bugworkaroundname": state.pbl_level.field[:], + "convection_fraction_bugworkaroundname": state.convection_fraction.field[:], + "surface_type_bugworkaroundname": state.surface_type.field[:], + "seed_convection_bugworkaroundname": state.seed_convection.field[:], + "land_fraction_bugworkaroundname": state.land_fraction.field[:], + "scalar_diffusivity_bugworkaroundname": state.scalar_diffusivity.field[:], + "buoyancy_bugworkaroundname": state.buoyancy.field[:], + "convective_precipitation_GF_bugworkaroundname": state.convective_precipitation_GF.field[:], + "convective_precipitation_RAS_bugworkaroundname": state.convective_precipitation_RAS.field[:], + "sensible_heat_flux_bugworkaroundname": state.sensible_heat_flux.field[:], + "total_water_flux_deep_convection_interface_bugworkaroundname": state.total_water_flux_deep_convection_interface.field[ + : + ], + "evaporation_bugworkaroundname": state.evaporation.field[:], + "sublimation_of_convective_precipitation_bugworkaroundname": state.sublimation_of_convective_precipitation.field[ + : + ], + "evaporation_of_convective_precipitation_bugworkaroundname": state.evaporation_of_convective_precipitation.field[ + : + ], + "ice_precip_flux_interface_bugworkaroundname": state.ice_precip_flux_interface.field[:], + "liquid_precip_flux_interface_bugworkaroundname": state.liquid_precip_flux_interface.field[:], + "convective_condensate_source_bugworkaroundname": state.convective_condensate_source.field[:], + "convective_condensate_grid_mean_bugworkaroundname": state.convective_condensate_grid_mean.field[ + : + ], + "entrainment_parameter_bugworkaroundname": state.entrainment_parameter.field[:], + "lateral_entrainment_rate_bugworkaroundname": state.lateral_entrainment_rate.field[:], + "lateral_entrainment_rate_shallow_bugworkaroundname": state.lateral_entrainment_rate_shallow.field[ + : + ], + "lateral_entrainment_rate_mid_bugworkaroundname": state.lateral_entrainment_rate_mid.field[:], + "lateral_entrainment_rate_deep_bugworkaroundname": state.lateral_entrainment_rate_deep.field[:], + "updraft_areal_fraction_bugworkaroundname": state.updraft_areal_fraction.field[:], + "updraft_vertical_velocity_bugworkaroundname": state.updraft_vertical_velocity.field[:], + "dtdt_shortwave_bugworkaroundname": state.dtdt_shortwave.field[:], + "dtdt_longwave_bugworkaroundname": state.dtdt_longwave.field[:], + "dspecific_humiditydt_pbl_bugworkaroundname": state.dspecific_humiditydt_pbl.field[:], + "dtdt_pbl_bugworkaroundname": state.dtdt_pbl.field[:], + "dtdt_from_dynamics_bugworkaroundname": state.dtdt_from_dynamics.field[:], + "dvapordt_from_dynamics_bugworkaroundname": state.dvapordt_from_dynamics.field[:], + "sigma_mid_bugworkaroundname": state.sigma_mid.field[:], + "sigma_deep_bugworkaroundname": state.sigma_deep.field[:], + "total_precipitable_water_initial_bugworkaroundname": state.total_precipitable_water_initial.field[ + : + ], + "saturation_total_precipitable_water_initial_bugworkaroundname": state.saturation_total_precipitable_water_initial.field[ + : + ], + "dvapordt_deep_convection_bugworkaroundname": state.dvapordt_deep_convection.field[:], + "dtdt_deep_convection_bugworkaroundname": state.dtdt_deep_convection.field[:], + "dudt_deep_convection_bugworkaroundname": state.dudt_deep_convection.field[:], + "dvdt_deep_convection_bugworkaroundname": state.dvdt_deep_convection.field[:], + "dliquiddt_deep_convection_bugworkaroundname": state.dliquiddt_deep_convection.field[:], + "dicedt_deep_convection_bugworkaroundname": state.dicedt_deep_convection.field[:], + "dcloudfractiondt_deep_convection_bugworkaroundname": state.dcloudfractiondt_deep_convection.field[ + : + ], + "pressure_shallow_convective_cloud_top_bugworkaroundname": state.pressure_shallow_convective_cloud_top.field[ + : + ], + "pressure_mid_convective_cloud_top_bugworkaroundname": state.pressure_mid_convective_cloud_top.field[ + : + ], + "pressure_deep_convective_cloud_top_bugworkaroundname": state.pressure_deep_convective_cloud_top.field[ + : + ], + "mass_flux_shallow_bugworkaroundname": state.mass_flux_shallow.field[:], + "mass_flux_mid_bugworkaroundname": state.mass_flux_mid.field[:], + "mass_flux_deep_updraft_bugworkaroundname": state.mass_flux_deep_updraft.field[:], + "mass_flux_deep_updraft_interface_bugworkaroundname": state.mass_flux_deep_updraft_interface.field[ + : + ], + "mass_flux_deep_updraft_detrained_bugworkaroundname": state.mass_flux_deep_updraft_detrained.field[ + : + ], + "mass_flux_deep_downdraft_bugworkaroundname": state.mass_flux_deep_downdraft.field[:], + "mass_flux_cloud_base_bugworkaroundname": state.mass_flux_cloud_base.field[:], + "mass_flux_cloud_base_shallow_bugworkaroundname": state.mass_flux_cloud_base_shallow.field[:], + "mass_flux_cloud_base_mid_bugworkaroundname": state.mass_flux_cloud_base_mid.field[:], + "mass_flux_cloud_base_deep_bugworkaroundname": state.mass_flux_cloud_base_deep.field[:], + "total_cumulative_mass_flux_interface_bugworkaroundname": state.total_cumulative_mass_flux_interface.field[ + : + ], + "total_detraining_mass_flux_bugworkaroundname": state.total_detraining_mass_flux.field[:], + "convection_code_shallow_bugworkaroundname": state.convection_code_shallow.field[:], + "convection_code_mid_bugworkaroundname": state.convection_code_mid.field[:], + "convection_code_deep_bugworkaroundname": state.convection_code_deep.field[:], + "cloud_workfunction_0_bugworkaroundname": state.cloud_workfunction_0.field[:], + "cloud_workfunction_1_bugworkaroundname": state.cloud_workfunction_1.field[:], + "cloud_workfunction_2_bugworkaroundname": state.cloud_workfunction_2.field[:], + "cloud_workfunction_3_bugworkaroundname": state.cloud_workfunction_3.field[:], + "cloud_workfunction_1_pbl_bugworkaroundname": state.cloud_workfunction_1_pbl.field[:], + "cloud_workfunction_1_cin_bugworkaroundname": state.cloud_workfunction_1_cin.field[:], + "pbl_time_scale_bugworkaroundname": state.pbl_time_scale.field[:], + "cape_removal_time_scale_bugworkaroundname": state.cape_removal_time_scale.field[:], + "lighting_density_bugworkaroundname": state.lightning_density.field[:], + "convection_tracer_bugworkaroundname": state.convection_tracer.field[:], + } + + return outputs diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Setup.py new file mode 100644 index 000000000..a7d4ae2ba --- /dev/null +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/convection/GF_2020/translate_GF2020_Setup.py @@ -0,0 +1,650 @@ +import numpy as np +from f90nml import Namelist + +from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing.savepoint import DataLoader +from ndsl.stencils.testing.translate import TranslateFortranData2Py +from pyMoist.constants import NUMBER_OF_TRACERS +from pyMoist.convection.GF_2020.config import GF2020Config +from pyMoist.convection.GF_2020.cumulus_parameterization.constants import NUMBER_OF_PLUMES +from pyMoist.convection.GF_2020.cumulus_parameterization.state import GF2020CumulusParameterizationState +from pyMoist.convection.GF_2020.locals import GF2020Locals +from pyMoist.convection.GF_2020.setup import GF2020Setup +from pyMoist.convection.GF_2020.state import GF2020State +from pyMoist.convection_tracers import ConvectionTracers +from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable + + +class TranslateGF2020_Setup(TranslateFortranData2Py): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "latitude_bugworkaroundname": {}, + "longitude_bugworkaroundname": {}, + "p_interface_bugworkaroundname": {}, + "t_bugworkaroundname": {}, + "u_bugworkaroundname": {}, + "v_bugworkaroundname": {}, + "w_bugworkaroundname": {}, + "omega_bugworkaroundname": {}, + "t_2m_bugworkaroundname": {}, + "specific_humidity_2m_bugworkaroundname": {}, + "t_surface_bugworkaroundname": {}, + "specific_humidity_surface_bugworkaroundname": {}, + "vapor_bugworkaroundname": {}, + "convective_liquid_bugworkaroundname": {}, + "convective_ice_bugworkaroundname": {}, + "large_scale_liquid_bugworkaroundname": {}, + "large_scale_ice_bugworkaroundname": {}, + "convective_cloud_fraction_bugworkaroundname": {}, + "large_scale_cloud_fraction_bugworkaroundname": {}, + "p_interface_timestep_start_bugworkaroundname": {}, + "t_timestep_start_bugworkaroundname": {}, + "u_timestep_start_bugworkaroundname": {}, + "v_timestep_start_bugworkaroundname": {}, + "vapor_timestep_start_bugworkaroundname": {}, + "geopotential_height_interface_bugworkaroundname": {}, + "geopotential_height_surface_bugworkaroundname": {}, + "area_bugworkaroundname": {}, + "pbl_level_bugworkaroundname": {}, + "convection_fraction_bugworkaroundname": {}, + "surface_type_bugworkaroundname": {}, + "seed_convection_bugworkaroundname": {}, + "land_fraction_bugworkaroundname": {}, + "scalar_diffusivity_bugworkaroundname": {}, + "buoyancy_bugworkaroundname": {}, + "convective_precipitation_GF_bugworkaroundname": {}, + "convective_precipitation_RAS_bugworkaroundname": {}, + "sensible_heat_flux_bugworkaroundname": {}, + "total_water_flux_deep_convection_interface_bugworkaroundname": {}, + "evaporation_bugworkaroundname": {}, + "convective_condensate_source_bugworkaroundname": {}, + "convective_condensate_grid_mean_bugworkaroundname": {}, + "entrainment_parameter_bugworkaroundname": {}, + "lateral_entrainment_rate_bugworkaroundname": {}, + "lateral_entrainment_rate_shallow_bugworkaroundname": {}, + "lateral_entrainment_rate_mid_bugworkaroundname": {}, + "lateral_entrainment_rate_deep_bugworkaroundname": {}, + "updraft_areal_fraction_bugworkaroundname": {}, + "updraft_vertical_velocity_bugworkaroundname": {}, + "dtdt_shortwave_bugworkaroundname": {}, + "dtdt_longwave_bugworkaroundname": {}, + "dspecific_humiditydt_pbl_bugworkaroundname": {}, + "dtdt_pbl_bugworkaroundname": {}, + "dtdt_from_dynamics_bugworkaroundname": {}, + "dvapordt_from_dynamics_bugworkaroundname": {}, + "sigma_mid_bugworkaroundname": {}, + "sigma_deep_bugworkaroundname": {}, + "total_precipitable_water_initial_bugworkaroundname": {}, + "saturation_total_precipitable_water_initial_bugworkaroundname": {}, + "dvapordt_deep_convection_bugworkaroundname": {}, + "dtdt_deep_convection_bugworkaroundname": {}, + "dudt_deep_convection_bugworkaroundname": {}, + "dvdt_deep_convection_bugworkaroundname": {}, + "pressure_shallow_convective_cloud_top_bugworkaroundname": {}, + "pressure_mid_convective_cloud_top_bugworkaroundname": {}, + "pressure_deep_convective_cloud_top_bugworkaroundname": {}, + "mass_flux_shallow_bugworkaroundname": {}, + "mass_flux_mid_bugworkaroundname": {}, + "mass_flux_deep_updraft_bugworkaroundname": {}, + "mass_flux_deep_updraft_interface_bugworkaroundname": {}, + "mass_flux_deep_updraft_detrained_bugworkaroundname": {}, + "mass_flux_deep_downdraft_bugworkaroundname": {}, + "mass_flux_cloud_base_bugworkaroundname": {}, + "mass_flux_cloud_base_shallow_bugworkaroundname": {}, + "mass_flux_cloud_base_mid_bugworkaroundname": {}, + "mass_flux_cloud_base_deep_bugworkaroundname": {}, + "convection_code_shallow_bugworkaroundname": {}, + "convection_code_mid_bugworkaroundname": {}, + "convection_code_deep_bugworkaroundname": {}, + "cloud_workfunction_0_bugworkaroundname": {}, + "cloud_workfunction_1_bugworkaroundname": {}, + "cloud_workfunction_2_bugworkaroundname": {}, + "cloud_workfunction_3_bugworkaroundname": {}, + "cloud_workfunction_1_pbl_bugworkaroundname": {}, + "cloud_workfunction_1_cin_bugworkaroundname": {}, + "pbl_time_scale_bugworkaroundname": {}, + "cape_removal_time_scale_bugworkaroundname": {}, + "lighting_density_bugworkaroundname": {}, + "convection_tracer_bugworkaroundname": {}, + } + + # NOTE disabled fields are nan in fortran - zero in python, disabled so the test passes + self.out_vars: dict = { + # fields saved midway through GF2020 fortran + "internal_lateral_entrainment_rate": {}, + "local_edge_height_above_surface": {}, + "local_layer_height_above_surface": {}, + "local_p": {}, + "local_p_kappa": {}, + "local_th": {}, + "local_mass": {}, + "local_modified_area": {}, + "local_vertical_velocity": {}, + "local_dz": {}, + "local_air_density": {}, + "local_scalar_diffusivity": {}, + "local_vapor_current": {}, + # CumulusParameterization state - input + "t_excess": {}, + "vapor_excess": {}, + "grid_scale_forcing_t": {}, + "grid_scale_forcing_vapor": {}, + "subgrid_scale_forcing_t": {}, + "subgrid_scale_forcing_vapor": {}, + "seed_convection": {}, + # "saturation_water_vapor": {}, + "ocean_fraction": {}, + "convection_fraction": {}, + "surface_type": {}, + "lateral_entrainment_rate": {}, + "last_error_code": {}, + # CumulusParameterization state - output + "dtdt": {}, + "dvapordt": {}, + "dcloudicedt": {}, + "dudt": {}, + "dvdt": {}, + "dnliquiddt": {}, + "dnicedt": {}, + "dbuoyancydt": {}, + # "dconvectiveicedt": {}, + # "dlargescaleicedt": {}, + # "dconvectiveliquiddt": {}, + # "dlargescaleliquiddt": {}, + # "dconvectivecloudfractiondt": {}, + # "dlargescalecloudfractiondt": {}, + "error_code": {}, + "downdraft_origin_level": {}, + "lcl_level": {}, + "updraft_origin_level": {}, + "updraft_lfc_level": {}, + "cloud_top_level": {}, + "kstabi": {}, + "kstabm": {}, + "precip": {}, + "cloud_base_mass_flux_modified": {}, + "epsilon_forced": {}, + "total_normalized_integrated_condensate_forced": {}, + "scale_dependence_factor": {}, + "p_cloud_levels_forced": {}, + "entrainment_rate": {}, + "mass_entrainment_updraft_forced": {}, + "mass_entrainment_downdraft_forced": {}, + "mass_detrainment_updraft_forced": {}, + "mass_detrainment_downdraft_forced": {}, + "normalized_massflux_updraft_forced": {}, + "normalized_massflux_downdraft_forced": {}, + "condensate_to_fall_forced": {}, + "evaporate_in_downdraft_forced": {}, + "cloud_liquid_after_rain_forced": {}, + "t_updraft": {}, + "convective_cloud_fraction_output": {}, + "cloud_workfunction_0": {}, + "cloud_workfunction_1": {}, + "cloud_workfunction_2": {}, + "cloud_workfunction_3": {}, + "cloud_workfunction_1_pbl": {}, + "cloud_workfunction_1_cin": {}, + "cape_removal_time_scale": {}, + "pbl_time_scale": {}, + "lightning_density": {}, + "evaporation_sublimation_tendency": {}, + "convective_precip_flux": {}, + "t_perturbation": {}, + # CumulusParameterization state - input-output + "grid_length": {}, + "pbl_level": {}, + "ccn": {}, + "air_density": {}, + "omega": {}, + "topography_height_no_negative": {}, + "sensible_heat_flux": {}, + "latent_heat_flux": {}, + "longitude_degrees": {}, + "latitude_degrees": {}, + "t_old": {}, + "vapor_old": {}, + "t_modified_by_advection": {}, + "vapor_modified_by_advection": {}, + "geopotential_height_forced": {}, + "p_forced": {}, + "p_surface": {}, + "t_surface": {}, + "u": {}, + "v": {}, + "w": {}, + "mass": {}, + "convective_scale_velocity": {}, + "buoyancy_excess": {}, + # "large_scale_ice": {}, + # "convective_ice": {}, + # "large_scale_liquid": {}, + # "convective_liquid": {}, + # "large_scale_cloud_fraction": {}, + # "convective_cloud_fraction": {}, + "chemistry_tracers": {}, + "chemistry_tracers_output": {}, + } + + def extra_data_load(self, data_loader: DataLoader): + self.constants = data_loader.load("GF2020-constants") + self.convection_tracers_input = data_loader.load("GF2020_ConvectionTracers") + + def compute(self, inputs): + config = GF2020Config(SINGLE_COLUMN_MODE=False, **self.constants) + + # initialize GF2020 state + state = GF2020State.zeros(self.quantity_factory) + + # initialize GF2020 CumulusParameterization state + cumulus_parameterization_state = GF2020CumulusParameterizationState.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": NUMBER_OF_PLUMES, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # fill GF2020 state with input data + state.latitude.field[:] = inputs["latitude_bugworkaroundname"] + state.longitude.field[:] = inputs["longitude_bugworkaroundname"] + state.p_interface.field[:] = inputs["p_interface_bugworkaroundname"] + state.t.field[:] = inputs["t_bugworkaroundname"] + state.u.field[:] = inputs["u_bugworkaroundname"] + state.v.field[:] = inputs["v_bugworkaroundname"] + state.w.field[:] = inputs["w_bugworkaroundname"] + state.omega.field[:] = inputs["omega_bugworkaroundname"] + state.t_2m.field[:] = inputs["t_2m_bugworkaroundname"] + state.specific_humidity_2m.field[:] = inputs["specific_humidity_2m_bugworkaroundname"] + state.t_surface.field[:] = inputs["t_surface_bugworkaroundname"] + state.specific_humidity_surface.field[:] = inputs["specific_humidity_surface_bugworkaroundname"] + state.vapor.field[:] = inputs["vapor_bugworkaroundname"] + state.convective_liquid.field[:] = inputs["convective_liquid_bugworkaroundname"] + state.convective_ice.field[:] = inputs["convective_ice_bugworkaroundname"] + state.large_scale_liquid.field[:] = inputs["large_scale_liquid_bugworkaroundname"] + state.large_scale_ice.field[:] = inputs["large_scale_ice_bugworkaroundname"] + state.convective_cloud_fraction.field[:] = inputs["convective_cloud_fraction_bugworkaroundname"] + state.large_scale_cloud_fraction.field[:] = inputs["large_scale_cloud_fraction_bugworkaroundname"] + state.p_interface_timestep_start.field[:] = inputs["p_interface_timestep_start_bugworkaroundname"] + state.t_timestep_start.field[:] = inputs["t_timestep_start_bugworkaroundname"] + state.u_timestep_start.field[:] = inputs["u_timestep_start_bugworkaroundname"] + state.v_timestep_start.field[:] = inputs["v_timestep_start_bugworkaroundname"] + state.vapor_timestep_start.field[:] = inputs["vapor_timestep_start_bugworkaroundname"] + state.geopotential_height_interface.field[:] = inputs[ + "geopotential_height_interface_bugworkaroundname" + ] + state.geopotential_height_surface.field[:] = inputs["geopotential_height_surface_bugworkaroundname"] + state.area.field[:] = inputs["area_bugworkaroundname"] + state.pbl_level.field[:] = inputs["pbl_level_bugworkaroundname"] - 1 + state.convection_fraction.field[:] = inputs["convection_fraction_bugworkaroundname"] + state.surface_type.field[:] = inputs["surface_type_bugworkaroundname"] + state.seed_convection.field[:] = inputs["seed_convection_bugworkaroundname"] + state.land_fraction.field[:] = inputs["land_fraction_bugworkaroundname"] + state.scalar_diffusivity.field[:] = inputs["scalar_diffusivity_bugworkaroundname"] + state.buoyancy.field[:] = inputs["buoyancy_bugworkaroundname"] + state.convective_precipitation_GF.field[:] = inputs["convective_precipitation_GF_bugworkaroundname"] + state.convective_precipitation_RAS.field[:] = inputs["convective_precipitation_RAS_bugworkaroundname"] + state.sensible_heat_flux.field[:] = inputs["sensible_heat_flux_bugworkaroundname"] + state.total_water_flux_deep_convection_interface.field[:] = inputs[ + "total_water_flux_deep_convection_interface_bugworkaroundname" + ] + state.evaporation.field[:] = inputs["evaporation_bugworkaroundname"] + state.convective_condensate_source.field[:] = inputs["convective_condensate_source_bugworkaroundname"] + state.convective_condensate_grid_mean.field[:] = inputs[ + "convective_condensate_grid_mean_bugworkaroundname" + ] + state.entrainment_parameter.field[:] = inputs["entrainment_parameter_bugworkaroundname"] + state.lateral_entrainment_rate.field[:] = inputs["lateral_entrainment_rate_bugworkaroundname"] + state.lateral_entrainment_rate_shallow.field[:] = inputs[ + "lateral_entrainment_rate_shallow_bugworkaroundname" + ] + state.lateral_entrainment_rate_mid.field[:] = inputs["lateral_entrainment_rate_mid_bugworkaroundname"] + state.lateral_entrainment_rate_deep.field[:] = inputs[ + "lateral_entrainment_rate_deep_bugworkaroundname" + ] + state.updraft_areal_fraction.field[:] = inputs["updraft_areal_fraction_bugworkaroundname"] + state.updraft_vertical_velocity.field[:] = inputs["updraft_vertical_velocity_bugworkaroundname"] + state.dtdt_shortwave.field[:] = inputs["dtdt_shortwave_bugworkaroundname"] + state.dtdt_longwave.field[:] = inputs["dtdt_longwave_bugworkaroundname"] + state.dspecific_humiditydt_pbl.field[:] = inputs["dspecific_humiditydt_pbl_bugworkaroundname"] + state.dtdt_pbl.field[:] = inputs["dtdt_pbl_bugworkaroundname"] + state.dtdt_from_dynamics.field[:] = inputs["dtdt_from_dynamics_bugworkaroundname"] + state.dvapordt_from_dynamics.field[:] = inputs["dvapordt_from_dynamics_bugworkaroundname"] + state.sigma_mid.field[:] = inputs["sigma_mid_bugworkaroundname"] + state.sigma_deep.field[:] = inputs["sigma_deep_bugworkaroundname"] + state.total_precipitable_water_initial.field[:] = inputs[ + "total_precipitable_water_initial_bugworkaroundname" + ] + state.saturation_total_precipitable_water_initial.field[:] = inputs[ + "saturation_total_precipitable_water_initial_bugworkaroundname" + ] + state.dvapordt_deep_convection.field[:] = inputs["dvapordt_deep_convection_bugworkaroundname"] + state.dtdt_deep_convection.field[:] = inputs["dtdt_deep_convection_bugworkaroundname"] + state.dudt_deep_convection.field[:] = inputs["dudt_deep_convection_bugworkaroundname"] + state.dvdt_deep_convection.field[:] = inputs["dvdt_deep_convection_bugworkaroundname"] + state.pressure_shallow_convective_cloud_top.field[:] = inputs[ + "pressure_shallow_convective_cloud_top_bugworkaroundname" + ] + state.pressure_mid_convective_cloud_top.field[:] = inputs[ + "pressure_mid_convective_cloud_top_bugworkaroundname" + ] + state.pressure_deep_convective_cloud_top.field[:] = inputs[ + "pressure_deep_convective_cloud_top_bugworkaroundname" + ] + state.mass_flux_shallow.field[:] = inputs["mass_flux_shallow_bugworkaroundname"] + state.mass_flux_mid.field[:] = inputs["mass_flux_mid_bugworkaroundname"] + state.mass_flux_deep_updraft.field[:] = inputs["mass_flux_deep_updraft_bugworkaroundname"] + state.mass_flux_deep_updraft_interface.field[:] = inputs[ + "mass_flux_deep_updraft_interface_bugworkaroundname" + ] + state.mass_flux_deep_updraft_detrained.field[:] = inputs[ + "mass_flux_deep_updraft_detrained_bugworkaroundname" + ] + state.mass_flux_deep_downdraft.field[:] = inputs["mass_flux_deep_downdraft_bugworkaroundname"] + state.mass_flux_cloud_base.field[:] = inputs["mass_flux_cloud_base_bugworkaroundname"] + state.mass_flux_cloud_base_shallow.field[:] = inputs["mass_flux_cloud_base_shallow_bugworkaroundname"] + state.mass_flux_cloud_base_mid.field[:] = inputs["mass_flux_cloud_base_mid_bugworkaroundname"] + state.mass_flux_cloud_base_deep.field[:] = inputs["mass_flux_cloud_base_deep_bugworkaroundname"] + state.convection_code_shallow.field[:] = inputs["convection_code_shallow_bugworkaroundname"] + state.convection_code_mid.field[:] = inputs["convection_code_mid_bugworkaroundname"] + state.convection_code_deep.field[:] = inputs["convection_code_deep_bugworkaroundname"] + state.cloud_workfunction_0.field[:] = inputs["cloud_workfunction_0_bugworkaroundname"] + state.cloud_workfunction_1.field[:] = inputs["cloud_workfunction_1_bugworkaroundname"] + state.cloud_workfunction_2.field[:] = inputs["cloud_workfunction_2_bugworkaroundname"] + state.cloud_workfunction_3.field[:] = inputs["cloud_workfunction_3_bugworkaroundname"] + state.cloud_workfunction_1_pbl.field[:] = inputs["cloud_workfunction_1_pbl_bugworkaroundname"] + state.cloud_workfunction_1_cin.field[:] = inputs["cloud_workfunction_1_cin_bugworkaroundname"] + state.pbl_time_scale.field[:] = inputs["pbl_time_scale_bugworkaroundname"] + state.cape_removal_time_scale.field[:] = inputs["cape_removal_time_scale_bugworkaroundname"] + state.lightning_density.field[:] = inputs["lighting_density_bugworkaroundname"] + state.convection_tracer.field[:] = inputs["convection_tracer_bugworkaroundname"] + + # initialize GF2020 locals + locals = GF2020Locals.zeros( + self.quantity_factory, + data_dimensions={ + "plumes": 3, + "convection_tracers": config.NUMBER_OF_TRACERS, + }, + ) + + # initialize convection tracers + convection_tracers = ConvectionTracers.ones( + self.quantity_factory, + data_dimensions={ + "convection_tracers": config.NUMBER_OF_TRACERS, + "size_three_dimension": 3, + "size_four_dimension": 4, + }, + ) + + convection_tracers.tracers.field[:] = np.moveaxis(self.convection_tracers_input["tracers"], 0, 3) + convection_tracers.vect_hcts.field[:] = self.convection_tracers_input["vect_hcts"] + convection_tracers.kc_scal.field[:] = self.convection_tracers_input["kc_scal"] + convection_tracers.fscav.field[:] = self.convection_tracers_input["fscav"] + convection_tracers.convfaci2g.field[:] = self.convection_tracers_input["convfaci2g"] + convection_tracers.retfactor.field[:] = self.convection_tracers_input["retfactor"] + convection_tracers.liq_and_gas.field[:] = self.convection_tracers_input["liq_and_gas"] + convection_tracers.online_cldliq.field[:] = self.convection_tracers_input["online_cldliq"] + convection_tracers.online_vud.field[:] = self.convection_tracers_input["online_vud"] + convection_tracers.ftemp_threshold.field[:] = self.convection_tracers_input["ftemp_threshold"] + convection_tracers.use_gcc_washout.field[:] = self.convection_tracers_input["use_gcc_washout"] + convection_tracers.use_gocart.field[:] = self.convection_tracers_input["use_gocart"] + convection_tracers.is_wetdep.field[:] = self.convection_tracers_input["is_wetdep"] + + saturation_tables = SaturationVaporPressureTable(self.stencil_factory.backend) + + code = GF2020Setup( + stencil_factory=self.stencil_factory, + quantity_factory=self.quantity_factory, + config=config, + saturation_tables=saturation_tables, + ) + + code( + state=state, + locals=locals, + cumulus_parameterization_state=cumulus_parameterization_state, + convection_tracers=convection_tracers, + scm_stop=False, + ) + + # fill output dictionary for testing + + # collapse plume dim for chemistry_tracers_output + # NOTE ideally this has no numpy dependency + chemistry_tracers_output_5d_reordered = ( + cumulus_parameterization_state.input_output.chemistry_tracers_output.field[:, :, :, [2, 0, 1], :] + ) + grid_size = self.stencil_factory.grid_indexing.get_shape([X_DIM, Y_DIM, Z_DIM]) + chemistry_tracers_output_4d = np.full( + [grid_size[0], grid_size[1], grid_size[2], NUMBER_OF_PLUMES * NUMBER_OF_TRACERS], np.nan + ) + for plume in range(NUMBER_OF_PLUMES): + chemistry_tracers_output_4d[ + :, + :, + :, + plume * config.NUMBER_OF_TRACERS : plume * config.NUMBER_OF_TRACERS + + config.NUMBER_OF_TRACERS, + ] = chemistry_tracers_output_5d_reordered[:, :, :, plume, :] + + # fill top level of a few fields with nans to make test pass + # these levels are NEVER read, so they don't need to be tested anyway + cumulus_parameterization_state.input_output.t_old.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.vapor_old.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.air_density.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.t_modified_by_advection.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.vapor_modified_by_advection.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.geopotential_height_forced.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.p_forced.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.u.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.v.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.w.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.mass.field[:, :, -1] = np.nan + cumulus_parameterization_state.input_output.buoyancy_excess.field[:, :, -1] = np.nan + + outputs = { + # GF2020 locals + "internal_lateral_entrainment_rate": state.lateral_entrainment_rate.field[:], + "local_edge_height_above_surface": locals.derived_state.edge_height_above_surface.field[:], + "local_layer_height_above_surface": locals.derived_state.layer_height_above_surface.field[:], + "local_p": locals.derived_state.p.field[:], + "local_p_kappa": locals.derived_state.p_kappa.field[:], + "local_th": locals.derived_state.th.field[:], + "local_mass": locals.derived_state.mass.field[:], + "local_modified_area": locals.derived_state.modified_area.field[:], + "local_vertical_velocity": locals.derived_state.vertical_velocity.field[:], + "local_dz": locals.derived_state.dz.field[:], + "local_air_density": locals.derived_state.air_density.field[:], + "local_scalar_diffusivity": np.moveaxis(locals.derived_state.scalar_diffusivity.field[:], 2, 0), + "local_vapor_current": np.moveaxis(locals.flipped_copy.vapor_current.field[:], 2, 0), + # GF2020 CumulusParameterization fields + # input fields + "t_excess": cumulus_parameterization_state.input.t_excess.field[:], + "vapor_excess": cumulus_parameterization_state.input.vapor_excess.field[:], + "grid_scale_forcing_t": cumulus_parameterization_state.input.grid_scale_forcing_t.field[:], + "grid_scale_forcing_vapor": cumulus_parameterization_state.input.grid_scale_forcing_vapor.field[ + : + ], + "subgrid_scale_forcing_t": cumulus_parameterization_state.input.subgrid_scale_forcing_t.field[:], + "subgrid_scale_forcing_vapor": cumulus_parameterization_state.input.subgrid_scale_forcing_vapor.field[ + : + ], + "seed_convection": cumulus_parameterization_state.input.seed_convection.field[:], + "saturation_water_vapor": cumulus_parameterization_state.input.saturation_water_vapor.field[:], + "ocean_fraction": cumulus_parameterization_state.input.ocean_fraction.field[:], + "convection_fraction": cumulus_parameterization_state.input.convection_fraction.field[:], + "surface_type": cumulus_parameterization_state.input.surface_type.field[:], + "lateral_entrainment_rate": cumulus_parameterization_state.input.lateral_entrainment_rate.field[ + : + ], + "last_error_code": cumulus_parameterization_state.input.last_error_code.field[:], + # output fields + "dtdt": cumulus_parameterization_state.output.dtdt.field[:, :, :, [2, 0, 1]], + "dvapordt": cumulus_parameterization_state.output.dvapordt.field[:, :, :, [2, 0, 1]], + "dcloudicedt": cumulus_parameterization_state.output.dcloudicedt.field[:, :, :, [2, 0, 1]], + "dudt": cumulus_parameterization_state.output.dudt.field[:, :, :, [2, 0, 1]], + "dvdt": cumulus_parameterization_state.output.dvdt.field[:, :, :, [2, 0, 1]], + "dnliquiddt": cumulus_parameterization_state.output.dnliquiddt.field[:, :, :, [2, 0, 1]], + "dnicedt": cumulus_parameterization_state.output.dnicedt.field[:, :, :, [2, 0, 1]], + "dbuoyancydt": cumulus_parameterization_state.output.dbuoyancydt.field[:, :, :, [2, 0, 1]], + "dconvectiveicedt": cumulus_parameterization_state.output.dconvectiveicedt.field[ + :, :, :, [2, 0, 1] + ], + "dlargescaleicedt": cumulus_parameterization_state.output.dlargescaleicedt.field[ + :, :, :, [2, 0, 1] + ], + "dconvectiveliquiddt": cumulus_parameterization_state.output.dconvectiveliquiddt.field[ + :, :, :, [2, 0, 1] + ], + "dlargescaleliquiddt": cumulus_parameterization_state.output.dlargescaleliquiddt.field[ + :, :, :, [2, 0, 1] + ], + "dconvectivecloudfractiondt": cumulus_parameterization_state.output.dconvectivecloudfractiondt.field[ + :, :, :, [2, 0, 1] + ], + "dlargescalecloudfractiondt": cumulus_parameterization_state.output.dlargescalecloudfractiondt.field[ + :, :, :, [2, 0, 1] + ], + "error_code": cumulus_parameterization_state.output.error_code.field[:, :, [2, 0, 1]], + "downdraft_origin_level": cumulus_parameterization_state.output.downdraft_origin_level.field[ + :, :, [2, 0, 1] + ], + "lcl_level": cumulus_parameterization_state.output.lcl_level.field[:, :, [2, 0, 1]], + "updraft_origin_level": cumulus_parameterization_state.output.updraft_origin_level.field[ + :, :, [2, 0, 1] + ], + "updraft_lfc_level": cumulus_parameterization_state.output.updraft_lfc_level.field[ + :, :, [2, 0, 1] + ], + "cloud_top_level": cumulus_parameterization_state.output.cloud_top_level.field[:, :, [2, 0, 1]], + "kstabi": cumulus_parameterization_state.output.kstabi.field[:, :, [2, 0, 1]], + "kstabm": cumulus_parameterization_state.output.kstabm.field[:, :, [2, 0, 1]], + "precip": cumulus_parameterization_state.output.precip.field[:, :, [2, 0, 1]], + "cloud_base_mass_flux_modified": cumulus_parameterization_state.output.cloud_base_mass_flux_modified.field[ + :, :, [2, 0, 1] + ], + "epsilon_forced": cumulus_parameterization_state.output.epsilon_forced.field[:, :, [2, 0, 1]], + "total_normalized_integrated_condensate_forced": cumulus_parameterization_state.output.total_normalized_integrated_condensate_forced.field[ + :, :, [2, 0, 1] + ], + "scale_dependence_factor": cumulus_parameterization_state.output.scale_dependence_factor.field[ + :, :, [2, 0, 1] + ], + "p_cloud_levels_forced": cumulus_parameterization_state.output.p_cloud_levels_forced.field[ + :, :, :, [2, 0, 1] + ], + "entrainment_rate": cumulus_parameterization_state.output.entrainment_rate.field[ + :, :, :, [2, 0, 1] + ], + "mass_entrainment_updraft_forced": cumulus_parameterization_state.output.mass_entrainment_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_entrainment_downdraft_forced": cumulus_parameterization_state.output.mass_entrainment_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_detrainment_updraft_forced": cumulus_parameterization_state.output.mass_detrainment_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "mass_detrainment_downdraft_forced": cumulus_parameterization_state.output.mass_detrainment_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "normalized_massflux_updraft_forced": cumulus_parameterization_state.output.normalized_massflux_updraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "normalized_massflux_downdraft_forced": cumulus_parameterization_state.output.normalized_massflux_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "condensate_to_fall_forced": cumulus_parameterization_state.output.condensate_to_fall_forced.field[ + :, :, :, [2, 0, 1] + ], + "evaporate_in_downdraft_forced": cumulus_parameterization_state.output.evaporate_in_downdraft_forced.field[ + :, :, :, [2, 0, 1] + ], + "cloud_liquid_after_rain_forced": cumulus_parameterization_state.output.cloud_liquid_after_rain_forced.field[ + :, :, :, [2, 0, 1] + ], + "t_updraft": cumulus_parameterization_state.output.t_updraft.field[:, :, :, [2, 0, 1]], + "convective_cloud_fraction_output": cumulus_parameterization_state.output.convective_cloud_fraction.field[ + :, :, :, [2, 0, 1] + ], + "cloud_workfunction_0": cumulus_parameterization_state.output.cloud_workfunction_0.field[:], + "cloud_workfunction_1": cumulus_parameterization_state.output.cloud_workfunction_1.field[:], + "cloud_workfunction_2": cumulus_parameterization_state.output.cloud_workfunction_2.field[:], + "cloud_workfunction_3": cumulus_parameterization_state.output.cloud_workfunction_3.field[:], + "cloud_workfunction_1_pbl": cumulus_parameterization_state.output.cloud_workfunction_1_pbl.field[ + : + ], + "cloud_workfunction_1_cin": cumulus_parameterization_state.output.cloud_workfunction_1_cin.field[ + : + ], + "cape_removal_time_scale": cumulus_parameterization_state.output.cape_removal_time_scale.field[:], + "pbl_time_scale": cumulus_parameterization_state.output.pbl_time_scale.field[:], + "lightning_density": cumulus_parameterization_state.output.lightning_density.field[:], + "evaporation_sublimation_tendency": cumulus_parameterization_state.output.evaporation_sublimation_tendency.field[ + : + ], + "convective_precip_flux": cumulus_parameterization_state.output.convective_precip_flux.field[:], + "t_perturbation": cumulus_parameterization_state.output.t_perturbation.field[:], + # input/output fields + "grid_length": cumulus_parameterization_state.input_output.grid_length.field[:], + "pbl_level": cumulus_parameterization_state.input_output.pbl_level.field[:] + 1, + "ccn": cumulus_parameterization_state.input_output.ccn.field[:], + "air_density": cumulus_parameterization_state.input_output.air_density.field[:], + "omega": cumulus_parameterization_state.input_output.omega.field[:], + "topography_height_no_negative": cumulus_parameterization_state.input_output.topography_height_no_negative.field[ + : + ], + "sensible_heat_flux": cumulus_parameterization_state.input_output.sensible_heat_flux.field[:], + "latent_heat_flux": cumulus_parameterization_state.input_output.latent_heat_flux.field[:], + "longitude_degrees": cumulus_parameterization_state.input_output.longitude_degrees.field[:], + "latitude_degrees": cumulus_parameterization_state.input_output.latitude_degrees.field[:], + "t_old": cumulus_parameterization_state.input_output.t_old.field[:], + "vapor_old": cumulus_parameterization_state.input_output.vapor_old.field[:], + "t_modified_by_advection": cumulus_parameterization_state.input_output.t_modified_by_advection.field[ + : + ], + "vapor_modified_by_advection": cumulus_parameterization_state.input_output.vapor_modified_by_advection.field[ + : + ], + "geopotential_height_forced": cumulus_parameterization_state.input_output.geopotential_height_forced.field[ + : + ], + "p_forced": cumulus_parameterization_state.input_output.p_forced.field[:], + "p_surface": cumulus_parameterization_state.input_output.p_surface.field[:], + "t_surface": cumulus_parameterization_state.input_output.t_surface.field[:], + "u": cumulus_parameterization_state.input_output.u.field[:], + "v": cumulus_parameterization_state.input_output.v.field[:], + "w": cumulus_parameterization_state.input_output.w.field[:], + "mass": cumulus_parameterization_state.input_output.mass.field[:], + "convective_scale_velocity": cumulus_parameterization_state.input_output.convective_scale_velocity.field[ + : + ], + "buoyancy_excess": cumulus_parameterization_state.input_output.buoyancy_excess.field[:], + "large_scale_ice": cumulus_parameterization_state.input_output.large_scale_ice.field[:], + "convective_ice": cumulus_parameterization_state.input_output.convective_ice.field[:], + "large_scale_liquid": cumulus_parameterization_state.input_output.large_scale_liquid.field[:], + "convective_liquid": cumulus_parameterization_state.input_output.convective_liquid.field[:], + "large_scale_cloud_fraction": cumulus_parameterization_state.input_output.large_scale_cloud_fraction.field[ + : + ], + "convective_cloud_fraction": cumulus_parameterization_state.input_output.convective_cloud_fraction.field[ + : + ], + "chemistry_tracers": cumulus_parameterization_state.input_output.chemistry_tracers.field[:], + "chemistry_tracers_output": chemistry_tracers_output_4d, + } + + return outputs