使用scipy与xarray计算CDF：gamma拟合及apply_ufunc调用失败排查

我需要对多维DataArray（test_data）沿时间维度拟合Gamma分布、计算CDF后转换为正态分布的PPF。使用numpy.apply_along_axis可以正常运行，但调用xarray.apply_ufunc时抛出错误，无法正常执行。

示例代码 #

import xarray as xr
import numpy as np
import pandas as pd
import scipy.stats as st

# 生成随机数据
random_data = np.random.randint(1, 100, (100, 50, 50))
# 生成时间维度
time_dim = pd.date_range("2000-01-01", periods=len(random_data), freq='MS')
# 创建DataArray
test_data = xr.DataArray(random_data, dims=("time", "y", "x"))
test_data["time"] = time_dim

# 定义拟合Gamma分布并转换为正态PPF的函数
def fit_gamma(a):
    if np.isnan(a).all():
        return np.array(a, dtype=np.float32)
    else:
        filter_nan = a[~np.isnan(a)]
        shape, loc, scale = st.gamma.fit(filter_nan, scale=np.std(filter_nan))
        cdf = st.gamma.cdf(a, shape, loc=loc, scale=scale)
        ppf = st.norm.ppf(cdf)
        return np.array(ppf, dtype=np.float32)

# numpy版本可以正常运行
result_numpy = np.apply_along_axis(fit_gamma, 0, test_data.values)

# xarray版本报错
result = xr.apply_ufunc(fit_gamma, test_data, input_core_dims=[["time"]], vectorize=True)

报错信息 #

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
TypeError: only size-1 arrays can be converted to Python scalars

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_6860/3359275878.py in <module>
     23 # result_numpy=np.apply_along_axis(fit_gamma,0, test_data.values)
     24 # test for xarray ufunc, but failed to work
---> 25 result=xr.apply_ufunc(fit_gamma, test_data, input_core_dims=[["time"]], vectorize=True)

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\xarray\core\computation.py in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
   1163     # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
   1164     elif any(isinstance(a, DataArray) for a in args):
-> 1165         return apply_dataarray_vfunc(
   1166             variables_vfunc,
   1167             *args,

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\xarray\core\computation.py in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    288 
    289     data_vars = [getattr(a, "variable", a) for a in args]
-> 290     result_var = func(*data_vars)
    291 
    292     if signature.num_outputs > 1:

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\xarray\core\computation.py in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    731             )
    732 
-> 733     result_data = func(*input_data)
    734 
    735     if signature.num_outputs == 1:

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\numpy\lib\function_base.py in __call__(self, *args, **kwargs)
   2327             vargs.extend([kwargs[_n] for _n in names])
   2328 
-> 2329         return self._vectorize_call(func=func, args=vargs)
   2330 
   2331     def _get_ufunc_and_otypes(self, func, args):

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\numpy\lib\function_base.py in _vectorize_call(self, func, args)
   2401         """Vectorized call to `func` over positional `args`."""
   2402         if self.signature is not None:
-> 2403             res = self._vectorize_call_with_signature(func, args)
   2404         elif not args:
   2405             res = func()

c:\users\hava_tu\appdata\local\programs\python\python38\lib\site-packages\numpy\lib\function_base.py in _vectorize_call_with_signature(self, func, args)
   2461 
   2462             for output, result in zip(outputs, results):
-> 2463                 output[index] = result
   2464 
   2465         if outputs is None:

ValueError: setting an array element with a sequence.

问题原因与解决方案 #

错误原因 #

设置vectorize=True后，xarray会调用numpy的vectorize功能，将你的函数当作逐元素处理的函数——也就是期望函数接收单个标量输入，返回单个标量输出。但你的fit_gamma函数是接收一维数组（时间维度的序列），并返回同样长度的一维数组，这就导致numpy尝试把单个标量传给函数，函数返回数组后无法匹配标量的输出要求，最终报错。

修正代码 #

只需要两个调整：

1. 移除vectorize=True，因为我们的函数已经是处理一维数组的逻辑，不需要自动向量化
2. 添加output_core_dims=[["time"]]，告诉xarray：函数的输出也保留time这个核心维度，这样xarray能正确对齐非核心维度（y、x）

修正后的apply_ufunc调用代码：

result = xr.apply_ufunc(
    fit_gamma, 
    test_data, 
    input_core_dims=[["time"]], 
    output_core_dims=[["time"]]
)

验证结果 #

可以检查xarray结果和numpy结果是否一致：

# 验证结果是否一致
print(np.allclose(result.values, result_numpy))  # 输出 True

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