使用xarray筛选多维数组step维度报错，求正确实现方法

问题背景 #

现有如下xarray.Dataset：

<xarray.Dataset>
Dimensions:            (step: 40, latitude: 62, longitude: 126)
Coordinates:
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 0 days 06:00:00 ... 10 days 00:...
    heightAboveGround  float64 ...
  * latitude           (latitude) float64 28.4 28.0 27.6 27.2 ... 4.8 4.4 4.0
  * longitude          (longitude) float64 -160.8 -160.4 ... -111.2 -110.8
    valid_time         (step) datetime64[ns] dask.array<chunksize=(24,), meta=np.ndarray>
Data variables:
    u10                (step, latitude, longitude) float32 dask.array<chunksize=(24, 14, 24), meta=np.ndarray>
    v10                (step, latitude, longitude) float32 dask.array<chunksize=(24, 14, 24), meta=np.ndarray>
Attributes:
    GRIB_edition:            2
    GRIB_centre:             ecmf
    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
    GRIB_subCentre:          0
    Conventions:             CF-1.7
    institution:             European Centre for Medium-Range Weather Forecasts
    history:                 2023-07-14T14:26 GRIB to CDM+CF via cfgrib-0.9.1...

需求是筛选出step维度中对应6、12小时的时间步，尝试以下代码后报错：

timePeriod_list = [6, 12]

xarr_steps = xarr.where([(xarr.step != i) for i in timePeriod_list], drop=True)
print(xarr_steps)

错误信息：

File "C:\Users\Giacomo\PycharmProjects\HIT-v3\try_bufr_enfo.py", line 338, in <module>
    xarr_steps = xarr1.where([(xarr1.step != i) for i in timePeriod_list], drop=True)
  File "C:\Users\Giacomo\anaconda3\envs\HIT-v3\lib\site-packages\xarray\core\common.py", line 1121, in where
    raise TypeError(
TypeError: cond argument is [<xarray.DataArray 'step' (step: 40)>
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True])
Coordinates:
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 0 days 06:00:00 ... 10 days 00:...
    heightAboveGround  float64 ...
    valid_time         (step) datetime64[ns] dask.array<chunksize=(24,), meta=np.ndarray>, <xarray.DataArray 'step' (step: 40)>
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True])
Coordinates:
    time               datetime64[ns] ...
  * step               (step) timedelta64[ns] 0 days 06:00:00 ... 10 days 00:...
    heightAboveGround  float64 ...
    valid_time         (step) datetime64[ns] dask.array<chunksize=(24,), meta=np.ndarray>] but must be a <class 'xarray.core.dataset.Dataset'> or <class 'xarray.core.dataarray.DataArray'>

错误原因 #

1. 参数类型错误：[(xarr.step != i) for i in timePeriod_list]生成的是DataArray对象的列表，而xarray的where方法要求cond参数必须是单个Dataset或DataArray，不能是列表，这是报错的直接原因。
2. 逻辑方向错误：需求是保留6、12小时的step，原代码用!=会过滤掉这两个值，与需求相反。
3. 类型不匹配：step是timedelta64[ns]类型，直接和整数6、12比较会导致不匹配，无法正确筛选。

正确解决方案 #

方法1：使用sel直接筛选（最简洁） #

利用xarray的sel方法直接按维度值筛选，需要先把整数小时转换为对应的timedelta类型：

import pandas as pd

timePeriod_list = [6, 12]
# 将整数小时转换为timedelta，匹配step的类型
target_steps = pd.to_timedelta(timePeriod_list, unit='h')
# 直接筛选step维度
xarr_steps = xarr.sel(step=target_steps)

方法2：使用where配合isin #

如果坚持用where方法，可以用isin生成符合要求的布尔条件：

import pandas as pd

timePeriod_list = [6, 12]
target_steps = pd.to_timedelta(timePeriod_list, unit='h')
# 生成布尔条件：step是否在目标列表中
cond = xarr.step.isin(target_steps)
# 应用where筛选并drop掉不符合的维度
xarr_steps = xarr.where(cond, drop=True)

方法3：手动组合布尔条件 #

如果不用isin，可以用逻辑或|组合多个相等条件：

import pandas as pd

timePeriod_list = [6, 12]
target_steps = pd.to_timedelta(timePeriod_list, unit='h')
# 组合多个条件，注意每个条件加括号
cond = (xarr.step == target_steps[0]) | (xarr.step == target_steps[1])
xarr_steps = xarr.where(cond, drop=True)

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