Merge pull request #35 from NOAA-GFDL/develop

Release 2024.04.00

ndsl/grid/eta.py

@@ -1,6 +1,7 @@
 import math
 import os
 from dataclasses import dataclass
+from typing import Optional, Tuple
 
 import numpy as np
 import xarray as xr
@@ -28,8 +29,25 @@ class HybridPressureCoefficients:
     bk: np.ndarray
 
 
+def _load_ak_bk_from_file(eta_file: str) -> Tuple[np.ndarray, np.ndarray]:
+    if eta_file == "None":
+        raise ValueError("eta file not specified")
+    if not os.path.isfile(eta_file):
+        raise ValueError("file " + eta_file + " does not exist")
+
+    # read file into ak, bk arrays
+    data = xr.open_dataset(eta_file)
+    ak = data["ak"].values
+    bk = data["bk"].values
+
+    return ak, bk
+
+
 def set_hybrid_pressure_coefficients(
-    km: int, eta_file: str
+    km: int,
+    eta_file: str,
+    ak_data: Optional[np.ndarray] = None,
+    bk_data: Optional[np.ndarray] = None,
 ) -> HybridPressureCoefficients:
     """
     Sets the coefficients describing the hybrid pressure coordinates.
@@ -44,25 +62,19 @@ def set_hybrid_pressure_coefficients(
     Returns:
         a HybridPressureCoefficients dataclass
     """
-
-    if eta_file == "None":
-        raise ValueError("eta file not specified")
-    if not os.path.isfile(eta_file):
-        raise ValueError("file " + eta_file + " does not exist")
-
-    # read file into ak, bk arrays
-    data = xr.open_dataset(eta_file)
-    ak = data["ak"].values
-    bk = data["bk"].values
+    if ak_data is None or bk_data is None:
+        ak, bk = _load_ak_bk_from_file(eta_file)
+    else:
+        ak, bk = ak_data, bk_data
 
     # check size of ak and bk array is km+1
     if ak.size - 1 != km:
-        raise ValueError(f"size of ak array is not equal to km={km}")
+        raise ValueError(f"size of ak array {ak.size} is not equal to km+1={km+1}")
     if bk.size - 1 != km:
-        raise ValueError(f"size of bk array is not equal to km={km}")
+        raise ValueError(f"size of bk array {ak.size} is not equal to km+1={km+1}")
 
     # check that the eta values computed from ak and bk are monotonically increasing
-    eta, etav = check_eta(ak, bk)
+    eta, etav = _check_eta(ak, bk)
 
     if not np.all(eta[:-1] <= eta[1:]):
         raise ValueError("ETA values are not monotonically increasing")
