91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
"""Download NaPTAN data and extract railway/metro station POIs."""
|
|
|
|
import argparse
|
|
import io
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
import polars as pl
|
|
|
|
NAPTAN_CSV_URL = "https://naptan.api.dft.gov.uk/v1/access-nodes?dataFormat=csv"
|
|
|
|
|
|
STOP_TYPES = {
|
|
"AIR": "Airport",
|
|
"FTD": "Ferry",
|
|
"RSE": "Rail station",
|
|
"BCT": "Bus stop",
|
|
"BCE": "Bus station",
|
|
"TXR": "Taxi rank",
|
|
"TMU": "Metro or Tram stop",
|
|
"MET": "Metro or Tram stop",
|
|
}
|
|
|
|
|
|
def download_naptan(output: Path) -> None:
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"Downloading NaPTAN data from {NAPTAN_CSV_URL}")
|
|
with urllib.request.urlopen(NAPTAN_CSV_URL) as resp:
|
|
raw = resp.read()
|
|
|
|
print(f"Downloaded {len(raw) / (1024 * 1024):.1f} MB")
|
|
|
|
df = (
|
|
pl.read_csv(io.BytesIO(raw), infer_schema_length=0)
|
|
.with_columns(
|
|
pl.col("Latitude").cast(pl.Float64, strict=False),
|
|
pl.col("Longitude").cast(pl.Float64, strict=False),
|
|
)
|
|
.drop_nulls(subset=["Latitude", "Longitude"])
|
|
.filter(pl.col("StopType").is_in(list(STOP_TYPES.keys())))
|
|
.select(
|
|
pl.col("ATCOCode").alias("id"),
|
|
pl.col("CommonName").alias("name"),
|
|
pl.col("StopType").replace(STOP_TYPES).alias("category"),
|
|
pl.col("Latitude").alias("lat"),
|
|
pl.col("Longitude").alias("lng"),
|
|
pl.col("NptgLocalityCode").alias("locality"),
|
|
)
|
|
)
|
|
|
|
before = len(df)
|
|
|
|
# Deduplicate: one record per name+category+locality
|
|
# (merges entrances, bus stop pairs on opposite sides of the road, etc.)
|
|
has_loc = df.filter(
|
|
pl.col("locality").is_not_null() & (pl.col("locality") != "")
|
|
)
|
|
no_loc = df.filter(
|
|
pl.col("locality").is_null() | (pl.col("locality") == "")
|
|
)
|
|
cols = ["id", "name", "category", "lat", "lng"]
|
|
deduped = has_loc.group_by("name", "category", "locality").agg(
|
|
pl.col("id").first(),
|
|
pl.col("lat").mean(),
|
|
pl.col("lng").mean(),
|
|
)
|
|
df = pl.concat([deduped.select(cols), no_loc.select(cols)])
|
|
|
|
print(f"Deduplicated {before:,} → {len(df):,} stops (by name+category+locality)")
|
|
|
|
df.write_parquet(output)
|
|
size_mb = output.stat().st_size / (1024 * 1024)
|
|
print(f"Wrote {output} ({size_mb:.1f} MB, {len(df):,} stations)")
|
|
|
|
counts = df.group_by("category").len().sort("len", descending=True)
|
|
for row in counts.iter_rows(named=True):
|
|
print(f" {row['category']}: {row['len']:,}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Download NaPTAN station data")
|
|
parser.add_argument(
|
|
"--output", type=Path, required=True, help="Output parquet file path"
|
|
)
|
|
args = parser.parse_args()
|
|
download_naptan(args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|