variables.py

from veros import runtime_settings


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