internship/swash/processing/reflex3s.py

import numpy as np
from scipy import fft
from scipy import signal as sgl
from scipy import optimize as opti


def reflex3s(eta, x, h, f):
    eta_psd = np.stack([fft.rfft(z) for z in eta])
    f_psd = fft.rfftfreq(eta.shape[1], 1 / f)

    eta_amp = np.abs(eta_psd) / (f_psd.size - 1)
    eta_phase = np.angle(eta_psd)

    g = 9.81
    k = np.asarray(
        [
            opti.root_scalar(
                f=lambda k: k * np.tanh(k) - (2 * np.pi * f) ** 2 / g * h,
                fprime=lambda k: np.tanh(k) + k * (1 - np.tanh(k) ** 2),
                x0=0.5,
            ).root
            / h
            for f in f_psd
        ]
    )

    s1 = 1 + np.exp(2j * k * (x[1] - x[0])) + np.exp(2j * k * (x[2] - x[0]))
    s2 = 1 + np.exp(-2j * k * (x[1] - x[0])) + np.exp(-2j * k * (x[2] - x[0]))
    s12 = 3
    s3 = (
        eta_amp[0] * np.exp(-1j * (eta_phase[0]))
        + eta_amp[1] * np.exp(-1j * (eta_phase[1] - k * (x[1] - x[0])))
        + eta_amp[2] * np.exp(-1j * (eta_phase[2] - k * (x[2] - x[0])))
    )
    s4 = (
        eta_amp[0] * np.exp(-1j * (eta_phase[0]))
        + eta_amp[1] * np.exp(-1j * (eta_phase[1] + k * (x[1] - x[0])))
        + eta_amp[2] * np.exp(-1j * (eta_phase[2] + k * (x[2] - x[0])))
    )
    s5 = s1 * s2 - s12**2

    ai = np.abs((s2 * s3 - s12 * s4) / s5)
    ar = np.abs((s1 * s4 - s12 * s3) / s5)

    return f_psd, ai, ar