From 189f135840f2dbfa4868d524faf6e59d156eba1d Mon Sep 17 00:00:00 2001 From: randogoth Date: Tue, 4 Mar 2025 15:26:36 +0000 Subject: [PATCH] slice and water --- slice.py | 94 ++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 81 insertions(+), 13 deletions(-) diff --git a/slice.py b/slice.py index 28ed8bb..42d4dfc 100644 --- a/slice.py +++ b/slice.py @@ -7,8 +7,14 @@ import json import cbor2 from pathlib import Path from shapely import wkb +from shapely.geometry import box from typing import List, Dict +# Configure logging +import logging +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) +logger = logging.getLogger(__name__) + def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]: """Generates all 4-character Geohashes that intersect a bounding box.""" geohashes = set() @@ -37,12 +43,12 @@ def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame): combined_data = combined_data.reset_index(drop=True) combined_data.to_parquet(file_path, index=False) - print(f"✅ Appended new data to {file_path}") + logger.info(f"✅ Appended new data to {file_path}") except Exception as e: - print(f"❌ Error while merging {file_path}: {e}") + logger.warning(f"❌ Error while merging {file_path}: {e}") else: new_data.to_parquet(file_path, index=False) - print(f"✅ Created new file: {file_path}") + logger.info(f"✅ Created new file: {file_path}") def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a GeoJSON file.""" @@ -51,7 +57,7 @@ def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path) with open(geojson_file, "w", encoding="utf-8") as f: json.dump(geojson_dict, f) - print(f"✅ Saved GeoJSON: {geojson_file}") + logger.info(f"✅ Saved GeoJSON: {geojson_file}") def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a CBOR file (compact binary format).""" @@ -60,16 +66,76 @@ def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): with open(cbor_file, "wb") as f: cbor2.dump(geojson_dict, f) - print(f"✅ Saved CBOR: {cbor_file}") + logger.info(f"✅ Saved CBOR: {cbor_file}") -def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: bool, export_geojson: bool, export_cbor: bool): - """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles.""" - input_path, output_path = Path(input_file), Path(output_dir) +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 add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame): + """Adds water polygons to a geohash tile.""" + logger.info(f"Processing {geohash_file}...") + + # Load existing geohash file + try: + existing_gdf = gpd.read_parquet(geohash_file) + except Exception as e: + logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}") + return + + existing_gdf = ensure_crs_consistency(existing_gdf) + + # Get the geohash bounding box + bbox = geohash_bbox(geohash_code) + + # Extract water polygons that overlap with this geohash + water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy() + + if water_in_tile.empty: + logger.info(f"No water found for {geohash_code}. Skipping water addition.") + return + + # 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(geohash_file, index=False) + logger.info(f"Updated geohash file {geohash_file} with water polygons.") + +def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str, export_parquet: bool, export_geojson: bool, export_cbor: bool): + """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and adds water polygons.""" + input_path, output_path, water_path = Path(input_file), Path(output_dir), Path(water_file) if not input_path.exists(): raise FileNotFoundError(f"File '{input_file}' not found.") + if not water_path.exists(): + raise FileNotFoundError(f"File '{water_file}' not found.") output_path.mkdir(parents=True, exist_ok=True) - print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") + logger.info(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") + + # Load water dataset + logger.info(f"Loading water dataset from {water_file}...") + water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path)) # Connect to DuckDB and load spatial extension con = duckdb.connect() @@ -129,6 +195,7 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: if export_parquet: geohash_file = output_path / f"{geohash_code}.parquet" append_to_parquet(geohash_file, gdf) + add_water_to_geohash(geohash_code, geohash_file, water_gdf) if export_geojson: save_as_geojson(geohash_code, gdf, output_path) @@ -137,12 +204,13 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: save_as_cbor(geohash_code, gdf, output_path) con.close() - print("🎉 Processing complete!") + logger.info("🎉 Processing complete!") def main(): - parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles with export format options.") + parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, add water polygons, and export in desired formats.") parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.") parser.add_argument("-o", "--output", required=True, help="Directory to save output files.") + parser.add_argument("-w", "--water", required=True, help="Path to the water.parquet file.") parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet") parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON") parser.add_argument("--cbor", action="store_true", help="Export as CBOR") @@ -153,9 +221,9 @@ def main(): args.parquet = True # Default to GeoParquet if no options are given try: - slice_and_split_geoparquet(args.input, args.output, args.parquet, args.geojson, args.cbor) + slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor) except Exception as e: - print(f"❌ Error: {e}") + logger.error(f"❌ Error: {e}") if __name__ == "__main__": main() \ No newline at end of file