Format python

This commit is contained in:
Andras Schmelczer 2026-01-31 13:07:09 +00:00
parent 85f5770e09
commit 4c258018c3
17 changed files with 348 additions and 248 deletions

View file

@ -73,29 +73,39 @@ class TestHaversineKmExpr:
def test_same_point(self):
"""Distance from a point to itself should be zero."""
df = pl.DataFrame({"lat": [51.5074], "lon": [-0.1278]})
result = df.select(haversine_km_expr("lat", "lon", 51.5074, -0.1278).alias("dist"))
result = df.select(
haversine_km_expr("lat", "lon", 51.5074, -0.1278).alias("dist")
)
assert result["dist"][0] == pytest.approx(0.0, abs=1e-10)
def test_known_distance_london_to_paris(self):
"""Test distance from London to Paris (~344 km)."""
df = pl.DataFrame({"lat": [51.5074], "lon": [-0.1278]})
result = df.select(haversine_km_expr("lat", "lon", 48.8566, 2.3522).alias("dist"))
result = df.select(
haversine_km_expr("lat", "lon", 48.8566, 2.3522).alias("dist")
)
assert result["dist"][0] == pytest.approx(344, rel=0.01)
def test_known_distance_new_york_to_london(self):
"""Test distance from New York to London (~5570 km)."""
df = pl.DataFrame({"lat": [40.7128], "lon": [-74.0060]})
result = df.select(haversine_km_expr("lat", "lon", 51.5074, -0.1278).alias("dist"))
result = df.select(
haversine_km_expr("lat", "lon", 51.5074, -0.1278).alias("dist")
)
assert result["dist"][0] == pytest.approx(5570, rel=0.01)
def test_multiple_points(self):
"""Test calculating distances from multiple points to a single destination."""
df = pl.DataFrame({
"lat": [51.5074, 48.8566, 40.7128], # London, Paris, NYC
"lon": [-0.1278, 2.3522, -74.0060],
})
df = pl.DataFrame(
{
"lat": [51.5074, 48.8566, 40.7128], # London, Paris, NYC
"lon": [-0.1278, 2.3522, -74.0060],
}
)
# Distance to Edinburgh
result = df.select(haversine_km_expr("lat", "lon", 55.9533, -3.1883).alias("dist"))
result = df.select(
haversine_km_expr("lat", "lon", 55.9533, -3.1883).alias("dist")
)
dists = result["dist"].to_numpy()
# All distances should be positive
@ -128,7 +138,9 @@ class TestHaversineConsistency:
# Polars version
df = pl.DataFrame({"lat": lats, "lon": lons})
polars_result = df.select(haversine_km_expr("lat", "lon", dest_lat, dest_lon).alias("dist"))
polars_result = df.select(
haversine_km_expr("lat", "lon", dest_lat, dest_lon).alias("dist")
)
polars_dists = polars_result["dist"].to_numpy()
# Should be identical (or at least very close due to floating point)