faster_clip.py

import xarray as xr
import rasterio
from rasterio.features import geometry_mask


def order_bounds(
    minx: float,
    miny: float,
    maxx: float,
    maxy: float,
    resolution_x: float,
    resolution_y: float,
) -> tuple[float, float, float, float]:
    """
    Make sure that the bounds are in the correct order (Yanked from rioxarray)
    """
    if resolution_y < 0:
        top = maxy
        bottom = miny
    else:
        top = miny
        bottom = maxy
    if resolution_x < 0:
        left = maxx
        right = minx
    else:
        left = minx
        right = maxx

    return left, bottom, right, top


def get_window_given_box(
    array: xr.DataArray,
    minx: float,
    miny: float,
    maxx: float,
    maxy: float,
) -> xr.DataArray:
    """Inspired by rioxarrays clip_box."""

    if array.rio.crs is None:
        raise ValueError("crs must be set")

    minx, miny, maxx, maxy = rasterio.warp.transform_bounds(
        src_crs=array.rio.crs,
        dst_crs=array.rio.crs,
        left=minx,
        bottom=miny,
        right=maxx,
        top=maxy,
    )

    resolution_x, resolution_y = array.rio.resolution()
    # make sure that if the coordinates are
    # in reverse order that it still works
    left, bottom, right, top = order_bounds(
        minx=minx,
        miny=miny,
        maxx=maxx,
        maxy=maxy,
        resolution_x=resolution_x,
        resolution_y=resolution_y,
    )

    window = rasterio.windows.from_bounds(
        left=np.array(left).item(),
        bottom=np.array(bottom).item(),
        right=np.array(right).item(),
        top=np.array(top).item(),
        transform=array.rio.transform(recalc=True),
    )

    return window


# go get a bounding box and dataset
box = ...
dset = ...
# assuming all variables share coordinates, avoid clipping for each variable
# as done in rioxarray https://github.com/corteva/rioxarray/blob/15b7fcc44f7fab3e7250a26fe80132e7a29a46a4/rioxarray/raster_dataset.py#L293 
window = get_window_given_box(array=dset["example_feature"], *box)
clipped = dset.rio.isel_window(window)


def get_masked_array_from_geometries(
    array: xr.DataArray, geometries, all_touched=False
):
    clip_mask_arr = geometry_mask(
        geometries=geometries,
        out_shape=(int(array.rio.height), int(array.rio.width)),
        transform=array.rio.transform(recalc=True),
        all_touched=all_touched,
    )
    clip_mask = xr.DataArray(
        clip_mask_arr,
        dims=(array.rio.y_dim, array.rio.x_dim),
    )

    return clip_mask


geometries = ...
array_mask = get_masked_array_from_geometries(array=xr.ones_like(dset['example_feature']), geometries=geometries])
# burn in nans for all variables
clipped = dset.where(array_mask)