[NDSL] UW Convection port

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore

@@ -9,4 +9,8 @@ test_data/
 .vscode
 test_data/
 sandbox/
+<<<<<<< HEAD
 *.mod
+=======
+*.mod
+>>>>>>> origin/dsl/QSat

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/compute_uwshcu.py

@@ -0,0 +1,116 @@
+import gt4py.cartesian.gtscript as gtscript
+from gt4py.cartesian.gtscript import computation, interval, PARALLEL
+import pyMoist.pyMoist_constants as constants
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.typing import FloatField, Float, IntField, Int 
+from ndsl import StencilFactory, QuantityFactory 
+
+def slope(field: FloatField, p0: FloatField, slope: FloatField):
+    with computation(PARALLEL), interval(0, 1):
+        value = (field[0, 0, 1] - field[0, 0, 0]) / (p0[0, 0, 1] - p0[0, 0, 0])
+        if value > 0.0:
+            slope = max(0.0, value)
+        else:
+            slope = min(0.0, value)
+    with computation(PARALLEL), interval(1, -1):
+        above_value = (field[0, 0, 1] - field[0, 0, 0]) / (p0[0, 0, 1] - p0[0, 0, 0])
+        below_value = (field[0, 0, 0] - field[0, 0, -1]) / (p0[0, 0, 0] - p0[0, 0, -1])
+        if above_value > 0.0:
+            slope = max(0.0, min(above_value, below_value))
+        else:
+            slope = min(0.0, max(above_value, below_value))
+    with computation(PARALLEL), interval(-1, None):
+        slope = slope[0, 0, -1]
+
+
+def compute_uwshcu(
+    dotransport: Int,           # Transport tracers [1 true]
+    exnifc0_in: FloatField,     # Exner function at interfaces 
+    pmid0_in: FloatField,       # Environmental pressure at midpoints [Pa] 
+    zmid0_in: FloatField,       # Environmental height at midpoints [m] 
+    exnmid0_in: FloatField,     # Exner function at midpoints
+    u0_in: FloatField,          # Environmental zonal wind [m/s] 
+    v0_in: FloatField,          # Environmental meridional wind [m/s] 
+    qv0_in: FloatField,         # Environmental specific humidity
+    ql0_in: FloatField,         # Environmental liquid water specific humidity 
+    qi0_in: FloatField,         # Environmental ice specific humidity
+    th0_in: FloatField,         # Environmental potential temperature [K]
+    tr0_inout: FloatField,      # Environmental tracers [ #, kg/kg ]
+    tr0: FloatField,
+    ssthl0: FloatField,
+    ssqt0: FloatField,
+    ssu0: FloatField,
+    ssv0: FloatField,
+    sstr0: FloatField
+):
+
+    '''
+    University of Washington Shallow Convection Scheme          
+
+    Described in Park and Bretherton. 2008. J. Climate :        
+
+    'The University of Washington shallow convection and         
+    moist turbulent schemes and their impact on climate         
+    simulations with the Community Atmosphere Model'            
+
+    Coded in CESM by Sungsu Park. Oct.2005. May.2008.                     
+
+    Coded in GEOS by Nathan Arnold. July 2016.                  
+
+    For general questions, email [email protected] or             
+    [email protected]    
+
+    For GEOS-specific questions, email [email protected]                                                         
+    '''
+
+    # Start Main Calculation
+
+    with computation(PARALLEL), interval(...):
+         #id_exit = False
+
+         zvir = Float(0.609) # r_H2O/r_air-1
+         pmid0 = pmid0_in
+         zmid0 = zmid0_in
+         u0 = u0_in
+         v0 = v0_in
+         qv0 = qv0_in
+         ql0 = ql0_in
+         qi0 = qi0_in
+
+         #if dotransport == 1:
+         #   tr0 = tr0_inout
+
+         '''
+         Compute basic thermodynamic variables directly from  
+         input variables for each column     
+         '''                 
+
+         # Compute internal environmental variables
+         exnmid0 = exnmid0_in
+         exnifc0 = exnifc0_in
+         t0 = th0_in * exnmid0
+         s0 = constants.gravity*zmid0 + constants.cpdry*t0
+         qt0 = qv0 + ql0 + qi0
+         thl0 = (t0 - constants.latent_heat_vaporization*ql0/constants.cpdry - constants.latent_heat_sublimation*qi0/constants.cpdry) / exnmid0
+         thvl0 = (1.0 + zvir*qt0) * thl0
+
+
+         # Compute slopes of environmental variables in each layer
+         ssthl0 = slope(thl0, pmid0)
+         ssqt0 = slope(qt0 , pmid0)
+         ssu0 = slope(u0  , pmid0)
+         ssv0 = slope(v0  , pmid0)
+
+         #if dotransport == 1:
+         #   sstr0 = slope(tr0, pmid0)
+
+
+
+
+
+
+
+
+
+
+

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/conden.py

