将numpy.polyfit应用于xarray数据集

据我所知（包括我自己），这已成为xarray用户中一个非常普遍的问题，并且与这个Github问题密切相关。通常，存在某些函数的NumPy实现（在您的情况下为

np.polyfit()

在地球科学的背景下，有 两种主要的用例 ，一种是简单的解决方案，而另一种则更为复杂：

（1）简单案例 ：

您有一个xr.DataArray

temp

(time, lat,lon)

(lat,lon)

.apply()

受到来自Ryan
Abernathy的这个要旨的启发：<3

# Example datada = xr.DataArray(np.random.randn(20, 180, 360),       dims=('time', 'lat', 'lon'),       coords={'time': np.linspace(0,19, 20),        'lat': np.linspace(-90,90,180),        'lon': np.linspace(0,359, 360)})# define a function to compute a linear trend of a timeseriesdef linear_trend(x):    pf = np.polyfit(x.time, x, 1)    # need to return an xr.DataArray for groupby    return xr.DataArray(pf[0])# stack lat and lon into a single dimension called allpointsstacked = da.stack(allpoints=['lat','lon'])# apply the function over allpoints to calculate the trend at each pointtrend = stacked.groupby('allpoints').apply(linear_trend)# unstack back to lat lon coordinatestrend_unstacked = trend.unstack('allpoints')

缺点： 这种方法对于较大的阵列会变得非常慢，并且很难轻易地使其他问题在本质上感觉非常相似。这导致我们…

（2）比较困难的情况 （以及OP的问题）：

您有一个xr.Dataset，其中包含变量

temp

height

(plev, time, lat,lon)

temp

height

(time, lat, lon)

解决此问题的最简单方法是使用xr.apply_ufunc（），它为您提供一定程度的矢量化和dask兼容性。（速度！）

# Example DataArraysda1 = xr.DataArray(np.random.randn(20, 20, 180, 360),        dims=('plev', 'time', 'lat', 'lon'),        coords={'plev': np.linspace(0,19, 20),         'time': np.linspace(0,19, 20),         'lat': np.linspace(-90,90,180),         'lon': np.linspace(0,359, 360)})# Create datasetds = xr.Dataset({'Temp': da1, 'Height': da1})

和以前一样，我们创建一个函数来计算所需的线性趋势：

def linear_trend(x, y):    pf = np.polyfit(x, y, 1)    return xr.DataArray(pf[0])

现在，我们可以用

xr.apply_ufunc()

temp

height

plev

%%timeslopes = xr.apply_ufunc(linear_trend,  ds.Height, ds.Temp,  vectorize=True,  input_core_dims=[['plev'], ['plev']],# reduce along 'plev'  )

但是，这种方法也很慢，并且像以前一样，对于较大的阵列无法很好地扩展。

CPU times: user 2min 44s, sys: 2.1 s, total: 2min 46sWall time: 2min 48s

加快速度：

为了加快计算速度，我们可以将

height

temp

dask.arrays

xr.DataArray.chunk()

dask=parallelized

apply_ufunc()

注意：您必须小心，不要沿用要应用回归的维度！

dask_height = ds.Height.chunk({'lat':10, 'lon':10, 'time': 10})dask_temp   = ds.Temp.chunk({'lat':10, 'lon':10, 'time': 10})dask_height<xarray.DataArray 'Height' (plev: 20, time: 20, lat: 180, lon: 360)>dask.array<xarray-<this-array>, shape=(20, 20, 180, 360), dtype=float64, chunksize=(20, 10, 10, 10), chunktype=numpy.ndarray>Coordinates:  * plev     (plev) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19  * time     (time) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19  * lat      (lat) float64 -90.0 -88.99 -87.99 -86.98 ... 86.98 87.99 88.99 90.0  * lon      (lon) int64 0 1 2 3 4 5 6 7 8 ... 352 353 354 355 356 357 358 359

现在，再次进行计算！

%%timeslopes_dask = xr.apply_ufunc(linear_trend,       dask_height, dask_temp,       vectorize=True,       dask='parallelized',       input_core_dims=[['plev'], ['plev']], # reduce along 'plev'       output_dtypes=['d'],       )CPU times: user 6.55 ms, sys: 2.39 ms, total: 8.94 msWall time: 9.24 ms

显着的速度！

希望这可以帮助！我学到很多东西试图回答它:)

最好

编辑 ：正如评论中指出的那样，要 真正 比较dask和非dask方法之间的处理时间，应使用：

%%timeslopes_dask.compute()

这为您提供了与非黄昏方法相当的计时时间。

但是，必须指出的是，对于使用气候分析中遇到的大型数据集，最好对数据进行 延迟
操作（即直到绝对需要之前才将其加载）。因此，我仍然建议使用dask方法，因为这样您就可以在输入数组上操作许多不同的过程，而每个过程只需要花费几个时间

ms