@@ -75,20 +87,18 @@ def set_hybrid_pressure_coefficients(
     else:
         raise ValueError("bk must contain at least one 0.")
 
-    pressure_data = HybridPressureCoefficients(ks, ptop, ak, bk)
-
-    return pressure_data
+    return HybridPressureCoefficients(ks, ptop, ak, bk)
 
 
-def vertical_coordinate(eta_value):
+def vertical_coordinate(eta_value) -> np.ndarray:
     """
     Equation (1) JRMS2006
     computes eta_v, the auxiliary variable vertical coordinate
     """
     return (eta_value - ETA_0) * math.pi * 0.5
 
 
-def compute_eta(ak, bk):
+def compute_eta(ak, bk) -> Tuple[np.ndarray, np.ndarray]:
     """
     Equation (1) JRMS2006
     eta is the vertical coordinate and eta_v is an auxiliary vertical coordinate
@@ -98,5 +108,5 @@ def compute_eta(ak, bk):
     return eta, eta_v
 
 
-def check_eta(ak, bk):
+def _check_eta(ak, bk):
     return compute_eta(ak, bk)

ndsl/grid/generation.py

@@ -1,7 +1,7 @@
 import dataclasses
 import functools
 import warnings
-from typing import Tuple
+from typing import Optional, Tuple
 
 import numpy as np
 
@@ -238,6 +238,8 @@ def __init__(
         deglat: float = 15.0,
         extdgrid: bool = False,
         eta_file: str = "None",
+        ak: Optional[np.ndarray] = None,
+        bk: Optional[np.ndarray] = None,
     ):
         self._grid_type = grid_type
         self._dx_const = dx_const
@@ -300,7 +302,7 @@ def __init__(
             self._ptop,
             self._ak,
             self._bk,
-        ) = self._set_hybrid_pressure_coefficients(eta_file)
+        ) = self._set_hybrid_pressure_coefficients(eta_file, ak, bk)
         self._ec1 = None
         self._ec2 = None
         self._ew1 = None
@@ -2147,7 +2149,12 @@ def _compute_area_c_cartesian(self):
         area_cgrid_64.data[:, :] = self._dx_const * self._dy_const
         return quantity_cast_to_model_float(self.quantity_factory, area_cgrid_64)
 
-    def _set_hybrid_pressure_coefficients(self, eta_file):
+    def _set_hybrid_pressure_coefficients(
+        self,
+        eta_file,
+        ak_data: Optional[np.ndarray] = None,
+        bk_data: Optional[np.ndarray] = None,
+    ):
         ks = self.quantity_factory.zeros(
             [],
             "",
@@ -2169,7 +2176,10 @@ def _set_hybrid_pressure_coefficients(self, eta_file):
             dtype=Float,
         )
         pressure_coefficients = eta.set_hybrid_pressure_coefficients(
-            self._npz, eta_file
+            self._npz,
+            eta_file,
+            ak_data,
+            bk_data,
         )
         ks = pressure_coefficients.ks
         ptop = pressure_coefficients.ptop

ndsl/stencils/basic_operations.py

@@ -0,0 +1,96 @@
+import gt4py.cartesian.gtscript as gtscript
+from gt4py.cartesian.gtscript import PARALLEL, computation, interval
+
+from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ
+
+
+def copy_defn(q_in: FloatField, q_out: FloatField):
+    """
+    Copy q_in to q_out.
+
+    Args:
+        q_in: input field
+        q_out: output field
+    """
+    with computation(PARALLEL), interval(...):
+        q_out = q_in
+
+
+def adjustmentfactor_stencil_defn(adjustment: FloatFieldIJ, q_out: FloatField):
+    """
+    Multiplies every element of q_out
+    by every element of the adjustment
+    field over the interval, replacing
+    the elements of q_out by the result
+    of the multiplication.
+
+    Args:
+        adjustment: adjustment field
+        q_out: output field
+    """
+    with computation(PARALLEL), interval(...):
+        q_out = q_out * adjustment
+
+
+def set_value_defn(q_out: FloatField, value: Float):
+    """
+    Sets every element of q_out to the
+    value specified by value argument.
+
+    Args:
+        q_out: output field
+        value: NDSL Float type
+    """
+    with computation(PARALLEL), interval(...):
+        q_out = value
+
+
+def adjust_divide_stencil(adjustment: FloatField, q_out: FloatField):
+    """
+    Divides every element of q_out
+    by every element of the adjustment
+    field over the interval, replacing
+    the elements of q_out by the result
+    of the multiplication.
+
+    Args:
+        adjustment: adjustment field
+        q_out: output field
+    """
+    with computation(PARALLEL), interval(...):
+        q_out = q_out / adjustment
+
+
+@gtscript.function
+def sign(a, b):
+    """
+    Defines asignb as the absolute value
+    of a, and checks if b is positive
+    or negative, assigning the analogus
+    sign value to asignb. asignb is returned
+
+    Args:
+        a: A number
+        b: A number
+    """
+    asignb = abs(a)
+    if b > 0:
+        asignb = asignb
+    else:
+        asignb = -asignb
+    return asignb
+
+
+@gtscript.function
+def dim(a, b):
+    """
+    Performs a check on the difference
+    between the values in arguments
+    a and b. The variable diff is set
+    to the difference between a and b
+    when the difference is positive,
+    otherwise it is set to zero. The
+    function returns the diff variable.
+    """
+    diff = a - b if a - b > 0 else 0
+    return diff

setup.py

@@ -19,7 +19,7 @@ def local_pkg(name: str, relative_path: str) -> str:
 requirements: List[str] = [
     local_pkg("gt4py", "external/gt4py"),
     local_pkg("dace", "external/dace"),
-    "mpi4py",
+    "mpi4py==3.1.5",
     "cftime",
     "xarray",
     "f90nml>=1.1.0",
@@ -50,6 +50,6 @@ def local_pkg(name: str, relative_path: str) -> str:
     packages=find_namespace_packages(include=["ndsl", "ndsl.*"]),
     include_package_data=True,
     url="https://github.com/NOAA-GFDL/NDSL",
-    version="2024.03.01",
+    version="2024.04.00-RC",
     zip_safe=False,
 )