From 6cdc1916422ec24977fb4bc4299dfcf272319efb Mon Sep 17 00:00:00 2001 From: Shane Elipot Date: Wed, 23 Aug 2023 16:31:50 -0400 Subject: [PATCH] Nd-signal (#228) * removed pandas * drop list support * nd for transform * typo correct? * nd input for transforms * augmented help * black * ndarray fix. added loose tests * lint fixes * Docstrings * Refactor tests so that they run * Compatibility with xr.DataArray * analytic_transform tests pass; rotary_transform tests TODO * fix even case * added xa (anomaly) and it passes the tests * rotary_transform tests fix --------- Co-authored-by: milancurcic Co-authored-by: Philippe Miron --- clouddrift/analysis.py | 2 +- clouddrift/signal.py | 137 +++++++++++++++++++++++++------------ tests/signal_tests.py | 152 ++++++++++++++++++----------------------- 3 files changed, 161 insertions(+), 130 deletions(-) diff --git a/clouddrift/analysis.py b/clouddrift/analysis.py index 438b31a5..3ff9bf47 100644 --- a/clouddrift/analysis.py +++ b/clouddrift/analysis.py @@ -720,7 +720,7 @@ def velocity_from_position( Coordinate system that x and y arrays are in; possible values are "spherical" (default) or "cartesian". difference_scheme : str, optional Difference scheme to use; possible values are "forward", "backward", and "centered". - time_axis : int, optional) + time_axis : int, optional Axis along which to differentiate (default is -1) Returns diff --git a/clouddrift/signal.py b/clouddrift/signal.py index 8bc983ca..fe4683af 100644 --- a/clouddrift/signal.py +++ b/clouddrift/signal.py @@ -5,13 +5,13 @@ import numpy as np from typing import Optional, Tuple, Union import xarray as xr -import pandas as pd import warnings def analytic_transform( - x: Union[list, np.ndarray, xr.DataArray, pd.Series], + x: Union[np.ndarray, xr.DataArray], boundary: Optional[str] = "mirror", + time_axis: Optional[int] = -1, ) -> np.ndarray: """Return the analytic part of a real-valued signal or of a complex-valued signal. To obtain the anti-analytic part of a complex-valued signal apply analytic_transform @@ -21,8 +21,12 @@ def analytic_transform( ---------- x : np.ndarray Real- or complex-valued signal - boundary : str, optional ["mirror", "zeros", "periodic"] optionally specifies the - boundary condition to be imposed at the edges of the time series. Default is "mirror". + boundary : str, optional + The boundary condition to be imposed at the edges of the time series. + Allowed values are "mirror", "zeros", and "periodic". + Default is "mirror". + time_axis : int, optional + Axis on which the time is defined (default is -1) Returns ------- @@ -41,57 +45,86 @@ def analytic_transform( >>> zp = analytic_transform(z) >>> zn = analytic_transform(np.conj(z)) - To specify that a periodic boundary condition should be used + To specify that a periodic boundary condition should be used: >>> x = np.random.rand(99) - >>> z = analytic_transform(x,boundary="periodic") + >>> z = analytic_transform(x, boundary="periodic") + + To specify that the time axis is along the first axis and apply + zero boundary conditions: + >>> x = np.random.rand(100, 99) + >>> z = analytic_transform(x, time_axis=0, boundary="zeros") Raises ------ ValueError If ``boundary not in ["mirror", "zeros", "periodic"]``. """ - # assume unidimensional input; add dimension option - m0 = len(x) + # time_axis must be in valid range + if time_axis < -1 or time_axis > len(x.shape) - 1: + raise ValueError( + f"time_axis ({time_axis}) is outside of the valid range ([-1," + f" {len(x.shape) - 1}])." + ) + + # Swap the axis to make the time axis last (fast-varying). + # np.swapaxes returns a view to the input array, so no copy is made. + if time_axis != -1 and time_axis != len(x.shape) - 1: + x_ = np.swapaxes(x, time_axis, -1) + else: + x_ = x + + # time dimension length + N = np.shape(x_)[-1] - # remove mean - x = x - np.mean(x) + # Subtract mean along time axis (-1); convert to np.array for compatibility + # with xarray.DataArray. + xa = x_ - np.array(np.mean(x_, axis=-1, keepdims=True)) # apply boundary conditions if boundary == "mirror": - x = np.concatenate((np.flip(x), x, np.flip(x))) + xa = np.concatenate((np.flip(xa, axis=-1), xa, np.flip(xa, axis=-1)), axis=-1) elif boundary == "zeros": - x = np.concatenate((np.zeros_like(x), x, np.zeros_like(x))) + xa = np.concatenate((np.zeros_like(xa), xa, np.zeros_like(xa)), axis=-1) elif boundary == "periodic": - x = np.concatenate((x, x, x)) + xa = np.concatenate((xa, xa, xa), axis=-1) else: raise ValueError("boundary must be one of 'mirror', 'align', or 'zeros'.") - if np.isrealobj(x): - z = 2 * np.fft.fft(x) + if np.isrealobj(xa): + # fft should be taken along last axis + z = 2 * np.fft.fft(xa) else: - z = np.fft.fft(x) + z = np.fft.fft(xa) - m = len(x) + # time dimension of extended time series + M = np.shape(xa)[-1] # zero negative frequencies - if m % 2 == 0: - z[int(m / 2 + 2) - 1 : int(m + 1) + 1] = 0 + if M % 2 == 0: + z[..., int(M / 2 + 2) - 1 : int(M + 1) + 1] = 0 + # divide Nyquist component by 2 in even case + z[..., int(M / 2 + 1) - 1] = z[..., int(M / 2 + 1) - 1] / 2 else: - z[int((m + 3) / 2) - 1 : int(m + 1) + 1] = 0 + z[..., int((M + 3) / 2) - 1 : int(M + 1) + 1] = 0 - # inverse Fourier transform + # inverse Fourier transform along last axis z = np.fft.ifft(z) # return central part - z = z[int(m0 + 1) - 1 : int(2 * m0 + 1) - 1] + z = z[..., int(N + 1) - 1 : int(2 * N + 1) - 1] - return z + # return after reorganizing the axes + if time_axis != -1 and time_axis != len(x.shape) - 1: + return np.swapaxes(z, time_axis, -1) + else: + return z def rotary_transform( - u: Union[list, np.ndarray, xr.DataArray, pd.Series], - v: Union[list, np.ndarray, xr.DataArray, pd.Series], + u: Union[np.ndarray, xr.DataArray], + v: Union[np.ndarray, xr.DataArray], boundary: Optional[str] = "mirror", + time_axis: Optional[int] = -1, ) -> Tuple[np.ndarray, np.ndarray]: """Return time-domain rotary components time series (zp,zn) from Cartesian components time series (u,v). Note that zp and zn are both analytic time series which implies that the complex-valued time series @@ -108,9 +141,13 @@ def rotary_transform( Real-valued signal, first Cartesian component (zonal, east-west) v : np.ndarray Real-valued signal, second Cartesian component (meridional, north-south) - boundary : str, optional ["mirror", "zeros", "periodic"] optionally specifies the - boundary condition to be imposed at the edges of the time series for the underlying analytic - transform. Default is "mirror". + boundary : str, optional + The boundary condition to be imposed at the edges of the time series. + Allowed values are "mirror", "zeros", and "periodic". + Default is "mirror". + time_axis : int, optional + The axis of the time array. Default is -1, which corresponds to the + last axis. Returns ------- @@ -123,8 +160,16 @@ def rotary_transform( -------- To obtain the rotary components of a real-valued signal: + >>> u = np.random.rand(99) + >>> v = np.random.rand(99) >>> zp, zn = rotary_transform(u,v) + To specify that the time axis is along the first axis, and apply + zero boundary conditions: + >>> u = np.random.rand(100,99) + >>> v = np.random.rand(100,99) + >>> zp, zn = rotary_transform(u,v,time_axis=0,boundary="zeros") + Raises ------ ValueError @@ -135,20 +180,28 @@ def rotary_transform( :func:`analytic_transform` """ # to implement: input one complex argument instead of two real arguments - # u and v arrays must have the same shape or list length. - if type(u) == list and type(v) == list: - if not len(u) == len(v): - raise ValueError("u and v must have the same length.") - else: - if not u.shape == v.shape: - raise ValueError("u and v must have the same shape.") - - muv = np.mean(u) + 1j * np.mean(v) - - if muv == xr.DataArray: + # u and v arrays must have the same shape. + if not u.shape == v.shape: + raise ValueError("u and v must have the same shape.") + + # time_axis must be in valid range + if time_axis < -1 or time_axis > len(u.shape) - 1: + raise ValueError( + f"time_axis ({time_axis}) is outside of the valid range ([-1," + f" {len(u.shape) - 1}])." + ) + + muv = np.mean(u, axis=time_axis, keepdims=True) + 1j * np.mean( + v, axis=time_axis, keepdims=True + ) + + if type(muv) == xr.DataArray: muv = muv.to_numpy() - up = analytic_transform(u, boundary=boundary) - vp = analytic_transform(v, boundary=boundary) + up = analytic_transform(u, boundary=boundary, time_axis=time_axis) + vp = analytic_transform(v, boundary=boundary, time_axis=time_axis) + + zp = 0.5 * (up + 1j * vp) + 0.5 * muv + zn = 0.5 * (up - 1j * vp) + 0.5 * np.conj(muv) - return 0.5 * (up + 1j * vp) + 0.5 * muv, 0.5 * (up - 1j * vp) + 0.5 * np.conj(muv) + return zp, zn diff --git a/tests/signal_tests.py b/tests/signal_tests.py index 911c5e4b..d257bbd8 100644 --- a/tests/signal_tests.py +++ b/tests/signal_tests.py @@ -4,96 +4,74 @@ ) import numpy as np import unittest -import pandas as pd import xarray as xr if __name__ == "__main__": unittest.main() -def test_analytic_transform_real_odd(self): - x = np.random.rand(99) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_analytic_transform_real_even(self): - x = np.random.rand(100) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_analytic_transform_imag_odd(self): - z = np.random.rand(99) + 1j * np.random.rand(99) - zp = analytic_transform(z) - zn = analytic_transform(np.conj(z)) - self.assertTrue(np.allclose(z - np.mean(z), zp + np.conj(zn))) - - -def test_analytic_transform_imag_even(self): - z = np.random.rand(100) + 1j * np.random.rand(100) - zp = analytic_transform(z) - zn = analytic_transform(np.conj(z)) - self.assertTrue(np.allclose(z - np.mean(z), zp + np.conj(zn))) - - -def test_analytic_transform_boundary(self): - x = np.random.rand(99) - z1 = analytic_transform(x, boundary="mirror") - z2 = analytic_transform(x, boundary="zeros") - z3 = analytic_transform(x, boundary="periodic") - self.assertTrue(np.allclose(x - np.mean(x), np.real(z1))) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z2))) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z3))) - - -def test_analytic_transform_list(self): - x = list(np.random.rand(99)) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_analytic_transform_array(self): - x = np.random.rand(99) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_analytic_transform_pandas(self): - x = pd.Series(data=np.random.rand(99)) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_analytic_transform_xarray(self): - x = xr.DataArray(data=np.random.rand(99)) - z = analytic_transform(x) - self.assertTrue(np.allclose(x - np.mean(x), np.real(z))) - - -def test_rotary_transform_array(self): - u = np.random.rand(99) - v = np.random.rand(99) - zp, zn = rotary_transform(u, v) - self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) - - -def test_rotary_transform_list(self): - u = list(np.random.rand(99)) - v = list(np.random.rand(99)) - zp, zn = rotary_transform(u, v) - self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) - - -def test_rotary_transform_pandas(self): - u = pd.Series(data=np.random.rand(99)) - v = pd.Series(data=np.random.rand(99)) - zp, zn = rotary_transform(u, v) - self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) - - -def test_rotary_transform_xarray(self): - u = xr.DataArray(data=np.random.rand(99)) - v = xr.DataArray(data=np.random.rand(99)) - zp, zn = rotary_transform(u, v) - self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) +class analytic_transform_tests(unittest.TestCase): + def test_real_odd(self): + x = np.random.rand(99) + z = analytic_transform(x) + self.assertTrue(np.allclose(x - np.mean(x), z.real)) + + def test_real_even(self): + x = np.random.rand(100) + z = analytic_transform(x) + self.assertTrue(np.allclose(x - np.mean(x), z.real)) + + def test_imag_odd(self): + z = np.random.rand(99) + 1j * np.random.rand(99) + zp = analytic_transform(z) + zn = analytic_transform(np.conj(z)) + self.assertTrue(np.allclose(z - np.mean(z, keepdims=True), zp + np.conj(zn))) + + def test_imag_even(self): + z = np.random.rand(100) + 1j * np.random.rand(100) + zp = analytic_transform(z) + zn = analytic_transform(np.conj(z)) + self.assertTrue(np.allclose(z - np.mean(z, keepdims=True), zp + np.conj(zn))) + + def test_boundary(self): + x = np.random.rand(99) + z1 = analytic_transform(x, boundary="mirror") + z2 = analytic_transform(x, boundary="zeros") + z3 = analytic_transform(x, boundary="periodic") + self.assertTrue(np.allclose(x - np.mean(x), z1.real)) + self.assertTrue(np.allclose(x - np.mean(x), z2.real)) + self.assertTrue(np.allclose(x - np.mean(x), z3.real)) + + def test_ndarray(self): + x = np.random.random((9, 11, 13)) + for n in range(3): + z = analytic_transform(x, time_axis=n) + self.assertTrue(np.allclose(x - np.mean(x, axis=n, keepdims=True), z.real)) + + def test_xarray(self): + x = xr.DataArray(data=np.random.random((9, 11, 13))) + for n in range(3): + z = analytic_transform(x, time_axis=n) + self.assertTrue( + np.allclose(x - np.array(np.mean(x, axis=n, keepdims=True)), z.real) + ) + + +class rotary_transform_tests(unittest.TestCase): + def test_array(self): + u = np.random.random((9)) + v = np.random.random((9)) + zp, zn = rotary_transform(u, v) + self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) + + def test_ndarray(self): + u = np.random.rand(99, 10, 101) + v = np.random.rand(99, 10, 101) + zp, zn = rotary_transform(u, v) + self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn))) + + def test_xarray(self): + u = xr.DataArray(data=np.random.rand(99)) + v = xr.DataArray(data=np.random.rand(99)) + zp, zn = rotary_transform(u, v) + self.assertTrue(np.allclose(u + 1j * v, zp + np.conj(zn)))