Aplicando numpy.polyfit al conjunto de datos xarray

¿Admite Xarray numerosas funciones de cálculo como polyfit? ¿O hay una manera eficiente de aplicar tales funciones a los conjuntos de datos?

Ejemplo: quiero calcular la pendiente de una línea ajustada a dos variables (Temperatura y Altura), para calcular una tasa de lapso. Tengo un conjunto de datos (abajo) con estas dos variables con dimensiones de (vertical, tiempo, xgrid_0, ygrid_0).

 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 ... 

Si extraigo la temperatura y la altura para un tiempo determinado, xgrid_0, ygrid_0; Puedo usar la función 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) 

Pero este es un enfoque lento e ineficiente. ¿Alguna sugerencia sobre una mejor manera de abordar esto?