Module ocean_science_utilities.wavephysics.balance.st4_wind_input

import numpy as np

from numba import jit  # type: ignore
from typing import Optional, TypedDict

from ocean_science_utilities.wavephysics.balance._numba_settings import numba_default
from ocean_science_utilities.wavephysics.balance.generation import WindGeneration
from ocean_science_utilities.wavephysics.balance.wam_tail_stress import (
    tail_stress_parametrization_wam,
)
from ocean_science_utilities.wavephysics.fluidproperties import (
    AIR,
    WATER,
    GRAVITATIONAL_ACCELERATION,
)
from ocean_science_utilities.wavetheory.lineardispersion import (
    inverse_intrinsic_dispersion_relation,
)


class ST4WaveGenerationParameters(TypedDict):
    gravitational_acceleration: float
    charnock_maximum_roughness: float
    charnock_constant: float
    air_density: float
    water_density: float
    vonkarman_constant: float
    wave_age_tuning_parameter: float
    growth_parameter_betamax: float
    elevation: float
    air_viscosity: float
    viscous_stress_parameter: float


class ST4WindInput(WindGeneration):
    name = "st4 generation"

    def __init__(self, parameters: Optional[ST4WaveGenerationParameters] = None):
        super(ST4WindInput, self).__init__(parameters)
        self._wind_source_term_function = _st4_wind_generation_point
        self._tail_stress_parametrization_function = tail_stress_parametrization_wam

    @staticmethod
    def default_parameters() -> ST4WaveGenerationParameters:
        return ST4WaveGenerationParameters(
            wave_age_tuning_parameter=0.006,
            growth_parameter_betamax=1.52,
            gravitational_acceleration=GRAVITATIONAL_ACCELERATION,
            charnock_maximum_roughness=np.inf,
            charnock_constant=0.01,
            air_density=AIR.density,
            water_density=WATER.density,
            vonkarman_constant=AIR.vonkarman_constant,
            elevation=10,
            air_viscosity=AIR.kinematic_viscosity,
            viscous_stress_parameter=0.0,
        )


# ----------------------------------------------------------------------------------------------------------------------
# ST4 Point wise implementation functions (apply to a single spatial point)
# ----------------------------------------------------------------------------------------------------------------------


@jit(**numba_default)
def _st4_wind_generation_point(
    variance_density: np.ndarray,
    wind,
    depth: float,
    roughness_length: float,
    spectral_grid,
    parameters,
    wind_source=None,
):
    """
    Implementation of the st4 wind growth rate term as proposed by Janssen 1991-
    and as described in the ST4 paper.

    :param variance_density: input numpy array of shape (nf,nd), with nf number of
        frequencies, nd number of directions, that contains variance
        densities (unit m**2/Hz/Degree).
    :param wind: Wind input tuple (speed, direction, type) that descripes the wind
        speed, direction and kind of wind input. The latter can be: "u10"
        or "friction_velocity".
    :param depth: Local water depth in meter.
    :param roughness_length: Roughness length in meter.
    :param spectral_grid: numba dictionary containing the spectral grid
        (frequencies/directions)
    :param parameters: numba dictionary containing the parameters of the st4 algorithm.
        See the st4 class for details.
    :param wind_source: optional numpy array that contains the return values.
        This helps to avoid reallocating memory if the routine is called often
        in successive fashion (e.g. during iterative solutions).

    :return: numpy array of shape (nf,nd) containing for each frequency/direction
        pair the wind input in (m**2/Hz/Degree/s)
    """

    # Get the spectral size
    number_of_frequencies, number_of_directions = variance_density.shape

    vonkarman_constant = parameters["vonkarman_constant"]
    radian_frequency = spectral_grid["radian_frequency"]

    wind_forcing, wind_direction_degrees, wind_forcing_type = wind

    cosine_mutual_angle_wind_waves = np.cos(
        (
            spectral_grid["radian_direction"]
            - wind_direction_degrees * np.pi / 180
            + np.pi
        )
        % (2 * np.pi)
        - np.pi
    )

    # Allocate the output variable
    if wind_source is None:
        wind_source = np.empty((number_of_frequencies, number_of_directions))

    # Preprocess winds to the correct conventions (U10->friction velocity if need be,
    # wind direction degrees->radians)
    if wind_forcing_type == "u10":
        friction_velocity = (
            wind_forcing
            * vonkarman_constant
            / np.log(parameters["elevation"] / roughness_length)
        )
    elif wind_forcing_type in ["ustar", "friction_velocity"]:
        friction_velocity = wind_forcing

    else:
        raise ValueError("Unknown wind input type")

    # precalculate the constant multiplication.
    constant_factor = (
        parameters["growth_parameter_betamax"]
        / vonkarman_constant**2
        * parameters["air_density"]
        / parameters["water_density"]
    )

    if depth == np.inf:
        wavenumber = radian_frequency**2 / parameters["gravitational_acceleration"]
    else:
        wavenumber = inverse_intrinsic_dispersion_relation(radian_frequency, depth)

    # wave age tuning times the directional cosine.
    cosine_wave_age_tuning = (
        parameters["wave_age_tuning_parameter"] * cosine_mutual_angle_wind_waves
    )

    # Loop over all frequencies/directions
    for frequency_index in range(number_of_frequencies):
        # Since wavenumber and wavespeed only depend on frequency we calculate those
        # outside of the direction loop.
        relative_speed = (
            wavenumber[frequency_index]
            * friction_velocity
            / radian_frequency[frequency_index]
        )

        for direction_index in range(number_of_directions):
            # If the wave direction has an along wind component we continue.
            if cosine_mutual_angle_wind_waves[direction_index] > 0:
                # Calculate the directional factor in the wind input formulation
                W = relative_speed * cosine_mutual_angle_wind_waves[direction_index]

                # The logarithm of the approximate expression for the dimensionless
                # critical height (Ardhuin 2010, eq 20).
                log_dimensionless_critical_height = np.log(
                    wavenumber[frequency_index] * roughness_length
                ) + vonkarman_constant / (W + cosine_wave_age_tuning[direction_index])

                # if the dimensionless critical height > 1 we assume there is no more
                # energy transfer (Janssen 1991), his equation 18.
                if log_dimensionless_critical_height > 0:
                    log_dimensionless_critical_height = 0

                growth_rate = (
                    constant_factor
                    * np.exp(log_dimensionless_critical_height)
                    * log_dimensionless_critical_height**4
                    * W**2
                    * radian_frequency[frequency_index]
                )

                wind_source[frequency_index, direction_index] = (
                    growth_rate * variance_density[frequency_index, direction_index]
                )

            else:
                wind_source[frequency_index, direction_index] = 0.0

    return wind_source