CA/scripts/resample_spatial_features.py

#!/usr/bin/env python3
"""
Resample Spatial Features (DEM and Population Density) to Road Network Nodes.

Uses bilinear interpolation to sample:
- DEM (elevation) from Datas/DEM/CJJJD_DEM.TIF
- Population density from Datas/landscan-hd-china-v1-assets/landscan-hd-china-v1.tif

Input:
  - processed/graph/node_metadata.parquet (from US-004): contains node_id, lat, lon

Output:
  - processed/graph/node_features.parquet: updated with elevation_m and pop_density
"""

import argparse
import numpy as np
import pandas as pd
import rasterio
from pathlib import Path
from rasterio.features import bounds
from rasterio.warp import transform


def load_nodes(node_path: str) -> pd.DataFrame:
    """Load node metadata from parquet."""
    df = pd.read_parquet(node_path)
    print(f"Loaded {len(df)} nodes from {node_path}")
    print(f"Columns: {list(df.columns)}")
    return df


def create_sample_nodes(n: int = 1000, seed: int = 42) -> pd.DataFrame:
    """Create sample nodes within Wuhan boundary for testing.

    Wuhan approximate bounding box:
    - Lon: 114.257 to 114.403
    - Lat: 30.573 to 30.699
    """
    np.random.seed(seed)

    # Wuhan bounding box
    lon_min, lon_max = 114.257, 114.403
    lat_min, lat_max = 30.573, 30.699

    nodes = pd.DataFrame({
        'node_id': range(n),
        'lon': np.random.uniform(lon_min, lon_max, n),
        'lat': np.random.uniform(lat_min, lat_max, n)
    })
    print(f"Created {n} sample nodes within Wuhan bounding box")
    return nodes


def sample_raster_bilinear(df: pd.DataFrame, raster_path: str, col_name: str) -> pd.DataFrame:
    """Sample raster values at node locations using bilinear interpolation.

    Args:
        df: DataFrame with 'lon' and 'lat' columns (WGS84/EPSG:4326)
        raster_path: Path to raster file
        col_name: Name of the column to create in df

    Returns:
        df with new column added
    """
    print(f"Sampling {col_name} from {raster_path}...")

    with rasterio.open(raster_path) as rast:
        # Get raster bounds and CRS
        bounds = rast.bounds
        raster_crs = rast.crs
        print(f"  Raster bounds: {bounds}")
        print(f"  Raster CRS: {raster_crs}")

        width = rast.width
        height = rast.height

        # Transform node coordinates to raster CRS if needed
        from_crs = "EPSG:4326"  # WGS84 lat/lon
        if raster_crs.to_string() != from_crs:
            from pyproj import Transformer
            transformer = Transformer.from_crs(from_crs, raster_crs.to_string(), always_xy=True)
            node_x, node_y = transformer.transform(df['lon'].values, df['lat'].values)
            print(f"  Transformed {len(node_x)} nodes to {raster_crs.to_string()}")
        else:
            node_x = df['lon'].values
            node_y = df['lat'].values

        # Compute fractional pixel coordinates
        x_frac = (node_x - bounds.left) / (bounds.right - bounds.left) * (width - 1)
        y_frac = (bounds.top - node_y) / (bounds.top - bounds.bottom) * (height - 1)

        # Get integer pixel indices
        x_int = np.floor(x_frac).astype(int)
        y_int = np.floor(y_frac).astype(int)

        # Clip to valid range
        x_int = np.clip(x_int, 0, width - 2)
        y_int = np.clip(y_int, 0, height - 2)

        # Get fractional offsets for bilinear weights
        x_f = x_frac - x_int
        y_f = y_frac - y_int

        # Bilinear interpolation weights
        w00 = (1 - x_f) * (1 - y_f)
        w10 = x_f * (1 - y_f)
        w01 = (1 - x_f) * y_f
        w11 = x_f * y_f

        # Read all data at once
        data = rast.read(1)

        # Get 4 neighboring pixel values
        v00 = data[y_int, x_int]
        v10 = data[y_int, x_int + 1]
        v01 = data[y_int + 1, x_int]
        v11 = data[y_int + 1, x_int + 1]

        # Bilinear interpolation
        values = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11

        # Handle nodata
        nodata = rast.nodata
        if nodata is not None:
            valid_mask = (values != nodata)
            nan_count = (~valid_mask).sum()
            if nan_count > 0:
                pct = nan_count / len(values) * 100
                print(f"  Warning: {nan_count} nodes ({pct:.1f}%) outside raster or nodata")
            values = np.where(valid_mask, values, np.nan)

        df[col_name] = values
        print(f"  Sampled {len(df)} points, mean={np.nanmean(values):.2f}, std={np.nanstd(values):.2f}")

    return df


def main():
    parser = argparse.ArgumentParser(
        description="Resample DEM and population density to road network nodes"
    )
    parser.add_argument(
        "--input-nodes",
        type=str,
        default="processed/graph/node_metadata.parquet",
        help="Input node metadata parquet (from US-004)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default="processed/graph/node_features.parquet",
        help="Output parquet path"
    )
    parser.add_argument(
        "--dem",
        type=str,
        default="Datas/DEM/CJJJD_DEM.TIF",
        help="DEM raster path"
    )
    parser.add_argument(
        "--pop",
        type=str,
        default="Datas/landscan-hd-china-v1-assets/landscan-hd-china-v1.tif",
        help="Population density raster path"
    )
    parser.add_argument(
        "--use-sample-nodes",
        action="store_true",
        help="Use sample nodes instead of input (for testing when node_metadata doesn't exist)"
    )
    parser.add_argument(
        "--sample-nodes-count",
        type=int,
        default=1000,
        help="Number of sample nodes to create"
    )

    args = parser.parse_args()

    # Load or create nodes
    if args.use_sample_nodes:
        nodes = create_sample_nodes(n=args.sample_nodes_count)
    else:
        if not Path(args.input_nodes).exists():
            print(f"ERROR: {args.input_nodes} not found.")
            print("  US-004 (road network construction) must be completed first.")
            print("  Or use --use-sample-nodes to test with synthetic nodes.")
            return 1
        nodes = load_nodes(args.input_nodes)

    # Ensure output directory exists
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)

    # Sample DEM
    dem_path = Path(args.dem)
    if not dem_path.exists():
        print(f"ERROR: DEM not found at {dem_path}")
        return 1
    nodes = sample_raster_bilinear(nodes, str(dem_path), "elevation_m")

    # Sample population density
    pop_path = Path(args.pop)
    if not pop_path.exists():
        print(f"ERROR: Population density raster not found at {pop_path}")
        return 1
    nodes = sample_raster_bilinear(nodes, str(pop_path), "pop_density")

    # Save output
    print(f"Saving to {args.output}")
    nodes.to_parquet(args.output, index=False)

    # Verification
    df_verify = pd.read_parquet(args.output)
    print(f"\n=== Verification ===")
    print(f"Output shape: {df_verify.shape}")
    print(f"Columns: {list(df_verify.columns)}")

    required_cols = ["elevation_m", "pop_density"]
    for col in required_cols:
        if col in df_verify.columns:
            valid = df_verify[col].notna().sum()
            print(f"  {col}: {valid}/{len(df_verify)} valid values")
            print(f"    mean={df_verify[col].mean():.2f}, min={df_verify[col].min():.2f}, max={df_verify[col].max():.2f}")
        else:
            print(f"  {col}: MISSING")

    return 0


if __name__ == "__main__":
    exit(main())