@@ -0,0 +1,169 @@
+import gt4py.cartesian.gtscript as gtscript
+from gt4py.cartesian.gtscript import computation, interval, PARALLEL, sin
+import pyMoist.radiation_coupling_constants as radconstants
+import pyMoist.pyMoist_constants as constants
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.typing import FloatField, Float, IntField 
+from ndsl import StencilFactory, QuantityFactory 
+from pyMoist.saturation.qsat import QSat, QSat_Float, FloatField_Extra_Dim
+from pyMoist.saturation.formulation import SaturationFormulation
+
+
+@gtscript.function
+def exnerfn(
+    p: Float,
+)-> Float:
+
+    return (p / 100000.0) ** (radconstants.MAPL_RDRY / radconstants.MAPL_CPDRY)
+
+
+@gtscript.function
+def ice_fraction(
+    temp: Float,
+    cnv_frc: Float,
+    srf_type: Float,
+):
+    # Anvil clouds
+    # Anvil-Convective sigmoidal function like figure 6(right)
+    # Sigmoidal functions Hu et al 2010, doi:10.1029/2009JD012384
+    if temp <= constants.JaT_ICE_ALL:
+        icefrct_c = 1.000
+    elif temp > constants.JaT_ICE_ALL and temp <= constants.JaT_ICE_MAX:
+        icefrct_c = sin(0.5 * radconstants.MAPL_PI * (1.00 - (temp - constants.JaT_ICE_ALL) / (constants.JaT_ICE_MAX - constants.JaT_ICE_ALL)))
+    else:
+        icefrct_c = 0.00
+    icefrct_c = max(min(icefrct_c,1.00),0.00) ** constants.aICEFRPWR
+    # Sigmoidal functions like figure 6b/6c of Hu et al 2010, doi:10.1029/2009JD012384
+    if srf_type == 2.0:
+        if temp <= constants.JiT_ICE_ALL:
+            icefrct_m = 1.000
+        elif temp > constants.JiT_ICE_ALL and temp <= constants.JiT_ICE_MAX:
+            icefrct_m = 1.00 - (temp - constants.JiT_ICE_ALL) / (constants.JiT_ICE_MAX - constants.JiT_ICE_ALL)
+        else:
+            icefrct_m = 0.00
+        icefrct_m = max(min(icefrct_m, 1.00), 0.00) ** constants.iICEFRPWR
+    elif srf_type > 1.0:
+        if temp <= constants.lT_ICE_ALL:
+           icefrct_m = 1.000
+        elif temp > constants.lT_ICE_ALL and temp <= constants.lT_ICE_MAX:
+           icefrct_m = sin(0.5 * radconstants.MAPL_PI * (1.00 - (temp - constants.lT_ICE_ALL) / (constants.lT_ICE_MAX - constants.lT_ICE_ALL)))
+        else:
+            icefrct_m = 0.00
+        icefrct_m = max(min(icefrct_m, 1.00), 0.00) ** constants.lICEFRPWR
+    else:
+        if temp <= constants.oT_ICE_ALL:
+           icefrct_m = 1.000
+        elif temp > constants.oT_ICE_ALL and temp <= constants.oT_ICE_MAX:
+           icefrct_m = sin(0.5 * radconstants.MAPL_PI * (1.00 - (temp - constants.oT_ICE_ALL) / (constants.oT_ICE_MAX - constants.oT_ICE_ALL)))
+        else:
+            icefrct_m = 0.00
+        icefrct_m = max(min(icefrct_m, 1.00), 0.00) ** constants.oICEFRPWR
+    ice_frac = icefrct_m * (1.0-cnv_frc) + icefrct_c * cnv_frc
+    return ice_frac
+
+
+def conden(
+    p: FloatField, 
+    thl: FloatField, 
+    qt: FloatField,
+    th: FloatField, 
+    qv: FloatField, 
+    ql: FloatField,
+    qi: FloatField, 
+    rvls: FloatField,
+    id_check: IntField,
+    ese: FloatField_Extra_Dim,
+    esw: FloatField_Extra_Dim,
+    esx: FloatField_Extra_Dim,
+):
+
+    with computation(PARALLEL), interval(...):
+
+        tc = thl * exnerfn(p)
+
+        icefrac = ice_fraction(tc, 0.0, 0.0)
+        nu = icefrac
+        leff = (1.0 - nu) * radconstants.MAPL_ALHL + nu * radconstants.MAPL_ALHS # Effective latent heat
+        temps = tc
+        ps = p
+        qs, _ = QSat_Float(ese, esw, esx, temps, ps / 100.0) # Saturation specific humidity
+        rvls = qs
+
+        if qs >= qt:      # no condensation
+            id_check = 0
+            qv = qt
+            qc = 0.0
+            ql = 0.0
+            qi = 0.0
+            th = thl
+        else:             # condensation
+            iteration = 0
+            while iteration < 10:
+                temps = temps + ((tc - temps) * radconstants.MAPL_CPDRY / leff + qt - rvls) / (radconstants.MAPL_CPDRY / leff + (constants.rdry / constants.rvap) * leff * rvls / (radconstants.MAPL_RDRY * temps * temps))
+                qs, _ = QSat_Float(ese, esw, esx, temps, ps / 100.0)
+                rvls = qs
+                iteration+=1
+            qc = max(qt - qs, 0.0)
+            qv = qt - qc
+            ql = qc * (1.0 - nu)
+            qi = nu * qc
+            th = temps / exnerfn(p)
+            if abs((temps - (leff / radconstants.MAPL_CPDRY) * qc) - tc) >= 1.0:
+                id_check = 1
+            else:
+                id_check = 0
+
+
+class Conden:
+    def __init__(
+        self,
+        stencil_factory: StencilFactory,
+        quantity_factory: QuantityFactory,
+        formulation: SaturationFormulation = SaturationFormulation.Staars,
+        use_table_lookup: bool = True,
+    ) -> None:
+
+        self.stencil_factory = stencil_factory
+        self.quantity_factory = quantity_factory
+
+        self._conden = self.stencil_factory.from_dims_halo(
+            func=conden,
+            compute_dims=[X_DIM, Y_DIM, Z_DIM],
+        )
+
+    def __call__(
+        self,
+        pifc0_test: FloatField, 
+        thl0bot_test: FloatField, 
+        qt0bot_test: FloatField,
+        thj_test: FloatField, 
+        qvj_test: FloatField, 
+        qlj_test: FloatField,
+        qij_test: FloatField, 
+        qse_test: FloatField,
+        id_check_test: IntField,
+        formulation: SaturationFormulation = SaturationFormulation.Staars,
+        use_table_lookup: bool = True,
+    ):
+
+        self.qsat = QSat(
+            self.stencil_factory,
+            self.quantity_factory,
+            formulation=formulation,
+            use_table_lookup=use_table_lookup,
+        )
+
+        self._conden(
+                pifc0_test, 
+                thl0bot_test, 
+                qt0bot_test, 
+                thj_test, 
+                qvj_test, 
+                qlj_test, 
+                qij_test, 
+                qse_test, 
+                id_check_test, 
+                self.qsat._ese, 
+                self.qsat._esw, 
+                self.qsat._esx,
+        )