如何使用shapefile屏蔽xarray数据集？

Question

我有一个shapefile，并且想使用github中的代码来掩盖我的xarray dataset。地图的图是正确的，并且显示了地图的完整范围，但是当我要用于遮罩xarrray数据集时，仅绘制了地图的西南方向。该代码有什么错误？

import geopandas
from rasterio import features
from affine import Affine
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
%matplotlib inline

def transform_from_latlon(lat, lon):
    lat = np.asarray(lat)
    lon = np.asarray(lon)
    trans = Affine.translation(lon[0], lat[0])
    scale = Affine.scale(lon[1] - lon[0], lat[1] - lat[0])
    return trans * scale

def rasterize(shapes, coords, latitude='latitude', longitude='longitude',
              fill=np.nan, **kwargs):
    """Rasterize a list of (geometry, fill_value) tuples onto the given
    xray coordinates. This only works for 1d latitude and longitude
    arrays.
    """
    transform = transform_from_latlon(coords[latitude], coords[longitude])
    out_shape = (len(coords[latitude]), len(coords[longitude]))
    raster = features.rasterize(shapes, out_shape=out_shape,
                                fill=fill, transform=transform,
                                dtype=float, **kwargs)
    spatial_coords = {latitude: coords[latitude], longitude: coords[longitude]}
    return xr.DataArray(raster, coords=spatial_coords, dims=(latitude, longitude))

# Open Khuzestan counties shapefile
states = geopandas.read_file('CountyPolygon.shp', encoding="windows-1256")

khz_admin = states.reset_index(drop=True)
state_ids = {k: i for i, k in enumerate(khz_admin.Name)}

shapes = zip(khz_admin.geometry, range(len(khz_admin)))

states.plot();

完成地图

ds = xr.open_dataset('file_2018-01-01_01.nc').chunk({'time': 30, 'Lon': 200, 'Lat': 200})

ds['states'] = rasterize(shapes, ds.coords, longitude='Lon', latitude='Lat')

不完整的地图

%matplotlib inline
ds.states.plot();

两个图一起显示：