Add binding to calculate distance between two geometries

src/accessors-geog.cpp

@@ -7,6 +7,8 @@ namespace py = pybind11;
 namespace s2geog = s2geography;
 using namespace spherely;
 
+const double EARTH_RADIUS_METERS = 6371.01 * 1000;
+
 PyObjectGeography centroid(PyObjectGeography a) {
     const auto& a_ptr = a.as_geog_ptr()->geog();
     auto s2_point = s2geog::s2_centroid(a_ptr);
@@ -30,7 +32,15 @@ PyObjectGeography convex_hull(PyObjectGeography a) {
     return PyObjectGeography::from_geog(std::move(geog_ptr));
 }
 
+double distance(PyObjectGeography a, PyObjectGeography b, double radius = EARTH_RADIUS_METERS) {
+    const auto& a_index = a.as_geog_ptr()->geog_index();
+    const auto& b_index = b.as_geog_ptr()->geog_index();
+    return s2geog::s2_distance(a_index, b_index) * radius;
+}
+
 void init_accessors(py::module& m) {
+    m.attr("EARTH_RADIUS_METERS") = py::float_(EARTH_RADIUS_METERS);
+
     m.def("centroid",
           py::vectorize(&centroid),
           py::arg("a"),
@@ -69,4 +79,23 @@ void init_accessors(py::module& m) {
             Geography object
 
     )pbdoc");
+
+    m.def("distance",
+          py::vectorize(&distance),
+          py::arg("a"),
+          py::arg("b"),
+          py::arg("radius") = EARTH_RADIUS_METERS,
+          R"pbdoc(
+        Calculate the distance between two geographies.
+
+        Parameters
+        ----------
+        a : :py:class:`Geography` or array_like
+            Geography object
+        b : :py:class:`Geography` or array_like
+            Geography object
+        radius : float, optional
+            Radius of Earth in meters, default 6,371,010
+
+    )pbdoc");
 }

src/spherely.pyi

@@ -13,6 +13,8 @@ from typing import (
 import numpy as np
 import numpy.typing as npt
 
+EARTH_RADIUS_METERS: float = ...
+
 class Geography:
     def __init__(self, *args, **kwargs) -> None: ...
     @property
@@ -87,6 +89,28 @@ class _VFunc_Nin2_Nout1(Generic[_NameType, _ScalarReturnType, _ArrayReturnDType]
         self, a: npt.ArrayLike, b: Geography
     ) -> npt.NDArray[_ArrayReturnDType]: ...
 
+class _VFunc_Nin2optradius_Nout1(
+    Generic[_NameType, _ScalarReturnType, _ArrayReturnDType]
+):
+    @property
+    def __name__(self) -> _NameType: ...
+    @overload
+    def __call__(
+        self, a: Geography, b: Geography, radius: float = ...
+    ) -> _ScalarReturnType: ...
+    @overload
+    def __call__(
+        self, a: npt.ArrayLike, b: npt.ArrayLike, radius: float = ...
+    ) -> npt.NDArray[_ArrayReturnDType]: ...
+    @overload
+    def __call__(
+        self, a: Geography, b: npt.ArrayLike, radius: float = ...
+    ) -> npt.NDArray[_ArrayReturnDType]: ...
+    @overload
+    def __call__(
+        self, a: npt.ArrayLike, b: Geography, radius: float = ...
+    ) -> npt.NDArray[_ArrayReturnDType]: ...
+
 # Geography properties
 
 get_dimensions: _VFunc_Nin1_Nout1[Literal["get_dimensions"], Geography, Any]
@@ -112,6 +136,7 @@ disjoint: _VFunc_Nin2_Nout1[Literal["disjoint"], bool, bool]
 centroid: _VFunc_Nin1_Nout1[Literal["centroid"], Geography, Point]
 boundary: _VFunc_Nin1_Nout1[Literal["boundary"], Geography, Geography]
 convex_hull: _VFunc_Nin1_Nout1[Literal["convex_hull"], Geography, Polygon]
+distance: _VFunc_Nin2optradius_Nout1[Literal["distance"], float, float]
 
 # temp (remove)
 

tests/test_accessors.py

@@ -1,8 +1,9 @@
 import numpy as np
+import pytest
 
 import spherely
 
-import pytest
+earth_radius_meters = 6_371_010
 
 
 @pytest.mark.parametrize(
@@ -81,3 +82,47 @@ def test_convex_hull(geog, expected) -> None:
     actual = actual[0]
     assert isinstance(actual, spherely.Polygon)
     assert spherely.equals(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "geog_a, geog_b, expected",
+    [
+        (
+            spherely.Point(0, 0),
+            spherely.Point(90, 0),
+            np.pi / 2 * spherely.EARTH_RADIUS_METERS,
+        ),
+        (
+            spherely.Point(90, 0),
+            spherely.Point(30, 90),
+            np.pi / 3 * spherely.EARTH_RADIUS_METERS,
+        ),
+        (
+            spherely.Polygon([(0, 0), (60, 30), (-30, 60)]),
+            spherely.Point(90, 0),
+            np.pi / 6 * spherely.EARTH_RADIUS_METERS,
+        ),
+    ],
+)
+def test_distance(geog_a, geog_b, expected) -> None:
+    # scalar
+    actual = spherely.distance(geog_a, geog_b)
+    assert isinstance(actual, float)
+    assert actual == pytest.approx(expected, 1e-9)
+
+    # array
+    actual = spherely.distance([geog_a], [geog_b])
+    assert isinstance(actual, np.ndarray)
+    actual = actual[0]
+    assert isinstance(actual, float)
+    assert actual == pytest.approx(expected, 1e-9)
+
+
+def test_distance_with_custom_radius() -> None:
+    actual = spherely.distance(
+        spherely.Point(90, 0),
+        spherely.Point(0, 0),
+        radius=1,
+    )
+    assert isinstance(actual, float)
+    assert actual == pytest.approx(np.pi / 2)