如何用Xarray非循环实现每个时间步提取不同子域

问题背景 #

处理海平面气压数据集MSLP_ds，数据集结构如下：

<xarray.Dataset>
Dimensions:    (time: 85, nj: 1602, ni: 1602)
Coordinates:
  latitude   (nj, ni) float64 ...
  longitude  (nj, ni) float64 ...
  ni         (ni) float64 9.086e+05 9.089e+05 ... 1.309e+06 1.309e+06
  nj         (nj) float64 9.086e+05 9.089e+05 ... 1.309e+06 1.309e+06
  time       (time) datetime64[ns] 2018-12-22T01:00:00 ... 2018-12-23T23:30:00
Data variables:
  MSLP       (time, nj, ni) float64 ...
Attributes:
  nco_openmp_thread_number:   1
  history:                    Wed Jul 12 13:38:25 2023: ncks -A -v MSLP CI2...
  history_of_appended_files:  Wed Jul 12 13:38:25 2023: Appended file CI250...
  NCO:                        "4.6.2"

需求：为每个时间步提取不同的方形子域，子域的四角索引由4个长度与time维度一致的numpy数组定义：indi_w（西边界ni索引）、indi_e（东边界ni索引）、indj_s（南边界nj索引）、indj_n（北边界nj索引）。

当前可提取单个时间步子域：

tt=70
MSLP_ds_sub = MSLP_ds['MSLP'].isel(nj=slice(indj_s[tt],indj_n[tt]),ni=slice(indi_w[tt],indi_e[tt]))

但尝试批量提取时报错：

MSLP_ds_mod = MSLP_ds['MSLP'].isel(nj=slice(indj_s,indj_n),ni=slice(indi_w,indi_e) )

错误信息：TypeError: only integer scalar arrays can be converted to a scalar index，原因是slice无法接受数组作为参数。

解决方案 #

方法1：基于Xarray的堆叠与高级索引（适合子域大小固定的场景） #

通过堆叠空间维度，构造每个时间步对应的全局索引，实现批量提取：

import numpy as np
import xarray as xr

# 假设所有时间步的子域大小相同，先计算子域尺寸
nj_size = indj_n[0] - indj_s[0]
ni_size = indi_e[0] - indi_w[0]

# 构造每个时间步的nj、ni索引网格
nj_indices = np.array([np.arange(s, n) for s, n in zip(indj_s, indj_n)])
ni_indices = np.array([np.arange(w, e) for w, e in zip(indi_w, indi_e)])

# 将nj和ni维度堆叠为单个grid维度
mslp_stacked = MSLP_ds['MSLP'].stack(grid=('nj', 'ni'))

# 将二维索引转换为grid维度的一维索引
grid_indices = np.ravel_multi_index(
    (nj_indices[:, :, np.newaxis], ni_indices[:, np.newaxis, :]),
    dims=(MSLP_ds.dims['nj'], MSLP_ds.dims['ni'])
).reshape(85, nj_size * ni_size)

# 提取子域并恢复空间维度
mslp_sub = mslp_stacked.isel(time=np.arange(85), grid=grid_indices).unstack('grid')

# 调整维度顺序为(time, nj, ni)
mslp_sub = mslp_sub.transpose('time', 'nj', 'ni')

方法2：基于NumPy的take_along（灵活处理子域，保留坐标信息） #

利用NumPy的广播索引提取数据，再转回Xarray结构：

# 获取原始数据的numpy数组
mslp_np = MSLP_ds['MSLP'].values

# 构造每个时间步的nj、ni索引（添加维度实现广播）
nj_idx = np.array([np.arange(s, n) for s, n in zip(indj_s, indj_n)])[:, :, np.newaxis]
ni_idx = np.array([np.arange(w, e) for w, e in zip(indi_w, indi_e)])[:, np.newaxis, :]

# 嵌套使用take_along提取子域
mslp_sub_np = np.take_along(np.take_along(mslp_np, nj_idx, axis=1), ni_idx, axis=2)

# 构造对应的坐标数组（每个时间步的经纬度子域）
lat_sub = np.array([
    np.take_along(np.take_along(MSLP_ds['latitude'].values, nj_idx[t], axis=0), ni_idx[t], axis=1)
    for t in range(85)
])
lon_sub = np.array([
    np.take_along(np.take_along(MSLP_ds['longitude'].values, nj_idx[t], axis=0), ni_idx[t], axis=1)
    for t in range(85)
])

# 转回Xarray Dataset
mslp_sub_ds = xr.Dataset(
    {
        'MSLP': (['time', 'nj_sub', 'ni_sub'], mslp_sub_np)
    },
    coords={
        'time': MSLP_ds['time'],
        'latitude': (['time', 'nj_sub', 'ni_sub'], lat_sub),
        'longitude': (['time', 'nj_sub', 'ni_sub'], lon_sub)
    }
)

注意事项 #

• 如果各时间步的子域大小不一致，无法生成统一形状的数组，此时只能通过列表推导结合xr.concat实现（本质是隐式循环，但效率优于显式循环）：

  mslp_sub_list = [
      MSLP_ds['MSLP'].isel(time=t, nj=slice(indj_s[t], indj_n[t]), ni=slice(indi_w[t], indi_e[t]))
      for t in range(85)
  ]
  mslp_sub_ds = xr.concat(mslp_sub_list, dim='time')
  

内容的提问来源于stack exchange，提问作者Clément Soufflet