将numpy.polyfit应用于xarray数据集

Question

Xarray是否支持numpy计算功能（例如polyfit）？还是有一种有效的方法将这些函数应用于数据集？

示例：我想计算拟合两个变量（温度和高度）的直线的斜率，以计算失效率。我有一个数据集（下面），具有这两个变量，维度为（垂直，时间，xgrid_0，ygrid_0）。

<xarray.Dataset>
Dimensions:    (PressLev: 7, time: 48, xgrid_0: 685, ygrid_0: 485)
Coordinates:
    gridlat_0  (ygrid_0, xgrid_0) float32 44.6896 44.6956 44.7015 44.7075 ...
    gridlon_0  (ygrid_0, xgrid_0) float32 -129.906 -129.879 -129.851 ...
  * ygrid_0    (ygrid_0) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 ...
  * xgrid_0    (xgrid_0) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 ...
  * time       (time) datetime64[ns] 2016-08-15T01:00:00 2016-08-15T02:00:00 ...
  * PressLev   (PressLev) int64 0 1 2 3 4 5 6
Data variables:
    Temperature       (PressLev, time, ygrid_0, xgrid_0) float64 289.4 289.4 289.4 ...
    Height       (PressLev, time, ygrid_0, xgrid_0) float64 85.23 85.13 84.98 ...

如果我提取给定时间的温度和高度，则xgrid_0，ygrid_0; 我可以使用numpy.polyfit函数。

ds_LR = ds.TMP_P0_L103_GST0 * 0 -9999 # Quick way to make dataarray with -9999 values but with correct dims/coords
for cts in np.arange(0,len(ds_UA.time)):
        for cx in ds_UA.xgrid_0.values:
                for cy in ds_UA.ygrid_0.values:
                        x_temp = ds_UA.Temperature[:,cts,cy,cx] # Grab the vertical profile of air temperature
                        y_hgt  = ds_UA.Height[:,cts,cy,cx] # Grab the vertical heights of air temperature values
                        s      = np.polyfit(y_hgt,x_temp,1) # Fit a line to the data
                        ds_LR[cts,cy,cx].values = s[0] # Grab the slope (first element)

但这是一种缓慢而低效的方法。对解决此问题的更好方法有何建议？

Answer 1

据我所知（包括我自己），这已成为xarray用户中一个非常普遍的问题，并且与这个Github问题密切相关。通常，存在某些函数的NumPy实现（在您的情况下为np.polyfit()），但尚不清楚如何最好地将此计算应用于每个网格单元（可能跨多个维度）。

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

（1）简单案例 ：

您有一个xr.DataArray temp，它是的一个函数，(time, lat, lon)并且您想在每个网格框中找到时间趋势。最简单的方法是将(lat, lon)坐标分组为一个新坐标，然后将该坐标分组，然后使用该.apply()方法。

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

# Example data
da = 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 timeseries
def 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 allpoints
stacked = da.stack(allpoints=['lat','lon'])
# apply the function over allpoints to calculate the trend at each point
trend = stacked.groupby('allpoints').apply(linear_trend)
# unstack back to lat lon coordinates
trend_unstacked = trend.unstack('allpoints')

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

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

您有一个xr.Dataset，其中包含变量temp和height，每个变量的功能，(plev, time, lat, lon)并且您希望找到每个点temp对height（回归率）的回归(time, lat, lon)。

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

# Example DataArrays
da1 = 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 dataset
ds = 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()倒退的两个DataArraystemp和height反对对方，沿plev尺寸！

%%time
slopes = 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 46s
Wall time: 2min 48s

加快速度：

为了加快计算速度，我们可以将height和转换temp为dask.arraysusing
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

现在，再次进行计算！

%%time
slopes_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 ms
Wall time: 9.24 ms

显着的速度！

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

最好

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

%%time
slopes_dask.compute()

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

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