Extarct utils

pipeline/utils/haversine.py

@@ -0,0 +1,36 @@
+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()