from veros import runtime_settings
[docs]
class Variable:
def __init__(
self,
name,
dims,
units="",
long_description="",
dtype=None,
time_dependent=True,
scale=1.0,
write_to_restart=False,
extra_attributes=None,
mask=None,
active=True,
initial=None,
):
self.name = name
self.dims = dims
self.units = units
self.long_description = long_description
self.dtype = dtype
self.time_dependent = time_dependent
self.scale = scale
self.write_to_restart = write_to_restart
self.active = active
self.initial = initial
self.get_mask = lambda settings, vs: None
if mask is not None:
if not callable(mask):
raise TypeError("mask argument has to be callable")
self.get_mask = mask
elif isinstance(dims, tuple):
if dims[:3] in DEFAULT_MASKS:
self.get_mask = DEFAULT_MASKS[dims[:3]]
elif dims[:2] in DEFAULT_MASKS:
self.get_mask = DEFAULT_MASKS[dims[:2]]
#: Additional attributes to be written in netCDF output
self.extra_attributes = extra_attributes or {}
def __repr__(self):
attr_str = []
for v in vars(self):
attr_str.append(f"{v}={getattr(self, v)}")
attr_str = ", ".join(attr_str)
return f"{self.__class__.__qualname__}({attr_str})"
# fill value for netCDF output (invalid floating data is replaced by this value)
FLOAT_FILL_VALUE = -1e18
#
XT = ("xt",)
XU = ("xu",)
YT = ("yt",)
YU = ("yu",)
ZT = ("zt",)
ZW = ("zw",)
T_HOR = ("xt", "yt")
U_HOR = ("xu", "yt")
V_HOR = ("xt", "yu")
ZETA_HOR = ("xu", "yu")
T_GRID = ("xt", "yt", "zt")
U_GRID = ("xu", "yt", "zt")
V_GRID = ("xt", "yu", "zt")
W_GRID = ("xt", "yt", "zw")
ZETA_GRID = ("xu", "yu", "zt")
TIMESTEPS = ("timesteps",)
ISLE = ("isle",)
TENSOR_COMP = ("tensor1", "tensor2")
# those are written to netCDF output by default
BASE_DIMENSIONS = XT + XU + YT + YU + ZT + ZW + ISLE
GHOST_DIMENSIONS = ("xt", "yt", "xu", "yu")
# these are the settings that are getting used to determine shapes
DIM_TO_SHAPE_VAR = {
"xt": "nx",
"xu": "nx",
"yt": "ny",
"yu": "ny",
"zt": "nz",
"zw": "nz",
"timesteps": 3,
"tensor1": 2,
"tensor2": 2,
"isle": 0,
}
DEFAULT_MASKS = {
T_HOR: lambda settings, vs: vs.maskT[:, :, -1],
U_HOR: lambda settings, vs: vs.maskU[:, :, -1],
V_HOR: lambda settings, vs: vs.maskV[:, :, -1],
ZETA_HOR: lambda settings, vs: vs.maskZ[:, :, -1],
T_GRID: lambda settings, vs: vs.maskT,
U_GRID: lambda settings, vs: vs.maskU,
V_GRID: lambda settings, vs: vs.maskV,
W_GRID: lambda settings, vs: vs.maskW,
ZETA_GRID: lambda settings, vs: vs.maskZ,
}
# custom mask for streamfunction
def _get_psi_mask(settings, vs):
if not settings.enable_streamfunction:
return vs.maskT[:, :, -1]
# eroded around the edges
return vs.maskZ[:, :, -1] | ~vs.isle_boundary_mask
def get_fill_value(dtype):
import numpy as onp
if onp.issubdtype(dtype, onp.floating):
return FLOAT_FILL_VALUE
return onp.iinfo(dtype).max
def get_shape(dimensions, grid, include_ghosts=True, local=True):
from veros.routines import CURRENT_CONTEXT
from veros.distributed import SCATTERED_DIMENSIONS
if grid is None:
return ()
px, py = runtime_settings.num_proc
grid_shapes = dict(dimensions)
if local and CURRENT_CONTEXT.is_dist_safe:
for pxi, dims in zip((px, py), SCATTERED_DIMENSIONS):
for dim in dims:
if dim not in grid_shapes:
continue
grid_shapes[dim] = grid_shapes[dim] // pxi
if include_ghosts:
for d in GHOST_DIMENSIONS:
if d in grid_shapes:
grid_shapes[d] += 4
shape = []
for grid_dim in grid:
if isinstance(grid_dim, int):
shape.append(grid_dim)
continue
if grid_dim not in grid_shapes:
raise ValueError(f"unrecognized dimension {grid_dim}")
shape.append(grid_shapes[grid_dim])
return tuple(shape)
def remove_ghosts(array, dims):
if dims is None:
# scalar
return array
ghost_mask = tuple(slice(2, -2) if dim in GHOST_DIMENSIONS else slice(None) for dim in dims)
return array[ghost_mask]
VARIABLES = {
# scalars
"tau": Variable(
"Index of current time step",
None,
"",
"Index of current time step",
dtype="int32",
initial=1,
write_to_restart=True,
),
"taup1": Variable(
"Index of next time step", None, "", "Index of next time step", dtype="int32", initial=2, write_to_restart=True
),
"taum1": Variable(
"Index of last time step", None, "", "Index of last time step", dtype="int32", initial=0, write_to_restart=True
),
"time": Variable(
"Current time",
None,
"",
"Current time",
write_to_restart=True,
),
"itt": Variable("Current iteration", None, "", "Current iteration", dtype="int32", initial=0),
# base variables
"dxt": Variable("Zonal T-grid spacing", XT, "m", "Zonal (x) spacing of T-grid point", time_dependent=False),
"dxu": Variable("Zonal U-grid spacing", XU, "m", "Zonal (x) spacing of U-grid point", time_dependent=False),
"dyt": Variable(
"Meridional T-grid spacing", YT, "m", "Meridional (y) spacing of T-grid point", time_dependent=False
),
"dyu": Variable(
"Meridional U-grid spacing", YU, "m", "Meridional (y) spacing of U-grid point", time_dependent=False
),
"zt": Variable(
"Vertical coordinate (T)",
ZT,
"m",
"Vertical coordinate",
time_dependent=False,
extra_attributes={"positive": "up"},
),
"zw": Variable(
"Vertical coordinate (W)",
ZW,
"m",
"Vertical coordinate",
time_dependent=False,
extra_attributes={"positive": "up"},
),
"dzt": Variable("Vertical spacing (T)", ZT, "m", "Vertical spacing", time_dependent=False),
"dzw": Variable("Vertical spacing (W)", ZW, "m", "Vertical spacing", time_dependent=False),
"cost": Variable("Metric factor (T)", YT, "1", "Metric factor for spherical coordinates", time_dependent=False),
"cosu": Variable("Metric factor (U)", YU, "1", "Metric factor for spherical coordinates", time_dependent=False),
"tantr": Variable("Metric factor", YT, "1", "Metric factor for spherical coordinates", time_dependent=False),
"coriolis_t": Variable(
"Coriolis frequency", T_HOR, "1/s", "Coriolis frequency at T grid point", time_dependent=False
),
"kbot": Variable(
"Index of deepest cell",
T_HOR,
"",
"Index of the deepest grid cell (counting from 1, 0 means all land)",
dtype="int32",
time_dependent=False,
),
"ht": Variable("Total depth (T)", T_HOR, "m", "Total depth of the water column", time_dependent=False),
"hu": Variable("Total depth (U)", U_HOR, "m", "Total depth of the water column", time_dependent=False),
"hv": Variable("Total depth (V)", V_HOR, "m", "Total depth of the water column", time_dependent=False),
"hur": Variable(
"Total depth (U), masked", U_HOR, "m", "Total depth of the water column (masked)", time_dependent=False
),
"hvr": Variable(
"Total depth (V), masked", V_HOR, "m", "Total depth of the water column (masked)", time_dependent=False
),
"beta": Variable(
"Change of Coriolis freq.", T_HOR, "1/(ms)", "Change of Coriolis frequency with latitude", time_dependent=False
),
"area_t": Variable("Area of T-box", T_HOR, "m^2", "Area of T-box", time_dependent=False),
"area_u": Variable("Area of U-box", U_HOR, "m^2", "Area of U-box", time_dependent=False),
"area_v": Variable("Area of V-box", V_HOR, "m^2", "Area of V-box", time_dependent=False),
"maskT": Variable(
"Mask for tracer points",
T_GRID,
"",
"Mask in physical space for tracer points",
time_dependent=False,
dtype="bool",
),
"maskU": Variable(
"Mask for U points",
U_GRID,
"",
"Mask in physical space for U points",
time_dependent=False,
dtype="bool",
),
"maskV": Variable(
"Mask for V points",
V_GRID,
"",
"Mask in physical space for V points",
time_dependent=False,
dtype="bool",
),
"maskW": Variable(
"Mask for W points",
W_GRID,
"",
"Mask in physical space for W points",
time_dependent=False,
dtype="bool",
),
"maskZ": Variable(
"Mask for Zeta points",
ZETA_GRID,
"",
"Mask in physical space for Zeta points",
time_dependent=False,
dtype="bool",
),
"rho": Variable(
"Density",
T_GRID + TIMESTEPS,
"kg/m^3",
"In-situ density anomaly, relative to the surface mean value of 1024 kg/m^3",
write_to_restart=True,
),
"prho": Variable(
"Potential density",
T_GRID,
"kg/m^3",
"Potential density anomaly, relative to the surface mean value of 1024 kg/m^3 "
"(identical to in-situ density anomaly for equation of state type 1, 2, and 4)",
),
"int_drhodT": Variable(
"Der. of dyn. enthalpy by temperature",
T_GRID + TIMESTEPS,
"kg / (m^2 deg C)",
"Partial derivative of dynamic enthalpy by temperature",
write_to_restart=True,
),
"int_drhodS": Variable(
"Der. of dyn. enthalpy by salinity",
T_GRID + TIMESTEPS,
"kg / (m^2 g / kg)",
"Partial derivative of dynamic enthalpy by salinity",
write_to_restart=True,
),
"Nsqr": Variable(
"Square of stability frequency",
W_GRID + TIMESTEPS,
"1/s^2",
"Square of stability frequency",
write_to_restart=True,
),
"Hd": Variable("Dynamic enthalpy", T_GRID + TIMESTEPS, "m^2/s^2", "Dynamic enthalpy", write_to_restart=True),
"dHd": Variable(
"Change of dyn. enth. by adv.",
T_GRID + TIMESTEPS,
"m^2/s^3",
"Change of dynamic enthalpy due to advection",
write_to_restart=True,
),
"temp": Variable("Temperature", T_GRID + TIMESTEPS, "deg C", "Conservative temperature", write_to_restart=True),
"dtemp": Variable(
"Temperature tendency",
T_GRID + TIMESTEPS,
"deg C/s",
"Conservative temperature tendency",
write_to_restart=True,
),
"salt": Variable("Salinity", T_GRID + TIMESTEPS, "g/kg", "Salinity", write_to_restart=True),
"dsalt": Variable("Salinity tendency", T_GRID + TIMESTEPS, "g/(kg s)", "Salinity tendency", write_to_restart=True),
"dtemp_vmix": Variable(
"Change of temp. by vertical mixing",
T_GRID,
"deg C/s",
"Change of temperature due to vertical mixing",
),
"dtemp_hmix": Variable(
"Change of temp. by horizontal mixing",
T_GRID,
"deg C/s",
"Change of temperature due to horizontal mixing",
),
"dsalt_vmix": Variable(
"Change of sal. by vertical mixing",
T_GRID,
"deg C/s",
"Change of salinity due to vertical mixing",
),
"dsalt_hmix": Variable(
"Change of sal. by horizontal mixing",
T_GRID,
"deg C/s",
"Change of salinity due to horizontal mixing",
),
"dtemp_iso": Variable(
"Change of temp. by isop. mixing",
T_GRID,
"deg C/s",
"Change of temperature due to isopycnal mixing plus skew mixing",
),
"dsalt_iso": Variable(
"Change of sal. by isop. mixing",
T_GRID,
"deg C/s",
"Change of salinity due to isopycnal mixing plus skew mixing",
),
"forc_temp_surface": Variable(
"Surface temperature flux",
T_HOR,
"m deg C/s",
"Surface temperature flux",
),
"forc_salt_surface": Variable(
"Surface salinity flux",
T_HOR,
"m g/s kg",
"Surface salinity flux",
),
"u": Variable("Zonal velocity", U_GRID + TIMESTEPS, "m/s", "Zonal velocity", write_to_restart=True),
"v": Variable("Meridional velocity", V_GRID + TIMESTEPS, "m/s", "Meridional velocity", write_to_restart=True),
"w": Variable("Vertical velocity", W_GRID + TIMESTEPS, "m/s", "Vertical velocity", write_to_restart=True),
"du": Variable(
"Zonal velocity tendency", U_GRID + TIMESTEPS, "m/s", "Zonal velocity tendency", write_to_restart=True
),
"dv": Variable(
"Meridional velocity tendency", V_GRID + TIMESTEPS, "m/s", "Meridional velocity tendency", write_to_restart=True
),
"du_cor": Variable("Change of u by Coriolis force", U_GRID, "m/s^2", "Change of u due to Coriolis force"),
"dv_cor": Variable("Change of v by Coriolis force", V_GRID, "m/s^2", "Change of v due to Coriolis force"),
"du_mix": Variable(
"Change of u by vertical mixing", U_GRID, "m/s^2", "Change of u due to implicit vertical mixing"
),
"dv_mix": Variable(
"Change of v by vertical mixing", V_GRID, "m/s^2", "Change of v due to implicit vertical mixing"
),
"du_adv": Variable("Change of u by advection", U_GRID, "m/s^2", "Change of u due to advection"),
"dv_adv": Variable("Change of v by advection", V_GRID, "m/s^2", "Change of v due to advection"),
"p_hydro": Variable("Hydrostatic pressure", T_GRID, "m^2/s^2", "Hydrostatic pressure"),
"kappaM": Variable("Vertical viscosity", T_GRID, "m^2/s", "Vertical viscosity"),
"kappaH": Variable("Vertical diffusivity", W_GRID, "m^2/s", "Vertical diffusivity"),
"surface_taux": Variable(
"Surface wind stress",
U_HOR,
"N/m^2",
"Zonal surface wind stress",
),
"surface_tauy": Variable(
"Surface wind stress",
V_HOR,
"N/m^2",
"Meridional surface wind stress",
),
"forc_rho_surface": Variable("Surface density flux", T_HOR, "kg / (m^2 s)", "Surface potential density flux"),
"ssh": Variable(
"Sea surface height",
T_HOR,
"m",
"Sea surface height",
active=lambda settings: not settings.enable_streamfunction,
),
"psi": Variable(
lambda settings: "Streamfunction" if settings.enable_streamfunction else "Surface pressure",
lambda settings: ZETA_HOR + TIMESTEPS if settings.enable_streamfunction else T_HOR + TIMESTEPS,
lambda settings: "m^3/s" if settings.enable_streamfunction else "m^2/s^2",
lambda settings: "Barotropic streamfunction" if settings.enable_streamfunction else "Surface pressure",
write_to_restart=True,
mask=_get_psi_mask,
),
"dpsi": Variable(
"Streamfunction tendency",
ZETA_HOR + TIMESTEPS,
"m^3/s^2",
"Streamfunction tendency",
write_to_restart=True,
active=lambda settings: settings.enable_streamfunction,
),
"land_map": Variable(
"Land map",
T_HOR,
"",
"Land map",
dtype="int32",
active=lambda settings: settings.enable_streamfunction,
),
"isle": Variable(
"Island number",
ISLE,
"",
"Island number",
dtype="int32",
active=lambda settings: settings.enable_streamfunction,
),
"psin": Variable(
"Boundary streamfunction",
ZETA_HOR + ISLE,
"m^3/s",
"Boundary streamfunction",
time_dependent=False,
mask=_get_psi_mask,
active=lambda settings: settings.enable_streamfunction,
),
"dpsin": Variable(
"Boundary streamfunction factor",
ISLE + TIMESTEPS,
"m^3/s^2",
"Boundary streamfunction factor",
write_to_restart=True,
active=lambda settings: settings.enable_streamfunction,
),
"line_psin": Variable(
"Boundary line integrals",
ISLE + ISLE,
"m^4/s^2",
"Boundary line integrals",
time_dependent=False,
active=lambda settings: settings.enable_streamfunction,
),
"isle_boundary_mask": Variable(
"Island boundary mask",
T_HOR,
"",
"Island boundary mask",
time_dependent=False,
dtype="bool",
active=lambda settings: settings.enable_streamfunction,
),
"line_dir_south_mask": Variable(
"Line integral mask",
T_HOR + ISLE,
"",
"Line integral mask",
time_dependent=False,
dtype="bool",
active=lambda settings: settings.enable_streamfunction,
),
"line_dir_north_mask": Variable(
"Line integral mask",
T_HOR + ISLE,
"",
"Line integral mask",
time_dependent=False,
dtype="bool",
active=lambda settings: settings.enable_streamfunction,
),
"line_dir_east_mask": Variable(
"Line integral mask",
T_HOR + ISLE,
"",
"Line integral mask",
time_dependent=False,
dtype="bool",
active=lambda settings: settings.enable_streamfunction,
),
"line_dir_west_mask": Variable(
"Line integral mask",
T_HOR + ISLE,
"",
"Line integral mask",
time_dependent=False,
dtype="bool",
active=lambda settings: settings.enable_streamfunction,
),
"K_gm": Variable("Skewness diffusivity", W_GRID, "m^2/s", "GM diffusivity, either constant or from EKE model"),
"K_iso": Variable("Isopycnal diffusivity", W_GRID, "m^2/s", "Along-isopycnal diffusivity"),
"K_diss_v": Variable(
"Dissipation of kinetic Energy",
W_GRID,
"m^2/s^3",
"Kinetic energy dissipation by vertical, rayleigh and bottom friction",
write_to_restart=True,
),
"K_diss_bot": Variable(
"Dissipation of kinetic Energy", W_GRID, "m^2/s^3", "Mean energy dissipation by bottom and rayleigh friction"
),
"K_diss_h": Variable(
"Dissipation of kinetic Energy", W_GRID, "m^2/s^3", "Kinetic energy dissipation by horizontal friction"
),
"K_diss_gm": Variable(
"Dissipation of mean energy",
W_GRID,
"m^2/s^3",
"Mean energy dissipation by GM (TRM formalism only)",
),
"P_diss_v": Variable(
"Dissipation of potential Energy", W_GRID, "m^2/s^3", "Potential energy dissipation by vertical diffusion"
),
"P_diss_nonlin": Variable(
"Dissipation of potential Energy",
W_GRID,
"m^2/s^3",
"Potential energy dissipation by nonlinear equation of state",
),
"P_diss_iso": Variable(
"Dissipation of potential Energy", W_GRID, "m^2/s^3", "Potential energy dissipation by isopycnal mixing"
),
"P_diss_skew": Variable(
"Dissipation of potential Energy", W_GRID, "m^2/s^3", "Potential energy dissipation by GM (w/o TRM)"
),
"P_diss_hmix": Variable(
"Dissipation of potential Energy", W_GRID, "m^2/s^3", "Potential energy dissipation by horizontal mixing"
),
"P_diss_adv": Variable(
"Dissipation of potential Energy", W_GRID, "m^2/s^3", "Potential energy dissipation by advection"
),
"P_diss_sources": Variable(
"Dissipation of potential Energy",
W_GRID,
"m^2/s^3",
"Potential energy dissipation by external sources (e.g. restoring zones)",
),
"u_wgrid": Variable("U on W grid", W_GRID, "m/s", "Zonal velocity interpolated to W grid points"),
"v_wgrid": Variable("V on W grid", W_GRID, "m/s", "Meridional velocity interpolated to W grid points"),
"w_wgrid": Variable("W on W grid", W_GRID, "m/s", "Vertical velocity interpolated to W grid points"),
"xt": Variable(
"Zonal coordinate (T)",
XT,
lambda settings: "degrees_east" if settings.coord_degree else "km",
"Zonal (x) coordinate of T-grid point",
time_dependent=False,
scale=lambda settings: 1 if settings.coord_degree else 1e-3,
),
"xu": Variable(
"Zonal coordinate (U)",
XU,
lambda settings: "degrees_east" if settings.coord_degree else "km",
"Zonal (x) coordinate of U-grid point",
time_dependent=False,
scale=lambda settings: 1 if settings.coord_degree else 1e-3,
),
"yt": Variable(
"Meridional coordinate (T)",
YT,
lambda settings: "degrees_north" if settings.coord_degree else "km",
"Meridional (y) coordinate of T-grid point",
time_dependent=False,
scale=lambda settings: 1 if settings.coord_degree else 1e-3,
),
"yu": Variable(
"Meridional coordinate (U)",
YU,
lambda settings: "degrees_north" if settings.coord_degree else "km",
"Meridional (y) coordinate of U-grid point",
time_dependent=False,
scale=lambda settings: 1 if settings.coord_degree else 1e-3,
),
"temp_source": Variable(
"Source of temperature",
T_GRID,
"K/s",
"Non-conservative source of temperature",
active=lambda settings: settings.enable_tempsalt_sources,
),
"salt_source": Variable(
"Source of salt",
T_GRID,
"g/(kg s)",
"Non-conservative source of salt",
active=lambda settings: settings.enable_tempsalt_sources,
),
"u_source": Variable(
"Source of zonal velocity",
U_GRID,
"m/s^2",
"Non-conservative source of zonal velocity",
active=lambda settings: settings.enable_momentum_sources,
),
"v_source": Variable(
"Source of meridional velocity",
V_GRID,
"m/s^2",
"Non-conservative source of meridional velocity",
active=lambda settings: settings.enable_momentum_sources,
),
"K_11": Variable(
"Isopycnal mixing coefficient",
T_GRID,
"m^2/s",
"Isopycnal mixing tensor component",
active=lambda settings: settings.enable_neutral_diffusion,
),
"K_22": Variable(
"Isopycnal mixing coefficient",
T_GRID,
"m^2/s",
"Isopycnal mixing tensor component",
active=lambda settings: settings.enable_neutral_diffusion,
),
"K_33": Variable(
"Isopycnal mixing coefficient",
T_GRID,
"m^2/s",
"Isopycnal mixing tensor component",
active=lambda settings: settings.enable_neutral_diffusion,
),
"Ai_ez": Variable(
"Isopycnal diffusion coefficient",
T_GRID + TENSOR_COMP,
"Vertical isopycnal diffusion coefficient on eastern face of T cell",
"1",
active=lambda settings: settings.enable_neutral_diffusion,
),
"Ai_nz": Variable(
"Isopycnal diffusion coefficient",
T_GRID + TENSOR_COMP,
"Vertical isopycnal diffusion coefficient on northern face of T cell",
"1",
active=lambda settings: settings.enable_neutral_diffusion,
),
"Ai_bx": Variable(
"Isopycnal diffusion coefficient",
T_GRID + TENSOR_COMP,
"Zonal isopycnal diffusion coefficient on bottom face of T cell",
"1",
active=lambda settings: settings.enable_neutral_diffusion,
),
"Ai_by": Variable(
"Isopycnal diffusion coefficient",
T_GRID + TENSOR_COMP,
"Meridional isopycnal diffusion coefficient on bottom face of T cell",
"1",
active=lambda settings: settings.enable_neutral_diffusion,
),
"B1_gm": Variable(
"Zonal component of GM streamfunction",
V_GRID,
"m^2/s",
"Zonal component of GM streamfunction",
active=lambda settings: settings.enable_skew_diffusion,
),
"B2_gm": Variable(
"Meridional component of GM streamfunction",
U_GRID,
"m^2/s",
"Meridional component of GM streamfunction",
active=lambda settings: settings.enable_skew_diffusion,
),
"r_bot_var_u": Variable(
"Bottom friction coeff.",
U_HOR,
"1/s",
"Zonal bottom friction coefficient",
active=lambda settings: settings.enable_bottom_friction_var,
),
"r_bot_var_v": Variable(
"Bottom friction coeff.",
V_HOR,
"1/s",
"Meridional bottom friction coefficient",
active=lambda settings: settings.enable_bottom_friction_var,
),
"kappa_gm": Variable(
"Vertical diffusivity",
W_GRID,
"m^2/s",
"Vertical diffusivity",
active=lambda settings: settings.enable_TEM_friction,
),
"tke": Variable(
"Turbulent kinetic energy",
W_GRID + TIMESTEPS,
"m^2/s^2",
"Turbulent kinetic energy",
write_to_restart=True,
active=lambda settings: settings.enable_tke,
),
"sqrttke": Variable(
"Square-root of TKE",
W_GRID,
"m/s",
"Square-root of TKE",
active=lambda settings: settings.enable_tke,
),
"dtke": Variable(
"Turbulent kinetic energy tendency",
W_GRID + TIMESTEPS,
"m^2/s^3",
"Turbulent kinetic energy tendency",
write_to_restart=True,
active=lambda settings: settings.enable_tke,
),
"Prandtlnumber": Variable(
"Prandtl number",
W_GRID,
"",
"Prandtl number",
active=lambda settings: settings.enable_tke,
),
"mxl": Variable(
"Mixing length",
W_GRID,
"m",
"Mixing length",
active=lambda settings: settings.enable_tke,
),
"forc_tke_surface": Variable(
"TKE surface flux",
T_HOR,
"m^3/s^3",
"TKE surface flux",
active=lambda settings: settings.enable_tke,
),
"tke_diss": Variable(
"TKE dissipation",
W_GRID,
"m^2/s^3",
"TKE dissipation",
active=lambda settings: settings.enable_tke,
),
"tke_surf_corr": Variable(
"Correction of TKE surface flux",
T_HOR,
"m^3/s^3",
"Correction of TKE surface flux",
active=lambda settings: settings.enable_tke,
),
"eke": Variable(
"meso-scale energy",
W_GRID + TIMESTEPS,
"m^2/s^2",
"meso-scale energy",
write_to_restart=True,
active=lambda settings: settings.enable_eke,
),
"deke": Variable(
"meso-scale energy tendency",
W_GRID + TIMESTEPS,
"m^2/s^3",
"meso-scale energy tendency",
write_to_restart=True,
active=lambda settings: settings.enable_eke,
),
"sqrteke": Variable(
"square-root of eke",
W_GRID,
"m/s",
"square-root of eke",
active=lambda settings: settings.enable_eke,
),
"L_rossby": Variable(
"Rossby radius",
T_HOR,
"m",
"Rossby radius",
active=lambda settings: settings.enable_eke,
),
"L_rhines": Variable(
"Rhines scale",
W_GRID,
"m",
"Rhines scale",
active=lambda settings: settings.enable_eke,
),
"eke_len": Variable(
"Eddy length scale",
T_GRID,
"m",
"Eddy length scale",
active=lambda settings: settings.enable_eke,
),
"eke_diss_iw": Variable(
"Dissipation of EKE to IW",
W_GRID,
"m^2/s^3",
"Dissipation of EKE to internal waves",
active=lambda settings: settings.enable_eke,
),
"eke_diss_tke": Variable(
"Dissipation of EKE to TKE",
W_GRID,
"m^2/s^3",
"Dissipation of EKE to TKE",
active=lambda settings: settings.enable_eke,
),
"E_iw": Variable(
"Internal wave energy",
W_GRID + TIMESTEPS,
"m^2/s^2",
"Internal wave energy",
write_to_restart=True,
active=lambda settings: settings.enable_idemix,
),
"dE_iw": Variable(
"Internal wave energy tendency",
W_GRID + TIMESTEPS,
"m^2/s^2",
"Internal wave energy tendency",
write_to_restart=True,
active=lambda settings: settings.enable_idemix,
),
"c0": Variable(
"Vertical IW group velocity",
W_GRID,
"m/s",
"Vertical internal wave group velocity",
active=lambda settings: settings.enable_idemix,
),
"v0": Variable(
"Horizontal IW group velocity",
W_GRID,
"m/s",
"Horizontal internal wave group velocity",
active=lambda settings: settings.enable_idemix,
),
"alpha_c": Variable(
"?",
W_GRID,
"?",
"?",
active=lambda settings: settings.enable_idemix,
),
"iw_diss": Variable(
"IW dissipation",
W_GRID,
"m^2/s^3",
"Internal wave dissipation",
active=lambda settings: settings.enable_idemix,
),
"forc_iw_surface": Variable(
"IW surface forcing",
T_HOR,
"m^3/s^3",
"Internal wave surface forcing",
time_dependent=False,
active=lambda settings: settings.enable_idemix,
),
"forc_iw_bottom": Variable(
"IW bottom forcing",
T_HOR,
"m^3/s^3",
"Internal wave bottom forcing",
time_dependent=False,
active=lambda settings: settings.enable_idemix,
),
}
def manifest_metadata(var_meta, settings):
"""Evaluate callable metadata fields given the current settings."""
from copy import copy
out = {}
for var_name, var_val in var_meta.items():
var_val = copy(var_val)
for attr, attr_val in vars(var_val).items():
if callable(attr_val) and attr != "get_mask":
setattr(var_val, attr, attr_val(settings))
out[var_name] = var_val
return out
def allocate(dimensions, grid, dtype=None, include_ghosts=True, local=True, fill=0):
from veros.core.operators import numpy as npx
if dtype is None:
dtype = runtime_settings.float_type
shape = get_shape(dimensions, grid, include_ghosts=include_ghosts, local=local)
out = npx.full(shape, fill, dtype=dtype)
if runtime_settings.backend == "numpy":
out.flags.writeable = False
return out