NDSL port of the Grell-Fritas 2020 convection parameterization scheme

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore

@@ -17,3 +17,4 @@ sandbox/
 .venv/
 site/
 test_data/
+_dacecache/

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  # --

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,
+            )

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

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
+            )

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