Refactor pairwise haversine function. (#218)

Supports our typical centroid array format, and units besides miles as well.

doc/devlog/2024-12-06.ipynb

@@ -133,10 +133,7 @@
     "\n",
     "# This is copied with some minor changes from Centroids' dispersal_kernel function\n",
     "phi = 40.0\n",
-    "distance = pairwise_haversine(\n",
-    "    centroids[\"longitude\"],\n",
-    "    centroids[\"latitude\"],\n",
-    ")  # (N,N) array\n",
+    "distance = pairwise_haversine(centroids)  # (N,N) array\n",
     "kernel = 1 / np.exp(distance / phi)\n",
     "\n",
     "# the diagonal of this matrix represents movement from location A to location A,\n",

epymorph/data/mm/centroids.py

@@ -53,7 +53,7 @@ def dispersal_kernel(self) -> NDArray[np.float64]:
         """
         centroid = self.data("centroid")
         phi = self.data("phi")
-        distance = pairwise_haversine(centroid["longitude"], centroid["latitude"])
+        distance = pairwise_haversine(centroid)
         return row_normalize(1 / np.exp(distance / phi))
 
     def evaluate(self, tick: Tick) -> NDArray[np.int64]:

epymorph/data/mm/sparsemod.py

@@ -54,7 +54,7 @@ def dispersal_kernel(self) -> NDArray[np.float64]:
         """
         centroid = self.data("centroid")
         phi = self.data("phi")
-        distance = pairwise_haversine(centroid["longitude"], centroid["latitude"])
+        distance = pairwise_haversine(centroid)
         return row_normalize(1 / np.exp(distance / phi))
 
     def evaluate(self, tick: Tick) -> NDArray[np.int64]:

epymorph/test/util_test.py

@@ -106,6 +106,76 @@ def test_check_ndarray_03(self):
             util.check_ndarray(arr, dtype=m.dtype(np.str_))
 
 
+class TestHaversine(unittest.TestCase):
+    test_coords = np.array(
+        [(0, 0), (1, 0), (0, 1)],
+        dtype=[("longitude", np.float64), ("latitude", np.float64)],
+    )
+
+    def test_01(self):
+        # Test miles
+        actual1 = util.pairwise_haversine(self.test_coords)
+        actual2 = util.pairwise_haversine(self.test_coords, units="miles")
+        actual3 = util.pairwise_haversine(self.test_coords, radius=3963.1906)
+        expected = np.array(
+            [
+                [0.0, 69.17, 69.17],
+                [69.17, 0.0, 97.82],
+                [69.17, 97.82, 0.0],
+            ],
+            dtype=np.float64,
+        )
+        np.testing.assert_array_almost_equal(expected, actual1, decimal=2)
+        np.testing.assert_array_almost_equal(expected, actual2, decimal=2)
+        np.testing.assert_array_almost_equal(expected, actual3, decimal=2)
+
+    def test_02(self):
+        # Test kilometers
+        actual1 = util.pairwise_haversine(self.test_coords, units="kilometers")
+        actual2 = util.pairwise_haversine(self.test_coords, radius=6378.1370)
+        expected = np.array(
+            [
+                [0.0, 111.32, 111.32],
+                [111.32, 0.0, 157.43],
+                [111.32, 157.43, 0.0],
+            ],
+            dtype=np.float64,
+        )
+        np.testing.assert_array_almost_equal(expected, actual1, decimal=2)
+        np.testing.assert_array_almost_equal(expected, actual2, decimal=2)
+
+    def test_03(self):
+        # Test custom radius
+        actual1 = util.pairwise_haversine(self.test_coords, radius=10)
+        expected = np.array(
+            [
+                [0.0, 0.17, 0.17],
+                [0.17, 0.0, 0.25],
+                [0.17, 0.25, 0.0],
+            ],
+            dtype=np.float64,
+        )
+        np.testing.assert_array_almost_equal(expected, actual1, decimal=2)
+
+    def test_04(self):
+        # Test tuple input
+        actual1 = util.pairwise_haversine((np.array([0, 1, 0]), np.array([0, 0, 1])))
+        expected = np.array(
+            [
+                [0.0, 69.17, 69.17],
+                [69.17, 0.0, 97.82],
+                [69.17, 97.82, 0.0],
+            ],
+            dtype=np.float64,
+        )
+        np.testing.assert_array_almost_equal(expected, actual1, decimal=2)
+
+    def test_05(self):
+        # Test bad input
+        with self.assertRaises(ValueError):
+            util.pairwise_haversine([1, 2, 3])  # type: ignore
+
+
 class TestProgress(unittest.TestCase):
     def test_zero_percent(self):
         self.assertEqual(progress(0), "|                    | 0% ")

epymorph/util.py

@@ -357,25 +357,73 @@ def mask(length: int, selection: slice | list[int]) -> NDArray[np.bool_]:
     return mask
 
 
-RADIUS_MI = 3959.87433  # radius of earth in mi
+# values from: https://www.themathdoctors.org/distances-on-earth-2-the-haversine-formula
+_EARTH_RADIUS = {
+    "miles": 3963.1906,
+    "kilometers": 6378.1370,
+}
 
 
 def pairwise_haversine(
-    longitudes: NDArray[np.float64], latitudes: NDArray[np.float64]
+    coordinates: NDArray | tuple[NDArray[np.float64], NDArray[np.float64]],
+    *,
+    units: Literal["miles", "kilometers"] = "miles",
+    radius: float | None = None,
 ) -> NDArray[np.float64]:
-    """Compute the distances in miles between all pairs of coordinates."""
+    """Compute the distances between all pairs of coordinates.
+
+    Parameters
+    ----------
+    coordinates : NDArray | tuple[NDArray, NDArray]
+        The coordinates, given in one of two forms: either a structured numpy array
+        with dtype `[("longitude", np.float64), ("latitude", np.float64)]` or a tuple
+        of two numpy arrays, the first containing longitudes and the second latitudes.
+        The coordinates must be given in degrees.
+    units : Literal["miles", "kilometers"] = "miles",
+        The units of distance to use for the result, unless radius is given.
+    radius : float, optional
+        The radius of the Earth to use in calculating the results. If not given,
+        we will use an appropriate value for the given `units`.
+        Since the value of radius implies the distance units being used, if you
+        specify `radius` the value of `units` is ignored.
+
+    Returns
+    -------
+    NDArray[np.float64] :
+        An NxN array of distances where N is the number of coordinates given,
+        representing the distance between each pair of coordinates. The output
+        maintains the same ordering of coordinates as the input.
+
+    Raises
+    ------
+    ValueError :
+        if coordinates are not given in an expected format
+    """
     # https://www.themathdoctors.org/distances-on-earth-2-the-haversine-formula
-    lng = np.radians(longitudes)
-    lat = np.radians(latitudes)
-    dlng = lng[:, np.newaxis] - lng[np.newaxis, :]
-    dlat = lat[:, np.newaxis] - lat[np.newaxis, :]
-    cos_lat = np.cos(lat)
+    if isinstance(coordinates, np.ndarray) and coordinates.dtype == np.dtype(
+        [("longitude", np.float64), ("latitude", np.float64)]
+    ):
+        lng = coordinates["longitude"]
+        lat = coordinates["latitude"]
+    elif isinstance(coordinates, tuple) and len(coordinates) == 2:
+        lng = coordinates[0]
+        lat = coordinates[1]
+    else:
+        err = "Unable to interpret the given `coordinates`."
+        raise ValueError(err)
+
+    lngrad = np.radians(lng)
+    latrad = np.radians(lat)
+    dlng = lngrad[:, np.newaxis] - lngrad[np.newaxis, :]
+    dlat = latrad[:, np.newaxis] - latrad[np.newaxis, :]
+    cos_lat = np.cos(latrad)
 
     a = (
         np.sin(dlat / 2.0) ** 2
         + (cos_lat[:, np.newaxis] * cos_lat[np.newaxis, :]) * np.sin(dlng / 2.0) ** 2
     )
-    return 2 * RADIUS_MI * np.arcsin(np.sqrt(a))
+    r = radius if radius is not None else _EARTH_RADIUS[units]
+    return 2 * r * np.arcsin(np.sqrt(a))
 
 
 def top(size: int, arr: NDArray) -> NDIndices: