Skip to content

Commit

Permalink
Nd-signal (#228)
Browse files Browse the repository at this point in the history 
* 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 <[email protected]>
Co-authored-by: Philippe Miron <[email protected]>
  • Loading branch information
@selipot @milancurcic
3 people authored Aug 23, 2023
1 parent 49828cc commit 6cdc191
Show file tree
Hide file tree
Showing 3 changed files with 161 additions and 130 deletions.
2 changes: 1 addition & 1 deletion clouddrift/analysis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
137 changes: 95 additions & 42 deletions clouddrift/signal.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
-------
Expand All @@ -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
Expand All @@ -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
-------
Expand All @@ -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
Expand All @@ -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
152 changes: 65 additions & 87 deletions tests/signal_tests.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)))

0 comments on commit 6cdc191

Please sign in to comment.