tke.py

from veros import veros_kernel, veros_routine, KernelOutput
from veros.variables import allocate
from veros.core import advection, utilities
from veros.core.operators import update, update_add, at, for_loop, numpy as npx


@veros_routine
def set_tke_diffusivities(state):
    vs = state.variables
    settings = state.settings

    if settings.enable_tke:
        tke_diff_out = set_tke_diffusivities_kernel(state)
        vs.update(tke_diff_out)
    else:
        vs.kappaM = update(vs.kappaM, at[...], settings.kappaM_0)
        vs.kappaH = npx.where(vs.Nsqr[..., vs.tau] < 0.0, 1.0, settings.kappaH_0)


@veros_kernel
def set_tke_diffusivities_kernel(state):
    """
    set vertical diffusivities based on TKE model
    """
    vs = state.variables
    settings = state.settings

    Rinumber = allocate(state.dimensions, ("xt", "yt", "zt"))

    vs.sqrttke = npx.sqrt(npx.maximum(0.0, vs.tke[:, :, :, vs.tau]))
    """
    calculate buoyancy length scale
    """
    vs.mxl = npx.sqrt(2) * vs.sqrttke / npx.sqrt(npx.maximum(1e-12, vs.Nsqr[:, :, :, vs.tau])) * vs.maskW

    """
    apply limits for mixing length
    """
    if settings.tke_mxl_choice == 1:
        """
        bounded by the distance to surface/bottom
        """
        vs.mxl = npx.minimum(
            npx.minimum(vs.mxl, -vs.zw[npx.newaxis, npx.newaxis, :] + vs.dzw[npx.newaxis, npx.newaxis, :] * 0.5),
            vs.ht[:, :, npx.newaxis] + vs.zw[npx.newaxis, npx.newaxis, :],
        )
        vs.mxl = npx.maximum(vs.mxl, settings.mxl_min)
    elif settings.tke_mxl_choice == 2:
        """
        bound length scale as in mitgcm/OPA code
        """
        nz = state.dimensions["zt"]

        def backwards_pass(kinv, mxl):
            k = nz - kinv - 1
            return update(mxl, at[:, :, k], npx.minimum(mxl[:, :, k], mxl[:, :, k + 1] + vs.dzt[k + 1]))

        vs.mxl = for_loop(1, nz, backwards_pass, vs.mxl)
        vs.mxl = update(vs.mxl, at[:, :, -1], npx.minimum(vs.mxl[:, :, -1], settings.mxl_min + vs.dzt[-1]))

        def forwards_pass(k, mxl):
            return update(mxl, at[:, :, k], npx.minimum(mxl[:, :, k], mxl[:, :, k - 1] + vs.dzt[k]))

        vs.mxl = for_loop(1, nz, forwards_pass, vs.mxl)
        vs.mxl = npx.maximum(vs.mxl, settings.mxl_min)
    else:
        raise ValueError("unknown mixing length choice in tke_mxl_choice")

    """
    calculate viscosity and diffusivity based on Prandtl number
    """
    vs.K_diss_v = utilities.enforce_boundaries(vs.K_diss_v, settings.enable_cyclic_x)
    vs.kappaM = update(vs.kappaM, at[...], npx.minimum(settings.kappaM_max, settings.c_k * vs.mxl * vs.sqrttke))
    Rinumber = update(
        Rinumber, at[...], vs.Nsqr[:, :, :, vs.tau] / npx.maximum(vs.K_diss_v / npx.maximum(1e-12, vs.kappaM), 1e-12)
    )
    if settings.enable_idemix:
        Rinumber = update(
            Rinumber,
            at[...],
            npx.minimum(
                Rinumber,
                vs.kappaM * vs.Nsqr[:, :, :, vs.tau] / npx.maximum(1e-12, vs.alpha_c * vs.E_iw[:, :, :, vs.tau] ** 2),
            ),
        )

    if settings.enable_Prandtl_tke:
        vs.Prandtlnumber = npx.maximum(1.0, npx.minimum(10, 6.6 * Rinumber))
    else:
        vs.Prandtlnumber = update(vs.Prandtlnumber, at[...], settings.Prandtl_tke0)

    vs.kappaH = npx.maximum(settings.kappaH_min, vs.kappaM / vs.Prandtlnumber)

    if settings.enable_kappaH_profile:
        # Correct diffusivity according to
        # Bryan, K., and L. J. Lewis, 1979:
        # A water mass model of the world ocean. J. Geophys. Res., 84, 2503–2517.
        # It mainly modifies kappaH within 20S - 20N deg. belt
        vs.kappaH = npx.maximum(
            vs.kappaH,
            (0.8 + 1.05 / settings.pi * npx.arctan((-vs.zw[npx.newaxis, npx.newaxis, :] - 2500.0) / 222.2)) * 1e-4,
        )

    vs.kappaM = npx.maximum(settings.kappaM_min, vs.kappaM)

    return KernelOutput(
        sqrttke=vs.sqrttke,
        mxl=vs.mxl,
        kappaM=vs.kappaM,
        kappaH=vs.kappaH,
        Prandtlnumber=vs.Prandtlnumber,
        K_diss_v=vs.K_diss_v,
    )


