81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
import polars as pl
|
|
|
|
from pipeline.transform.area_crime_averages import (
|
|
NATIONAL_AREA,
|
|
SCOPE_NATIONAL,
|
|
SCOPE_OUTCODE,
|
|
SCOPE_SECTOR,
|
|
compute_area_crime_averages,
|
|
)
|
|
|
|
_BURGLARY = "Burglary (/yr, 7y)"
|
|
_ROBBERY = "Robbery (/yr, 7y)"
|
|
|
|
|
|
def _postcodes() -> pl.LazyFrame:
|
|
return pl.LazyFrame(
|
|
{
|
|
"Postcode": ["E14 2DG", "E14 2AB", "E14 9XY", "M1 1AE", "E14 2ZZ"],
|
|
# E14 9XY has no usable crime data; E14 2AB lacks robbery; E14 2ZZ has
|
|
# crime but (below) no properties, so it must not weight any average.
|
|
_BURGLARY: [10.0, 20.0, None, 5.0, 100.0],
|
|
_ROBBERY: [2.0, None, None, 1.0, 50.0],
|
|
# An unrelated column proves only the crime columns are averaged.
|
|
"Median age": [40.0, 41.0, 42.0, 30.0, 99.0],
|
|
}
|
|
)
|
|
|
|
|
|
def _properties() -> pl.LazyFrame:
|
|
# Property rows per postcode become the weights (3 / 1 / 2 / 4). E14 2ZZ has
|
|
# none, so it is excluded entirely.
|
|
postcodes = ["E14 2DG"] * 3 + ["E14 2AB"] + ["E14 9XY"] * 2 + ["M1 1AE"] * 4
|
|
return pl.LazyFrame({"Postcode": postcodes})
|
|
|
|
|
|
def _row(df: pl.DataFrame, scope: str, area: str) -> dict:
|
|
matched = df.filter((pl.col("scope") == scope) & (pl.col("area") == area))
|
|
assert matched.height == 1, f"expected one {scope} row for {area!r}"
|
|
return matched.to_dicts()[0]
|
|
|
|
|
|
def test_property_weighted_means_skip_nulls():
|
|
result = compute_area_crime_averages(_postcodes(), _properties())
|
|
|
|
national = _row(result, SCOPE_NATIONAL, NATIONAL_AREA)
|
|
# Burglary: (10*3 + 20*1 + 5*4) / (3+1+4) = 70/8; E14 9XY null dilutes nothing.
|
|
assert national[_BURGLARY] == 8.75
|
|
# Robbery: (2*3 + 1*4) / (3+4) = 10/7; both null postcodes are excluded from
|
|
# the numerator AND the denominator.
|
|
assert national[_ROBBERY] == pl.Series([10.0 / 7.0]).cast(pl.Float32).item()
|
|
|
|
outcode = _row(result, SCOPE_OUTCODE, "E14")
|
|
assert outcode[_BURGLARY] == 12.5 # (10*3 + 20*1) / 4
|
|
assert outcode[_ROBBERY] == 2.0 # only E14 2DG has robbery (2 * 3 / 3)
|
|
|
|
|
|
def test_sector_aggregation_and_all_null_rows_dropped():
|
|
result = compute_area_crime_averages(_postcodes(), _properties())
|
|
|
|
sector = _row(result, SCOPE_SECTOR, "E14 2")
|
|
assert sector[_BURGLARY] == 12.5
|
|
assert sector[_ROBBERY] == 2.0
|
|
|
|
# E14 9XY has properties but no crime data at all, so its sector "E14 9" is
|
|
# all-null and must be dropped rather than reported as a known area.
|
|
assert result.filter(pl.col("area") == "E14 9").height == 0
|
|
|
|
|
|
def test_postcodes_without_properties_are_excluded():
|
|
result = compute_area_crime_averages(_postcodes(), _properties())
|
|
|
|
# E14 2ZZ carries crime values but no properties; including it would pull the
|
|
# E14 outcode burglary mean toward its 100.0. It must contribute nothing.
|
|
outcode = _row(result, SCOPE_OUTCODE, "E14")
|
|
assert outcode[_BURGLARY] == 12.5
|
|
|
|
|
|
def test_only_crime_columns_are_emitted():
|
|
result = compute_area_crime_averages(_postcodes(), _properties())
|
|
assert set(result.columns) == {"scope", "area", _BURGLARY, _ROBBERY}
|
|
assert result.schema[_BURGLARY] == pl.Float32
|