Use Salem to rasterize a shapefile with multiple geometries to an xarray dataset

multiple shapes to xarray raster.py

import datetime
import numpy as np
import os
import xarray as xr

wd = 'C:\\Users\\xxx\\'
outputs = wd+'outputs_test'

# =============================================================================
# Working version using salem
# =============================================================================
# # Clone and install salem, e.g.
# git clone https://github.com/fmaussion/salem/
# pip install .
# or
# pip install salem
# Docs: https://salem.readthedocs.io/en/stable/



# Load shapefile
shp_fn = wd+'India\\IND_adm1.shp' #
shapes = salem.read_shapefile(shp_fn,  cached=True)
print('Start '+shp_fn)
print(dt.now())
namestr = 'NAME_1' # This is the name of the column in the shapefile that will be used to distinguish between regions,
# e.g. states, countries, basins
tempshape = wd+'temp.shp' #this is the filename of a temporary shapefile written in and out. No need to change

#temporary netcdf that gives the dimensions of output file,
# should match the salem.Grid defined below
tempgrid = wd+'template360x720_05.nc' 

out_fn = shp_fn[:-3] #output filename

#Define grid (must match template below)
nx = 720
ny = 360
grid = salem.Grid(proj=salem.wgs84, x0y0=(-180, -90), nxny=(nx, ny), dxdy=(0.5, 0.5))

#Open grid that is same size, e.g. from template (or create empty xarray with desired dimensions)
template = xr.open_dataarray(tempgrid)
nd = xr.full_like(template, np.nan, dtype=float)
nd.attrs = {}
td = nd.copy(deep=True)

# Run loop through geometries
idxs = {} # Save the different geometries with a different value
for g, geo in enumerate(shapes.geometry):
    print(g)
    region = shapes.iloc[[g]]
    region.to_file(tempshape)
    mask_default = grid.region_of_interest(shape=tempshape)
    mask_default = mask_default.astype(float)
    mask_default[mask_default==1] = g
    mask_default[mask_default==0] = np.nan
    # Flip the mask (no fancy transforms here)
    td.values = np.flipud(mask_default)
    nd = nd.combine_first(td)
    idxs[str(g)] = str(region[namestr].iloc[0])
     
nd.attrs = idxs

    
# File saves out using same name as the input shapefile.
nd.to_netcdf(shp_fn[:-3]+'nc')
os.remove(tempshape)
nd.plot()