162 lines
5 KiB
Python
162 lines
5 KiB
Python
import numpy as np
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from pipeline.utils.poi_counts import count_pois_per_postcode, min_distance_per_postcode
|
|
|
|
POI_GROUPS = {
|
|
"restaurants": ["Restaurant", "Fast Food"],
|
|
"groceries": ["Supermarket"],
|
|
"parks": ["Park"],
|
|
"train_tube": ["Rail station", "Tube station"],
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def pois():
|
|
"""POIs clustered around two locations: central London and 10km away."""
|
|
return pl.DataFrame(
|
|
{
|
|
"lat": [51.5074, 51.5075, 51.5080, 51.5076, 51.5073, 51.60],
|
|
"lng": [-0.1278, -0.1280, -0.1275, -0.1279, -0.1277, -0.20],
|
|
"category": [
|
|
"Restaurant",
|
|
"Fast Food",
|
|
"Supermarket",
|
|
"Park",
|
|
"Rail station",
|
|
"Restaurant", # too far from any property
|
|
],
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def postcodes():
|
|
"""Two postcodes: one near central London, one far away."""
|
|
return pl.DataFrame(
|
|
{
|
|
"postcode": ["EC1A 1BB", "ZZ99 9ZZ"],
|
|
"lat": [51.5074, 55.0],
|
|
"lon": [-0.1278, -3.0],
|
|
}
|
|
)
|
|
|
|
|
|
def test_counts_pois_within_radius(postcodes, pois):
|
|
result = count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=2.0)
|
|
|
|
expected_cols = {f"{g}_2km" for g in POI_GROUPS}
|
|
assert expected_cols.issubset(set(result.columns))
|
|
|
|
# Result must be aligned to postcodes (2 rows)
|
|
assert len(result) == 2
|
|
|
|
ec1a = result.filter(pl.col("postcode") == "EC1A 1BB")
|
|
assert ec1a["restaurants_2km"][0] == 2 # Restaurant + Fast Food
|
|
assert ec1a["groceries_2km"][0] == 1 # Supermarket
|
|
assert ec1a["parks_2km"][0] == 1 # Park
|
|
assert ec1a["train_tube_2km"][0] == 1 # Rail station
|
|
|
|
# Far-away postcode should have zero counts
|
|
zz99 = result.filter(pl.col("postcode") == "ZZ99 9ZZ")
|
|
for group in POI_GROUPS:
|
|
assert zz99[f"{group}_2km"][0] == 0
|
|
|
|
|
|
def test_no_pois_returns_zeros(postcodes):
|
|
empty_pois = pl.DataFrame(
|
|
{
|
|
"lat": pl.Series([], dtype=pl.Float64),
|
|
"lng": pl.Series([], dtype=pl.Float64),
|
|
"category": pl.Series([], dtype=pl.String),
|
|
}
|
|
)
|
|
result = count_pois_per_postcode(
|
|
postcodes, empty_pois, groups=POI_GROUPS, radius_km=2.0
|
|
)
|
|
|
|
for group in POI_GROUPS:
|
|
col = f"{group}_2km"
|
|
assert col in result.columns
|
|
assert result[col].to_list() == [0, 0]
|
|
|
|
|
|
def test_custom_radius(pois):
|
|
"""A tiny radius should exclude POIs that are even slightly away."""
|
|
postcodes = pl.DataFrame(
|
|
{
|
|
"postcode": ["EC1A 1BB"],
|
|
"lat": [51.5074],
|
|
"lon": [-0.1278],
|
|
}
|
|
)
|
|
|
|
# 0.01 km = 10m — only the POI at the exact same location should match
|
|
result = count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=0.01)
|
|
# The Restaurant at (51.5074, -0.1278) is at distance 0
|
|
assert result["restaurants_0km"][0] >= 1
|
|
# POIs >100m away should not be counted
|
|
total = sum(result[f"{g}_0km"][0] for g in POI_GROUPS)
|
|
assert total <= 2 # at most the co-located POIs
|
|
|
|
|
|
def test_counts_pois_across_multiple_grid_cells_within_5km():
|
|
"""A POI around 4.8km away must not be dropped by grid candidate lookup."""
|
|
postcodes = pl.DataFrame(
|
|
{
|
|
"postcode": ["GRID 5KM"],
|
|
"lat": [51.5],
|
|
"lon": [0.049],
|
|
}
|
|
)
|
|
pois = pl.DataFrame(
|
|
{
|
|
"lat": [51.5, 51.5],
|
|
"lng": [0.1183, 0.1240],
|
|
"category": ["Park", "Park"],
|
|
}
|
|
)
|
|
|
|
result = count_pois_per_postcode(
|
|
postcodes,
|
|
pois,
|
|
groups={"parks": ["Park"]},
|
|
radius_km=5.0,
|
|
)
|
|
|
|
assert result["parks_5km"][0] == 1
|
|
|
|
|
|
def test_min_distance_finds_nearest(postcodes, pois):
|
|
"""min_distance_per_postcode returns distance to closest POI per group."""
|
|
result = min_distance_per_postcode(postcodes, pois, groups=POI_GROUPS)
|
|
|
|
assert len(result) == 2
|
|
|
|
ec1a = result.filter(pl.col("postcode") == "EC1A 1BB")
|
|
# Rail station is at (51.5073, -0.1277), postcode at (51.5074, -0.1278) — very close
|
|
assert ec1a["train_tube_nearest_km"][0] < 0.05 # within 50m
|
|
# Restaurant is co-located — distance ~0
|
|
assert ec1a["restaurants_nearest_km"][0] < 0.01
|
|
|
|
# Far-away postcode should still get the global nearest distance.
|
|
zz99 = result.filter(pl.col("postcode") == "ZZ99 9ZZ")
|
|
assert zz99["train_tube_nearest_km"][0] > 300
|
|
|
|
|
|
def test_min_distance_no_pois_returns_nan(postcodes):
|
|
"""With no POIs, all distances should be NaN."""
|
|
empty_pois = pl.DataFrame(
|
|
{
|
|
"lat": pl.Series([], dtype=pl.Float64),
|
|
"lng": pl.Series([], dtype=pl.Float64),
|
|
"category": pl.Series([], dtype=pl.String),
|
|
}
|
|
)
|
|
result = min_distance_per_postcode(
|
|
postcodes, empty_pois, groups={"train_tube": ["Rail station"]}
|
|
)
|
|
|
|
assert "train_tube_nearest_km" in result.columns
|
|
assert all(np.isnan(v) for v in result["train_tube_nearest_km"].to_list())
|