Implement plane_to_sphere (#186)

* Implement plane_to_sphere

* Add round-trip test

* Expand docstrings; fix for N-d arrays

* Update clouddrift/sphere.py

Co-authored-by: Philippe Miron <philippemiron@gmail.com>

* Update clouddrift/sphere.py

Co-authored-by: Philippe Miron <philippemiron@gmail.com>

* Fix tests; update docstrings

---------

Co-authored-by: Philippe Miron <philippemiron@gmail.com>

Author: milancurcic

clouddrift/sphere.py

@@ -105,6 +105,82 @@ def recast_lon180(lon: np.ndarray) -> np.ndarray:
     return recast_lon(lon, -180)
 
 
+def plane_to_sphere(
+    x: np.ndarray, y: np.ndarray, lon_origin: float = 0, lat_origin: float = 0
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Convert Cartesian coordinates on a plane to spherical coordinates.
+
+    The arrays of input zonal and meridional displacements ``x`` and ``y`` are
+    assumed to follow a contiguous trajectory. The spherical coordinate of each
+    successive point is determined by following a great circle path from the
+    previous point. The spherical coordinate of the first point is determined by
+    following a great circle path from the origin, by default (0, 0).
+
+    The output arrays have the same floating-point output type as the input.
+
+    If projecting multiple trajectories onto the same plane, use
+    :func:`apply_ragged` for highest accuracy.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        An N-d array of zonal displacements in meters
+    y : np.ndarray
+        An N-d array of meridional displacements in meters
+    lon_origin : float, optional
+        Origin longitude of the tangent plane in degrees, default 0
+    lat_origin : float, optional
+        Origin latitude of the tangent plane in degrees, default 0
+
+    Returns
+    -------
+    lon : np.ndarray
+        Longitude in degrees
+    lat : np.ndarray
+        Latitude in degrees
+
+    Examples
+    --------
+    >>> plane_to_sphere(np.array([0., 0.]), np.array([0., 1000.]))
+    (array([0.00000000e+00, 5.50062664e-19]), array([0.       , 0.0089832]))
+
+    You can also specify an origin longitude and latitude:
+
+    >>> plane_to_sphere(np.array([0., 0.]), np.array([0., 1000.]), lon_origin=1, lat_origin=0)
+    (array([1., 1.]), array([0.       , 0.0089832]))
+
+    Raises
+    ------
+    AttributeError
+        If ``x`` and ``y`` are not NumPy arrays
+
+    See Also
+    --------
+    :func:`sphere_to_plane`
+    """
+    lon = np.empty_like(x)
+    lat = np.empty_like(y)
+
+    # Cartesian distances between each point
+    dx = np.diff(x, prepend=0)
+    dy = np.diff(y, prepend=0)
+
+    distances = np.sqrt(dx**2 + dy**2)
+    bearings = np.arctan2(dy, dx)
+
+    # Compute spherical coordinates following great circles between each
+    # successive point.
+    lat[..., 0], lon[..., 0] = haversine.position_from_distance_and_bearing(
+        lat_origin, lon_origin, distances[..., 0], bearings[..., 0]
+    )
+    for n in range(1, lon.shape[-1]):
+        lat[..., n], lon[..., n] = haversine.position_from_distance_and_bearing(
+            lat[..., n - 1], lon[..., n - 1], distances[..., n], bearings[..., n]
+        )
+
+    return lon, lat
+
+
 def sphere_to_plane(
     lon: np.ndarray, lat: np.ndarray, lon_origin: float = 0, lat_origin: float = 0
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -116,8 +192,7 @@ def sphere_to_plane(
     The Cartesian coordinate of the first point is determined by following a
     great circle path from the origin, by default (0, 0).
 
-    This function uses 64-bit floats for all intermediate calculations,
-    regardless of the type of input arrays, to avoid loss of precision.
+    The output arrays have the same floating-point output type as the input.
 
     If projecting multiple trajectories onto the same plane, use
     :func:`apply_ragged` for highest accuracy.
@@ -135,29 +210,36 @@ def sphere_to_plane(
 
     Returns
     -------
-    Tuple[np.ndarray, np.ndarray]
-        x- and y-coordinates of the tangent plane
+    x : np.ndarray
+        x-coordinates on the tangent plane
+    y : np.ndarray
+        y-coordinates on the tangent plane
 
     Examples
     --------
     >>> sphere_to_plane(np.array([0., 1.]), np.array([0., 0.]))
     (array([     0.        , 111318.84502145]), array([0., 0.]))
 
-    You can also specify an x and y origin:
+    You can also specify an origin longitude and latitude:
 
     >>> sphere_to_plane(np.array([0., 1.]), np.array([0., 0.]), lon_origin=1, lat_origin=0)
     (array([-111318.84502145,       0.        ]),
      array([1.36326267e-11, 1.36326267e-11]))
 
     Raises
     ------
-    TypeError
+    AttributeError
         If ``lon`` and ``lat`` are not NumPy arrays
+
+    See Also
+    --------
+    :func:`plane_to_sphere`
     """
-    x = np.empty(lon.shape, dtype=np.float64)
-    y = np.empty(lat.shape, dtype=np.float64)
-    distances = np.empty(lon.shape, dtype=np.float64)
-    bearings = np.empty(lon.shape, dtype=np.float64)
+    x = np.empty_like(lon)
+    y = np.empty_like(lat)
+
+    distances = np.empty_like(x)
+    bearings = np.empty_like(x)
 
     # Distance and bearing of the starting point relative to the origin
     distances[0] = haversine.distance(lat_origin, lon_origin, lat[..., 0], lon[..., 0])

tests/sphere_tests.py

@@ -1,8 +1,15 @@
 from clouddrift import haversine
-from clouddrift.sphere import recast_lon, recast_lon180, recast_lon360, sphere_to_plane
+from clouddrift.sphere import (
+    recast_lon,
+    recast_lon180,
+    recast_lon360,
+    plane_to_sphere,
+    sphere_to_plane,
+)
 import unittest
 import numpy as np
 
+ONE_DEGREE_METERS = np.deg2rad(haversine.EARTH_RADIUS_METERS)
 
 if __name__ == "__main__":
     unittest.main()
@@ -79,6 +86,50 @@ def test_decimals(self):
         )
 
 
+class plane_to_sphere_tests(unittest.TestCase):
+    def test_simple(self):
+        lon, lat = plane_to_sphere(
+            np.array([0.0, ONE_DEGREE_METERS]), np.array([0.0, 0.0])
+        )
+        self.assertTrue(np.allclose(lon, np.array([0.0, 1.0])))
+        self.assertTrue(np.allclose(lat, np.array([0.0, 0.0])))
+
+        lon, lat = plane_to_sphere(
+            np.array([0.0, 0.0]), np.array([0.0, ONE_DEGREE_METERS])
+        )
+        self.assertTrue(np.allclose(lon, np.array([0.0, 0.0])))
+        self.assertTrue(np.allclose(lat, np.array([0.0, 1.0])))
+
+    def test_with_origin(self):
+        lon_origin = 5
+        lat_origin = 0
+
+        lon, lat = plane_to_sphere(
+            np.array([0.0, ONE_DEGREE_METERS]),
+            np.array([0.0, 0.0]),
+            lon_origin,
+            lat_origin,
+        )
+        self.assertTrue(np.allclose(lon, np.array([lon_origin, lon_origin + 1])))
+        self.assertTrue(np.allclose(lat, np.array([lat_origin, lat_origin])))
+
+        lon_origin = 0
+        lat_origin = 5
+
+        lon, lat = plane_to_sphere(
+            np.array([0.0, 0.0]),
+            np.array([0.0, ONE_DEGREE_METERS]),
+            lon_origin,
+            lat_origin,
+        )
+        self.assertTrue(np.allclose(lon, np.array([lon_origin, lon_origin])))
+        self.assertTrue(np.allclose(lat, np.array([lat_origin, lat_origin + 1])))
+
+    def test_scalar_raises_error(self):
+        with self.assertRaises(Exception):
+            plane_to_sphere(0, 0)
+
+
 class sphere_to_plane_tests(unittest.TestCase):
     def test_simple(self):
         x, y = sphere_to_plane(np.array([0.0, 1.0]), np.array([0.0, 0.0]))
@@ -97,8 +148,6 @@ def test_with_origin(self):
         lon_origin = 5
         lat_origin = 0
 
-        ONE_DEGREE_METERS = np.deg2rad(haversine.EARTH_RADIUS_METERS)
-
         x, y = sphere_to_plane(
             np.array([0.0, 1.0]), np.array([0.0, 0.0]), lon_origin, lat_origin
         )
@@ -152,6 +201,14 @@ def test_scalar_raises_error(self):
         with self.assertRaises(Exception):
             sphere_to_plane(0, 0)
 
-    def test_list_raises_error(self):
-        with self.assertRaises(Exception):
-            sphere_to_plane([0, 1], [0, 0])
+
+class sphere_to_plane_roundtrip(unittest.TestCase):
+    def test_roundtrip(self):
+        expected_lon = 2 * np.cumsum(np.random.random((100)))
+        expected_lat = np.cumsum(np.random.random((100)))
+
+        x, y = sphere_to_plane(expected_lon, expected_lat)
+        lon, lat = plane_to_sphere(x, y)
+
+        self.assertTrue(np.allclose(lon, expected_lon))
+        self.assertTrue(np.allclose(lat, expected_lat))