43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
import math
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
_EARTH_RADIUS_KM = 6371.0
|
|
|
|
|
|
def haversine_km(
|
|
lat1: np.ndarray, lon1: np.ndarray, lat2: float, lon2: float
|
|
) -> np.ndarray:
|
|
"""Compute haversine distance in km between arrays (lat1, lon1) and a single point (lat2, lon2)."""
|
|
lat1_rad = np.radians(lat1)
|
|
lon1_rad = np.radians(lon1)
|
|
lat2_rad = np.radians(lat2)
|
|
lon2_rad = np.radians(lon2)
|
|
dlat = lat2_rad - lat1_rad
|
|
dlon = lon2_rad - lon1_rad
|
|
a = (
|
|
np.sin(dlat / 2) ** 2
|
|
+ np.cos(lat1_rad) * np.cos(lat2_rad) * np.sin(dlon / 2) ** 2
|
|
)
|
|
c = 2 * np.arcsin(np.sqrt(a))
|
|
return _EARTH_RADIUS_KM * c
|
|
|
|
|
|
def haversine_km_expr(
|
|
lat_col: str, lon_col: str, dest_lat: float, dest_lon: float
|
|
) -> pl.Expr:
|
|
"""Polars expression computing haversine distance in km to a fixed point."""
|
|
dest_lat_rad = math.radians(dest_lat)
|
|
dest_lon_rad = math.radians(dest_lon)
|
|
|
|
lat_rad = pl.col(lat_col).radians()
|
|
lon_rad = pl.col(lon_col).radians()
|
|
|
|
dlat = pl.lit(dest_lat_rad) - lat_rad
|
|
dlon = pl.lit(dest_lon_rad) - lon_rad
|
|
|
|
a = (dlat / 2).sin() ** 2 + pl.lit(dest_lat_rad).cos() * lat_rad.cos() * (
|
|
dlon / 2
|
|
).sin() ** 2
|
|
return 2 * _EARTH_RADIUS_KM * a.sqrt().arcsin()
|