Merge pull request #5 from ecmwf/develop

Merge develop into main for release 0.0.2

earthkit/geo/__init__.py

@@ -22,11 +22,14 @@
     nearest_point_haversine,
     nearest_point_kdtree,
 )
+from earthkit.geo.figure import IFS_SPHERE, UNIT_SPHERE
 
 __all__ = [
     "GeoKDTree",
     "haversine_distance",
+    IFS_SPHERE,
     "nearest_point_haversine",
     "nearest_point_kdtree",
+    UNIT_SPHERE,
     "__version__",
 ]

earthkit/geo/constants.py

@@ -11,15 +11,5 @@
 Collection of constants in SI units.
 """
 
-R_earth = 6371229
-r"""Average radius of the Earth [:math:`m`]. See [IFS-CY47R3-PhysicalProcesses]_
- (Chapter 12)."""
-
-full_circle = 360
-r"""Full circle in degrees"""
-
 north = 90
 r"""Latitude of the north pole in degrees"""
-
-south = -90
-r"""Latitude of the south pole in degrees"""

earthkit/geo/distance.py

@@ -8,17 +8,17 @@
 #
 
 import numpy as np
-from scipy.spatial import KDTree
 
 from . import constants
+from .figure import IFS_SPHERE, UNIT_SPHERE
 
 
-def regulate_lat(lat):
+def _regulate_lat(lat):
     return np.where(np.abs(lat) > constants.north, np.nan, lat)
 
 
-def haversine_distance(p1, p2):
-    """Compute haversine distance between two (sets of) points on Earth.
+def haversine_distance(p1, p2, figure=IFS_SPHERE):
+    """Compute haversine distance between two (sets of) points on a spheroid.
 
     Parameters
     ----------
@@ -28,11 +28,14 @@ def haversine_distance(p1, p2):
     p2: pair of array-like
         Locations of the second points. The first item specifies the latitudes,
         the second the longitudes (degrees)
+    figure: :class:`geo.figure.Figure`, optional
+        Figure of the spheroid (default: :obj:`geo.figure.IFS_SPHERE`)
 
     Returns
     -------
     number or ndarray
-        Spherical distance on the surface in Earth (m)
+        Distance (m) on the surface of the spheroid defined by ``figure``.
+
 
     Either ``p1`` or ``p2`` must be a single point.
 
@@ -74,8 +77,8 @@ def haversine_distance(p1, p2):
     if lat1.size != 1 and lat2.size != 1:
         raise ValueError("haversine_distance: either p1 or p2 must be a single point")
 
-    lat1 = regulate_lat(lat1)
-    lat2 = regulate_lat(lat2)
+    lat1 = _regulate_lat(lat1)
+    lat2 = _regulate_lat(lat2)
 
     lat1, lon1, lat2, lon2 = map(np.deg2rad, [lat1, lon1, lat2, lon2])
     d_lon = lon2 - lon1
@@ -84,15 +87,14 @@ def haversine_distance(p1, p2):
     a = np.sqrt(
         np.sin(d_lat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(d_lon / 2) ** 2
     )
-    c = 2 * np.arcsin(a)
-    distance = constants.R_earth * c
+    distance = 2 * np.arcsin(a)
 
-    return distance
+    return figure.scale(distance)
 
 
-def nearest_point_haversine(ref_points, points):
-    """Find the index of the nearest point to all ``ref_points`` in a set of ``points`` using the
-       haversine distance formula.
+def nearest_point_haversine(ref_points, points, figure=IFS_SPHERE):
+    """Find the index of the nearest point to all ``ref_points`` in a set of
+      ``points`` using the haversine distance formula.
 
     Parameters
     ----------
@@ -102,14 +104,17 @@ def nearest_point_haversine(ref_points, points):
         Locations of the set of points from which the nearest to
         ``ref_points`` is to be found. The first item specifies the latitudes,
         the second the longitudes (degrees)
+    figure: :class:`geo.figure.Figure`, optional
+        Figure of the spheroid the returned
+        distances are computed on (default: :obj:`geo.figure.IFS_SPHERE`)
 
     Returns
     -------
     ndarray
-        Indices of the nearest points to ``ref_points`.
+        Indices of the nearest points to ``ref_points``.
     ndarray
         The distance (m) between the ``ref_points`` and the corresponding nearest
-        point in ``points``.
+        point in ``points`` on the surface of the spheroid defined by ``figure``.
 
     Examples
     --------
@@ -139,45 +144,52 @@ def nearest_point_haversine(ref_points, points):
     res_index = []
     res_distance = []
     for lat, lon in ref_points.T:
-        distance = haversine_distance((lat, lon), points).flatten()
+        distance = haversine_distance((lat, lon), points, figure=UNIT_SPHERE).flatten()
         index = np.nanargmin(distance)
         index = index[0] if isinstance(index, np.ndarray) else index
         res_index.append(index)
         res_distance.append(distance[index])
-    return (np.array(res_index), np.array(res_distance))
+    return (np.array(res_index), figure.scale(np.array(res_distance)))
 
 
-def ll_to_xyz(lat, lon):
+def _latlon_to_xyz(lat, lon):
+    """Works on the unit sphere."""
     lat = np.asarray(lat)
     lon = np.asarray(lon)
     lat = np.radians(lat)
     lon = np.radians(lon)
-    x = constants.R_earth * np.cos(lat) * np.cos(lon)
-    y = constants.R_earth * np.cos(lat) * np.sin(lon)
-    z = constants.R_earth * np.sin(lat)
+    x = np.cos(lat) * np.cos(lon)
+    y = np.cos(lat) * np.sin(lon)
+    z = np.sin(lat)
+
     return x, y, z
 
 
-def cordlength_to_arclength(chord_length):
+def _chordlength_to_arclength(chord_length):
     """
     Convert 3D (Euclidean) distance to great circle arc length
     https://en.wikipedia.org/wiki/Great-circle_distance
+    Works on the unit sphere.
     """
-    central_angle = 2.0 * np.arcsin(chord_length / (2.0 * constants.R_earth))
-    return constants.R_earth * central_angle
+    central_angle = 2.0 * np.arcsin(chord_length / 2.0)
+    return central_angle
 
 
-def arclength_to_cordlenght(arc_length):
+def _arclength_to_chordlength(arc_length):
     """
     Convert great circle arc length to 3D (Euclidean) distance
     https://en.wikipedia.org/wiki/Great-circle_distance
+    Works on the unit sphere.
     """
-    central_angle = arc_length / constants.R_earth
-    return np.sin(central_angle / 2) * 2.0 * constants.R_earth
+    central_angle = arc_length
+    return np.sin(central_angle / 2) * 2.0
 
 
 class GeoKDTree:
     def __init__(self, lats, lons):
+        """KDTree built from ``lats`` and ``lons``."""
+        from scipy.spatial import KDTree
+
         lats = np.asarray(lats).flatten()
         lons = np.asarray(lons).flatten()
 
@@ -187,31 +199,51 @@ def __init__(self, lats, lons):
             lats = lats[mask]
             lons = lons[mask]
 
-        x, y, z = ll_to_xyz(lats, lons)
+        x, y, z = _latlon_to_xyz(lats, lons)
         v = np.column_stack((x, y, z))
         self.tree = KDTree(v)
 
         # TODO: allow user to specify max distance
-        self.max_distance_arc = 10000 * 1000  # m
-        if self.max_distance_arc <= np.pi * constants.R_earth:
-            self.max_distance_cord = arclength_to_cordlenght(self.max_distance_arc)
+        self.max_distance_arc = np.pi / 4
+        if self.max_distance_arc <= np.pi:
+            self.max_distance_chord = _arclength_to_chordlength(self.max_distance_arc)
         else:
-            self.max_distance_cord = np.inf
-
-    def nearest_point(self, points):
-        lat, lon = points
-        x, y, z = ll_to_xyz(lat, lon)
+            self.max_distance_chord = np.inf
+
+    def nearest_point(self, ref_points, figure=IFS_SPHERE):
+        """Find the index of the nearest point to all ``ref_points``.
+
+        Parameters
+        ----------
+        ref_points: pair of array-like
+            Latitude and longitude coordinates of the reference point (degrees)
+        figure: :class:`geo.figure.Figure`, optional
+            Figure of the spheroid the returned
+            distances are computed on (default: :obj:`geo.figure.IFS_SPHERE`)
+
+        Returns
+        -------
+        ndarray
+            Indices of the nearest points to ``ref_points``.
+        ndarray
+            The distance (m) between the ``ref_points`` and the corresponding nearest
+            point in ``points``. Computed on the surface of the spheroid defined
+            by ``figure``.
+
+        """
+        lat, lon = ref_points
+        x, y, z = _latlon_to_xyz(lat, lon)
         points = np.column_stack((x, y, z))
 
         # find the nearest point
         distance, index = self.tree.query(
             points, distance_upper_bound=self.max_distance_arc
         )
 
-        return index, cordlength_to_arclength(distance)
+        return index, figure.scale(_chordlength_to_arclength(distance))
 
 
-def nearest_point_kdtree(ref_points, points):
+def nearest_point_kdtree(ref_points, points, figure=IFS_SPHERE):
     """Find the index of the nearest point to all ``ref_points`` in a set of ``points`` using a KDTree.
 
     Parameters
@@ -222,14 +254,18 @@ def nearest_point_kdtree(ref_points, points):
         Locations of the set of points from which the nearest to
         ``ref_points`` is to be found. The first item specifies the latitudes,
         the second the longitudes (degrees)
+    figure: :class:`geo.figure.Figure`, optional
+        Figure of the spheroid the returned
+        distances are computed on (default: :obj:`geo.figure.IFS_SPHERE`)
 
     Returns
     -------
     ndarray
-        Indices of the nearest points to ``ref_points`.
+        Indices of the nearest points to ``ref_points``.
     ndarray
         The distance (m) between the ``ref_points`` and the corresponding nearest
-        point in ``points``.
+        point in ``points``. Computed on the surface of the spheroid defined
+        by ``figure``.
 
     Examples
     --------
@@ -250,5 +286,5 @@ def nearest_point_kdtree(ref_points, points):
     """
     lats, lons = points
     tree = GeoKDTree(lats, lons)
-    index, distance = tree.nearest_point(ref_points)
+    index, distance = tree.nearest_point(ref_points, figure=figure)
     return index, distance

earthkit/geo/figure.py

@@ -0,0 +1,54 @@
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+
+class Figure:
+    """Figure representing the size and shape of a spheroid"""
+
+    pass
+
+
+class Sphere(Figure):
+    """Base class for a sphere. The radius is in metres."""
+
+    def __init__(self, radius):
+        self._radius = radius
+        if self._radius <= 0.0:
+            raise ValueError(f"Radius={self._radius} must be positive")
+
+    @property
+    def radius(self):
+        return self._radius
+
+    def scale(self, *args):
+        if len(args) > 1:
+            return tuple([x * self.radius for x in args])
+        else:
+            return args[0] * self.radius
+
+
+class UnitSphere(Sphere):
+    """Unit sphere (with the radius of 1 m)."""
+
+    def __init__(self):
+        super().__init__(1.0)
+
+    def scale(self, *args):
+        if len(args) > 1:
+            return args
+        else:
+            return args[0]
+
+
+IFS_SPHERE = Sphere(6371229.0)
+r"""Object of spherical Earth with radius=6371229 m as used in the IFS.
+See [IFS-CY47R3-PhysicalProcesses]_ (Chapter 12)."""
+
+UNIT_SPHERE = UnitSphere()
+r"""Object of unit sphere (with the radius of 1 m)."""

tests/distance/test_distance.py

@@ -13,6 +13,7 @@
 import pytest
 
 from earthkit.geo import (
+    UNIT_SPHERE,
     haversine_distance,
     nearest_point_haversine,
     nearest_point_kdtree,
@@ -103,6 +104,15 @@ def test_haversine_distance(dlon_ref, dlon_p, p_ref, d_ref):
     ), f"p_ref={p_ref},dlon_ref={dlon_ref},dlon_p={dlon_p}"
 
 
+def test_haversine_distance_with_radius():
+    p = (0, 90)
+    p_ref = (48.0, 20)
+    d_ref = 1.33989384590445
+
+    d = haversine_distance(p_ref, p, figure=UNIT_SPHERE)
+    assert np.allclose(d, d_ref)
+
+
 # lat is out of range
 @pytest.mark.parametrize("p_ref", [((90.00001, 0)), (((-90.00001, 0)))])
 def test_haversine_distance_invalid(p_ref):
@@ -134,6 +144,25 @@ def test_nearest_single_ref(mode, p_ref, index_ref, dist_ref):
     assert np.allclose(distance, dist_ref)
 
 
+@pytest.mark.parametrize("mode", ["haversine", "kdtree"])
+@pytest.mark.parametrize(
+    "p_ref,index_ref,dist_ref",
+    [
+        ((0, 0), 0, 0),
+        ((15, 22), 0, 0.46103882),
+        ((-50, -18), 8, 0.04174471),
+    ],
+)
+def test_nearest_single_ref_with_radius(mode, p_ref, index_ref, dist_ref):
+    lats = np.array([0.0, 0, 0, 0, 90, -90, 48, 48, -48, -48, np.nan])
+    lons = np.array([0, 90, -90, 180, 0, 0, 20, -20, -20, 20, 1.0])
+
+    meth = get_nearest_method(mode)
+    index, distance = meth(p_ref, (lats, lons), figure=UNIT_SPHERE)
+    assert np.allclose(index_ref, index)
+    assert np.allclose(distance, dist_ref)
+
+
 @pytest.mark.parametrize("mode", ["haversine", "kdtree"])
 def test_nearest_multi_ref_single_point(mode):
     p_ref = [(15, 44, 44), (22, 10, -10)]
@@ -167,7 +196,7 @@ def test_nearest_multi_ref_multi_points(mode):
 
 
 @pytest.mark.parametrize("mode", ["haversine", "kdtree"])
-def test_haversine_nearest_invalid(mode):
+def test_nearest_invalid(mode):
     # checks: p_ref must have a (2,) shape
     p_ref = ([15, 21], [22, 7], [44, 12])
     p = ([0.0, 0], [0, 90])