99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
from pathlib import Path
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from pipeline.transform.tree_density import (
|
|
STREET_TREE_COVERAGE_COL,
|
|
STREET_TREE_DENSITY_COL,
|
|
_coverage_percentile_expr,
|
|
_metric_columns,
|
|
_postcode_density_percentile_col,
|
|
_with_postcode_density_percentiles,
|
|
_write_street_rollups,
|
|
)
|
|
|
|
|
|
def test_coverage_percentile_expr_ranks_higher_coverage_higher() -> None:
|
|
df = pl.DataFrame({"coverage": [0.0, 5.0, 10.0, None]})
|
|
|
|
result = df.lazy().with_columns(
|
|
_coverage_percentile_expr("coverage", "percentile")
|
|
).collect()
|
|
|
|
assert result["percentile"].to_list() == [0.0, 50.0, 100.0, None]
|
|
|
|
|
|
def test_coverage_percentile_expr_uses_exact_scale_endpoints() -> None:
|
|
df = pl.DataFrame({"coverage": [0.0, 0.0, 5.0, 10.0, 10.0]})
|
|
|
|
result = df.lazy().with_columns(
|
|
_coverage_percentile_expr("coverage", "percentile")
|
|
).collect()
|
|
|
|
assert result["percentile"].to_list() == [0.0, 0.0, 50.0, 100.0, 100.0]
|
|
|
|
|
|
def test_street_rollup_percentiles_are_ranked_over_raw_street_coverage(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
radius_m = 50
|
|
density_col, area_col, count_col, height_col = _metric_columns(radius_m)
|
|
percentile_col = _postcode_density_percentile_col(radius_m)
|
|
|
|
postcode_metrics = _with_postcode_density_percentiles(
|
|
pl.DataFrame(
|
|
{
|
|
"postcode": ["AA1 1AA", "AA1 1AB", "AA1 1AC"],
|
|
density_col: [10.0, 30.0, 50.0],
|
|
area_col: [100.0, 300.0, 500.0],
|
|
count_col: [1, 3, 5],
|
|
height_col: [4.0, 6.0, 8.0],
|
|
}
|
|
),
|
|
radius_m,
|
|
)
|
|
|
|
price_paid = pl.DataFrame(
|
|
{
|
|
"postcode": ["AA1 1AA", "AA1 1AA", "AA1 1AB", "AA1 1AC"],
|
|
"paon": ["1", "2", "3", "4"],
|
|
"saon": ["", "", "", ""],
|
|
"street": ["Oak Road", "Oak Road", "Oak Road", "Elm Street"],
|
|
"locality": ["", "", "", ""],
|
|
"town_city": ["Test Town", "Test Town", "Test Town", "Test Town"],
|
|
"district": ["Test District"] * 4,
|
|
"county": ["Test County"] * 4,
|
|
"date_of_transfer": [
|
|
"2024-01-01",
|
|
"2024-01-02",
|
|
"2024-01-03",
|
|
"2024-01-04",
|
|
],
|
|
}
|
|
)
|
|
price_paid_path = tmp_path / "price-paid.parquet"
|
|
output_streets = tmp_path / "streets.parquet"
|
|
output_addresses = tmp_path / "addresses.parquet"
|
|
price_paid.write_parquet(price_paid_path)
|
|
|
|
_write_street_rollups(
|
|
postcode_metrics=postcode_metrics,
|
|
price_paid_path=price_paid_path,
|
|
output_streets=output_streets,
|
|
output_addresses=output_addresses,
|
|
radius_m=radius_m,
|
|
)
|
|
|
|
streets = pl.read_parquet(output_streets).sort("street")
|
|
addresses = pl.read_parquet(output_addresses)
|
|
|
|
assert streets["street"].to_list() == ["Elm Street", "Oak Road"]
|
|
assert streets[STREET_TREE_COVERAGE_COL].to_list() == pytest.approx([50.0, 16.7])
|
|
assert streets.select("street", STREET_TREE_DENSITY_COL).rows() == [
|
|
("Elm Street", 100.0),
|
|
("Oak Road", 0.0),
|
|
]
|
|
assert percentile_col in addresses.columns
|
|
assert STREET_TREE_COVERAGE_COL in addresses.columns
|
|
assert STREET_TREE_DENSITY_COL in addresses.columns
|