1
Fork 0
internship/swash/processing/reflex3s.py

48 lines
1.2 KiB
Python
Raw Normal View History

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):
#_welch = [sgl.welch(z, f) for z in eta]
#eta_psd = np.stack([_w[1] for _w in _welch])
#f_psd = _welch[0][0]
eta_psd = np.stack([fft.rfft(z) for z in eta])
f_psd = fft.rfftfreq(eta.shape[1])
eta_amp = np.abs(eta_psd) / eta_psd.shape[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.2,
).root for f in f_psd])
dx = np.roll(x, 1) - x
dphi = np.roll(eta_phase, 1, axis=0) - eta_phase
ai = np.sqrt(
eta_amp**2
+ np.roll(eta_amp, 1, axis=0)**2
- 2
* eta_amp
* np.roll(eta_amp, 1, axis=0)
* np.cos(dphi - k * dx[:, None])
/ (2 * np.abs(np.sin(k * dx[:, None])))
)
ar = np.sqrt(
eta_amp**2
+ np.roll(eta_amp, 1, axis=0)**2
- 2
* eta_amp
* np.roll(eta_amp, 1, axis=0)
* np.cos(dphi + k * dx[:, None])
/ (2 * np.abs(np.sin(k * dx[:, None])))
)
return f_psd, ar / ai