Move dict

This commit is contained in:
Andras Schmelczer 2026-02-01 08:48:07 +00:00
parent a3c1b6090e
commit ac45af8514
3 changed files with 19 additions and 16 deletions

View file

@ -8,6 +8,16 @@ import polars as pl
from pipeline.utils.poi_counts import _count_pois_per_postcode
# POI category groups for proximity counting
POI_GROUPS = {
"restaurants": ["Restaurant", "Fast Food"],
"groceries": ["Greengrocer", "Grocery Shop", "Supermarket", "Convenience Store"],
"parks": ["Park", "Garden", "Nature Reserve"],
"public_transport": ["Metro or Tram stop", "Rail station", "Bus stop", "Bus station"], # comes from naptan.py
}
def main():
parser = argparse.ArgumentParser(
description="Count POIs within radius per postcode"
@ -31,7 +41,7 @@ def main():
pois = pl.read_parquet(args.pois)
result = _count_pois_per_postcode(postcodes, pois, radius_km=2)
result = _count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=2)
result.write_parquet(args.output)
size_mb = args.output.stat().st_size / (1024 * 1024)