Ported and verified fluxbelowinv and made minor updates to other func…

…tions

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_alpha.py

@@ -15,15 +15,15 @@ def compute_alpha(
     #Subroutine to compute proportionality factor for 
     #implicit CIN calculation. 
 
-    x0 = 0.0
-    del_CIN8_f64 = del_CIN
-    ke8_f64 = ke
+    x0:f64 = f64(0.0)
+    del_CIN8_f64:f64 = f64(del_CIN)
+    ke8_f64:f64 = ke
     iteration = 0
     while iteration < 10:
         x1 = x0 - (exp(-x0*ke8_f64*del_CIN8_f64) - x0)/(-ke8_f64*del_CIN8_f64*exp(-x0*ke8_f64*del_CIN8_f64) - 1.0)
         x0 = x1
         iteration+=1
 
-    compute_alpha = x0
+    compute_alpha = f32(x0)
 
     return compute_alpha

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_mumin2.py

@@ -1,29 +1,11 @@
 import gt4py.cartesian.gtscript as gtscript
-from gt4py.cartesian.gtscript import computation, interval, PARALLEL, sqrt, exp
+from gt4py.cartesian.gtscript import computation, interval, PARALLEL, sqrt, exp, erfc
 import pyMoist.constants as constants
-from ndsl.constants import X_DIM, Y_DIM, Z_DIM
 from ndsl.dsl.typing import (
     Float,
     Int,
 )
-from ndsl import StencilFactory, QuantityFactory
 
-@gtscript.function
-def erfc(x: Float):
-    '''
-    Complimentary error function approximation borowed from 
-    Abramowits and Stegun's Handbook of Mathematical Functions
-    '''
-
-    # Constants for the approximation
-    p = 0.3275911
-    a1, a2, a3, a4, a5 = 0.254829592, -0.284496736, 1.421413741, -1.453152027, 1.061405429
-
-    # Approximation for the error function complement
-    t = 1 / (1 + p * x)
-    erf_approx = 1 - (a1 * t + a2 * t**2 + a3 * t**3 + a4 * t**4 + a5 * t**5) * exp(-x**2)
-
-    return erf_approx
 
 @gtscript.function
 def compute_mumin2(
@@ -35,19 +17,19 @@ def compute_mumin2(
     #Subroutine to compute critical 'mu' (normalized CIN) such 
     #that updraft fraction at the LCL is equal to 'rmaxfrac'.  
 
-    x0 = mulow
+    x0:f64 = mulow
     iteration = 0
     while iteration < 10:
-       ex = exp(-x0**2)
-       ef = erfc(x0) # Fortran error fraction function
-       exf = ex/ef
-       f  = 0.5*exf**2 - 0.5*(ex/2.0/rmaxfrax)**2 - (mulcl*2.5066/2.0)**2
-       fs = (2.*exf**2)*(exf/sqrt(constants.MAPL_PI)-x0) + (0.5*x0*ex**2)/(rmaxfrax**2)
-       x1 = x0 - f/fs     
+       ex:f64 = exp(-x0**2)
+       ef:f64 = erfc(x0) # Complimentary error fraction function
+       exf:f64 = ex/ef
+       f:f64  = f64(0.5)*exf**2 - f64(0.5)*(ex/f64(2.0)/rmaxfrax)**2 - (mulcl*f64(2.5066)/f64(2.0))**2
+       fs:f64 = (f64(2.0)*exf**2)*(exf/sqrt(constants.MAPL_PI)-x0) + (f64(0.5)*x0*ex**2)/(rmaxfrax**2)
+       x1:f64 = x0 - f/fs     
        x0 = x1
        iteration+=1
 
-    compute_mumin2 = x0
+    compute_mumin2 = f32(x0)
 
     return compute_mumin2
 

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_ppen.py

@@ -2,7 +2,6 @@
 from ndsl.dsl.typing import FloatField, Float
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
 from ndsl.stencils.testing.translate import TranslateFortranData2Py
-from pyMoist.saturation import QSat
 import gt4py.cartesian.gtscript as gtscript
 from gt4py.cartesian.gtscript import computation, PARALLEL, interval, exp, FORWARD
 import gt4py.cartesian.gtscript as gtscript
@@ -27,34 +26,34 @@ def compute_ppen(
     #non-zero fer(kpen).  
 
     # Buoyancy slope
-    SB = ( bogtop - bogbot ) / dpen
+    SB:f64 = ( bogtop - bogbot ) / dpen
 
     # Sign of slope, 'f' at x = 0
     # If 's00>0', 'w' increases with height.
-    s00 = bogbot / rho0j - D * wtwb
+    s00:f64 = bogbot / rho0j - D * wtwb
 
-    if D*dpen < 1.0e-4:
-        if s00 >= 0.0:
-            x0 = dpen       
+    if D*dpen < f64(1.0e-4):
+        if s00 >= f64(0.0):
+            x0:f64 = dpen       
         else:
-            x0 = max(0.0,min(dpen,-0.5*wtwb/s00))
+            x0:f64 = max(f64(0.0),min(dpen,f64(-0.5)*wtwb/s00))
     else:
-        if s00 >= 0.0:
-            x0 = dpen
+        if s00 >= f64(0.0):
+            x0:f64 = dpen
         else:
-            x0 = 0.0
+            x0:f64 = f64(0.0)
 
         iteration=0
         while iteration < 5:
-            aux  = min(max(-2.0*D*x0, -20.0), 20.0)
+            aux:f64  = min(max(f64(-2.0)*D*x0, -20.0), 20.0)
 
-            f  = exp(aux)*(wtwb-(bogbot-SB/(2.0*D))/(D*rho0j)) + (SB*x0+bogbot-SB/(2.0*D))/(D*rho0j)
-            fs = -2.0*D*exp(aux)*(wtwb-(bogbot-SB/(2.0*D))/(D*rho0j)) + (SB)/(D*rho0j)
+            f:f64  = exp(aux)*(wtwb-(bogbot-SB/(2.0*D))/(D*rho0j)) + (SB*x0+bogbot-SB/(2.0*D))/(D*rho0j)
+            fs:f64 = -2.0*D*exp(aux)*(wtwb-(bogbot-SB/(2.0*D))/(D*rho0j)) + (SB)/(D*rho0j)
 
-            x1 = x0 - f/fs     
-            x0 = x1
+            x1:f64 = x0 - f/fs     
+            x0:f64 = x1
             iteration+=1
 
-    compute_ppen = -max(0.0,min(dpen,x0))
+    compute_ppen = -max(f64(0.0),min(dpen,x0))
 
     return compute_ppen

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/fluxbelowinv.py

@@ -0,0 +1,135 @@
+import gt4py.cartesian.gtscript as gtscript
+from gt4py.cartesian.gtscript import computation, interval, PARALLEL, FORWARD, sqrt, exp
+import pyMoist.constants as constants
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.typing import (
+    Float,
+    Int,
+    IntField,
+    IntFieldIJ,
+    FloatFieldIJ,
+    FloatField
+)
+from ndsl import StencilFactory, QuantityFactory
+
+
+def fluxbelowinv(
+    cbmf: FloatFieldIJ,
+    ps0: FloatField,
+    kinv: IntFieldIJ,
+    dt: FloatFieldIJ,
+    xsrc: FloatFieldIJ,
+    xmean: FloatFieldIJ,
+    xtopin: FloatFieldIJ,
+    xbotin: FloatFieldIJ,
+    xflx: FloatField,
+    k_idx: FloatField,
+    ):
+    '''
+    Stencil to calculate turbulent fluxes at and below 'kinv-1' interfaces.
+    Check in the main program such that input 'cbmf' should not be zero.        
+    If the reconstructed inversion height does not go down below the 'kinv-1' 
+    interface, then turbulent flux at 'kinv-1' interface  is simply a product 
+    of 'cmbf' and 'qtsrc-xbot' where 'xbot' is the value at the top interface 
+    of 'kinv-1' layer. This flux is linearly interpolated down to the surface 
+    assuming turbulent fluxes at surface are zero. If reconstructed inversion 
+    height goes down below the 'kinv-1' interface, subsidence warming &drying 
+    measured by 'xtop-xbot', where  'xtop' is the value at the base interface 
+    of 'kinv+1' layer, is added ONLY to the 'kinv-1' layer, using appropriate 
+    mass weighting ( rpinv and rcbmf, or rr = rpinv / rcbmf ) between current 
+    and next provisional time step. Also impose a limiter to enforce outliers 
+    of thermodynamic variables in 'kinv' layer  to come back to normal values 
+    at the next step.  
+
+    '''
+    with computation(FORWARD), interval(...):
+        xflx[0, 0, 1] = 0.0 
+
+    with computation(PARALLEL), interval(...):
+        xflx = 0.0
+        k_below=kinv-1
+        dp = ps0.at(K=k_below) - ps0.at(K=kinv)    
+
+        xbot = xbotin
+        xtop = xtopin
+
+        # Compute reconstructed inversion height
+        xtop_ori = xtop
+        xbot_ori = xbot
+        rcbmf = ( cbmf * constants.MAPL_GRAV * dt ) / dp                  # Can be larger than 1 : 'OK'      
+
+        if xbot >= xtop:
+            rpeff = ( xmean - xtop ) / max(  1.e-20, xbot - xtop ) 
+        else:
+            rpeff = ( xmean - xtop ) / min( -1.e-20, xbot - xtop ) 
+
+        rpeff = min( max(0.0,rpeff), 1.0 )                # As of this, 0<= rpeff <= 1   
+        if rpeff == 0.0 or rpeff == 1.0:
+            xbot = xmean
+            xtop = xmean
+
+        #Below two commented-out lines are the old code replacing the above 'if' block.   
+        # if(rpeff.eq.1) xbot = xmean
+        # if(rpeff.eq.0) xtop = xmean 
+
+        rr       = rpeff / rcbmf
+        pinv     = ps0.at(K=k_below) - rpeff * dp             # "pinv" before detraining mass
+        pinv_eff = ps0.at(K=k_below) + ( rcbmf - rpeff ) * dp # Effective "pinv" after detraining mass
+
+        # Compute turbulent fluxes.                                               
+        # Below two cases exactly converges at 'kinv-1' interface when rr = 1. 
+        if k_idx <= k_below:
+            xflx = cbmf * ( xsrc - xbot ) * ( ps0.at(K=0) - ps0 ) / ( ps0.at(K=0) - pinv )
+        if k_idx == k_below:
+            if rr <= 1:
+                xflx = xflx - ( 1.0 - rr ) * cbmf * ( xtop_ori - xbot_ori )
+
+
+class FluxBelowInv:
+    def __init__(
+        self,
+        stencil_factory: StencilFactory,
+        quantity_factory: QuantityFactory,
+    ) -> None:
+
+        self.stencil_factory = stencil_factory
+        self.quantity_factory = quantity_factory
+        grid_indexing = stencil_factory.grid_indexing
+        self._FluxBelowInv = self.stencil_factory.from_dims_halo(
+            func=fluxbelowinv,
+            compute_dims=[X_DIM, Y_DIM, Z_DIM],
+        )
+
+        self._k_idx = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
+        for k in range(0, self._k_idx.view[:].shape[2]):
+            self._k_idx.view[:, :, k] = k
+
+    def __call__(
+        self,
+        cbmf: FloatFieldIJ,
+        ps0: FloatField,
+        kinv: IntFieldIJ,
+        dt: FloatFieldIJ,
+        xsrc: FloatFieldIJ,
+        xmean: FloatFieldIJ,
+        xtopin: FloatFieldIJ,
+        xbotin: FloatFieldIJ,
+        xflx: FloatField,
+    ):
+
+        self._FluxBelowInv(
+            cbmf=cbmf,
+            ps0=ps0,
+            kinv=kinv,
+            dt=dt,
+            xsrc=xsrc,
+            xmean=xmean,
+            xtopin=xtopin,
+            xbotin=xbotin,
+            xflx=xflx,
+            k_idx=self._k_idx,
+        )
+
+
+
+

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/getbuoy.py

@@ -1,11 +1,8 @@
 import gt4py.cartesian.gtscript as gtscript
-from gt4py.cartesian.gtscript import computation, interval, PARALLEL, log
 import pyMoist.constants as constants
-from ndsl.constants import X_DIM, Y_DIM, Z_DIM
 from ndsl.dsl.typing import (
     Float,
 )
-from ndsl import StencilFactory, QuantityFactory
 from pyMoist.UW.compute_uwshcu import exnerfn
 
 

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/translate_compute_alpha.py

@@ -1,7 +1,5 @@
 from ndsl import Namelist,StencilFactory
-from ndsl.dsl.typing import FloatField, Float
 from ndsl.stencils.testing.translate import TranslateFortranData2Py
-import gt4py.cartesian.gtscript as gtscript
 from gt4py.cartesian.gtscript import computation, PARALLEL, interval
 import gt4py.cartesian.gtscript as gtscript
 from ndsl.dsl.typing import (

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/translate_compute_mumin2.py

@@ -1,13 +1,13 @@
 from ndsl import Namelist,StencilFactory
 from ndsl.stencils.testing.translate import TranslateFortranData2Py
-from gt4py.cartesian.gtscript import computation, PARALLEL, interval, exp, sqrt
+from gt4py.cartesian.gtscript import computation, PARALLEL, interval
 import gt4py.cartesian.gtscript as gtscript
 from ndsl.dsl.typing import (
     Float,
     FloatField
 )
 import pyMoist.constants as constants
-from pyMoist.UW.compute_mumin2 import compute_mumin2, erfc
+from pyMoist.UW.compute_mumin2 import compute_mumin2
 
 def harness_stencil(
     mulcl: FloatField,
@@ -45,18 +45,18 @@ def __init__(self, grid, namelist: Namelist, stencil_factory: StencilFactory):
         self.stencil = stencil_factory.from_origin_domain(
             func=harness_stencil,
             origin=(0, 0, 0),
-            domain=(1, 1, 2),
+            domain=(1, 1, 10),
         )
 
     def compute(self, inputs):
-        mulcl = self.quantity_factory._numpy.empty((1, 1, 2), dtype=Float)
+        mulcl = self.quantity_factory._numpy.empty((1, 1, 10), dtype=Float)
         mulcl[0, 0, :] = inputs["mulcl"]
-        rmaxfrax = self.quantity_factory._numpy.empty((1, 1, 2), dtype=Float)
+        rmaxfrax = self.quantity_factory._numpy.empty((1, 1, 10), dtype=Float)
         rmaxfrax[0, 0, :] = inputs["rmaxfrax"]
-        mulow = self.quantity_factory._numpy.empty((1, 1, 2), dtype=Float)
+        mulow = self.quantity_factory._numpy.empty((1, 1, 10), dtype=Float)
         mulow[0, 0, :] = inputs["mu"]
 
-        mumin2 = self.quantity_factory._numpy.empty((1, 1, 2), dtype=Float)
+        mumin2 = self.quantity_factory._numpy.empty((1, 1, 10), dtype=Float)
 
         self.stencil(mulcl, rmaxfrax, mulow, mumin2)
 

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/translate_compute_ppen.py

@@ -1,12 +1,7 @@
-from ndsl import Namelist, Quantity, StencilFactory
-from ndsl.dsl.typing import FloatField, Float
-from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl import Namelist, Quantity
 from ndsl.stencils.testing.translate import TranslateFortranData2Py
-from pyMoist.saturation import QSat
+from gt4py.cartesian.gtscript import computation, PARALLEL, interval
 import gt4py.cartesian.gtscript as gtscript
-from gt4py.cartesian.gtscript import computation, PARALLEL, interval, exp, FORWARD
-import gt4py.cartesian.gtscript as gtscript
-from ndsl.constants import X_DIM, Y_DIM, Z_DIM
 from ndsl.dsl.typing import (
     FloatFieldIJ,
     Float,
@@ -57,25 +52,25 @@ def __init__(self, grid, namelist: Namelist, stencil_factory: StencilFactory):
         self.stencil = stencil_factory.from_origin_domain(
             func=harness_stencil,
             origin=(0, 0, 0),
-            domain=(1, 1, 107),
+            domain=(1, 1, 120),
         )
 
     def compute(self, inputs):
-        wtwb = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        wtwb = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         wtwb[0, 0, :] = inputs["wtwb"]
-        drage = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        drage = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         drage[0, 0, :] = inputs["drage"]
-        bogbot = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        bogbot = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         bogbot[0, 0, :] = inputs["bogbot"]
-        bogtop = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        bogtop = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         bogtop[0, 0, :] = inputs["bogtop"]
-        rhomid0j = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        rhomid0j = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         rhomid0j[0, 0, :] = inputs["rhomid0j"]
-        dp0 = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        dp0 = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
         dp0[0, 0, :] = inputs["dp0"]
 
 
-        ppen = self.quantity_factory._numpy.empty((1, 1, 107), dtype=Float)
+        ppen = self.quantity_factory._numpy.empty((1, 1, 120), dtype=Float)
 
         self.stencil(wtwb, drage, bogbot, bogtop, rhomid0j, dp0, ppen)