Add OSM road benchmark (pyogrio / Geofabrik extracts)

.gitignore

@@ -22,6 +22,7 @@ src/snkit/_version.py
 
 # data
 *.gpkg
+.benchmark-data/
 *.shp
 *.prj
 *.dbf

benchmarks/README.md

@@ -0,0 +1,23 @@
+# OpenStreetMap road benchmark
+
+This benchmark addresses issue #61 using the current Geofabrik extracts for
+Oxfordshire and metropolitan France. It reads every way in GDAL's OSM `lines`
+layer whose `highway` field is set. It then independently times creating unique
+endpoint nodes, assigning node/edge IDs, and assigning edge topology IDs.
+
+## Run
+
+```console
+uv sync --extra benchmark
+uv run python benchmarks/benchmark_osm_roads.py all \
+  --output benchmarks/osm-results.json
+```
+
+Downloads can be retried safely after a failure (incomplete downloads are
+discarded), cached under `.benchmark-data`, and excluded from the timings. The
+JSON output records exact source URLs, PBF sizes, platform and Python version
+alongside node/edge counts and timings.
+
+For an already downloaded or otherwise supplied extract, invoke the Python API
+with `benchmark(area, pbf_path)`. This is useful on compute hosts where data is
+staged separately.

benchmarks/benchmark_osm_roads.py

@@ -1,105 +1,126 @@
-"""Benchmark snkit topology operations on OpenStreetMap road networks.
+"""Benchmark snkit topology operations on OpenStreetMap highway ways.
 
-Issue #61: measure performance on realistically large road networks.
+The input files are the current Geofabrik Oxfordshire and metropolitan-France
+extracts.  Downloads are cached in ``.benchmark-data`` and are deliberately not
+included in any measured interval.
 
-The benchmark expects Geofabrik .osm.pbf extracts. By default it downloads:
-- Oxfordshire, cropped from the England extract (small case)
-- France (France Metropolitaine, i.e. mainland France; large case)
+Run both cases, writing machine-readable results, with::
 
-Run with:
-    python benchmarks/benchmark_osm_roads.py oxfordshire
-    python benchmarks/benchmark_osm_roads.py france
-
-The OSM parsing time is reported separately from the snkit timings. The snkit
-benchmark starts with highway way geometries, then measures creation of endpoint
-nodes, ID assignment and topology assignment.
+    python benchmarks/benchmark_osm_roads.py all --output benchmarks/osm-results.json
 """
 
 from __future__ import annotations
 
 import argparse
 import json
+import platform
+from collections.abc import Callable
 from pathlib import Path
 from time import perf_counter
+from typing import Any, TypeVar
+from urllib.request import urlopen
 
 import geopandas as gpd
-from pyrosm import OSM, get_data, get_data_by_geocoding
 
 import snkit
 
 
 DATA_DIR = Path(".benchmark-data")
+EXTRACTS = {
+    "oxfordshire": "https://download.geofabrik.de/europe/united-kingdom/england/oxfordshire-latest.osm.pbf",
+    # Geofabrik describes this extract as France métropolitaine.
+    "france": "https://download.geofabrik.de/europe/france-latest.osm.pbf",
+}
+T = TypeVar("T")
+
+
+def download(area: str, data_dir: Path = DATA_DIR) -> Path:
+    """Return a cached extract, downloading it atomically when absent."""
+    url = EXTRACTS[area]
+    data_dir.mkdir(parents=True, exist_ok=True)
+    destination = data_dir / url.rsplit("/", 1)[-1]
+    if destination.exists():
+        return destination
+
+    partial = destination.with_suffix(destination.suffix + ".part")
+    print(f"download: {url}", flush=True)
+    try:
+        with urlopen(url) as response, partial.open("wb") as output:  # noqa: S310
+            while chunk := response.read(1024 * 1024):
+                output.write(chunk)
+        partial.replace(destination)
+    except BaseException:
+        partial.unlink(missing_ok=True)
+        raise
+    return destination
+
+
+def read_highway_ways(pbf: Path) -> gpd.GeoDataFrame:
+    """Read highway geometries from GDAL's OSM ``lines`` layer."""
+    edges = gpd.read_file(
+        pbf,
+        layer="lines",
+        columns=[],
+        where="highway IS NOT NULL",
+        engine="pyogrio",
+    )
+    # OSM relations can contribute non-linear geometry to the lines layer.
+    edges = edges.loc[edges.geometry.notna() & edges.geom_type.isin(["LineString", "MultiLineString"])]
+    return edges.reset_index(drop=True)
 
 
-def get_pbf(area: str) -> Path:
-    DATA_DIR.mkdir(exist_ok=True)
-    if area == "oxfordshire":
-        # Download the covering Geofabrik extract and crop to the administrative
-        # area returned by geocoding. This keeps the small benchmark reproducible.
-        return Path(get_data_by_geocoding("Oxfordshire, United Kingdom", directory=DATA_DIR))
-    if area == "france":
-        # Geofabrik's France extract contains France Metropolitaine.
-        return Path(get_data("france", directory=DATA_DIR))
-    raise ValueError(area)
-
-
-def timed(label, fn):
+def timed(label: str, function: Callable[[], T]) -> tuple[T, float]:
     start = perf_counter()
