diff --git a/ndsl/grid/eta.py b/ndsl/grid/eta.py index 19663fde..90db8c4a 100644 --- a/ndsl/grid/eta.py +++ b/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,12 +87,10 @@ 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 @@ -88,7 +98,7 @@ def vertical_coordinate(eta_value): 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) diff --git a/ndsl/grid/generation.py b/ndsl/grid/generation.py index 12275d7d..275db563 100644 --- a/ndsl/grid/generation.py +++ b/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 diff --git a/ndsl/stencils/basic_operations.py b/ndsl/stencils/basic_operations.py new file mode 100644 index 00000000..b46123a3 --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index d81523d1..bb24006a 100644 --- a/setup.py +++ b/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, ) diff --git a/tests/test_basic_operations.py b/tests/test_basic_operations.py new file mode 100644 index 00000000..0d707240 --- /dev/null +++ b/tests/test_basic_operations.py @@ -0,0 +1,212 @@ +import numpy as np + +from ndsl import ( + CompilationConfig, + DaceConfig, + DaCeOrchestration, + GridIndexing, + Quantity, + RunMode, + StencilConfig, + StencilFactory, +) +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.stencils import basic_operations as basic + + +nx = 20 +ny = 20 +nz = 79 +nhalo = 0 +backend = "numpy" + +dace_config = DaceConfig( + communicator=None, backend=backend, orchestration=DaCeOrchestration.Python +) + +compilation_config = CompilationConfig( + backend=backend, + rebuild=True, + validate_args=True, + format_source=False, + device_sync=False, + run_mode=RunMode.BuildAndRun, + use_minimal_caching=False, +) + +stencil_config = StencilConfig( + compare_to_numpy=False, + compilation_config=compilation_config, + dace_config=dace_config, +) + +grid_indexing = GridIndexing( + domain=(nx, ny, nz), + n_halo=nhalo, + south_edge=True, + north_edge=True, + west_edge=True, + east_edge=True, +) + +stencil_factory = StencilFactory(config=stencil_config, grid_indexing=grid_indexing) + + +class Copy: + def __init__(self, stencil_factory: StencilFactory): + grid_indexing = stencil_factory.grid_indexing + self._copy_stencil = stencil_factory.from_origin_domain( + basic.copy_defn, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + def __call__( + self, + f_in: FloatField, + f_out: FloatField, + ): + self._copy_stencil(f_in, f_out) + + +class AdjustmentFactor: + def __init__(self, stencil_factory: StencilFactory): + grid_indexing = stencil_factory.grid_indexing + self._adjustmentfactor_stencil = stencil_factory.from_origin_domain( + basic.adjustmentfactor_stencil_defn, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + def __call__( + self, + factor: FloatFieldIJ, + f_out: FloatField, + ): + self._adjustmentfactor_stencil(factor, f_out) + + +class SetValue: + def __init__(self, stencil_factory: StencilFactory): + grid_indexing = stencil_factory.grid_indexing + self._set_value_stencil = stencil_factory.from_origin_domain( + basic.set_value_defn, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + def __call__( + self, + f_out: FloatField, + value: Float, + ): + self._set_value_stencil(f_out, value) + + +class AdjustDivide: + def __init__(self, stencil_factory: StencilFactory): + grid_indexing = stencil_factory.grid_indexing + self._adjust_divide_stencil = stencil_factory.from_origin_domain( + basic.adjust_divide_stencil, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + def __call__( + self, + factor: FloatField, + f_out: FloatField, + ): + self._adjust_divide_stencil(factor, f_out) + + +def test_copy(): + copy = Copy(stencil_factory) + + infield = Quantity( + data=np.zeros([20, 20, 79]), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + outfield = Quantity( + data=np.ones([20, 20, 79]), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + copy(f_in=infield.data, f_out=outfield.data) + + assert (infield.data == outfield.data).all() + + +def test_adjustmentfactor(): + adfact = AdjustmentFactor(stencil_factory) + + factorfield = Quantity( + data=np.full(shape=[20, 20], fill_value=2.0), + dims=[X_DIM, Y_DIM], + units="m", + ) + + outfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=2.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + testfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=4.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + adfact(factor=factorfield.data, f_out=outfield.data) + assert (outfield.data == testfield.data).all() + + +def test_setvalue(): + setvalue = SetValue(stencil_factory) + + outfield = Quantity( + data=np.zeros(shape=[20, 20, 79]), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + testfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=2.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + setvalue(f_out=outfield.data, value=2.0) + + assert (outfield.data == testfield.data).all() + + +def test_adjustdivide(): + addiv = AdjustDivide(stencil_factory) + + factorfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=2.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + outfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=2.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + testfield = Quantity( + data=np.full(shape=[20, 20, 79], fill_value=1.0), + dims=[X_DIM, Y_DIM, Z_DIM], + units="m", + ) + + addiv(factor=factorfield.data, f_out=outfield.data) + + assert (outfield.data == testfield.data).all()