Files
CA/scripts/generate_grid.py

101 lines
3.3 KiB
Python
Raw Permalink Normal View History

#!/usr/bin/env python3
"""Generate 100x100m resolution grid index for Wuhan city, China."""
import geopandas as gpd
import pandas as pd
import numpy as np
from shapely.geometry import box
from pathlib import Path
def generate_wuhan_grid(
boundary_path: str = "Datas/武汉市.geojson",
output_dir: str = "processed",
grid_size: float = 100,
) -> tuple[gpd.GeoDataFrame, pd.DataFrame]:
"""Generate 100m resolution grid covering Wuhan boundary."""
print(f"Loading Wuhan boundary from {boundary_path}...")
wuhan = gpd.read_file(boundary_path)
bounds = wuhan.total_bounds
print(f"Wuhan bounds: minx={bounds[0]:.4f}, miny={bounds[1]:.4f}, maxx={bounds[2]:.4f}, maxy={bounds[3]:.4f}")
minx, miny, maxx, maxy = bounds
cell_size_deg = grid_size / 111000.0
print(f"Creating grid with {grid_size}m cells (vectorized)...")
x_coords = np.arange(minx, maxx, cell_size_deg)
y_coords = np.arange(miny, maxy, cell_size_deg)
print(f" Grid dimensions: {len(x_coords)} x {len(y_coords)}")
x_grid, y_grid = np.meshgrid(x_coords, y_coords)
x_flat = x_grid.flatten()
y_flat = y_grid.flatten()
print(f" Total cells in bounding box: {len(x_flat)}")
minxs = x_flat
minys = y_flat
maxxs = minxs + cell_size_deg
maxys = minys + cell_size_deg
geometries = [box(mx, my, Mx, My) for mx, my, Mx, My in zip(minxs, minys, maxxs, maxys)]
cells = np.arange(len(geometries))
rows = cells // len(x_coords)
cols = cells % len(x_coords)
print(" Building GeoDataFrame...")
grid_gdf = gpd.GeoDataFrame({
'row': rows,
'col': cols,
'geometry': geometries
}, crs="EPSG:4326")
print("Filtering to cells intersecting Wuhan boundary...")
wuhan_union = wuhan.unary_union
mask = grid_gdf.intersects(wuhan_union)
grid_gdf = grid_gdf[mask].copy().reset_index(drop=True)
print(f"Cells within Wuhan boundary: {len(grid_gdf)}")
grid_gdf['grid_id'] = [f"r{r}_c{c}" for r, c in zip(grid_gdf['row'], grid_gdf['col'])]
centroids = grid_gdf.geometry.centroid
grid_gdf['center_lon'] = centroids.x
grid_gdf['center_lat'] = centroids.y
grid_gdf['polygon'] = grid_gdf.geometry.apply(lambda g: g.wkt)
parquet_df = grid_gdf[['grid_id', 'center_lon', 'center_lat', 'row', 'col', 'polygon']].copy()
return grid_gdf, parquet_df
def main():
output_dir = Path("processed")
output_dir.mkdir(parents=True, exist_ok=True)
grid_gdf, parquet_df = generate_wuhan_grid()
geojson_path = output_dir / "grid_100m_index.geojson"
print(f"Exporting to GeoJSON: {geojson_path}")
grid_gdf.to_file(geojson_path, driver="GeoJSON")
print(f" Exported {len(grid_gdf)} features")
parquet_path = output_dir / "grid_100m_index.parquet"
print(f"Exporting to Parquet: {parquet_path}")
parquet_df.to_parquet(parquet_path, index=False)
print(f" Exported {len(parquet_df)} rows")
print("\n=== Grid Summary ===")
print(f"Total grid cells: {len(grid_gdf)}")
print(f"Bounds: {grid_gdf.total_bounds}")
print(f"Grid ID format example: {grid_gdf['grid_id'].iloc[0]}")
print(f"Center coordinate range:")
print(f" Lon: {parquet_df['center_lon'].min():.4f} to {parquet_df['center_lon'].max():.4f}")
print(f" Lat: {parquet_df['center_lat'].min():.4f} to {parquet_df['center_lat'].max():.4f}")
if __name__ == "__main__":
main()