-    result = fn()
+    result = function()
     elapsed = perf_counter() - start
-    print(f"{label}: {elapsed:.3f} s")
+    print(f"{label}: {elapsed:.3f} s", flush=True)
     return result, elapsed
 
 
-def benchmark(area: str) -> dict:
-    pbf = get_pbf(area)
-    size_mb = pbf.stat().st_size / 1024**2
-    print(f"dataset: {area}")
-    print(f"pbf: {pbf} ({size_mb:.1f} MiB)")
-
-    # France is several GB: use pyrosm's bounded-memory streaming reader.
-    engine = "out_of_core" if area == "france" else "in_memory"
-    osm = OSM(str(pbf), engine=engine, workers="auto" if area == "france" else None)
+def benchmark(area: str, pbf: Path | None = None) -> dict[str, Any]:
+    pbf = pbf or download(area)
+    size_mib = pbf.stat().st_size / 1024**2
+    print(f"\ndataset: {area}\npbf: {pbf} ({size_mib:.1f} MiB)", flush=True)
 
-    edges, parse_seconds = timed(
-        "read highway ways",
-        lambda: osm.get_network(network_type="driving"),
-    )
-    # snkit operates on undirected geometries here; duplicate directional rows
-    # from OSM are unnecessary for the topology benchmark.
-    edges = gpd.GeoDataFrame(edges[["geometry"]].copy(), geometry="geometry", crs=edges.crs)
-    edges = edges.drop_duplicates(subset="geometry").reset_index(drop=True)
+    edges, read_seconds = timed("read highway ways", lambda: read_highway_ways(pbf))
     network = snkit.Network(edges=edges)
-
-    network, endpoints_seconds = timed(
-        "add endpoint nodes", lambda: snkit.network.add_endpoints(network)
-    )
-    network, ids_seconds = timed("add ids", lambda: snkit.network.add_ids(network))
-    network, topology_seconds = timed(
-        "add topology ids", lambda: snkit.network.add_topology(network)
-    )
+    network, endpoints_seconds = timed("create endpoint nodes", lambda: snkit.network.add_endpoints(network))
+    network, ids_seconds = timed("add IDs", lambda: snkit.network.add_ids(network))
+    network, topology_seconds = timed("add topology IDs", lambda: snkit.network.add_topology(network))
 
     result = {
         "area": area,
-        "pbf_mib": round(size_mb, 1),
+        "source": EXTRACTS[area],
+        "pbf_mib": round(size_mib, 1),
         "edges": len(network.edges),
         "nodes": len(network.nodes),
         "seconds": {
-            "osm_read_highway_ways": parse_seconds,
-            "snkit_add_endpoints": endpoints_seconds,
-            "snkit_add_ids": ids_seconds,
-            "snkit_add_topology": topology_seconds,
+            "read_highway_ways": read_seconds,
+            "create_endpoint_nodes": endpoints_seconds,
+            "add_ids": ids_seconds,
+            "add_topology_ids": topology_seconds,
         },
     }
-    print(json.dumps(result, indent=2))
+    print(json.dumps(result, indent=2), flush=True)
     return result
 
 
 def main() -> None:
     parser = argparse.ArgumentParser()
-    parser.add_argument("area", choices=["oxfordshire", "france"])
+    parser.add_argument("area", choices=[*EXTRACTS, "all"])
+    parser.add_argument("--data-dir", type=Path, default=DATA_DIR)
+    parser.add_argument("--output", type=Path)
     args = parser.parse_args()
-    benchmark(args.area)
+
+    areas = list(EXTRACTS) if args.area == "all" else [args.area]
+    results = [benchmark(area, download(area, args.data_dir)) for area in areas]
+    report = {
+        "system": {"platform": platform.platform(), "python": platform.python_version()},
+        "results": results,
+    }
+    if args.output:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -27,6 +27,7 @@ dev = ["mypy", "nbstripout", "pre-commit", "pytest", "pytest-cov", "ruff"]
 docs = ["myst-parser", "sphinx"]
 networkx = ["networkx>=3.0"]
 igraph = ["python-igraph>=1.0"]
+benchmark = ["pyogrio>=0.7"]
 
 [project.urls]
 Homepage = "https://snkit.readthedocs.io/en/latest/"