Turbulence kinetic energy calculation #3756

def tke(u, v, w, perturbation=False, axis=-1):
    r"""Compute turbulence kinetic energy.

    Compute the turbulence kinetic energy (e) from the time series of the
    velocity components.

    Parameters
    ----------
    u : array-like
        The wind component along the x-axis
    v : array-like
        The wind component along the y-axis
    w : array-like
        The wind component along the z-axis
    perturbation : bool, optional
                   True if the `u`, `v`, and `w` components of wind speed
                   supplied to the function are perturbation velocities.
                   If False, perturbation velocities will be calculated by
                   removing the mean value from each component.

    Returns
    -------
    array-like
        The corresponding turbulence kinetic energy value

    Other Parameters
    ----------------
    axis : int
           The index of the time axis. Default is -1

    See Also
    --------
    get_perturbation : Used to compute perturbations if `perturbation`
                       is False.

    Notes
    -----
    Turbulence Kinetic Energy is computed as:

    .. math:: e = 0.5 \sqrt{\overline{u^{\prime2}} +
                            \overline{v^{\prime2}} +
                            \overline{w^{\prime2}}},

    where the velocity components

    .. math:: u^{\prime}, v^{\prime}, u^{\prime}

    are perturbation velocities. For more information on the subject, please
    see [Garratt1994]_.

    """
    if not perturbation:
        u = get_perturbation(u, axis=axis)
        v = get_perturbation(v, axis=axis)
        w = get_perturbation(w, axis=axis)

    u_cont = np.mean(u**2, axis=axis)
    v_cont = np.mean(v**2, axis=axis)
    w_cont = np.mean(w**2, axis=axis)

    return 0.5 * np.sqrt(u_cont + v_cont + w_cont)