logo标签列表

应用numpy.polyfit到xarray数据集

pythonnumpypython-xarray
9

Xarray是否支持像polyfit这样的numpy计算函数？或者有没有一种有效的方法将这些函数应用于数据集？

例如：我想计算适合两个变量（温度和高度）的线的斜率，以计算温度随高度变化的速率。我有一个数据集（如下），其中包含这两个变量，维度为（垂直、时间、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)

但这是一种缓慢而低效的方法。有没有更好的方法来解决这个问题？
numpy.polyfit期望x坐标的一维值和y坐标值的集合，从而形成2D输入。因此，您可以将多维数组重塑以满足此要求，然后选择所需元素并再次重塑回原始形状。 - mdurant
我认为答案在.reduce()中，但你需要玩弄kwargs。 - claude
2
我想回来并引导您使用xr.apply_ufunc() - 我昨天尝试过。http://xarray.pydata.org/en/stable/generated/xarray.apply_ufunc.html 它不一定比循环更快。（在我的情况下不是）特别是如果您使用vectorize=True（我必须使用）。我最终选择使用chain.from_iterable(zip(*dataset))并进行循环。但是，您必须将xgrid和ygrid作为前两个维度（我使用了转置方法）。我会尽快写出答案，但您可以尝试一下。 - claude
2个回答
12

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

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

(1) 简单情况：

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

受Ryan Abernathy这个Gist启发：<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的问题）：

你有一个带有变量temp和height的xr.Dataset，它们都是由（plev，time，lat，lon）的函数组成的，并且你想要找到每个（time，lat，lon）点上temp对height（即递减率）的回归。

绕过这个最简单的方法是使用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()对temp和height这两个DataArray进行回归分析，并沿着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

加速计算：

为了加快计算速度，我们可以使用xr.DataArray.chunk()将height和temp数据转换为dask.arrays。这将把我们的数据分成小的可管理的块，然后我们可以使用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方法相当的墙时。

然而，重要的是指出，在处理气候分析中遇到的大型数据集时，“惰性”操作数据（即不加载数据直到绝对需要）更加优先。因此，我仍建议使用dask方法，因为您可以在输入数组上操作多个不同的进程，每个进程只需要几个ms，最后只需等待几分钟即可得到完成的产品。:)

1
为了能够比较两种情况的处理时间，对于第二种情况，您必须使用slopes_dask.compute()，是吗？ - CamiloEr
1
是的，你说得对！而且.compute()部分确实需要大约一分钟左右的时间，但是对于大多数工作流程，您希望尽可能地延迟操作数据，直到最后可能的时刻，因此在我看来，这仍然是一个可以接受的比较。不过我会添加一个编辑，这是一个很好的观点。 :) - Andrew Williams
1
@CamiloEr 请看我上面的编辑，.compute() 方法所需的时间与之前大致相同，但是如果您进行惰性操作，则可以在将数据加载到内存之前对数据集执行多个操作。而将数据加载到内存需要一些时间，但是这个时间不会（在一阶）因为您首先对数据集执行一次、两次或三次惰性操作而显著增加。 - Andrew Williams
网页内容由stack overflow 提供, 点击上面的
可以查看英文原文，
原文链接