将 xarray 与自定义函数一起使用并重新采样答案

【问题标题】：Use xarray with custom function and resample将 xarray 与自定义函数一起使用并重新采样
【发布时间】：2020-02-25 03:32:23
【问题描述】：

我正在尝试获取一个数组并使用自定义函数对其进行重新采样。来自这篇文章：Apply function along time dimension of XArray

def special_mean(x, drop_min=False):
    s = np.sum(x)
    n = len(x)
    if drop_min:
        s = s - x.min()
    n -= 1
    return s/n

是一个示例 sample_mean。

我有一个数据集是：

<xarray.Dataset>
Dimensions:  (lat: 100, lon: 130, time: 7305)
Coordinates:
  * lon      (lon) float32 -99.375 -99.291664 -99.208336 ... -88.708336 -88.625
  * lat      (lat) float32 49.78038 49.696426 49.61247 ... 41.552795 41.46884
    lev      float32 1.0
  * time     (time) datetime64[ns] 2040-01-01 2040-01-02 ... 2059-12-31
Data variables:
    tmin     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
    tmax     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
    prec     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
    relh     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
    wspd     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
    rads     (time, lat, lon) float32 dask.array<chunksize=(366, 100, 130), meta=np.ndarray>
Attributes:
    history:  Fri Jun 14 10:32:22 2019: ncatted -a _FillValue,,o,d,9e+20 IBIS...

然后我应用一个重新采样：

data.resample(time='1MS').map(special_mean)


<xarray.Dataset>
Dimensions:  (time: 240)
Coordinates:
  * time     (time) datetime64[ns] 2040-01-01 2040-02-01 ... 2059-12-01
    lev      float32 1.0
Data variables:
    tmin     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    tmax     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    prec     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    relh     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    wspd     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>
    rads     (time) float32 dask.array<chunksize=(1,), meta=np.ndarray>

如何执行此功能，以便我可以保留'lon'和'lat'坐标，就像做时一样

data.resample(time='1MS').mean()

【问题讨论】：

标签： dask python-xarray

【解决方案1】：

这是一个如何使用 xr.apply_ufunc() 的示例。

import xarray as xr
data = xr.tutorial.open_dataset('air_temperature')

def special_mean(x, drop_min=False):
    s = np.sum(x)
    n = len(x)
    if drop_min:
        s = s - x.min()
    n -= 1
    return s/n

def special_func(data):
    return xr.apply_ufunc(special_mean, data, input_core_dims=[["time"]], 
            kwargs={'drop_min': True}, dask = 'allowed', vectorize = True)

data.resample(time='1MS').apply(special_func)

<xarray.Dataset>
Dimensions:  (lat: 25, lon: 53, time: 24)
Coordinates:
  * time     (time) datetime64[ns] 2013-01-01 2013-02-01 ... 2014-12-01
  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0
  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0
Data variables:
    air      (time, lat, lon) float64 244.6 244.7 244.7 ... 297.7 297.7 297.7

【讨论】：

如文档中所示，对于我的大小数据集，vectorize=True 非常慢。你知道我怎样才能让它更快一点吗？
另外，有没有办法将超出范围的变量传递给 special_func。用例将是 drop_min only 1 vs drop_min 两次或类似的东西。
@blueduckyy 尝试使用 .chunk() 将您的数据转换为 dask 数组，然后将参数切换为 dask='parallelized' :) 它将允许 ufunc 使用 dask 对您的数据进行惰性操作，您可以然后稍后使用da.compute() 将数据加载到内存中。看看我的另一个答案，看看有没有帮助。 :) stackoverflow.com/questions/38960903/…

【解决方案2】：

我怀疑你可以用apply_ufunc 方法做你想做的事。

（尽管作为免责声明我不太了解 Xarray API。）

【讨论】：

怎么样？我尝试过使用 ufunc，但我不确定我是否只是不知道它是如何工作的？去哪儿了？