Rasterize a vector dataset based on a reference netcdf grid

rasterize_onto_netcdf.py

import geopandas
import numpy
import xarray
from affine import Affine
from rasterio.features import rasterize

def rasterize_vector_onto_netcdf(aoi_path, netcdf_path, target_filepath):
    """Rasterize a vector dataset based on a reference netcdf grid.
    
    This example assumes,
    1. the vector and netcdf have a common coordinate system.
    2. the netcdf has dimensions named 'lon' and 'lat'.
    3. the netcdf has regular grid spacing (distance between points does not vary).
    """
    with xarray.open_dataset(netcdf_path) as dataset:
        coords = dataset.coords

    width = coords['lon'][1] - coords['lon'][0]
    height = coords['lat'][1] - coords['lat'][0]  # should be negative
    translation = Affine.translation(
        coords['lon'][0] - width / 2, coords['lat'][0] + numpy.abs(height) / 2)
    scale = Affine.scale(width, height)
    transform = translation * scale

    aoi_df = geopandas.read_file(aoi_path)
    out_shape = (len(coords['lat']), len(coords['lon']))
    # rasterize defaults to value of 1 where geometry is present, 0 elsewhere
    raster = rasterize(
        aoi_df.geometry,
        out_shape=out_shape,
        transform=transform,
        dtype=int,
        all_touched=True)

    da = xarray.DataArray(
        raster,
        coords={
            'lon': coords['lon'],
            'lat': coords['lat']
        },
        dims=('lat', 'lon'))
    
    ds = xarray.Dataset({'aoi': da})
    ds.to_netcdf(target_filepath)