@veros_routine
def integrate_tke(state):
    vs = state.variables
    tke_out = integrate_tke_kernel(state)
    vs.update(tke_out)


@veros_kernel
def integrate_tke_kernel(state):
    """
    integrate Tke equation on W grid with surface flux boundary condition
    """
    vs = state.variables
    settings = state.settings

    conditional_outputs = {}

    flux_east = allocate(state.dimensions, ("xt", "yt", "zt"))
    flux_north = allocate(state.dimensions, ("xt", "yt", "zt"))
    flux_top = allocate(state.dimensions, ("xt", "yt", "zt"))

    dt_tke = settings.dt_mom  # use momentum time step to prevent spurious oscillations

    """
    Sources and sinks by vertical friction, vertical mixing, and non-conservative advection
    """
    forc = vs.K_diss_v - vs.P_diss_v - vs.P_diss_adv

    """
    store transfer due to vertical mixing from dyn. enthalpy by non-linear eq.of
    state either to TKE or to heat
    """
    if not settings.enable_store_cabbeling_heat:
        forc = forc - vs.P_diss_nonlin

    """
    transfer part of dissipation of EKE to TKE
    """
    if settings.enable_eke:
        forc = forc + vs.eke_diss_tke

    if settings.enable_idemix:
        """
        transfer dissipation of internal waves to TKE
        """
        forc = forc + vs.iw_diss
        """
        store bottom friction either in TKE or internal waves
        """
        if settings.enable_store_bottom_friction_tke:
            forc = forc + vs.K_diss_bot

    else:  # short-cut without idemix
        if settings.enable_eke:
            forc = forc + vs.eke_diss_iw

        else:  # and without EKE model
            if settings.enable_store_cabbeling_heat:
                forc = forc + vs.K_diss_gm + vs.K_diss_h - vs.P_diss_skew - vs.P_diss_hmix - vs.P_diss_iso
            else:
                forc = forc + vs.K_diss_gm + vs.K_diss_h - vs.P_diss_skew

        forc = forc + vs.K_diss_bot

    """
    vertical mixing and dissipation of TKE
    """
    _, water_mask, edge_mask = utilities.create_water_masks(vs.kbot[2:-2, 2:-2], settings.nz)

    a_tri, b_tri, c_tri, d_tri, delta = (
        allocate(state.dimensions, ("xt", "yt", "zt"))[2:-2, 2:-2, :] for _ in range(5)
    )

    delta = update(
        delta,
        at[:, :, :-1],
        dt_tke
        / vs.dzt[npx.newaxis, npx.newaxis, 1:]
        * settings.alpha_tke
        * 0.5
        * (vs.kappaM[2:-2, 2:-2, :-1] + vs.kappaM[2:-2, 2:-2, 1:]),
    )

    a_tri = update(a_tri, at[:, :, 1:-1], -delta[:, :, :-2] / vs.dzw[npx.newaxis, npx.newaxis, 1:-1])
    a_tri = update(a_tri, at[:, :, -1], -delta[:, :, -2] / (0.5 * vs.dzw[-1]))

    b_tri = update(
        b_tri,
        at[:, :, 1:-1],
        1
        + (delta[:, :, 1:-1] + delta[:, :, :-2]) / vs.dzw[npx.newaxis, npx.newaxis, 1:-1]
        + dt_tke * settings.c_eps * vs.sqrttke[2:-2, 2:-2, 1:-1] / vs.mxl[2:-2, 2:-2, 1:-1],
    )
    b_tri = update(
        b_tri,
        at[:, :, -1],
        1
        + delta[:, :, -2] / (0.5 * vs.dzw[-1])
        + dt_tke * settings.c_eps / vs.mxl[2:-2, 2:-2, -1] * vs.sqrttke[2:-2, 2:-2, -1],
    )
    b_tri_edge = (
        1
        + delta / vs.dzw[npx.newaxis, npx.newaxis, :]
        + dt_tke * settings.c_eps / vs.mxl[2:-2, 2:-2, :] * vs.sqrttke[2:-2, 2:-2, :]
    )

    c_tri = update(c_tri, at[:, :, :-1], -delta[:, :, :-1] / vs.dzw[npx.newaxis, npx.newaxis, :-1])

    d_tri = update(d_tri, at[...], vs.tke[2:-2, 2:-2, :, vs.tau] + dt_tke * forc[2:-2, 2:-2, :])
    d_tri = update_add(d_tri, at[:, :, -1], dt_tke * vs.forc_tke_surface[2:-2, 2:-2] / (0.5 * vs.dzw[-1]))

    sol = utilities.solve_implicit(a_tri, b_tri, c_tri, d_tri, water_mask, b_edge=b_tri_edge, edge_mask=edge_mask)
    vs.tke = update(vs.tke, at[2:-2, 2:-2, :, vs.taup1], npx.where(water_mask, sol, vs.tke[2:-2, 2:-2, :, vs.taup1]))

    """
    store tke dissipation for diagnostics
    """
    vs.tke_diss = settings.c_eps / vs.mxl * vs.sqrttke * vs.tke[:, :, :, vs.taup1]

    """
    Add TKE if surface density flux drains TKE in uppermost box
    """
    mask = vs.tke[2:-2, 2:-2, -1, vs.taup1] < 0.0
    vs.tke_surf_corr = update(
        vs.tke_surf_corr,
        at[2:-2, 2:-2],
        npx.where(mask, -vs.tke[2:-2, 2:-2, -1, vs.taup1] * 0.5 * vs.dzw[-1] / dt_tke, 0.0),
    )
    vs.tke = update(vs.tke, at[2:-2, 2:-2, -1, vs.taup1], npx.maximum(0.0, vs.tke[2:-2, 2:-2, -1, vs.taup1]))

    if settings.enable_tke_hor_diffusion:
        """
        add tendency due to lateral diffusion
        """
        flux_east = update(
            flux_east,
            at[:-1, :, :],
            settings.K_h_tke
            * (vs.tke[1:, :, :, vs.tau] - vs.tke[:-1, :, :, vs.tau])
            / (vs.cost[npx.newaxis, :, npx.newaxis] * vs.dxu[:-1, npx.newaxis, npx.newaxis])
            * vs.maskU[:-1, :, :],
        )

        flux_north = update(
            flux_north,
            at[:, :-1, :],
            settings.K_h_tke
            * (vs.tke[:, 1:, :, vs.tau] - vs.tke[:, :-1, :, vs.tau])
            / vs.dyu[npx.newaxis, :-1, npx.newaxis]
            * vs.maskV[:, :-1, :]
            * vs.cosu[npx.newaxis, :-1, npx.newaxis],
        )
        flux_north = update(flux_north, at[:, -1, :], 0.0)

        vs.tke = update_add(
            vs.tke,
            at[2:-2, 2:-2, :, vs.taup1],
            dt_tke
            * vs.maskW[2:-2, 2:-2, :]
            * (
                (flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
                / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dxt[2:-2, npx.newaxis, npx.newaxis])
                + (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
                / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dyt[npx.newaxis, 2:-2, npx.newaxis])
            ),
        )

    """
    add tendency due to advection
    """
    if settings.enable_tke_superbee_advection:
        flux_east, flux_north, flux_top = advection.adv_flux_superbee_wgrid(state, vs.tke[:, :, :, vs.tau])

    if settings.enable_tke_upwind_advection:
        flux_east, flux_north, flux_top = advection.adv_flux_upwind_wgrid(state, vs.tke[:, :, :, vs.tau])

    if settings.enable_tke_superbee_advection or settings.enable_tke_upwind_advection:
        vs.dtke = update(
            vs.dtke,
            at[2:-2, 2:-2, :, vs.tau],
            vs.maskW[2:-2, 2:-2, :]
            * (
                -(flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :])
                / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dxt[2:-2, npx.newaxis, npx.newaxis])
                - (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :])
                / (vs.cost[npx.newaxis, 2:-2, npx.newaxis] * vs.dyt[npx.newaxis, 2:-2, npx.newaxis])
            ),
        )
        vs.dtke = update_add(vs.dtke, at[:, :, 0, vs.tau], -flux_top[:, :, 0] / vs.dzw[0])
        vs.dtke = update_add(
            vs.dtke, at[:, :, 1:-1, vs.tau], -(flux_top[:, :, 1:-1] - flux_top[:, :, :-2]) / vs.dzw[1:-1]
        )
        vs.dtke = update_add(
            vs.dtke, at[:, :, -1, vs.tau], -(flux_top[:, :, -1] - flux_top[:, :, -2]) / (0.5 * vs.dzw[-1])
        )

        """
        Adam Bashforth time stepping
        """
        vs.tke = update_add(
            vs.tke,
            at[:, :, :, vs.taup1],
            settings.dt_tracer
            * (
                (1.5 + settings.AB_eps) * vs.dtke[:, :, :, vs.tau]
                - (0.5 + settings.AB_eps) * vs.dtke[:, :, :, vs.taum1]
            ),
        )

        conditional_outputs.update(dtke=vs.dtke)

    return KernelOutput(tke=vs.tke, tke_surf_corr=vs.tke_surf_corr, tke_diss=vs.tke_diss, **conditional_outputs)