110 lines
No EOL
4.1 KiB
Python
110 lines
No EOL
4.1 KiB
Python
import os
|
|
import argparse
|
|
import pandas as pd
|
|
import geopandas as gpd
|
|
import geohash
|
|
import logging
|
|
from shapely.geometry import box
|
|
|
|
# Configure logging
|
|
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def load_geohash_files(folder):
|
|
"""Finds all existing geohash parquet files in the folder."""
|
|
geohash_files = {}
|
|
for filename in os.listdir(folder):
|
|
if filename.endswith(".parquet"):
|
|
geohash_code = filename[:-8] # Remove '.parquet' extension
|
|
geohash_files[geohash_code] = os.path.join(folder, filename)
|
|
logger.info(f"Found {len(geohash_files)} geohash parquet files in {folder}.")
|
|
return geohash_files
|
|
|
|
def geohash_bbox(geohash_code):
|
|
"""Returns a bounding box polygon for a geohash."""
|
|
bbox = geohash.bbox(geohash_code)
|
|
return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"])
|
|
|
|
def ensure_crs_consistency(gdf):
|
|
"""Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches."""
|
|
if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326":
|
|
gdf = gdf.to_crs("EPSG:4326")
|
|
return gdf
|
|
|
|
def process_geohash_files(water_gdf, geohash_files):
|
|
"""Processes each geohash file: adds water polygons from water.parquet."""
|
|
# Precompute bounding boxes for all geohash codes
|
|
geohash_bboxes = {geohash_code: geohash_bbox(geohash_code) for geohash_code in geohash_files.keys()}
|
|
|
|
# Spatial index for faster spatial queries
|
|
water_sindex = water_gdf.sindex
|
|
|
|
for geohash_code, file_path in geohash_files.items():
|
|
logger.info(f"Processing {file_path}...")
|
|
|
|
# Load existing geohash file
|
|
try:
|
|
existing_gdf = gpd.read_parquet(file_path)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to read {file_path}. Skipping. Error: {e}")
|
|
continue
|
|
|
|
existing_gdf = ensure_crs_consistency(existing_gdf)
|
|
|
|
# Get the precomputed bounding box for this geohash
|
|
bbox = geohash_bboxes[geohash_code]
|
|
|
|
# Use spatial index to find intersecting water polygons
|
|
possible_matches_index = list(water_sindex.intersection(bbox.bounds))
|
|
possible_matches = water_gdf.iloc[possible_matches_index]
|
|
water_in_tile = possible_matches[possible_matches.intersects(bbox)].copy()
|
|
|
|
if water_in_tile.empty:
|
|
logger.info(f"No water found for {geohash_code}. Skipping water addition.")
|
|
continue
|
|
|
|
# Clip water polygons to the geohash tile boundary
|
|
water_in_tile["geometry"] = water_in_tile.intersection(bbox)
|
|
|
|
# Assign "water" metadata
|
|
water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile)
|
|
|
|
# Ensure columns match before merging
|
|
for col in ["feature_id", "tags"]:
|
|
if col not in existing_gdf.columns:
|
|
existing_gdf[col] = None
|
|
if col not in water_in_tile.columns:
|
|
water_in_tile[col] = None
|
|
|
|
# Merge updated water polygons into geohash file
|
|
updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326")
|
|
|
|
# Save back to parquet
|
|
updated_gdf.to_parquet(file_path, index=False)
|
|
logger.info(f"Updated geohash file {file_path} with water polygons.")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Add water polygons to geohash parquet files from a global water dataset.")
|
|
parser.add_argument("water", help="Path to the water.parquet file.")
|
|
parser.add_argument("geohash_folder", help="Path to the folder containing <geohash>.parquet files.")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Load water dataset
|
|
logger.info(f"Loading water dataset from {args.water}...")
|
|
water_gdf = ensure_crs_consistency(gpd.read_parquet(args.water))
|
|
|
|
# Get existing geohash parquet files
|
|
geohash_files = load_geohash_files(args.geohash_folder)
|
|
|
|
if not geohash_files:
|
|
logger.warning("No geohash parquet files found. Exiting.")
|
|
return
|
|
|
|
# Process geohash files
|
|
process_geohash_files(water_gdf, geohash_files)
|
|
|
|
logger.info("Processing complete. All water polygons added.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |