Custom implementation of reflex3s (not working)

swash/processing/reflex3s.py

@@ -0,0 +1,47 @@
+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