perfect-postcode/pipeline/utils/poi_counts.py

"""Count POIs within a radius of properties, optimized via postcode deduplication."""

import tempfile

import numpy as np
import polars as pl

from .haversine import haversine_km


def _count_pois_per_postcode(
    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.
    Uses spatial grid with vectorized distance calculations.
    """
    print(f"Counting POIs within {radius_km}km per postcode...")

    n_postcodes = len(postcodes_df)
    n_pois = len(pois)
    print(f"  {n_postcodes:,} postcodes, {n_pois:,} POIs")

    # Build spatial grid for POIs (0.05 degree cells ~5.5km)
    grid_size = 0.05
    print("  Building POI spatial grid...")

    # Convert to numpy arrays
    poi_lats = pois["lat"].to_numpy()
    poi_lngs = pois["lng"].to_numpy()
    poi_cats = pois["category"].to_numpy()

    # Compute grid coordinates for all POIs
    poi_grid_lats = np.floor(poi_lats / grid_size).astype(np.int32)
    poi_grid_lngs = np.floor(poi_lngs / grid_size).astype(np.int32)

    # Build grid cell lookup using numpy indexing
    poi_grid = {}
    for i in range(n_pois):
        key = (poi_grid_lats[i], poi_grid_lngs[i])
        if key not in poi_grid:
            poi_grid[key] = []
        poi_grid[key].append(i)

    # Convert grid values to numpy arrays for faster indexing
    for key in poi_grid:
        poi_grid[key] = np.array(poi_grid[key], dtype=np.int32)

    print(f"  POI grid has {len(poi_grid):,} occupied cells")

    # Pre-compute category masks
    category_masks = {}
    for group, categories in groups.items():
        mask = np.isin(poi_cats, categories)
        category_masks[group] = mask
        print(f"  {group}: {mask.sum():,} POIs")

    # Extract postcode coordinates as numpy arrays
    pc_lats = postcodes_df["lat"].to_numpy()
    pc_lons = postcodes_df["lon"].to_numpy()
    pc_codes = postcodes_df["postcode"].to_list()

    # Initialize result arrays
    result_counts = {
        group: np.zeros(n_postcodes, dtype=np.int32) for group in groups
    }

    # Process in batches with progress
    batch_size = 50000
    n_batches = (n_postcodes + batch_size - 1) // batch_size

    print(f"  Processing {n_postcodes:,} postcodes in {n_batches} batches...")

    for batch_idx in range(n_batches):
        start_idx = batch_idx * batch_size
        end_idx = min(start_idx + batch_size, n_postcodes)

        if batch_idx % 5 == 0:
            print(
                f"  Batch {batch_idx + 1}/{n_batches}: postcodes {start_idx:,} - {end_idx:,}"
            )

        # Process batch
        for i in range(start_idx, end_idx):
            pc_lat = pc_lats[i]
            pc_lon = pc_lons[i]

            # Find grid cells to check (3x3 grid)
            grid_lat = int(np.floor(pc_lat / grid_size))
            grid_lng = int(np.floor(pc_lon / grid_size))

            # Collect nearby POI indices
            nearby_indices = []
            for dlat in [-1, 0, 1]:
                for dlng in [-1, 0, 1]:
                    cell_key = (grid_lat + dlat, grid_lng + dlng)
                    if cell_key in poi_grid:
                        nearby_indices.append(poi_grid[cell_key])

            if not nearby_indices:
                continue

            # Concatenate all nearby POI indices
            nearby = np.concatenate(nearby_indices)

            # Vectorized distance calculation for all nearby POIs
            distances = haversine_km(poi_lats[nearby], poi_lngs[nearby], pc_lat, pc_lon)

            # Filter by radius
            within_mask = distances <= radius_km
            within_indices = nearby[within_mask]

            if len(within_indices) == 0:
                continue

            # Count by category group using pre-computed masks
            for group, cat_mask in category_masks.items():
                result_counts[group][i] = cat_mask[within_indices].sum()

    # Build result dataframe
    result_data = {"postcode": pc_codes}
    for group in groups:
        result_data[f"{group}_{int(radius_km)}km"] = result_counts[group]

    result = pl.DataFrame(result_data)
    print("  Completed POI counting")
    return result


def count_pois_within_radius(
    properties: pl.DataFrame, pois: pl.DataFrame, radius_km: float = 2.0
) -> dict[str, pl.Series]:
    """
    Count POIs within radius for properties, optimized by deduplicating postcodes.

    Returns dict of {column_name: count_series} aligned to properties dataframe.
    """
    # Get unique postcodes with coordinates
    print("Deduplicating postcodes...")
    unique_postcodes = properties.select(["postcode", "lat", "lon"]).unique(
        subset=["postcode"]
    )

    print(
        f"  {len(properties):,} properties → {len(unique_postcodes):,} unique postcodes"
    )

    # Count POIs per postcode
    postcode_counts = _count_pois_per_postcode(unique_postcodes, pois, radius_km)

    print("  Writing postcode counts to temp file...")
    with tempfile.NamedTemporaryFile(suffix=".parquet") as tmp:
        tmp_path = tmp.name
        postcode_counts.write_parquet(tmp_path)

        # Join using lazy evaluation
        print("  Joining counts back to properties (lazy)...")
        count_cols = [f"{group}_{int(radius_km)}km" for group in POI_GROUPS]

        # Convert properties to lazy frame, join, then collect
        result_lazy = (
            properties.lazy()
            .select("postcode")
            .join(pl.scan_parquet(tmp_path), on="postcode", how="left")
            .select(count_cols)
            .fill_null(0)
        )

        result_df = result_lazy.collect(engine="streaming")

        return {col: result_df[col] for col in count_cols}