diff --git a/water.py b/water.py index 053ec9f..13cf13f 100644 --- a/water.py +++ b/water.py @@ -5,7 +5,8 @@ import pygeohash as pgh import cbor2 import pandas as pd import logging -from shapely.geometry import box, Polygon +from shapely.geometry import box +from shapely.ops import unary_union from tqdm import tqdm # Configure logging @@ -19,13 +20,40 @@ def geohash_bbox(geohash_code): lat_max, lon_max = lat + lat_err, lon + lon_err return box(lon_min, lat_min, lon_max, lat_max) -def generate_4digit_geohashes(): - """Generates all possible 4-digit geohashes covering the world.""" - base_geohashes = "0123456789bcdefghjkmnpqrstuvwxyz" - return [a + b + c + d for a in base_geohashes for b in base_geohashes for c in base_geohashes for d in base_geohashes] +def get_existing_geohashes(output_dir): + """Returns a set of existing geohash parquet files in the output directory.""" + existing_geohashes = set() + for filename in os.listdir(output_dir): + if filename.endswith(".parquet"): + geohash_code = filename[:-8] # Remove '.parquet' extension + existing_geohashes.add(geohash_code) + logger.info(f"Found {len(existing_geohashes)} existing geohash files in {output_dir}.") + return existing_geohashes + +def subtract_existing_islands(water_gdf, existing_gdf): + """Subtracts intersecting land polygons (islands) from water polygons.""" + if existing_gdf.empty: + return water_gdf # No islands to subtract + + # Collect all existing non-water geometries (islands or coastlines) + land_geometries = existing_gdf[existing_gdf["tags"].apply(lambda tags: tags and tags.get("natural") != "water")] + + if land_geometries.empty: + return water_gdf # No islands to subtract + + # Merge islands into a single geometry + islands_union = unary_union(land_geometries.geometry) + + # Subtract islands from the water polygons + water_gdf["geometry"] = water_gdf.geometry.difference(islands_union) + + # Remove empty geometries resulting from subtraction + water_gdf = water_gdf[~water_gdf.geometry.is_empty] + + return water_gdf def intersect_geohash_tiles(water_polygons, coastline, geohashes): - """Finds and slices water polygons into 4-digit geohash tiles along the coastline.""" + """Finds and slices water polygons into geohash tiles along the coastline.""" results = {} logger.info("Processing geohashes for intersection with water polygons...") @@ -39,29 +67,55 @@ def intersect_geohash_tiles(water_polygons, coastline, geohashes): # Clip with geohash bounding box clipped_water = gpd.clip(water_subset, ghash_poly) - + # Remove areas already covered by coastline coastline_subset = coastline[coastline.intersects(ghash_poly)] if not coastline_subset.empty: clipped_water = clipped_water.overlay(coastline_subset, how="difference") + # Add water-specific tags if not clipped_water.empty: + clipped_water["tags"] = [{"natural": "water", "water": "sea"}] * len(clipped_water) + clipped_water = clipped_water[["tags", "geometry"]] # Drop unnecessary columns results[ghash] = clipped_water + logger.debug(f"Geohash {ghash} contains {len(clipped_water)} water polygons.") logger.info(f"Completed geohash processing. {len(results)} geohash tiles contain water polygons.") return results -def save_data(geohash_tiles, output_dir, export_cbor, export_geojson): - """Saves geohash-sliced water polygons in requested formats.""" - os.makedirs(output_dir, exist_ok=True) - logger.info(f"Saving data to {output_dir}...") - - for ghash, gdf in tqdm(geohash_tiles.items(), desc="Saving geohashes"): +def append_to_existing_parquet(geohash_tiles, output_dir): + """Appends sliced water polygons to existing parquet files after subtracting islands.""" + for ghash, gdf in tqdm(geohash_tiles.items(), desc="Appending data"): parquet_path = os.path.join(output_dir, f"{ghash}.parquet") - gdf.to_parquet(parquet_path) - logger.debug(f"Saved {ghash}.parquet") + # Ensure input water polygons are in EPSG:4326 + gdf = gdf.to_crs("EPSG:4326") + + if os.path.exists(parquet_path): + try: + existing_gdf = gpd.read_parquet(parquet_path) + + # Ensure the existing file is in EPSG:4326 before merging + if existing_gdf.crs != "EPSG:4326": + existing_gdf = existing_gdf.to_crs("EPSG:4326") + + # Subtract islands (land features) from the new water polygons + gdf = subtract_existing_islands(gdf, existing_gdf) + + # Merge the existing and new data + gdf = pd.concat([existing_gdf, gdf], ignore_index=True) + logger.debug(f"Appended new water polygons to {ghash}.parquet") + except Exception as e: + logger.warning(f"Failed to read {parquet_path}. Overwriting instead. Error: {e}") + + # Save the updated file + gdf.to_parquet(parquet_path) + logger.debug(f"Updated {ghash}.parquet with new data.") + +def save_optional_formats(geohash_tiles, output_dir, export_cbor, export_geojson): + """Saves geohash-sliced water polygons in CBOR and GeoJSON formats.""" + for ghash, gdf in tqdm(geohash_tiles.items(), desc="Saving optional formats"): if export_cbor: cbor_path = os.path.join(output_dir, f"{ghash}.cbor") with open(cbor_path, "wb") as f: @@ -73,10 +127,8 @@ def save_data(geohash_tiles, output_dir, export_cbor, export_geojson): gdf.to_file(geojson_path, driver="GeoJSON") logger.debug(f"Saved {ghash}.geojson") - logger.info("Data saving process completed.") - def main(): - parser = argparse.ArgumentParser(description="Slice water polygons along the coastline into 4-digit geohash tiles.") + parser = argparse.ArgumentParser(description="Slice water polygons along the coastline into geohash tiles.") parser.add_argument("coastline", help="Path to the coastline Parquet file.") parser.add_argument("water", help="Path to the water polygons Parquet file.") parser.add_argument("output_dir", help="Directory to save the output files.") @@ -102,18 +154,27 @@ def main(): logger.info(f"Loaded coastline dataset with {len(coastline)} features.") logger.info(f"Loaded water polygons dataset with {len(water_polygons)} features.") - # Generate 4-digit geohashes - logger.info("Generating all possible 4-digit geohashes...") - geohashes = generate_4digit_geohashes() - logger.info(f"Generated {len(geohashes)} geohash tiles.") + # Get existing geohashes from the output directory + existing_geohashes = get_existing_geohashes(args.output_dir) - # Intersect water polygons with geohash tiles - geohash_tiles = intersect_geohash_tiles(water_polygons, coastline, geohashes) + if not existing_geohashes: + logger.error("No existing geohash parquet files found in the output directory. Exiting.") + return - # Save results - save_data(geohash_tiles, args.output_dir, args.cbor, args.geojson) + # Intersect water polygons with only the existing geohash tiles + geohash_tiles = intersect_geohash_tiles(water_polygons, coastline, existing_geohashes) - logger.info("Processing complete. All files saved.") + if not geohash_tiles: + logger.info("No new data to append. Exiting.") + return + + # Append results to existing parquet files + append_to_existing_parquet(geohash_tiles, args.output_dir) + + # Save optional formats + save_optional_formats(geohash_tiles, args.output_dir, args.cbor, args.geojson) + + logger.info("Processing complete. All files updated.") if __name__ == "__main__": main()