diff --git a/src/accessors-geog.cpp b/src/accessors-geog.cpp index 172b786..927a54a 100644 --- a/src/accessors-geog.cpp +++ b/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(¢roid), 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"); } diff --git a/src/spherely.pyi b/src/spherely.pyi index 425280c..f01e74f 100644 --- a/src/spherely.pyi +++ b/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) diff --git a/tests/test_accessors.py b/tests/test_accessors.py index 205823a..5f64c28 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -1,9 +1,8 @@ import numpy as np +import pytest import spherely -import pytest - @pytest.mark.parametrize( "geog, expected", @@ -81,3 +80,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)