dopecarpet/lookup.py
2025-03-05 18:15:40 +00:00

61 lines
No EOL
2.5 KiB
Python

import json
import click
import pygeohash as pgh
import pyarrow.parquet as pq
import geopandas as gpd
from shapely.geometry import Point, LineString, MultiLineString
from pyproj import Geod
import numpy as np
@click.command()
@click.argument("latitude", type=float)
@click.argument("longitude", type=float)
@click.option("--distance", type=float, default=10, help="Distance threshold in meters for nearby lines.")
@click.option("--parquet-dir", type=str, default="geohash", help="Directory containing geohash Parquet files.")
def geohash_lookup(latitude, longitude, distance, parquet_dir):
"""
Finds polygons containing the given point and lines within a threshold (in meters).
"""
# Compute geohash (precision can be adjusted)
geohash = pgh.encode(latitude, longitude, precision=4) # Adjust precision as needed
parquet_file = f"{parquet_dir}/{geohash}.parquet"
try:
# Load the Parquet file directly into a GeoDataFrame
df = gpd.read_parquet(parquet_file)
except FileNotFoundError:
click.echo(json.dumps({"error": f"No Parquet file found for geohash {geohash}."}))
return
point = Point(longitude, latitude)
geod = Geod(ellps="WGS84") # Accurate geodetic distance calculations
found_geometries = []
# Vectorized distance calculation for LineStrings and MultiLineStrings
def calculate_min_distance(geom):
if isinstance(geom, LineString):
return min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in geom.coords)
elif isinstance(geom, MultiLineString):
return min(min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in line.coords) for line in geom.geoms)
return np.inf
for row in df.itertuples(index=False):
geom = row.geometry # Directly use the Shapely geometry object
metadata = {}
if isinstance(row.tags, dict):
metadata = {k: v for k, v in row.tags.items() if v is not None} # Remove empty values
if geom.contains(point):
found_geometries.append({"type": "polygon", **metadata})
elif isinstance(geom, (LineString, MultiLineString)):
min_distance = calculate_min_distance(geom)
if min_distance <= distance:
found_geometries.append({"type": "line" if isinstance(geom, LineString) else "multi_line", **metadata})
# Output JSON result
click.echo(json.dumps(found_geometries, indent=2, ensure_ascii=False)) # Ensure proper Unicode display
if __name__ == "__main__":
geohash_lookup()