Add postcode mapping

pipeline/transform/postcode_boundaries.py

@@ -0,0 +1,715 @@
+"""Generate postcode boundary polygons from OA boundaries, INSPIRE parcels, and UPRN data.
+
+Produces per-district GeoJSON files compatible with the Rust server's postcode loader.
+Each postcode gets a polygon (or MultiPolygon) guaranteed to be contained within its
+Output Area(s), with 100% OA coverage and no overlaps between postcodes within an OA.
+
+Algorithm per OA:
+  1. Single-postcode OA → entire OA polygon assigned to that postcode
+  2. Multi-postcode OA:
+     a. Assign INSPIRE parcels to postcodes via UPRN point-in-polygon majority vote
+     b. Union INSPIRE parcels per postcode, clip to OA → "claimed" area
+     c. Distribute remaining (unclaimed) OA area via Voronoi of UPRN points
+     d. Final polygon = claimed + Voronoi share
+
+Memory-efficient design (<12GB total):
+  - INSPIRE polygons stored as raw coordinate bytes in parquet; Shapely objects built
+    lazily per-OA via numpy bbox pre-filter (~100-500 candidates at a time)
+  - UPRNs kept as sorted polars DataFrame with offset dict (Arrow storage, ~1.2GB)
+  - OA processing runs sequentially (no multiprocess INSPIRE duplication)
+
+Output format: {output}/units/{DISTRICT}.geojson with properties.postcodes and
+properties.mapit_code fields matching server-rs/src/data/postcodes.rs expectations.
+"""
+
+import argparse
+import ctypes
+import gc
+import json
+import sqlite3
+import zipfile
+from collections import Counter, defaultdict
+from pathlib import Path
+from xml.etree.ElementTree import iterparse
+
+import numpy as np
+import polars as pl
+from pyproj import Transformer
+from scipy.spatial import Voronoi
+from shapely import STRtree, make_valid, wkb
+from shapely.geometry import MultiPolygon, Polygon
+from shapely.ops import unary_union
+from tqdm import tqdm
+
+# ---------------------------------------------------------------------------
+# GeoPackage helpers
+# ---------------------------------------------------------------------------
+
+
+def _release_memory() -> None:
+    """Force Python + glibc to release freed memory back to the OS."""
+    gc.collect()
+    try:
+        ctypes.CDLL("libc.so.6").malloc_trim(0)
+    except OSError:
+        pass
+
+
+_ENVELOPE_SIZES = {0: 0, 1: 32, 2: 48, 3: 48, 4: 64}
+
+
+def parse_gpkg_geometry(blob: bytes):
+    """Extract a Shapely geometry from a GeoPackage binary blob."""
+    flags = blob[3]
+    envelope_type = (flags >> 1) & 0x07
+    header_size = 8 + _ENVELOPE_SIZES[envelope_type]
+    return wkb.loads(blob[header_size:])
+
+
+def load_oa_boundaries(gpkg_path: Path) -> dict[str, Polygon | MultiPolygon]:
+    """Load OA boundary polygons from a GeoPackage. Geometry is already in BNG."""
+    print("Loading OA boundaries...")
+
+    conn = sqlite3.connect(str(gpkg_path))
+    cur = conn.cursor()
+    cur.execute("SELECT OA21CD, SHAPE FROM OA_2021_EW_BGC_V2")
+
+    oa_geoms: dict[str, Polygon | MultiPolygon] = {}
+    for oa_code, blob in cur:
+        geom = parse_gpkg_geometry(bytes(blob))
+        if geom.geom_type == "MultiPolygon" and len(geom.geoms) == 1:
+            geom = geom.geoms[0]
+        oa_geoms[oa_code] = geom
+
+    conn.close()
+    print(f"  Loaded {len(oa_geoms)} OA boundaries")
+    return oa_geoms
+
+
+# ---------------------------------------------------------------------------
+# UPRN loading (memory-efficient: sorted polars DataFrame + offset dict)
+# ---------------------------------------------------------------------------
+
+
+def load_uprns(uprn_path: Path) -> tuple[pl.DataFrame, dict[str, tuple[int, int]]]:
+    """Load UPRNs as a sorted polars DataFrame with OA offset lookup.
+
+    Returns (df, offsets) where offsets[oa_code] = (start_row, end_row).
+    Peak ~5GB during sort, steady state ~1.5GB (Arrow columnar with compact strings).
+    """
+    import tempfile
+
+    print("Loading UPRN lookup...")
+
+    # Sort via streaming sink to avoid polars doubling memory during in-memory sort
+    with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+        tmp_path = Path(tmp.name)
+    (
+        pl.scan_parquet(uprn_path)
+        .select("GRIDGB1E", "GRIDGB1N", "PCDS", "OA21CD")
+        .filter(~pl.col("OA21CD").str.starts_with("S"))
+        .filter(pl.col("GRIDGB1E").is_not_null() & pl.col("GRIDGB1N").is_not_null())
+        .filter(pl.col("PCDS").is_not_null() & (pl.col("PCDS") != ""))
+        .with_columns(pl.col("PCDS").str.strip_chars())
+        .sort("OA21CD")
+        .sink_parquet(tmp_path)
+    )
+    _release_memory()
+
+    # Read the sorted data — only one copy in memory (~2GB)
+    df = pl.read_parquet(tmp_path)
+    tmp_path.unlink()
+    n = len(df)
+    print(f"  Loaded {n:,} UPRNs (England & Wales)")
+
+    # Compute OA group offsets using polars (avoids 37M Python string creation)
+    boundary_df = (
+        df.lazy()
+        .with_row_index("_i")
+        .filter(pl.col("OA21CD") != pl.col("OA21CD").shift(1))
+        .select("_i", "OA21CD")
+        .collect()
+    )
+    starts_list = boundary_df["_i"].to_list()
+    oa_list = boundary_df["OA21CD"].to_list()
+    del boundary_df
+    offsets: dict[str, tuple[int, int]] = {}
+    for j in range(len(starts_list)):
+        end = starts_list[j + 1] if j + 1 < len(starts_list) else n
+        offsets[oa_list[j]] = (starts_list[j], end)
+    del starts_list, oa_list
+
+    # Drop OA column (no longer needed) to save ~400MB
+    df = df.select("GRIDGB1E", "GRIDGB1N", "PCDS")
+    _release_memory()
+
+    print(f"  Grouped into {len(offsets)} OAs")
+    return df, offsets
+
+
+def get_oa_uprns(
+    df: pl.DataFrame, offsets: dict[str, tuple[int, int]], oa_code: str
+) -> tuple[np.ndarray, list[str]]:
+    """Get UPRN coordinates and postcodes for a single OA.
+
+    Returns (points_nx2, postcodes_list).
+    """
+    s, e = offsets[oa_code]
+    sub = df[s:e]
+    points = np.column_stack(
+        [
+            sub["GRIDGB1E"].to_numpy(),
+            sub["GRIDGB1N"].to_numpy(),
+        ]
+    )
+    postcodes = sub["PCDS"].to_list()
+    return points, postcodes
+
+
+# ---------------------------------------------------------------------------
+# INSPIRE GML parsing and caching
+# ---------------------------------------------------------------------------
+
+_GML_NS = "{http://www.opengis.net/gml/3.2}"
+_LR_NS = "{www.landregistry.gov.uk}"
+
+
+def parse_inspire_zip(zip_path: Path) -> list[np.ndarray]:
+    """Parse a single INSPIRE ZIP → list of Nx2 coordinate arrays (easting, northing)."""
+    results = []
+    with zipfile.ZipFile(zip_path) as zf:
+        gml_names = [n for n in zf.namelist() if n.endswith(".gml")]
+        if not gml_names:
+            return results
+        with zf.open(gml_names[0]) as f:
+            for event, elem in iterparse(f, events=("end",)):
+                if elem.tag != f"{_LR_NS}PREDEFINED":
+                    continue
+                pos_list = elem.find(f".//{_GML_NS}posList")
+                if pos_list is not None and pos_list.text:
+                    vals = pos_list.text.split()
+                    n = len(vals) // 2
+                    if n >= 3:
+                        coords = np.array(vals, dtype=np.float64).reshape(n, 2)
+                        results.append(coords)
+                elem.clear()
+    return results
+
+
+def cache_inspire(inspire_dir: Path, cache_dir: Path) -> None:
+    """Parse all INSPIRE ZIPs and cache as memory-mappable binary files.
+
+    Processes ZIPs sequentially to keep memory under control (~2GB peak).
+    Each ZIP's polygons are streamed to disk immediately after parsing.
+
+    Writes three files:
+      - inspire_bboxes.npy: float64 array (N, 4) of [min_e, min_n, max_e, max_n]
+      - inspire_offsets.npy: int64 array (N, 2) of [byte_offset, n_points]
+      - inspire_coords.bin: flat binary of all float64 coordinate pairs
+    """
+    zip_files = sorted(inspire_dir.glob("*.zip"))
+    print(
+        f"Parsing {len(zip_files)} INSPIRE ZIP files (sequential, streaming to disk)..."
+    )
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # Pre-allocate arrays for bboxes and offsets (grow if needed)
+    capacity = 25_000_000
+    bboxes = np.empty((capacity, 4), dtype=np.float64)
+    offsets = np.empty((capacity, 2), dtype=np.int64)
+    count = 0
+    byte_offset = 0
+
+    coords_path = cache_dir / "inspire_coords.bin"
+    with open(coords_path, "wb") as cf:
+        for zip_path in tqdm(zip_files, desc="INSPIRE ZIPs", unit="file"):
+            for coords in parse_inspire_zip(zip_path):
+                if count >= capacity:
+                    capacity = int(capacity * 1.5)
+                    bboxes.resize((capacity, 4), refcheck=False)
+                    offsets.resize((capacity, 2), refcheck=False)
+                bboxes[count, 0] = coords[:, 0].min()
+                bboxes[count, 1] = coords[:, 1].min()
+                bboxes[count, 2] = coords[:, 0].max()
+                bboxes[count, 3] = coords[:, 1].max()
+                offsets[count, 0] = byte_offset
+                offsets[count, 1] = len(coords)
+                raw = coords.astype(np.float64).tobytes()
+                cf.write(raw)
+                byte_offset += len(raw)
+                count += 1
+
+    # Trim to actual size and save
+    bboxes = bboxes[:count]
+    offsets = offsets[:count]
+    np.save(cache_dir / "inspire_bboxes.npy", bboxes)
+    np.save(cache_dir / "inspire_offsets.npy", offsets)
+    size_mb = byte_offset / (1024 * 1024)
+    print(f"  Cached {count:,} INSPIRE polygons (coords: {size_mb:.0f} MB)")
+
+
+def _inspire_cache_exists(cache_dir: Path) -> bool:
+    return all(
+        (cache_dir / f).exists()
+        for f in ("inspire_bboxes.npy", "inspire_offsets.npy", "inspire_coords.bin")
+    )
+
+
+def load_inspire(
+    cache_dir: Path,
+) -> tuple[np.ndarray, np.ndarray, np.memmap]:
+    """Load INSPIRE cache → (bboxes, offsets, coords_mmap).
+
+    Memory usage: ~1.1GB (bboxes ~777MB + offsets ~290MB, coords memory-mapped).
+    """
+    print(f"Loading INSPIRE cache from {cache_dir}...")
+    bboxes = np.load(cache_dir / "inspire_bboxes.npy")
+    offsets = np.load(cache_dir / "inspire_offsets.npy")
+    coords_mmap = np.memmap(
+        cache_dir / "inspire_coords.bin", dtype=np.float64, mode="r"
+    )
+    print(
+        f"  Loaded {len(bboxes):,} INSPIRE polygon bboxes (~{bboxes.nbytes // (1024 * 1024)} MB)"
+    )
+    print(f"  Offsets: ~{offsets.nbytes // (1024 * 1024)} MB, coords: memory-mapped")
+    return bboxes, offsets, coords_mmap
+
+
+def get_inspire_candidates(
+    oa_bounds: tuple[float, float, float, float],
+    bboxes: np.ndarray,
+    offsets: np.ndarray,
+    coords_mmap: np.memmap,
+) -> list[Polygon]:
+    """Get INSPIRE polygons overlapping an OA via bbox pre-filter.
+
+    Builds Shapely objects only for matches (typically 10-500 per OA).
+    Reads coordinate data on-demand from memory-mapped file.
+    """
+    min_e, min_n, max_e, max_n = oa_bounds
+
+    # Vectorized bbox overlap test
+    mask = (
+        (bboxes[:, 2] >= min_e)
+        & (bboxes[:, 0] <= max_e)
+        & (bboxes[:, 3] >= min_n)
+        & (bboxes[:, 1] <= max_n)
+    )
+    idxs = np.where(mask)[0]
+    if len(idxs) == 0:
+        return []
+
+    # Build Shapely polygons only for candidates (coords from mmap)
+    candidates = []
+    for i in idxs:
+        byte_offset = offsets[i, 0]
+        n_pts = offsets[i, 1]
+        float_offset = byte_offset // 8  # float64 = 8 bytes
+        coords = coords_mmap[float_offset : float_offset + n_pts * 2].reshape(-1, 2)
+        poly = Polygon(coords)
+        if not poly.is_valid:
+            poly = make_valid(poly)
+            if poly.geom_type == "MultiPolygon":
+                poly = max(poly.geoms, key=lambda g: g.area)
+            elif poly.geom_type != "Polygon":
+                continue
+        if not poly.is_empty:
+            candidates.append(poly)
+    return candidates
+
+
+# ---------------------------------------------------------------------------
+# Voronoi computation
+# ---------------------------------------------------------------------------
+
+
+def compute_voronoi_regions(
+    points: np.ndarray, postcodes: list[str], boundary: Polygon | MultiPolygon
+) -> dict[str, Polygon | MultiPolygon]:
+    """Compute Voronoi regions for points, clipped to boundary, grouped by postcode."""
+    if len(points) == 0:
+        return {}
+    if len(points) == 1:
+        return {postcodes[0]: boundary}
+
+    # Deduplicate points per postcode (flats at same coords)
+    seen: dict[tuple[float, float], str] = {}
+    unique_pts = []
+    unique_pcs = []
+    for i in range(len(points)):
+        key = (points[i, 0], points[i, 1])
+        if key not in seen:
+            seen[key] = postcodes[i]
+            unique_pts.append(points[i])
+            unique_pcs.append(postcodes[i])
+
+    if len(unique_pts) == 1:
+        return {unique_pcs[0]: boundary}
+
+    pts = np.array(unique_pts)
+    min_e, min_n = pts.min(axis=0)
+    max_e, max_n = pts.max(axis=0)
+    span = max(max_e - min_e, max_n - min_n, 100)
+
+    dummy = np.array(
+        [
+            [min_e - span * 10, min_n - span * 10],
+            [max_e + span * 10, min_n - span * 10],
+            [min_e - span * 10, max_n + span * 10],
+            [max_e + span * 10, max_n + span * 10],
+        ]
+    )
+    all_points = np.vstack([pts, dummy])
+
+    try:
+        vor = Voronoi(all_points)
+    except Exception:
+        return {unique_pcs[0]: boundary}
+
+    n_real = len(pts)
+    pc_polys: dict[str, list[Polygon]] = defaultdict(list)
+
+    for i in range(n_real):
+        region_idx = vor.point_region[i]
+        region = vor.regions[region_idx]
+        if -1 in region or len(region) < 3:
+            continue
+        vertices = vor.vertices[region]
+        poly = Polygon(vertices)
+        if not poly.is_valid:
+            poly = make_valid(poly)
+        clipped = poly.intersection(boundary)
+        if not clipped.is_empty:
+            pc_polys[unique_pcs[i]].append(clipped)
+
+    result = {}
+    for pc, parts in pc_polys.items():
+        merged = unary_union(parts)
+        if not merged.is_empty:
+            result[pc] = merged
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Per-OA processing
+# ---------------------------------------------------------------------------
+
+
+def process_oa(
+    oa_geom: Polygon | MultiPolygon,
+    points: np.ndarray,
+    postcodes: list[str],
+    inspire_candidates: list[Polygon],
+) -> list[tuple[str, Polygon | MultiPolygon]]:
+    """Process a single OA → list of (postcode, geometry) fragments."""
+    unique_pcs = set(postcodes)
+    if len(unique_pcs) == 1:
+        return [(next(iter(unique_pcs)), oa_geom)]
+
+    # Try INSPIRE-based assignment
+    claimed: dict[str, Polygon | MultiPolygon] = {}
+
+    if inspire_candidates:
+        cand_tree = STRtree(inspire_candidates)
+
+        from shapely import points as shp_points
+
+        uprn_pts = shp_points(points)
+        pt_idx, cand_idx = cand_tree.query(uprn_pts, predicate="intersects")
+
+        # Majority vote per candidate polygon
+        cand_postcodes: dict[int, list[str]] = defaultdict(list)
+        for pi, ci in zip(pt_idx, cand_idx):
+            cand_postcodes[ci].append(postcodes[pi])
+
+        pc_inspire_polys: dict[str, list[Polygon]] = defaultdict(list)
+        for ci, pc_list in cand_postcodes.items():
+            winner = Counter(pc_list).most_common(1)[0][0]
+            pc_inspire_polys[winner].append(inspire_candidates[ci])
+
+        for pc, polys in pc_inspire_polys.items():
+            merged = unary_union(polys)
+            clipped = merged.intersection(oa_geom)
+            if not clipped.is_empty:
+                if not clipped.is_valid:
+                    clipped = make_valid(clipped)
+                claimed[pc] = clipped
+
+    # Compute remaining area
+    if claimed:
+        all_claimed = unary_union(list(claimed.values()))
+        if not all_claimed.is_valid:
+            all_claimed = make_valid(all_claimed)
+        remaining = oa_geom.difference(all_claimed)
+        if not remaining.is_valid:
+            remaining = make_valid(remaining)
+    else:
+        remaining = oa_geom
+
+    # Distribute remaining area via Voronoi
+    if not remaining.is_empty and remaining.area > 0.01:
+        voronoi_result = compute_voronoi_regions(points, postcodes, remaining)
+    else:
+        voronoi_result = {}
+
+    # Combine claimed + voronoi
+    result: dict[str, list] = defaultdict(list)
+    for pc, geom in claimed.items():
+        result[pc].append(geom)
+    for pc, geom in voronoi_result.items():
+        result[pc].append(geom)
+
+    fragments = []
+    for pc, parts in result.items():
+        merged = unary_union(parts)
+        if not merged.is_empty:
+            if not merged.is_valid:
+                merged = make_valid(merged)
+            fragments.append((pc, merged))
+
+    return fragments
+
+
+# ---------------------------------------------------------------------------
+# Output: merge fragments and write GeoJSON
+# ---------------------------------------------------------------------------
+
+_to_wgs84 = None
+
+
+def _get_to_wgs84():
+    global _to_wgs84
+    if _to_wgs84 is None:
+        _to_wgs84 = Transformer.from_crs("EPSG:27700", "EPSG:4326", always_xy=True)
+    return _to_wgs84
+
+
+def to_wgs84_geojson(
+    geom: Polygon | MultiPolygon, tolerance: float = 1.0
+) -> dict | None:
+    """Simplify geometry in BNG, convert to WGS84, return GeoJSON dict."""
+    if geom.is_empty:
+        return None
+
+    simplified = geom.simplify(tolerance, preserve_topology=True)
+    if simplified.is_empty:
+        return None
+
+    transformer = _get_to_wgs84()
+
+    def transform_ring(coords):
+        xs, ys = zip(*coords)
+        lons, lats = transformer.transform(list(xs), list(ys))
+        return list(zip(lons, lats))
+
+    def transform_polygon(poly):
+        exterior = transform_ring(poly.exterior.coords)
+        holes = [transform_ring(h.coords) for h in poly.interiors]
+        return [exterior] + holes
+
+    if simplified.geom_type == "Polygon":
+        return {
+            "type": "Polygon",
+            "coordinates": transform_polygon(simplified),
+        }
+    elif simplified.geom_type == "MultiPolygon":
+        return {
+            "type": "MultiPolygon",
+            "coordinates": [transform_polygon(p) for p in simplified.geoms],
+        }
+    elif simplified.geom_type == "GeometryCollection":
+        polys = [
+            g for g in simplified.geoms if g.geom_type in ("Polygon", "MultiPolygon")
+        ]
+        if not polys:
+            return None
+        return to_wgs84_geojson(unary_union(polys), tolerance=0)
+    return None
+
+
+def merge_fragments(
+    all_fragments: list[tuple[str, Polygon | MultiPolygon]],
+) -> dict[str, Polygon | MultiPolygon]:
+    """Merge cross-OA fragments for postcodes spanning multiple OAs."""
+    by_postcode: dict[str, list] = defaultdict(list)
+    for pc, geom in all_fragments:
+        by_postcode[pc].append(geom)
+
+    merged = {}
+    for pc, parts in by_postcode.items():
+        combined = unary_union(parts)
+        if combined.is_empty:
+            continue
+        if not combined.is_valid:
+            combined = make_valid(combined)
+        # Close tiny gaps between adjacent OA boundary edges (float mismatches)
+        if combined.geom_type == "MultiPolygon":
+            combined = combined.buffer(1.0).buffer(-1.0)
+            if not combined.is_valid:
+                combined = make_valid(combined)
+        # Postcodes are contiguous delivery routes — keep only the largest
+        # polygon; small detached fragments are algorithm artifacts
+        if combined.geom_type == "MultiPolygon":
+            combined = max(combined.geoms, key=lambda g: g.area)
+        merged[pc] = combined
+    return merged
+
+
+def write_district_geojson(
+    postcodes: dict[str, Polygon | MultiPolygon], output_dir: Path
+) -> int:
+    """Group postcodes by district, write GeoJSON files. Returns file count."""
+    units_dir = output_dir / "units"
+    units_dir.mkdir(parents=True, exist_ok=True)
+
+    by_district: dict[str, list[tuple[str, Polygon | MultiPolygon]]] = defaultdict(list)
+    for pc, geom in postcodes.items():
+        parts = pc.split()
+        district = parts[0] if parts else pc[:4]
+        by_district[district].append((pc, geom))
+
+    file_count = 0
+    for district, entries in tqdm(
+        sorted(by_district.items()), desc="Writing GeoJSON", unit="file"
+    ):
+        features = []
+        for pc, geom in sorted(entries, key=lambda x: x[0]):
+            geojson_geom = to_wgs84_geojson(geom)
+            if geojson_geom is None:
+                continue
+            mapit_code = pc.replace(" ", "")
+            features.append(
+                {
+                    "type": "Feature",
+                    "geometry": geojson_geom,
+                    "properties": {
+                        "postcodes": pc,
+                        "mapit_code": mapit_code,
+                    },
+                }
+            )
+
+        if not features:
+            continue
+
+        collection = {"type": "FeatureCollection", "features": features}
+        out_path = units_dir / f"{district}.geojson"
+        with open(out_path, "w") as f:
+            json.dump(collection, f, separators=(",", ":"))
+        file_count += 1
+
+    return file_count
+
+
+# ---------------------------------------------------------------------------
+# Main orchestration
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Generate postcode boundary polygons from OA + INSPIRE + UPRN data"
+    )
+    parser.add_argument("--uprn", type=Path, required=True, help="UPRN lookup parquet")
+    parser.add_argument(
+        "--oa-boundaries", type=Path, required=True, help="OA boundaries GeoPackage"
+    )
+    parser.add_argument(
+        "--inspire", type=Path, required=True, help="INSPIRE ZIP directory"
+    )
+    parser.add_argument("--output", type=Path, required=True, help="Output directory")
+    parser.add_argument(
+        "--limit", type=int, default=0, help="Process only first N OAs (0=all)"
+    )
+    args = parser.parse_args()
+
+    # Phase 1: Load all data
+    print("=" * 60)
+    print("Phase 1: Loading data")
+    print("=" * 60)
+
+    oa_geoms = load_oa_boundaries(args.oa_boundaries)
+    uprn_df, uprn_offsets = load_uprns(args.uprn)
+
+    # Phase 2: Parse/load INSPIRE
+    print()
+    print("=" * 60)
+    print("Phase 2: INSPIRE data")
+    print("=" * 60)
+
+    inspire_cache_dir = args.output / "inspire_cache"
+    if not _inspire_cache_exists(inspire_cache_dir):
+        cache_inspire(args.inspire, inspire_cache_dir)
+    inspire_bboxes, inspire_offsets, inspire_coords = load_inspire(inspire_cache_dir)
+
+    # Phase 3: Process OAs
+    print()
+    print("=" * 60)
+    print("Phase 3: Processing OAs")
+    print("=" * 60)
+
+    # Build work list — precompute which OAs are single vs multi-postcode
+    oa_codes_with_data = sorted(set(oa_geoms.keys()) & set(uprn_offsets.keys()))
+    skipped_no_uprn = len(oa_geoms) - len(oa_codes_with_data)
+    skipped_no_boundary = len(uprn_offsets) - len(oa_codes_with_data)
+
+    if args.limit > 0:
+        oa_codes_with_data = oa_codes_with_data[: args.limit]
+
+    print(f"  OAs with UPRNs + boundaries: {len(oa_codes_with_data)}")
+    print(f"  Skipped (no UPRNs): {skipped_no_uprn}")
+    print(f"  Skipped (no boundary): {skipped_no_boundary}")
+
+    all_fragments: list[tuple[str, Polygon | MultiPolygon]] = []
+    single_count = 0
+    multi_count = 0
+
+    for oa_code in tqdm(
+        oa_codes_with_data,
+        desc="Processing OAs",
+        unit="OA",
+        smoothing=0.01,
+        miniters=100,
+    ):
+        oa_geom = oa_geoms[oa_code]
+        points, postcodes = get_oa_uprns(uprn_df, uprn_offsets, oa_code)
+
+        if len(set(postcodes)) == 1:
+            # Fast path: entire OA = one postcode
+            all_fragments.append((postcodes[0], oa_geom))
+            single_count += 1
+            continue
+
+        # Get INSPIRE candidates via bbox pre-filter
+        candidates = get_inspire_candidates(
+            oa_geom.bounds, inspire_bboxes, inspire_offsets, inspire_coords
+        )
+
+        fragments = process_oa(oa_geom, points, postcodes, candidates)
+        all_fragments.extend(fragments)
+        multi_count += 1
+
+    print(f"\n  Single-postcode OAs (fast path): {single_count}")
+    print(f"  Multi-postcode OAs (INSPIRE+Voronoi): {multi_count}")
+    print(f"  Total fragments: {len(all_fragments)}")
+
+    # Free data no longer needed
+    del oa_geoms, uprn_df, uprn_offsets
+    del inspire_bboxes, inspire_offsets, inspire_coords
+
+    # Phase 4: Merge and write
+    print()
+    print("=" * 60)
+    print("Phase 4: Merging fragments and writing GeoJSON")
+    print("=" * 60)
+
+    merged = merge_fragments(all_fragments)
+    print(f"  Merged into {len(merged)} unique postcodes")
+
+    file_count = write_district_geojson(merged, args.output)
+    print(f"\n  Wrote {file_count} district GeoJSON files to {args.output / 'units'}")
+    print("Done!")
+
+
+if __name__ == "__main__":
+    main()