如何优化使用xarray和pymannkendall计算大文件趋势的代码？

原问题描述 #

我正尝试使用xarray模块和pymannkendall计算大文件的趋势。
我的数据集维度如下：

dimensions:
    time = UNLIMITED ; // (20 currently)
    lon = 3502 ;
    lat = 3724 ;

我目前用于计算趋势的代码如下：

import numpy as np
import xarray as xr
import pymannkendall as mk

# Slope-Trend Function
def calc_slope(ds):
    # Slope for Theil-Sen
    output_slope = []
    for lat in np.arange(len(ds.lat.values)):
        for lon in np.arange(len(ds.lon.values)):
            try:
                slope = mk.hamed_rao_modification_test(ds[:, lat, lon]).slope
            except:
                slope = np.nan  # 处理存在缺测值的情况
            output_slope.append(slope)

    output_slope = np.copy(output_slope).reshape(ds.lat.size, ds.lon.size)
    return xr.DataArray(output_slope, dims=('lat', 'lon'), coords={
        'lat': ds.lat, 'lon': ds.lon})

### 打开数据集
ds = xr.open_dataset('data.nc')

### 提取变量
ts = ds.ts

### 计算趋势斜率
trend = calc_slope(ts)

这段代码在处理小文件时可以正常运行，但处理大文件时运算时间长达数小时甚至数天。请问是否有其他方案可以在xarray中实现该趋势计算？

优化方案 #

你当前代码性能差的核心原因是使用了Python层面的嵌套循环遍历所有格点，3502*3724超过1300万个格点的循环开销极大，也没有利用CPU多核心并行能力，处理大尺寸空间数据时必然耗时极长。可以通过以下方案优化：

方案1：使用xarray.apply_ufunc + Dask并行（最易实现，性能提升明显） #

利用xarray的向量化执行能力代替手动循环，同时通过Dask分块并行利用多核心算力，还能避免大文件直接加载到内存导致的溢出问题。

import numpy as np
import xarray as xr
import pymannkendall as mk
from dask.diagnostics import ProgressBar

# 单时间序列MK斜率计算函数
def calc_mk_slope(ts_1d):
    try:
        return mk.hamed_rao_modification_test(ts_1d).slope
    except:
        return np.nan

# 打开数据集时按空间维度分块，time维度完整保留
ds = xr.open_dataset('data.nc', chunks={'lat': 100, 'lon': 100, 'time': -1})
ts = ds.ts

# 批量应用计算函数到所有格点
trend = xr.apply_ufunc(
    calc_mk_slope,
    ts,
    input_core_dims=[['time']],
    output_dtypes=[np.float32],
    vectorize=True,
    dask='parallelized'
)

# 触发计算并显示进度
with ProgressBar():
    trend = trend.compute()

该方案无需改动核心计算逻辑，即可实现数十倍的性能提升。

方案2：Numba预编译计算逻辑（极致性能优化） #

pymannkendall内置的检验函数为纯Python实现，如果你需要进一步压缩计算时间，可以用Numba自行实现hamed_rao修正的MK检验逻辑，编译为机器码后执行，速度还能再提升3-10倍。

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