Inertial modifications (#310)

* name changes

* xarray conversion

* add optional argument

* expanded example

Author: selipot

clouddrift/kinematics.py

@@ -19,10 +19,11 @@
 from clouddrift.wavelet import morse_logspace_freq, morse_wavelet, wavelet_transform
 
 
-def inertial_oscillations_from_positions(
+def inertial_oscillation_from_position(
     longitude: np.ndarray,
     latitude: np.ndarray,
-    relative_bandwidth: float,
+    relative_bandwidth: Optional[float] = None,
+    wavelet_duration: Optional[float] = None,
     time_step: Optional[float] = 3600.0,
     relative_vorticity: Optional[Union[float, np.ndarray]] = 0.0,
 ) -> np.ndarray:
@@ -31,19 +32,27 @@ def inertial_oscillations_from_positions(
     This function acts by performing a time-frequency analysis of horizontal displacements
     with analytic Morse wavelets. It extracts the portion of the wavelet transform signal
     that follows the inertial frequency (opposite of Coriolis frequency) as a function of time,
-    potentially shifted in frequency by a measure of relative vorticity.
+    potentially shifted in frequency by a measure of relative vorticity. The result is a pair
+    of zonal and meridional relative displacements in meters.
+
+    This function is equivalent to a bandpass filtering of the horizontal displacements. The characteristics
+    of the filter are defined by the relative bandwidth of the wavelet transform or by the duration of the wavelet,
+    see the parameters below.
 
     Parameters
     ----------
     longitude : array-like
         Longitude sequence. Unidimensional array input.
     latitude : array-like
         Latitude sequence. Unidimensional array input.
-    relative_bandwidth : float
+    relative_bandwidth : float, optional
         Bandwidth of the frequency-domain equivalent filter for the extraction of the inertial
         oscillations; a number less or equal to one which is a fraction of the inertial frequency.
         A value of 0.1 leads to a bandpass filter equivalent of +/- 10 percent of the inertial frequency.
-    time_step : float
+    wavelet_duration : float, optional
+        Duration of the wavelet, or inverse of the relative bandwidth, which can be passed instead of the
+        relative bandwidth.
+    time_step : float, optional
         The constant time interval between data points in seconds. Default is 3600.
     relative_vorticity: Optional, float or array-like
         Relative vorticity adding to the local Coriolis frequency. If "f" is the Coriolis
@@ -63,28 +72,44 @@ def inertial_oscillations_from_positions(
     To extract displacements from inertial oscillations from sequences of longitude
     and latitude values, equivalent to bandpass around 20 percent of the local inertial frequency:
 
-    >>> xhat, yhat = extract_inertial_from_position(longitude, latitude, 0.2)
+    >>> xhat, yhat = inertial_oscillation_from_position(longitude, latitude, relative_bandwidth=0.2)
+
+    The same result can be obtained by specifying the wavelet duration instead of the relative bandwidth:
+
+    >>> xhat, yhat = inertial_oscillation_from_position(longitude, latitude, wavelet_duration=5)
 
     Next, the residual positions from the inertial displacements can be obtained with another function:
 
-    >>> residual_longitudes, residual_latitudes = residual_positions_from_displacements(longitude, latitude, xhat, yhat)
+    >>> residual_longitudes, residual_latitudes = residual_position_from_displacement(longitude, latitude, xhat, yhat)
 
     Raises
     ------
     ValueError
         If longitude and latitude arrays do not have the same shape.
         If relative_vorticity is an array and does not have the same shape as longitude and latitude.
         If time_step is not a float.
+        If both relative_bandwidth and wavelet_duration are specified.
+        If neither relative_bandwidth nor wavelet_duration are specified.
         If the absolute value of relative_bandwidth is not in the range (0,1].
+        If the wavelet duration is not greater than or equal to 1.
 
     See Also
     --------
-    :func:`residual_positions_from_displacements`, `wavelet_transform`, `morse_wavelet`
+    :func:`residual_position_from_displacement`, `wavelet_transform`, `morse_wavelet`
 
     """
     if longitude.shape != latitude.shape:
         raise ValueError("longitude and latitude arrays must have the same shape.")
 
+    if relative_bandwidth is not None and wavelet_duration is not None:
+        raise ValueError(
+            "Only one of 'relative_bandwidth' and 'wavelet_duration' can be specified"
+        )
+    elif relative_bandwidth is None and wavelet_duration is None:
+        raise ValueError(
+            "One of 'relative_bandwidth' and 'wavelet_duration' must be specified"
+        )
+
     # length of data sequence
     data_length = longitude.shape[0]
 
@@ -95,14 +120,20 @@ def inertial_oscillations_from_positions(
             raise ValueError(
                 "relative_vorticity must be a float or the same shape as longitude and latitude."
             )
+    if relative_bandwidth is not None:
+        if not 0 < np.abs(relative_bandwidth) <= 1:
+            raise ValueError("relative_bandwidth must be in the (0, 1]) range")
 
-    if not 0 < np.abs(relative_bandwidth) <= 1:
-        raise ValueError("relative_bandwidth must be in the (0, 1]) range")
+    if wavelet_duration is not None:
+        if not wavelet_duration >= 1:
+            raise ValueError("wavelet_duration must be greater than or equal to 1")
 
     # wavelet parameters are gamma and beta
     gamma = 3  # symmetric wavelet
     density = 16  # results relative insensitive to this parameter
-    wavelet_duration = 1 / np.abs(relative_bandwidth)  # P parameter
+    # calculate beta from wavelet duration or from relative bandwidth
+    if relative_bandwidth is not None:
+        wavelet_duration = 1 / np.abs(relative_bandwidth)  # P parameter
     beta = wavelet_duration**2 / gamma
 
     if isinstance(latitude, xr.DataArray):
@@ -187,9 +218,9 @@ def inertial_oscillations_from_positions(
     return xhat, yhat
 
 
-def residual_positions_from_displacements(
-    longitude: Union[float, np.ndarray],
-    latitude: Union[float, np.ndarray],
+def residual_position_from_displacement(
+    longitude: Union[float, np.ndarray, xr.DataArray],
+    latitude: Union[float, np.ndarray, xr.DataArray],
     x: Union[float, np.ndarray],
     y: Union[float, np.ndarray],
 ) -> Union[Tuple[float], Tuple[np.ndarray]]:
@@ -202,9 +233,9 @@ def residual_positions_from_displacements(
 
     Parameters
     ----------
-    longitude : float or np.ndarray
+    longitude : float or array-like
         Longitude in degrees.
-    latitude : float or np.ndarray
+    latitude : float or array-like
         Latitude in degrees.
     x : float or np.ndarray
         Zonal displacement in meters.
@@ -224,9 +255,15 @@ def residual_positions_from_displacements(
     circumference of the Earth from original position (longitude,latitude) = (1,0):
 
     >>> from clouddrift.sphere import EARTH_RADIUS_METERS
-    >>> residual_positions_from_displacements(1,0,2 * np.pi * EARTH_RADIUS_METERS / 360,0)
+    >>> residual_position_from_displacement(1,0,2 * np.pi * EARTH_RADIUS_METERS / 360,0)
     (0.0, 0.0)
     """
+    # convert to numpy arrays to insure consistent outputs
+    if isinstance(longitude, xr.DataArray):
+        longitude = longitude.to_numpy()
+    if isinstance(latitude, xr.DataArray):
+        latitude = latitude.to_numpy()
+
     latitudehat = 180 / np.pi * y / EARTH_RADIUS_METERS
     longitudehat = (
         180 / np.pi * x / (EARTH_RADIUS_METERS * np.cos(np.radians(latitude)))

tests/kinematics_tests.py

@@ -1,8 +1,8 @@
 from clouddrift.kinematics import (
-    inertial_oscillations_from_positions,
+    inertial_oscillation_from_position,
     position_from_velocity,
     velocity_from_position,
-    residual_positions_from_displacements,
+    residual_position_from_displacement,
 )
 from clouddrift.sphere import EARTH_RADIUS_METERS, coriolis_frequency
 from clouddrift.raggedarray import RaggedArray
@@ -79,22 +79,22 @@ def sample_ragged_array() -> RaggedArray:
     return ra
 
 
-class inertial_oscillations_from_positions_tests(unittest.TestCase):
+class inertial_oscillation_from_position_tests(unittest.TestCase):
     def setUp(self):
         self.INPUT_SIZE = 1440
         self.longitude = np.linspace(0, 45, self.INPUT_SIZE)
         self.latitude = np.linspace(40, 45, self.INPUT_SIZE)
         self.relative_vorticity = 0.1 * coriolis_frequency(self.latitude)
 
     def test_result_has_same_size_as_input(self):
-        x, y = inertial_oscillations_from_positions(
+        x, y = inertial_oscillation_from_position(
             self.longitude, self.latitude, relative_bandwidth=0.10, time_step=3600
         )
         self.assertTrue(np.all(self.longitude.shape == x.shape))
         self.assertTrue(np.all(self.longitude.shape == y.shape))
 
     def test_relative_vorticity(self):
-        x, y = inertial_oscillations_from_positions(
+        x, y = inertial_oscillation_from_position(
             self.longitude,
             self.latitude,
             relative_bandwidth=0.10,
@@ -105,7 +105,7 @@ def test_relative_vorticity(self):
         self.assertTrue(np.all(self.longitude.shape == y.shape))
 
     def test_time_step(self):
-        x, y = inertial_oscillations_from_positions(
+        x, y = inertial_oscillation_from_position(
             self.longitude,
             self.latitude,
             relative_bandwidth=0.10,
@@ -140,21 +140,32 @@ def test_simulation_case(self):
         )
         x_expected = x_expected - np.mean(x_expected)
         y_expected = y_expected - np.mean(y_expected)
-        xhat, yhat = inertial_oscillations_from_positions(
+        xhat1, yhat1 = inertial_oscillation_from_position(
             lon1, lat1, relative_bandwidth=0.10, time_step=3600
         )
-        xhat = xhat - np.mean(xhat)
-        yhat = yhat - np.mean(yhat)
+        xhat1 = xhat1 - np.mean(xhat1)
+        yhat1 = yhat1 - np.mean(yhat1)
+        xhat2, yhat2 = inertial_oscillation_from_position(
+            lon1, lat1, wavelet_duration=10, time_step=3600
+        )
+        xhat2 = xhat2 - np.mean(xhat2)
+        yhat2 = yhat2 - np.mean(yhat2)
         m = 10
         self.assertTrue(
-            np.allclose(xhat[m * 24 : -m * 24], x_expected[m * 24 : -m * 24], atol=20)
+            np.allclose(xhat2[m * 24 : -m * 24], x_expected[m * 24 : -m * 24], atol=20)
+        )
+        self.assertTrue(
+            np.allclose(yhat2[m * 24 : -m * 24], y_expected[m * 24 : -m * 24], atol=20)
+        )
+        self.assertTrue(
+            np.allclose(xhat2[m * 24 : -m * 24], xhat1[m * 24 : -m * 24], atol=20)
         )
         self.assertTrue(
-            np.allclose(yhat[m * 24 : -m * 24], y_expected[m * 24 : -m * 24], atol=20)
+            np.allclose(yhat2[m * 24 : -m * 24], yhat1[m * 24 : -m * 24], atol=20)
         )
 
 
-class residual_positions_from_displacements_tests(unittest.TestCase):
+class residual_position_from_displacement_tests(unittest.TestCase):
     def setUp(self):
         self.INPUT_SIZE = 100
         self.lon = np.rad2deg(
@@ -165,14 +176,14 @@ def setUp(self):
         self.y = np.random.rand(self.INPUT_SIZE)
 
     def test_result_has_same_size_as_input(self):
-        lon, lat = residual_positions_from_displacements(
+        lon, lat = residual_position_from_displacement(
             self.lon, self.lat, self.x, self.y
         )
         self.assertTrue(np.all(self.lon.shape == lon.shape))
         self.assertTrue(np.all(self.lon.shape == lat.shape))
 
     def test_simple_case(self):
-        lon, lat = residual_positions_from_displacements(
+        lon, lat = residual_position_from_displacement(
             1, 0, 2 * np.pi * EARTH_RADIUS_METERS / 360, 0
         )
         self.assertTrue(np.isclose((np.mod(lon, 360), lat), (0, 0)).all())