dask 性能沿轴应用

Question

我正在尝试使用 dask 在大型高分辨率海洋模型数据集上计算随时间推移的线性趋势。

我按照这个例子 (Applying a function along an axis of a dask array) 发现apply_along_axis 的语法更容易。

我目前正在使用dask.array.apply_along_axis 在一维数组上包装一个 numpy 函数，然后将生成的 dask 数组打包到一个 xarray Dataarray 中。使用top -u <username> 表明计算不是并行执行的（~100% cpu 使用）。

我应该期待map_blocks 有更好的表现吗？或者对如何提高apply_along_axis的性能有什么建议吗？ 非常感谢任何提示。

import numpy as np
from scipy import optimize
import xarray as xr
import dask.array as dsa

def _lin_trend(y):
    x = np.arange(len(y))
    return np.polyfit(x, y, 1)



def linear_trend(da, dim, name='parameter'):
    da = da.copy()
    axis_num = da.get_axis_num(dim)

    dims = list(da.dims)
    dims[axis_num] = name
    coords = da.rename({dim:name}).coords
    coords[name] = ['slope', 'intercept']

    dsk = da.data
    dsk_trend = dsa.apply_along_axis(_lin_trend,0,dsk)
    out = xr.DataArray(dsk_trend, dims=dims, coords=coords)
    return out

Answer 1

我一直在使用 xarray 的 apply_ufunc 做类似的事情（需要 xarray v0.10 或更高版本）。这可能比在 dask 中使用 apply_along_axis 函数更容易管理。

import xarray as xr
import numpy as np
from scipy import stats

def _calc_slope(x, y):
    '''wrapper that returns the slop from a linear regression fit of x and y'''
    slope = stats.linregress(x, y)[0]  # extract slope only
    return slope


def linear_trend(obj):
    time_nums = xr.DataArray(obj['time'].values.astype(np.float),
                             dims='time',
                             coords={'time': obj['time']},
                             name='time_nums')
    trend = xr.apply_ufunc(_calc_slope, time_nums, obj,
                           vectorize=True,
                           input_core_dims=[['time'], ['time']],
                           output_core_dims=[[]],
                           output_dtypes=[np.float],
                           dask='parallelized')

    return trend

解决您关于性能为何不符合预期的问题。这可能有多种原因。你的 dask 数组是如何分块的？您使用的是哪个 dask 调度程序？在我更好地了解您的配置后，我会更新答案的第二部分？

Answer 2

我认为最终性能受到我正在处理的文件系统的限制。不过，为了回答您的问题，我的数据集具有以下形状：

<xarray.Dataset>
Dimensions:         (st_edges_ocean: 51, st_ocean: 50, time: 101, xt_ocean: 3600, yt_ocean: 2700)
Coordinates:
  * xt_ocean        (xt_ocean) float64 -279.9 -279.8 -279.7 -279.6 -279.5 ...
  * yt_ocean        (yt_ocean) float64 -81.11 -81.07 -81.02 -80.98 -80.94 ...
  * st_ocean        (st_ocean) float64 5.034 15.1 25.22 35.36 45.58 55.85 ...
  * st_edges_ocean  (st_edges_ocean) float64 0.0 10.07 20.16 30.29 40.47 ...
  * time            (time) float64 3.634e+04 3.671e+04 3.707e+04 3.744e+04 ...

所以它相当大，需要很长时间才能从磁盘读取。我已经重新分块了，所以时间维度是一个块

dask.array<concatenate, shape=(101, 50, 2700, 3600), dtype=float64, 
chunksize=(101, 1, 270, 3600)>

这对性能没有太大影响（该功能仍然需要大约 20 小时才能完成（包括读取和写入磁盘）。我目前只是及时分块，例如

dask.array<concatenate, shape=(101, 50, 2700, 3600), dtype=float64, 
chunksize=(1, 1, 2700, 3600)>

我对这两种方法的相对性能很感兴趣，并在我的笔记本电脑上进行了测试。

import xarray as xr
import numpy as np
from scipy import stats
import dask.array as dsa

slope = 10
intercept = 5
t = np.arange(250)
x = np.arange(10)
y = np.arange(500)
z = np.arange(200)
chunks = {'x':10, 'y':10}

noise = np.random.random([len(x), len(y), len(z), len(t)])
ones = np.ones_like(noise)
time = ones*t
data = (time*slope+intercept)+noise
da = xr.DataArray(data, dims=['x', 'y', 'z', 't'],
                 coords={'x':('x', x),
                        'y':('y', y),
                        'z':('z', z),
                        't':('t', t)})
da = da.chunk(chunks)
da

我现在定义了一组私有函数（使用 linregress 和 polyfit 来计算时间序列的斜率），以及使用 dask.apply_along 和 xarray.apply_ufunc 的不同实现。

def _calc_slope_poly(y):
    """ufunc to be used by linear_trend"""
    x = np.arange(len(y))
    return np.polyfit(x, y, 1)[0]


def _calc_slope(y):
    '''returns the slop from a linear regression fit of x and y'''
    x = np.arange(len(y))
    return stats.linregress(x, y)[0]

def linear_trend_along(da, dim):
    """computes linear trend over 'dim' from the da.
       Slope and intercept of the least square fit are added to a new
       DataArray which has the dimension 'name' instead of 'dim', containing
       slope and intercept for each gridpoint
    """
    da = da.copy()
    axis_num = da.get_axis_num(dim)
    trend = dsa.apply_along_axis(_calc_slope, axis_num, da.data)
    return trend

def linear_trend_ufunc(obj, dim):
    trend = xr.apply_ufunc(_calc_slope, obj,
                           vectorize=True,
                           input_core_dims=[[dim]],
                           output_core_dims=[[]],
                           output_dtypes=[np.float],
                           dask='parallelized')

    return trend

def linear_trend_ufunc_poly(obj, dim):
    trend = xr.apply_ufunc(_calc_slope_poly, obj,
                           vectorize=True,
                           input_core_dims=[[dim]],
                           output_core_dims=[[]],
                           output_dtypes=[np.float],
                           dask='parallelized')

    return trend

def linear_trend_along_poly(da, dim):
    """computes linear trend over 'dim' from the da.
       Slope and intercept of the least square fit are added to a new
       DataArray which has the dimension 'name' instead of 'dim', containing
       slope and intercept for each gridpoint
    """
    da = da.copy()
    axis_num = da.get_axis_num(dim)
    trend = dsa.apply_along_axis(_calc_slope_poly, axis_num, da.data)
    return trend

trend_ufunc = linear_trend_ufunc(da, 't')
trend_ufunc_poly = linear_trend_ufunc_poly(da, 't')
trend_along = linear_trend_along(da, 't')
trend_along_poly = linear_trend_along_poly(da, 't')

计算时间似乎表明apply_along 方法可能稍微快一些。不过，使用 polyfit 代替 linregress 似乎有相当大的影响。我不知道为什么这要快得多，但也许这对你很感兴趣。

%%timeit 
print(trend_ufunc[1,1,1].data.compute())

每个循环 4.89 秒 ± 180 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）

%%timeit 
trend_ufunc_poly[1,1,1].compute()

每个循环 2.74 秒 ± 182 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）

%%timeit 
trend_along[1,1,1].compute()

每个循环 4.58 秒 ± 193 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）

%%timeit
trend_along_poly[1,1,1].compute()

每个循环 2.64 秒 ± 65 毫秒（平均值 ± 标准偏差，7 次运行，每个循环 1 个）