dopecarpet/water.py
2025-03-04 08:07:21 +00:00

119 lines
4.8 KiB
Python

import os
import argparse
import geopandas as gpd
import pygeohash as pgh
import cbor2
import pandas as pd
import logging
from shapely.geometry import box, Polygon
from tqdm import tqdm
# Configure logging
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)
def geohash_bbox(geohash_code):
"""Returns a Shapely Polygon representing the bounding box of a geohash."""
lat, lon, lat_err, lon_err = pgh.decode_exactly(geohash_code)
lat_min, lon_min = lat - lat_err, lon - lon_err
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 intersect_geohash_tiles(water_polygons, coastline, geohashes):
"""Finds and slices water polygons into 4-digit geohash tiles along the coastline."""
results = {}
logger.info("Processing geohashes for intersection with water polygons...")
for ghash in tqdm(geohashes, desc="Processing geohashes"):
ghash_poly = geohash_bbox(ghash)
# Select only water polygons that intersect with the geohash tile
water_subset = water_polygons[water_polygons.intersects(ghash_poly)]
if water_subset.empty:
continue
# 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")
if not clipped_water.empty:
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"):
parquet_path = os.path.join(output_dir, f"{ghash}.parquet")
gdf.to_parquet(parquet_path)
logger.debug(f"Saved {ghash}.parquet")
if export_cbor:
cbor_path = os.path.join(output_dir, f"{ghash}.cbor")
with open(cbor_path, "wb") as f:
cbor2.dump(gdf.to_dict(), f)
logger.debug(f"Saved {ghash}.cbor")
if export_geojson:
geojson_path = os.path.join(output_dir, f"{ghash}.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.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.")
parser.add_argument("--cbor", action="store_true", help="Export as CBOR.")
parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON.")
parser.add_argument("--verbose", action="store_true", help="Enable detailed logging.")
args = parser.parse_args()
# Set logging level
if args.verbose:
logger.setLevel(logging.DEBUG)
logger.debug("Verbose logging enabled.")
logger.info("Loading input datasets...")
coastline = gpd.read_parquet(args.coastline)
water_polygons = gpd.read_parquet(args.water)
# Ensure the geometries are valid
coastline = coastline[coastline.geometry.is_valid]
water_polygons = water_polygons[water_polygons.geometry.is_valid]
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.")
# Intersect water polygons with geohash tiles
geohash_tiles = intersect_geohash_tiles(water_polygons, coastline, geohashes)
# Save results
save_data(geohash_tiles, args.output_dir, args.cbor, args.geojson)
logger.info("Processing complete. All files saved.")
if __name__ == "__main__":
main()