如何为xarray数据集每个格点的时序值计算百分位数并保留原维度

我有一个包含单个变量的xarray数据集，维度为time、latitude、longitude，部分格点存在NaN值，数据集结构如下：

<xarray.Dataset>
Dimensions:    (time = 6300, latitude: 300, longitude: 360)
Coordinates:
  * latitude   (latitude) float64 49.62 49.88 50.12 50.38 ... 70.88 71.12 71.38
  * longitude  (longitude) float64 -9.875 -9.625 -9.375 ... 39.38 39.62 39.88
  * time       (time) datetime64[ns] 1950-06-01 1950-06-02 ... 2018-08-31
Data variables:
    precip (time, latitude, longitude) float32 dask.array<chunksize=(6300, 300, 360)

需要为数据集中的每个值计算百分位数：沿时间轴计算（每个格点用自身的时间序列），最终要得到和原数据集维度一致的结果，目标结构如下：

<xarray.Dataset>
Dimensions:    (time = 6300, latitude: 300, longitude: 360)
Coordinates:
  * latitude   (latitude) float64 49.62 49.88 50.12 50.38 ... 70.88 71.12 71.38
  * longitude  (longitude) float64 -9.875 -9.625 -9.375 ... 39.38 39.62 39.88
  * time       (time) datetime64[ns] 1950-06-01 1950-06-02 ... 2018-08-31
Data variables:
    precip_percentile (time, latitude, longitude) float32 dask.array<chunksize=(6300, 300, 360)

我尝试用以下代码计算，但结果丢失了time维度，仅保留lat和lon：

def percentileofscore_weak(x):
    return stats.percentileofscore(x, x, kind='rank')

# Apply percentileofscore_weak along the time axis using apply_ufunc
percentiles = xr.apply_ufunc(
    percentileofscore_weak,
    mean_month,
    input_core_dims=[['time']],
    output_core_dims=[[]],
    dask='parallelized',  # Enable parallelization for large datasets
    dask_gufunc_kwargs={'allow_rechunk': False}
)

请问如何修改才能生成符合目标结构的数据集？

问题出在output_core_dims的设置和函数的输出逻辑上，修改后的代码如下：

1. 处理NaN的百分位数计算函数 #

import scipy.stats as stats
import numpy as np
import xarray as xr

def percentileofscore_with_nan(x):
    # 过滤当前格点时间序列中的NaN值
    valid_vals = x[~np.isnan(x)]
    # 若该格点无有效数据，返回全NaN的时间序列
    if len(valid_vals) == 0:
        return np.full_like(x, np.nan)
    # 对每个时间点的值，计算其在有效序列中的百分位数
    return np.array([stats.percentileofscore(valid_vals, val, kind='rank') for val in x])

2. 使用apply_ufunc正确计算 #

# 假设原数据集名为ds，变量为precip
percentile_array = xr.apply_ufunc(
    percentileofscore_with_nan,
    ds['precip'],
    input_core_dims=[['time']],
    output_core_dims=[['time']],  # 保留time维度，因为输出和输入时间序列长度一致
    vectorize=True,  # 对每个(lat, lon)格点独立处理
    dask='parallelized',
    output_dtypes=[np.float32],
    dask_gufunc_kwargs={'allow_rechunk': True}
)

# 将结果封装为目标结构的数据集
result_ds = xr.Dataset(
    {'precip_percentile': percentile_array},
    coords=ds.coords
)

修改说明 #

• output_core_dims调整：原代码中output_core_dims=[[]]会压缩time维度，改为[['time']]后，xarray会保留该维度，匹配输入的时间序列长度。
• NaN处理：新增的函数会过滤每个格点的NaN值，避免无效计算，同时对全NaN格点返回对应长度的NaN序列。
• vectorize参数：确保函数对每个(lat, lon)格点的时间序列单独执行，正确处理多维度广播。

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