Initial plotting with xarray and GeoPandas

spatial_plotting.py

import geopandas as gpd
import xarray as xr
from pathlib import Path
import matplotlib.pyplot as plt


def load_latlon_points(data_dir: Path) -> gpd.GeoSeries:
    static = xr.open_dataset(data_dir / "camels_static.nc")
    d = static[["gauge_lat", "gauge_lon"]].to_dataframe()
    points = gpd.GeoSeries(
        gpd.points_from_xy(d["gauge_lon"], d["gauge_lat"]),
        index=d.index,
        crs="epsg:4326",
    )
    points.name = "geometry"
    return points

  
  
  if __name__ == "__main__":
    data_dir = Path("path/to/data")
    ds = xr.open_dataset(data_dir / "ALL_dynamic_ds.nc")
    points = load_latlon_points(data_dir)
    
    # ENSURE THAT BOTH datatypes are integers! (otherwise don't match)
    points.index = [int(sid) for sid in points.index]
    ds["station_id"] = ds["station_id"].astype(int)
    
    # plot just the location
    f, ax = plt.subplots(figsize=(5, 8))
    points.plot(ax=ax)
    
    # color the points by mean precipitaiton
    precip = ds["precipitation"].mean(dim="time").to_dataframe()
    gdf = gpd.GeoDataFrame(precip.join(points))
    
    f, ax = plt.subplots(figsize=(5, 8))
    gdf.plot("precipitation", ax=ax, legend=True, cmap="viridis")
    plt.show()