Position from velocity (#166)

* Get position from bearing and distance

* Basic implementation works

* Change wording in docstrings

* Fix argument name

* Fix argument name, again

* Test simple roundtrip for three integration schemes; centered fails

* Allow N-dim

* Apply Haversine in centered differences over half-points

* Loosen tolerance for centered scheme

* Refactor for arbitrary time axis; test not passing yet

* time_axis argument in test

* Allow scalar in recast_lon

* Recast to -180-180 in haversine position function

* Update docstrings

* Bump version to 0.12.0

* Don't recast longitude in haversine functions

* Complete tests

* Docstrings

* Update clouddrift/analysis.py

* Update clouddrift/analysis.py

* Update clouddrift/analysis.py

* Update clouddrift/analysis.py

* Update clouddrift/analysis.py

* Allow uppercase in `integration_scheme` values.

* Allow uppercase in `coord_system` values.

Author: milancurcic

clouddrift/analysis.py

@@ -5,7 +5,7 @@
 from concurrent import futures
 from datetime import timedelta
 import warnings
-from clouddrift.haversine import distance, bearing
+from clouddrift.haversine import distance, bearing, position_from_distance_and_bearing
 from clouddrift.dataformat import unpack_ragged
 
 
@@ -375,6 +375,207 @@ def segment(
         return np.concatenate(segment_sizes)
 
 
+def position_from_velocity(
+    u: np.ndarray,
+    v: np.ndarray,
+    time: np.ndarray,
+    x_origin: float,
+    y_origin: float,
+    coord_system: Optional[str] = "spherical",
+    integration_scheme: Optional[str] = "forward",
+    time_axis: Optional[int] = -1,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute positions from arrays of velocities and time and a pair of origin
+    coordinates.
+
+    The units of the result are degrees if ``coord_system == "spherical"`` (default).
+    If ``coord_system == "cartesian"``, the units of the result are equal to the
+    units of the input velocities multiplied by the units of the input time.
+    For example, if the input velocities are in meters per second and the input
+    time is in seconds, the units of the result will be meters.
+
+    Integration scheme can take one of three values:
+
+        1. "forward" (default): integration from x[i] to x[i+1] is performed
+            using the velocity at x[i].
+        2. "backward": integration from x[i] to x[i+1] is performed using the
+            velocity at x[i+1].
+        3. "centered": integration from x[i] to x[i+1] is performed using the
+            arithmetic average of the velocities at x[i] and x[i+1]. Note that
+            this method introduces some error due to the averaging.
+
+    u, v, and time can be multi-dimensional arrays. If the time axis, along
+    which the finite differencing is performed, is not the last one (i.e.
+    x.shape[-1]), use the ``time_axis`` optional argument to specify along which
+    axis should the differencing be done. ``x``, ``y``, and ``time`` must have
+    the same shape.
+
+    This function will not do any special handling of longitude ranges. If the
+    integrated trajectory crosses the antimeridian (dateline) in either direction, the
+    longitude values will not be adjusted to stay in any specific range such
+    as [-180, 180] or [0, 360]. If you need your longitudes to be in a specific
+    range, recast the resulting longitude from this function using the function
+    :func:`clouddrift.sphere.recast_lon`.
+
+    Parameters
+    ----------
+    u : np.ndarray
+        An array of eastward velocities.
+    v : np.ndarray
+        An array of northward velocities.
+    time : np.ndarray
+        An array of time values.
+    x_origin : float
+        Origin x-coordinate or origin longitude.
+    y_origin : float
+        Origin y-coordinate or origin latitude.
+    coord_system : str, optional
+        The coordinate system of the input. Can be "spherical" or "cartesian".
+        Default is "spherical".
+    integration_scheme : str, optional
+        The difference scheme to use for computing the position. Can be
+        "forward" or "backward". Default is "forward".
+    time_axis : int, optional
+        The axis of the time array. Default is -1, which corresponds to the
+        last axis.
+
+    Returns
+    -------
+    x : np.ndarray
+        An array of zonal displacements or longitudes.
+    y : np.ndarray
+        An array of meridional displacements or latitudes.
+
+    Examples
+    --------
+
+    Simple integration on a plane, using the forward scheme by default:
+
+    >>> import numpy as np
+    >>> from clouddrift.analysis import position_from_velocity
+    >>> u = np.array([1., 2., 3., 4.])
+    >>> v = np.array([1., 1., 1., 1.])
+    >>> time = np.array([0., 1., 2., 3.])
+    >>> x, y = position_from_velocity(u, v, time, 0, 0, coord_system="cartesian")
+    >>> x
+    array([0., 1., 3., 6.])
+    >>> y
+    array([0., 1., 2., 3.])
+
+    As above, but using centered scheme:
+
+    >>> x, y = position_from_velocity(u, v, time, 0, 0, coord_system="cartesian", integration_scheme="centered")
+    >>> x
+    array([0., 1.5, 4., 7.5])
+    >>> y
+    array([0., 1., 2., 3.])
+
+    Simple integration on a sphere (default):
+
+    >>> u = np.array([1., 2., 3., 4.])
+    >>> v = np.array([1., 1., 1., 1.])
+    >>> time = np.array([0., 1., 2., 3.]) * 1e5
+    >>> x, y = position_from_velocity(u, v, time, 0, 0)
+    >>> x
+    array([0.        , 0.89839411, 2.69584476, 5.39367518])
+    >>> y
+    array([0.        , 0.89828369, 1.79601515, 2.69201609])
+
+    Integrating across the antimeridian (dateline) by default does not
+    recast the resulting longitude:
+
+    >>> u = np.array([1., 1.])
+    >>> v = np.array([0., 0.])
+    >>> time = np.array([0, 1e5])
+    >>> x, y = position_from_velocity(u, v, time, 179.5, 0)
+    >>> x
+    array([179.5      , 180.3983205])
+    >>> y
+    array([0., 0.])
+
+    Use the ``clouddrift.sphere.recast_lon`` function to recast the longitudes
+    to the desired range:
+
+    >>> from clouddrift.sphere import recast_lon
+    >>> recast_lon(x, -180)
+    array([ 179.5      , -179.6016795])
+
+    Raises
+    ------
+    ValueError
+        If the input arrays do not have the same shape.
+        If the time axis is outside of the valid range ([-1, N-1]).
+        If the input coordinate system is not "spherical" or "cartesian".
+        If the input integration scheme is not "forward", "backward", or "centered"
+
+    See Also
+    --------
+    :func:`velocity_from_position`
+    """
+    # Positions and time arrays must have the same shape.
+    if not u.shape == v.shape == time.shape:
+        raise ValueError("u, v, and time 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(x.shape) - 1}])."
+        )
+
+    # Nominal order of axes on input, i.e. (0, 1, 2, ..., N-1)
+    target_axes = list(range(len(u.shape)))
+
+    # If time_axis is not the last one, transpose the inputs
+    if time_axis != -1 and time_axis < len(u.shape) - 1:
+        target_axes.append(target_axes.pop(target_axes.index(time_axis)))
+
+    # Reshape the inputs to ensure the time axis is last (fast-varying)
+    u_ = np.transpose(u, target_axes)
+    v_ = np.transpose(v, target_axes)
+    time_ = np.transpose(time, target_axes)
+
+    x = np.zeros(u_.shape, dtype=u.dtype)
+    y = np.zeros(v_.shape, dtype=v.dtype)
+
+    dt = np.diff(time_)
+
+    if integration_scheme.lower() == "forward":
+        x[..., 1:] = np.cumsum(u_[..., :-1] * dt, axis=-1)
+        y[..., 1:] = np.cumsum(v_[..., :-1] * dt, axis=-1)
+    elif integration_scheme.lower() == "backward":
+        x[..., 1:] = np.cumsum(u_[1:] * dt, axis=-1)
+        y[..., 1:] = np.cumsum(v_[1:] * dt, axis=-1)
+    elif integration_scheme.lower() == "centered":
+        x[..., 1:] = np.cumsum(0.5 * (u_[..., :-1] + u_[..., 1:]) * dt, axis=-1)
+        y[..., 1:] = np.cumsum(0.5 * (v_[..., :-1] + v_[..., 1:]) * dt, axis=-1)
+    else:
+        raise ValueError(
+            'integration_scheme must be "forward", "backward", or "centered".'
+        )
+
+    if coord_system.lower() == "cartesian":
+        x += x_origin
+        y += y_origin
+    elif coord_system.lower() == "spherical":
+        dx = np.diff(x)
+        dy = np.diff(y)
+        distances = np.sqrt(dx**2 + dy**2)
+        bearings = np.arctan2(dy, dx)
+        x[..., 0], y[..., 0] = x_origin, y_origin
+        for n in range(distances.shape[-1]):
+            y[..., n + 1], x[..., n + 1] = position_from_distance_and_bearing(
+                y[..., n], x[..., n], distances[..., n], bearings[..., n]
+            )
+    else:
+        raise ValueError('coord_system must be "spherical" or "cartesian".')
+
+    if target_axes == list(range(len(u.shape))):
+        return x, y
+    else:
+        return np.transpose(x, target_axes), np.transpose(y, target_axes)
+
+
 def velocity_from_position(
     x: np.ndarray,
     y: np.ndarray,
@@ -394,9 +595,9 @@ def velocity_from_position(
     units of seconds, the resulting velocity is in the units of meters per
     second. Otherwise, if coord_system == "cartesian", the units of the
     resulting velocity correspond to the units of the input. For example,
-    if Easting and Northing are in the units of kilometers and time is in
-    the units of hours, the resulting velocity is in the units of kilometers
-    per hour.
+    if zonal and meridional displacements are in the units of kilometers and
+    time is in the units of hours, the resulting velocity is in the units of
+    kilometers per hour.
 
     x, y, and time can be multi-dimensional arrays. If the time axis, along
     which the finite differencing is performed, is not the last one (i.e.
@@ -421,16 +622,39 @@ def velocity_from_position(
     case of a centered difference scheme, the start and end boundary points are
     evaluated using the forward and backward difference scheme, respectively.
 
-    Args:
-        x (array_like): An N-d array of x-positions (longitude in degrees or easting in any unit)
-        y (array_like): An N-d array of y-positions (latitude in degrees or northing in any unit)
-        time (array_like): An N-d array of times as floating point values (in any unit)
-        coord_system (str, optional): 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): Axis along which to differentiate (default is -1)
+    Parameters
+    ----------
+    x : array_like
+        An N-d array of x-positions (longitude in degrees or zonal displacement in any unit)
+    y : array_like
+        An N-d array of y-positions (latitude in degrees or meridional displacement in any unit)
+    time : array_like
+        An N-d array of times as floating point values (in any unit)
+    coord_system : str, optional
+        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)
+        Axis along which to differentiate (default is -1)
 
-    Returns:
-        out (Tuple[xr.DataArray[float], xr.DataArray[float]]): Arrays of x- and y-velocities
+    Returns
+    -------
+    u : np.ndarray
+        Zonal velocity
+    v : np.ndarray
+        Meridional velocity
+
+    Raises
+    ------
+    ValueError
+        If x, y, and time do not have the same shape.
+        If time_axis is outside of the valid range.
+        If coord_system is not "spherical" or "cartesian".
+        If difference_scheme is not "forward", "backward", or "centered".
+
+    See Also
+    --------
+    :function:`position_from_velocity`
     """
 
     # Positions and time arrays must have the same shape.
@@ -536,10 +760,14 @@ def velocity_from_position(
 
         elif coord_system == "spherical":
             # Inner values
-            distances = distance(y_[..., :-2], x_[..., :-2], y_[..., 2:], x_[..., 2:])
-            bearings = bearing(y_[..., :-2], x_[..., :-2], y_[..., 2:], x_[..., 2:])
-            dx[..., 1:-1] = distances * np.cos(bearings) / 2
-            dy[..., 1:-1] = distances * np.sin(bearings) / 2
+            y1 = (y_[..., :-2] + y_[..., 1:-1]) / 2
+            x1 = (x_[..., :-2] + x_[..., 1:-1]) / 2
+            y2 = (y_[..., 2:] + y_[..., 1:-1]) / 2
+            x2 = (x_[..., 2:] + x_[..., 1:-1]) / 2
+            distances = distance(y1, x1, y2, x2)
+            bearings = bearing(y1, x1, y2, x2)
+            dx[..., 1:-1] = distances * np.cos(bearings)
+            dy[..., 1:-1] = distances * np.sin(bearings)
 
             # Boundary values
             distance1 = distance(y_[..., 0], x_[..., 0], y_[..., 1], x_[..., 1])

clouddrift/haversine.py

@@ -1,4 +1,5 @@
 import numpy as np
+from typing import Tuple
 import xarray as xr
 
 EARTH_RADIUS_METERS = 6.3781e6
@@ -68,7 +69,7 @@ def bearing(
     θ = atan2(cos φ1 ⋅ sin φ2 - sin φ1 ⋅ cos φ2 ⋅ cos Δλ, sin Δλ ⋅ cos φ2)
 
     where (φ, λ) is (lat, lon) and θ is bearing, all in radians.
-    Bearing is defined as zero toward East and increasing counter-clockwise.
+    Bearing is defined as zero toward East and positive counterclockwise.
 
     Args:
         lat1 (array_like): Latitudes of the first set of points, in degrees
@@ -109,3 +110,53 @@ def bearing(
     )
 
     return theta
+
+
+def position_from_distance_and_bearing(
+    lat: float, lon: float, distance: float, bearing: float
+) -> Tuple[float, float]:
+    """Return elementwise new position in degrees from arrays of latitude and
+    longitude in degrees, distance in meters, and bearing in radians, based on
+    the spherical law of cosines.
+
+    The formula is:
+
+    φ2 = asin( sin φ1 ⋅ cos δ + cos φ1 ⋅ sin δ ⋅ cos θ )
+    λ2 = λ1 + atan2( sin θ ⋅ sin δ ⋅ cos φ1, cos δ − sin φ1 ⋅ sin φ2 )
+
+    where (φ, λ) is (lat, lon) and θ is bearing, all in radians.
+    Bearing is defined as zero toward East and positive counterclockwise.
+
+    Parameters
+    ----------
+    lat : float
+        Latitude of the first set of points, in degrees
+    lon : float
+        Longitude of the first set of points, in degrees
+    distance : array_like
+        Distance in meters
+    bearing : array_like
+        Bearing angles in radians
+
+    Returns
+    -------
+    lat2 : array_like
+        Latitudes of the second set of points, in degrees, in the range [-90, 90]
+    lon2 : array_like
+        Longitudes of the second set of points, in degrees, in the range [-180, 180]
+    """
+    lat_rad = np.deg2rad(lat)
+    lon_rad = np.deg2rad(lon)
+
+    distance_rad = distance / EARTH_RADIUS_METERS
+
+    lat2_rad = np.arcsin(
+        np.sin(lat_rad) * np.cos(distance_rad)
+        + np.cos(lat_rad) * np.sin(distance_rad) * np.sin(bearing)
+    )
+    lon2_rad = lon_rad + np.arctan2(
+        np.cos(bearing) * np.sin(distance_rad) * np.cos(lat_rad),
+        np.cos(distance_rad) - np.sin(lat_rad) * np.sin(lat2_rad),
+    )
+
+    return np.rad2deg(lat2_rad), np.rad2deg(lon2_rad)

clouddrift/sphere.py

@@ -7,14 +7,14 @@ def recast_lon(lon: np.ndarray, lon0: Optional[float] = -180) -> np.ndarray:
 
     Parameters
     ----------
-    lon : np.ndarray
+    lon : np.ndarray or float
         An N-d array of longitudes in degrees
     lon0 : float, optional
         Starting longitude of the recasted range (default -180).
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or float
         Converted longitudes in the range `[lon0, lon0+360]`
 
     Examples
@@ -43,9 +43,6 @@ def recast_lon(lon: np.ndarray, lon0: Optional[float] = -180) -> np.ndarray:
     --------
     :func:`recast_lon360`, :func:`recast_lon180`
     """
-    if np.isscalar(lon):
-        lon = np.array([lon])
-
     return np.mod(lon - lon0, 360) + lon0