如何用argmax结果在xarray中选择对应植被类型层级？

我尝试用ERA-Interim植被覆盖比例数据确定优势植被类型，涉及4个变量：低植被覆盖比例、高植被覆盖比例、低植被类型、高植被类型。

数据预处理代码 #

ERAI_high_veg_frac = ERAI_ds[['CVH_GDS4_SFC_S123']]
ERAI_high_veg_frac = ERAI_high_veg_frac.rename({'CVH_GDS4_SFC_S123':'vegetation_cover'})
ERAI_high_veg_frac = ERAI_high_veg_frac.sel(initial_time0_hours=slice('2010-06-01','2010-06-01'),drop=True)

ERAI_low_veg_frac = ERAI_ds[['CVL_GDS4_SFC_S123']]
ERAI_low_veg_frac = ERAI_low_veg_frac.rename({'CVL_GDS4_SFC_S123':'vegetation_cover'})
ERAI_low_veg_frac = ERAI_low_veg_frac.sel(initial_time0_hours=slice('2010-06-01','2010-06-01'),drop=True)

ERAI_high_veg_type = ERAI_ds[['TVH_GDS4_SFC_S123']]
ERAI_high_veg_type = ERAI_high_veg_type.rename({'TVH_GDS4_SFC_S123':'vegetation_type'})
ERAI_high_veg_type = ERAI_high_veg_type.sel(initial_time0_hours=slice('2010-06-01','2010-06-01'),drop=True)

ERAI_low_veg_type = ERAI_ds[['TVL_GDS4_SFC_S123']]
ERAI_low_veg_type = ERAI_low_veg_type.rename({'TVL_GDS4_SFC_S123':'vegetation_type'})
ERAI_low_veg_type = ERAI_low_veg_type.sel(initial_time0_hours=slice('2010-06-01','2010-06-01'),drop=True)

确定优势植被覆盖的索引 #

ERAI_vegfrac = xr.concat([ERAI_high_veg_frac,ERAI_low_veg_frac],pd.Index(['High Veg Frac','Low Veg Frac'],name='veg_cover'),fill_value=np.nan)
ERAI_domvegfrac = ERAI_vegfrac.max('veg_cover',skipna='True')
ERAI_domvegidx = ERAI_vegfrac.argmax('veg_cover',skipna='True')

ERAI_domvegidx = ERAI_domvegidx['vegetation_cover']

生成的ERAI_domvegidx结构：

<xarray.DataArray 'vegetation_cover' (initial_time0_hours: 1, g4_lat_1: 256, g4_lon_2: 512)>
dask.array<nanarg_agg-aggregate, shape=(1, 256, 512), dtype=int64, chunksize=(1, 256, 512), chunktype=numpy.ndarray>
Coordinates:
  * initial_time0_hours  (initial_time0_hours) datetime64[ns] 2010-06-01
  * g4_lon_2             (g4_lon_2) float32 0.0 0.7031 1.406 ... 358.6 359.3
  * g4_lat_1             (g4_lat_1) float32 89.46 88.77 88.07 ... -88.77 -89.46

合并植被类型数据 #

ERAI_vegtype = xr.concat([ERAI_high_veg_type,ERAI_low_veg_type],pd.Index(['High Veg Type','Low Veg Type'],name='veg_type'),fill_value=np.nan)
ERAI_vegtype = ERAI_vegtype['vegetation_type']

生成的ERAI_vegtype结构：

<xarray.DataArray 'vegetation_type' (veg_type: 2, initial_time0_hours: 1,g4_lat_1: 256, g4_lon_2: 512)>
dask.array<concatenate, shape=(2, 1, 256, 512), dtype=float32, chunksize=(1, 1, 256, 512), chunktype=numpy.ndarray>
Coordinates:
  * initial_time0_hours  (initial_time0_hours) datetime64[ns] 2010-06-01
  * g4_lon_2             (g4_lon_2) float32 0.0 0.7031 1.406 ... 358.6 359.3
  * g4_lat_1             (g4_lat_1) float32 89.46 88.77 88.07 ... -88.77 -89.46
  * veg_type             (veg_type) object 'High Veg Type' 'Low Veg Type'

报错代码与信息 #

尝试用ERAI_domvegidx的0/1值选择对应植被类型层级：

ERAI_domveg = ERAI_vegtype.isel(veg_type=ERAI_domvegidx)

返回错误：

TypeError: unexpected indexer type for VectorizedIndexer: dask.array<nanarg_agg-aggregate, shape=(1, 256, 512), dtype=int64, chunksize=(1, 256, 512), chunktype=numpy.ndarray>

报错核心原因是isel不支持直接传入Dask数组作为向量化索引，以下三种方法可解决问题：

方法1：先计算Dask数组转为本地数组 #

如果数据量不大，可先将索引数组加载到本地内存：

# 计算Dask数组得到本地numpy数组
ERAI_domvegidx_computed = ERAI_domvegidx.compute()
# 使用计算后的索引选择植被类型
ERAI_domveg = ERAI_vegtype.isel(veg_type=ERAI_domvegidx_computed)

方法2：使用xarray.take_along_axis（推荐，支持Dask） #

适合大数据量场景，无需加载全量数据到内存：

# 扩展索引维度，匹配veg_type的轴位置
ERAI_domvegidx_expanded = ERAI_domvegidx.expand_dims('veg_type', axis=0)
# 按索引选择对应植被类型并压缩多余维度
ERAI_domveg = xr.take_along_axis(ERAI_vegtype, ERAI_domvegidx_expanded, axis=0).squeeze('veg_type')

方法3：使用where条件筛选 #

通过条件判断分别选择高/低植被类型再合并：

# 选择高植被类型的区域
high_veg = ERAI_vegtype.sel(veg_type='High Veg Type').where(ERAI_domvegidx == 0)
# 选择低植被类型的区域
low_veg = ERAI_vegtype.sel(veg_type='Low Veg Type').where(ERAI_domvegidx == 1)
# 合并结果
ERAI_domveg = high_veg.combine_first(low_veg)

内容的提问来源于stack exchange，提问作者arctic_climate_science