Format python
This commit is contained in:
parent
85f5770e09
commit
4c258018c3
17 changed files with 348 additions and 248 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue