將 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 的問題）：

您有一個帶有變量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()沿着plev維度對temp和height的兩個 DataArray 進行回歸！

%%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 。 這將我們的數據拆分成小的、可管理的塊，然后我們可以使用它們在我們的apply_ufunc()使用dask=parallelized並行化我們的計算。

NB 你必須小心不要沿着你應用回歸的維度分塊！

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 ，那么只有在最后你才需要等待幾分鍾才能完成產品出來。 :)

Answer 2

僅供參考，從 v0.16.0 開始，xarray 已將 polyfit 實現為Dataset和DataArray以及相關聯的xarray.polyval函數的一種方法：

https://xarray.pydata.org/en/stable/generated/xarray.Dataset.polyfit.html

https://xarray.pydata.org/en/stable/generated/xarray.DataArray.polyfit.html

https://xarray.pydata.org/en/stable/generated/xarray.polyval.html

https://xarray.pydata.org/en/stable/whats-new.html#v0-16-0-2020-07-11