diff --git a/ndsl/constants.py b/ndsl/constants.py index eb89bd17..3269672b 100644 --- a/ndsl/constants.py +++ b/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 diff --git a/ndsl/initialization/allocator.py b/ndsl/initialization/allocator.py index 869086f6..a6b3d30c 100644 --- a/ndsl/initialization/allocator.py +++ b/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, diff --git a/ndsl/namelist.py b/ndsl/namelist.py index 8df5c207..a9249573 100644 --- a/ndsl/namelist.py +++ b/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,6 +122,8 @@ 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 @@ -126,7 +131,8 @@ class NamelistDefaults: 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""" diff --git a/ndsl/stencils/basic_operations.py b/ndsl/stencils/basic_operations.py index b46123a3..4d706858 100644 --- a/ndsl/stencils/basic_operations.py +++ b/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): """ diff --git a/ndsl/stencils/testing/translate.py b/ndsl/stencils/testing/translate.py index b980ef7a..75e0f0e7 100644 --- a/ndsl/stencils/testing/translate.py +++ 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( diff --git a/ndsl/stencils/tridiag.py b/ndsl/stencils/tridiag.py new file mode 100644 index 00000000..e1fe50c0 --- /dev/null +++ b/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]