Skip to content

Commit

Permalink
Add binding to calculate distance between two geometries
Browse files Browse the repository at this point in the history
  • Loading branch information
gahjelle committed Aug 30, 2024
1 parent 64d14e0 commit 895efde
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 1 deletion.
29 changes: 29 additions & 0 deletions src/accessors-geog.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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"),
Expand Down Expand Up @@ -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
Radius of Earth in meters, default 6,371,010
)pbdoc");
}
25 changes: 25 additions & 0 deletions src/spherely.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand All @@ -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)

Expand Down
47 changes: 46 additions & 1 deletion tests/test_accessors.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import numpy as np
import pytest

import spherely

import pytest
earth_radius_meters = 6_371_010


@pytest.mark.parametrize(
Expand Down Expand Up @@ -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)

0 comments on commit 895efde

Please sign in to comment.