pbf_to_tile_cache.py

#!/usr/bin/env python3
"""Convert a Geofabrik OSM PBF extract into the per-tile Overpass-JSON
cache format expected by build_standard_track_context_v1.py.

For each tile (0.25 degree by default), emit a JSON file at:
  <out_cache>/tiles/<TILE_ID>.json

where TILE_ID matches the convention used by the build script:
    f"{tile_lat:+08.3f}_{tile_lon:+09.3f}"

The JSON structure mirrors the Overpass API output that build_*.context.py's
_parse_ways() consumes — list of {"type": "way", "id": ..., "nodes": [...],
"tags": {"natural":"coastline"} | {"natural":"water"} |
{"man_made":"pier"|"breakwater"|"groyne"|"quay"}, "geometry": [{"lat":..,
"lon":..}, ...]} entries inside payload["elements"].

This script extracts only the way geometry types that the build script
queries via Overpass — keeping the cache size small and the parse fast.

Usage
-----
python pbf_to_tile_cache.py \
    --pbf <path.osm.pbf> \
    --out-cache <env/osm_cache or alt path> \
    --bbox south,west,north,east \
    [--tiles-list <file with one TILE_ID per line>] \
    [--tile-deg 0.25]

Either --bbox (full coverage) or --tiles-list (selective) MUST be given.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Iterable

import osmium


def tile_id_for_point(lat: float, lon: float, tile_deg: float) -> str:
    tile_lat = math.floor(lat / tile_deg) * tile_deg
    tile_lon = math.floor(lon / tile_deg) * tile_deg
    return f"{tile_lat:+08.3f}_{tile_lon:+09.3f}"


# Tags accepted, matching _make_overpass_query() in build_standard_track_context_v1.py
def is_target_way(tags: dict) -> bool:
    if not tags:
        return False
    nat = tags.get("natural", "")
    mm  = tags.get("man_made", "")
    if nat == "coastline":
        return True
    if nat == "water":
        return True
    if mm in ("pier", "breakwater", "groyne", "quay"):
        return True
    return False


class WayCollector(osmium.SimpleHandler):
    """Collect target ways with geometry into per-tile buckets.

    osmium gives us node coordinates by attaching a location cache before
    parsing. The SimpleHandler base class supports `apply_file(filename,
    locations=True)` which fills the cache automatically.
    """

    def __init__(self, tile_deg: float, accept_tiles: set | None = None,
                 bbox: tuple | None = None):
        super().__init__()
        self.tile_deg = tile_deg
        self.accept_tiles = accept_tiles
        self.bbox = bbox  # (south, west, north, east) or None
        self.tile_payload: dict[str, list[dict]] = defaultdict(list)
        self.way_count = 0
        self.kept_count = 0
        self.t0 = time.time()

    def way(self, w):
        self.way_count += 1
        if self.way_count % 100000 == 0:
            sys.stderr.write(
                f"[pbf] ways scanned={self.way_count:,}  kept={self.kept_count:,}  "
                f"elapsed={time.time()-self.t0:.0f}s\n")
        tags = {tag.k: tag.v for tag in w.tags}
        if not is_target_way(tags):
            return
        try:
            coords = []
            for n in w.nodes:
                # osmium WayNode -> location
                loc = n.location
                if loc.valid():
                    coords.append((loc.lat, loc.lon))
        except osmium.InvalidLocationError:
            return  # unresolved node coords; skip
        if len(coords) < 2:
            return
        # bbox filter
        if self.bbox is not None:
            s, ww, n, ee = self.bbox
            lats = [c[0] for c in coords]; lons = [c[1] for c in coords]
            if max(lats) < s or min(lats) > n or max(lons) < ww or min(lons) > ee:
                return
        # Group by tile: assign way to every tile that any node falls in
        node_tiles = set()
        for lat, lon in coords:
            node_tiles.add(tile_id_for_point(lat, lon, self.tile_deg))
        if self.accept_tiles is not None:
            node_tiles &= self.accept_tiles
            if not node_tiles:
                return
        elem = {
            "type": "way",
            "id":   w.id,
            "tags": tags,
            "geometry": [{"lat": float(lat), "lon": float(lon)} for lat, lon in coords],
        }
        for tid in node_tiles:
            self.tile_payload[tid].append(elem)
        self.kept_count += 1


def write_tile_jsons(payload_map: dict[str, list[dict]], out_root: Path) -> int:
    tiles_dir = out_root / "tiles"
    tiles_dir.mkdir(parents=True, exist_ok=True)
    n = 0
    for tid, elements in payload_map.items():
        path = tiles_dir / f"{tid}.json"
        payload = {
            "version": 0.6,
            "generator": "pbf_to_tile_cache.py",
            "osm3s": {
                "copyright":
                    "The data included in this document is from www.openstreetmap.org. "
                    "Available under the Open Database License (ODbL).",
            },
            "elements": elements,
        }
        path.write_text(json.dumps(payload))
        n += 1
    return n


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--pbf",     type=Path, required=True)
    p.add_argument("--out-cache", type=Path, required=True,
                   help="Destination root; tiles/ subdir will be created.")
    p.add_argument("--bbox",    type=str, default=None,
                   help="south,west,north,east — restrict to this area.")
    p.add_argument("--tiles-list", type=Path, default=None,
                   help="One TILE_ID per line; only emit these tiles. "
                        "If absent emit all tiles intersected by ways within --bbox.")
    p.add_argument("--tile-deg", type=float, default=0.25)
    args = p.parse_args()

    accept_tiles = None
    if args.tiles_list:
        accept_tiles = {ln.strip() for ln in args.tiles_list.read_text().splitlines() if ln.strip()}
        print(f"[pbf] accept_tiles loaded: {len(accept_tiles):,}", flush=True)
    bbox = None
    if args.bbox:
        s, ww, n, ee = [float(x) for x in args.bbox.split(",")]
        bbox = (s, ww, n, ee)
        print(f"[pbf] bbox=({s},{ww},{n},{ee})", flush=True)

    coll = WayCollector(args.tile_deg, accept_tiles=accept_tiles, bbox=bbox)
    print(f"[pbf] parsing {args.pbf}  size={args.pbf.stat().st_size/(1<<20):.0f} MiB", flush=True)
    # locations=True attaches the node location cache (in-RAM dense index by id)
    # idx="sparse_mem_array" is more compact for partial PBFs; use sparse_mmap_array
    # for big planet files.
    coll.apply_file(str(args.pbf), locations=True, idx="sparse_mem_array")
    print(f"[pbf] DONE  ways={coll.way_count:,}  kept={coll.kept_count:,}  "
          f"tiles_with_data={len(coll.tile_payload):,}", flush=True)

    written = write_tile_jsons(coll.tile_payload, args.out_cache)
    print(f"[pbf] wrote {written:,} tile JSON files under {args.out_cache}/tiles/", flush=True)


if __name__ == "__main__":
    main()