Move dict

pipeline/utils/__init__.py

@@ -1,7 +1,7 @@
 from .download import download, extract_zip
 from .fuzzy_join import fuzzy_join_on_postcode
 from .haversine import haversine_km, haversine_km_expr
-from .poi_counts import POI_GROUPS, count_pois_within_radius
+from .poi_counts import count_pois_within_radius
 
 __all__ = [
     "download",
@@ -9,6 +9,5 @@ __all__ = [
     "fuzzy_join_on_postcode",
     "haversine_km",
     "haversine_km_expr",
-    "POI_GROUPS",
     "count_pois_within_radius",
 ]

pipeline/utils/poi_counts.py

@@ -1,6 +1,5 @@
 """Count POIs within a radius of properties, optimized via postcode deduplication."""
 
-import os
 import tempfile
 
 import numpy as np
@@ -8,17 +7,12 @@ import polars as pl
 
 from .haversine import haversine_km
 
-# 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": ["Station", "Stop", "Bus Station"],
-}
-
 
 def _count_pois_per_postcode(
-    postcodes_df: pl.DataFrame, pois: pl.DataFrame, radius_km: float = 2.0
+    postcodes_df: pl.DataFrame,
+    pois: pl.DataFrame,
+    groups: dict[str, list[str]],
+    radius_km: float = 2.0,
 ) -> pl.DataFrame:
     """
     For each unique postcode, count POIs within radius_km by category group.
@@ -59,7 +53,7 @@ def _count_pois_per_postcode(
 
     # Pre-compute category masks
     category_masks = {}
-    for group, categories in POI_GROUPS.items():
+    for group, categories in groups.items():
         mask = np.isin(poi_cats, categories)
         category_masks[group] = mask
         print(f"  {group}: {mask.sum():,} POIs")
@@ -71,7 +65,7 @@ def _count_pois_per_postcode(
 
     # Initialize result arrays
     result_counts = {
-        group: np.zeros(n_postcodes, dtype=np.int32) for group in POI_GROUPS
+        group: np.zeros(n_postcodes, dtype=np.int32) for group in groups
     }
 
     # Process in batches with progress
@@ -128,7 +122,7 @@ def _count_pois_per_postcode(
 
     # Build result dataframe
     result_data = {"postcode": pc_codes}
-    for group in POI_GROUPS:
+    for group in groups:
         result_data[f"{group}_{int(radius_km)}km"] = result_counts[group]
 
     result = pl.DataFrame(result_data)