如何基于位置相关的时间切片筛选经纬度降水数据

现有数据 #

逐日降水数据集 #

<xarray.Dataset>
Dimensions:    (longitude: 243, latitude: 260, time: 1826)
Coordinates:
  * longitude  (longitude) float32 166.4 166.5 166.5 166.6 ... 178.4 178.5 178.5
  * latitude   (latitude) float32 -34.38 -34.42 -34.47 ... -47.22 -47.28 -47.33
  * time       (time) datetime64[ns] 2006-01-01 2006-01-02 ... 2010-12-31
Data variables:
    rain       (time, latitude, longitude) float32 dask.array<chunksize=(1826, 1, 1), meta=np.ndarray>

起始时段DataArray（start_da） #

每个经纬度单元格对应每年的时段起始日期，时间维度对应年份标记（如2006-08-30对应2007年的时段）：

<xarray.DataArray (time: 4, longitude: 2, latitude: 2)>
array([[[datetime.date(2007, 3, 23), datetime.date(2007, 4, 26)],
        [datetime.date(2007, 5, 18), datetime.date(2007, 4, 20)]],

       [[datetime.date(2008, 3, 25), datetime.date(2008, 5, 6)],
        [datetime.date(2008, 5, 30), datetime.date(2008, 4, 25)]],

       [[datetime.date(2009, 3, 28), datetime.date(2009, 5, 4)],
        [datetime.date(2009, 5, 29), datetime.date(2009, 4, 23)]],

       [[datetime.date(2010, 3, 18), datetime.date(2010, 4, 16)],
        [datetime.date(2010, 5, 5), datetime.date(2010, 4, 9)]]],
      dtype=object)
Coordinates:
  * time       (time) datetime64[ns] 2006-08-30 2007-08-30 2008-08-29 2009-08-30
  * longitude  (longitude) float32 172.4 172.5
  * latitude   (latitude) float32 -41.38 -41.42

结束时段DataArray（end_da） #

对应每个单元格每年的时段结束日期：

<xarray.DataArray (time: 4, longitude: 2, latitude: 2)>
array([[[datetime.date(2007, 5, 10), datetime.date(2007, 8, 7)],
        [datetime.date(2007, 9, 23), datetime.date(2007, 7, 8)]],

       [[datetime.date(2008, 5, 17), datetime.date(2008, 8, 19)],
        [datetime.date(2007, 9, 29), datetime.date(2008, 7, 24)]],

       [[datetime.date(2009, 5, 19), datetime.date(2009, 8, 22)],
        [datetime.date(2009, 9, 27), datetime.date(2009, 7, 21)]],

       [[datetime.date(2010, 5, 2), datetime.date(2010, 7, 1)],
        [datetime.date(2010, 8, 24), datetime.date(2010, 6, 13)]]],
      dtype=object)
Coordinates:
  * time       (time) datetime64[ns] 2006-09-30 2007-09-30 2008-09-29 2009-09-30
  * longitude  (longitude) float32 172.4 172.5
  * latitude   (latitude) float32 -41.38 -41.42
    fstar      <U4 '2820'

需求 #

计算每个经纬度单元格每年对应时段内的降水总和，即对每个单元格的每一年，用start_da和end_da的日期切片降水数据后求和。

尝试过的方法及报错 #

尝试用slice直接索引时段，但因start_da/end_da是多维数组而非标量，报错：

rain_in_period = rain_ds.sel(time=slice(start_data_array, end_data_array)).rain.sum().compute()

报错信息：

ValueError: cannot use non-scalar arrays in a slice for xarray indexing: <xarray.DataArray (longitude: 243, latitude: 260)>

遍历或枚举索引的方案要么无效，要么速度极慢且仍报错。

解决方案：用xarray.apply_ufunc实现逐单元格时段求和 #

核心思路是利用apply_ufunc对每个经纬度单元格的时间序列，结合对应的时段起止日期进行切片求和，同时保留原维度结构。

步骤1：统一日期类型 #

先把start_da和end_da的datetime.date类型转为datetime64[ns]，确保和降水数据集的时间坐标类型一致：

import numpy as np
import xarray as xr

# 转换日期类型
start_da = start_da.apply(lambda x: np.datetime64(x))
end_da = end_da.apply(lambda x: np.datetime64(x))

步骤2：定义时段求和函数 #

定义一个函数，输入单格的降水时间序列、该格的所有起始/结束日期，返回每一年的时段降水总和：

def compute_period_sum(rain_ts, starts, ends):
    # rain_ts: 一维数组，单格的逐日降水
    # starts: 一维数组，该格各年的起始日期
    # ends: 一维数组，该格各年的结束日期
    time_axis = rain_ts.coords['time'].values
    sums = []
    for s, e in zip(starts, ends):
        # 筛选该时段内的降水并求和
        mask = (time_axis >= s) & (time_axis <= e)
        sums.append(rain_ts[mask].sum().item())
    return np.array(sums)

步骤3：用apply_ufunc批量计算 #

指定核心维度为time，对经纬度维度进行广播计算：

# 确保降水数据的维度顺序和时段数组匹配
rain_da = rain_ds['rain'].transpose('time', 'latitude', 'longitude')

# 应用ufunc
period_rain_sum = xr.apply_ufunc(
    compute_period_sum,
    rain_da,
    start_da.transpose('time', 'latitude', 'longitude'),
    end_da.transpose('time', 'latitude', 'longitude'),
    input_core_dims=[['time'], ['time'], ['time']],
    output_core_dims=[['year_time']],  # 输出的时间维度对应原时段数组的time维度
    vectorize=True,  # 自动遍历非核心维度（latitude, longitude）
    dask='parallelized',  # 支持dask并行计算，适合大数据
    output_dtypes=[float]
)

# 重命名输出的时间维度，匹配原时段数组的时间坐标
period_rain_sum = period_rain_sum.rename({'year_time': 'time'})
period_rain_sum = period_rain_sum.assign_coords(time=start_da.coords['time'])

说明 #

• vectorize=True会自动遍历经纬度的每个单元格，避免手动循环，效率更高；
• dask='parallelized'保留原数据的dask分块，适合处理大尺度降水数据；
• 最终结果的维度为(time:4, latitude:260, longitude:243)，每个值对应对应年份、经纬度单元格的时段降水总和。

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