Features to support surface physics

ndsl/constants.py

@@ -158,8 +158,10 @@ class ConstantVersions(Enum):
     raise RuntimeError("Constant selector failed, bad code.")
 
 SECONDS_PER_DAY = Float(86400.0)
-SBC = 5.670400e-8
+SBC = Float(5.670400e-8)
 """Stefan-Boltzmann constant (W/m^2/K^4)"""
+RHO_H2O = Float(1000.0)
+"""Density of water in kg/m^3"""
 CV_AIR = CP_AIR - RDGAS
 """Heat capacity of dry air at constant volume"""
 RDG = -RDGAS / GRAV

ndsl/initialization/allocator.py

@@ -112,6 +112,28 @@ def from_array(
         base.data[:] = base.np.asarray(data)
         return base
 
+    def from_compute_array(
+        self,
+        data: np.ndarray,
+        dims: Sequence[str],
+        units: str,
+        allow_mismatch_float_precision: bool = False,
+    ):
+        """
+        Create a Quantity from a numpy array.
+
+        That numpy array must correspond to the correct shape and extent
+        of the compute domain for the given dims.
+        """
+        base = self.zeros(
+            dims=dims,
+            units=units,
+            dtype=data.dtype,
+            allow_mismatch_float_precision=allow_mismatch_float_precision,
+        )
+        base.view[:] = base.np.asarray(data)
+        return base
+
     def _allocate(
         self,
         allocator: Callable,

ndsl/namelist.py

@@ -111,6 +111,9 @@ class NamelistDefaults:
     tice = 273.16  # set tice = 165. to turn off ice - phase phys (kessler emulator)
     alin = 842.0  # "a" in lin1983
     clin = 4.8  # "c" in lin 1983, 4.8 -- > 6. (to enhance ql -- > qs)
+    mom4ice = False
+    lsm = 1
+    redrag = False
     isatmedmf = 0  # which version of satmedmfvdif to use
     dspheat = False  # flag for tke dissipative heating
     xkzm_h = 1.0  # background vertical diffusion for heat q over ocean
@@ -119,14 +122,17 @@ class NamelistDefaults:
     xkzm_ml = 1.0  # background vertical diffusion for momentum over land
     xkzm_hi = 1.0  # background vertical diffusion for heat q over ice
     xkzm_mi = 1.0  # background vertical diffusion for momentum over ice
+    xkzm_ho = 1.0  # background vertical diffusion for heat q over ocean
+    xkzm_mo = 1.0  # background vertical diffusion for momentum over ocean
     xkzm_s = 1.0  # sigma threshold for background mom. diffusion
     xkzm_lim = 0.01  # background vertical diffusion limit
     xkzminv = 0.15  # diffusivity in inversion layers
     xkgdx = 25.0e3  # background vertical diffusion threshold
     rlmn = 30.0  # lower-limiter on asymtotic mixing length in satmedmfdiff
     rlmx = 300.0  # upper-limiter on asymtotic mixing length in satmedmfdiff
     do_dk_hb19 = False  # flag for using hb19 background diff formula in satmedmfdiff
-    cap_k0_land = False  # flag for applying limiter on background diff in inversion layer over land in satmedmfdiff
+    cap_k0_land = True  # flag for applying limter on background diff in inversion layer over land in satmedmfdiff
+    lsoil = 4  # Number of soil levels in land surface model
 
     @classmethod
     def as_dict(cls):
@@ -309,6 +315,9 @@ class Namelist:
     tice: float = NamelistDefaults.tice
     alin: float = NamelistDefaults.alin
     clin: float = NamelistDefaults.clin
+    mom4ice: bool = NamelistDefaults.mom4ice
+    lsm: int = NamelistDefaults.lsm
+    redrag: bool = NamelistDefaults.redrag
     isatmedmf: int = NamelistDefaults.isatmedmf
     dspheat: bool = NamelistDefaults.dspheat
     xkzm_h: float = NamelistDefaults.xkzm_h
@@ -317,6 +326,8 @@ class Namelist:
     xkzm_ml: float = NamelistDefaults.xkzm_ml
     xkzm_hi: float = NamelistDefaults.xkzm_hi
     xkzm_mi: float = NamelistDefaults.xkzm_mi
+    xkzm_ho: float = NamelistDefaults.xkzm_ho
+    xkzm_mo: float = NamelistDefaults.xkzm_mo
     xkzm_s: float = NamelistDefaults.xkzm_s
     xkzm_lim: float = NamelistDefaults.xkzm_lim
     xkzminv: float = NamelistDefaults.xkzminv
@@ -490,6 +501,7 @@ class Namelist:
     nf_omega: int = NamelistDefaults.nf_omega
     fv_sg_adj: int = NamelistDefaults.fv_sg_adj
     n_sponge: int = NamelistDefaults.n_sponge
+    lsoil: int = NamelistDefaults.lsoil
     daily_mean: bool = False
     """Flag to replace cosz with daily mean value in physics"""
 

ndsl/stencils/basic_operations.py

@@ -61,6 +61,23 @@ def adjust_divide_stencil(adjustment: FloatField, q_out: FloatField):
         q_out = q_out / adjustment
 
 
+def average_in(
+    q_out: FloatField,
+    adjustement: FloatField,
+):
+    """
+    Averages every element of q_out
+    with every element of the adjustment
+    field, overwriting q_out.
+
+    Args:
+        adjustment: adjustment field
+        q_out: output field
+    """
+    with computation(PARALLEL), interval(...):
+        q_out = (q_out + adjustement) * 0.5
+
+
 @gtscript.function
 def sign(a, b):
     """

ndsl/stencils/testing/translate.py

@@ -214,6 +214,9 @@ def make_storage_data_input_vars(
             inputs_out[p] = inputs_in[p]
         for d, info in storage_vars.items():
             serialname = info["serialname"] if "serialname" in info else d
+            index_variable = (
+                info["index_variable"] if "index_variable" in info else False
+            )
             self.update_info(info, inputs_in)
             if "kaxis" in info:
                 inputs_in[serialname] = np.moveaxis(
@@ -231,6 +234,8 @@ def make_storage_data_input_vars(
 
             dummy_axes = info.get("dummy_axes", None)
             axis = info.get("axis", 2)
+            if index_variable:
+                inputs_in[serialname] -= 1
             inputs_out[d] = self.make_storage_data(
                 np.squeeze(inputs_in[serialname]),
                 istart=istart,
@@ -258,9 +263,16 @@ def slice_output(self, inputs, out_data=None) -> dict[str, Any]:
             info = self.out_vars[var]
             self.update_info(info, inputs)
             serialname = info["serialname"] if "serialname" in info else var
+            index_variable = (
+                info["index_variable"] if "index_variable" in info else False
+            )
             ds = self.grid.default_domain_dict()
             ds.update(info)
             data_result = as_numpy(out_data[var])
+            if index_variable:
+                if isinstance(data_result, dict):
+                    raise TypeError(f"Variable {serialname} is a 4D dict, not an index")
+                data_result += 1
             if isinstance(data_result, dict):
                 names_4d = info.get("names_4d", utils.tracer_variables)
                 var4d = np.zeros(

ndsl/stencils/tridiag.py

@@ -0,0 +1,89 @@
+from gt4py.cartesian.gtscript import BACKWARD, FORWARD, computation, interval
+
+from ndsl.dsl.typing import BoolFieldIJ, FloatField
+
+
+def tridiag_solve(
+    a: FloatField,
+    b: FloatField,
+    c: FloatField,
+    d: FloatField,
+    x: FloatField,
+    delta: FloatField,
+):
+    """
+    This stencil solves a square, k x k tridiagonal matrix system
+    with coefficients a, b, and c, and vectors p and d using the Thomas algorithm:
+    ! ###                                            ### ###  ###   ###  ###!
+    ! #b(0), c(0),  0  ,  0  ,  0  ,   . . .  ,    0   # # x(0) #   # d(0) #!
+    ! #a(1), b(1), c(1),  0  ,  0  ,   . . .  ,    0   # # x(1) #   # d(1) #!
+    ! # 0  , a(2), b(2), c(2),  0  ,   . . .  ,    0   # # x(2) #   # d(2) #!
+    ! # 0  ,  0  , a(3), b(3), c(3),   . . .  ,    0   # # x(3) #   # d(3) #!
+    ! # 0  ,  0  ,  0  , a(4), b(4),   . . .  ,    0   # # x(4) #   # d(4) #!
+    ! # .                                          .   # #  .   # = #   .  #!
+    ! # .                                          .   # #  .   #   #   .  #!
+    ! # .                                          .   # #  .   #   #   .  #!
+    ! # 0  , . . . , 0 , a(k-2), b(k-2), c(k-2),   0   # #x(k-3)#   #d(k-3)#!
+    ! # 0  , . . . , 0 ,   0   , a(k-1), b(k-1), c(k-1)# #x(k-2)#   #d(k-2)#!
+    ! # 0  , . . . , 0 ,   0   ,   0   ,  a(k) ,  b(k) # #x(k-1)#   #d(k-1)#!
+    ! ###                                            ### ###  ###   ###  ###!
+
+    Args:
+        a (in): lower-diagonal matrix coefficients
+        b (in): diagonal matrix coefficients
+        c (in): upper-diagonal matrix coefficients
+        d (in): Result vector
+        x (out): The vector to solve for
+        delta (out): d post-pivot
+    """
+    with computation(FORWARD):  # Forward sweep
+        with interval(0, 1):
+            x = c / b
+            delta = d / b
+        with interval(1, None):
+            x = c / (b - a * x[0, 0, -1])
+            delta = (d - a * delta[0, 0, -1]) / (b - a * x[0, 0, -1])
+    with computation(BACKWARD):  # Reverse sweep
+        with interval(-1, None):
+            x = delta
+        with interval(0, -1):
+            x = delta - x * x[0, 0, 1]
+
+
+def masked_tridiag_solve(
+    a: FloatField,
+    b: FloatField,
+    c: FloatField,
+    d: FloatField,
+    x: FloatField,
+    delta: FloatField,
+    mask: BoolFieldIJ,
+):
+    """
+    Same as tridiag_solve but restricted to a subset of horizontal points
+
+    Args:
+        a (in): lower-diagonal matrix coefficients
+        b (in): diagonal matrix coefficients
+        c (in): upper-diagonal matrix coefficients
+        d (in): Result vector
+        mask (in): Columns to execute the stencil on
+        x (out): The vector to solve for
+        delta (out): d post-pivot
+    """
+    with computation(FORWARD):  # Forward sweep
+        with interval(0, 1):
+            if mask:
+                x = c / b
+                delta = d / b
+        with interval(1, None):
+            if mask:
+                x = c / (b - a * x[0, 0, -1])
+                delta = (d - a * delta[0, 0, -1]) / (b - a * x[0, 0, -1])
+    with computation(BACKWARD):  # Reverse sweep
+        with interval(-1, None):
+            if mask:
+                x = delta
+        with interval(0, -1):
+            if mask:
+                x = delta - x * x[0, 0, 1]