xarray 对于性能关键代码来说太慢了答案

【问题标题】：xarray too slow for performance critical codexarray 对于性能关键代码来说太慢了
【发布时间】：2020-02-22 21:47:26
【问题描述】：

我计划在我正在编写的一些数字密集型科学代码中广泛使用 xarray。到目前为止，它使代码非常优雅，但我认为我将不得不放弃它，因为性能成本太高了。

这是一个示例，它使用 xarray（具有多种索引方案）和 numpy.我使用了 num_comp=2 和 num_x=10000：

Line #      Hits     Time   Per Hit   % Time  Line Contents
 4                                           @profile
 5                                           def xr_timing(num_comp, num_x):
 6         1         4112   4112.0     10.1      da1 = xr.DataArray(np.random.random([num_comp, num_x]).astype(np.float32), dims=['component', 'x'], coords={'component': ['a', 'b'], 'x': np.linspace(0, 1, num_x)})
 7         1          438    438.0      1.1      da2 = da1.copy()
 8         1         1398   1398.0      3.4      da2[:] = np.random.random([num_comp, num_x]).astype(np.float32)
 9         1         7148   7148.0     17.6      da3 = da1.isel(component=0).drop('component') * da2.isel(component=0).drop('component')
10         1         6298   6298.0     15.5      da4 = da1[dict(component=0)].drop('component') * da2[dict(component=0)].drop('component')
11         1         7541   7541.0     18.6      da5 = da1.sel(component='a').drop('component') * da2.sel(component='a').drop('component')
12         1         7184   7184.0     17.7      da6 = da1.loc[dict(component='a')].drop('component') * da2.loc[dict(component='a')].drop('component')
13         1         6479   6479.0     16.0      da7 = da1[0, :].drop('component') * da2[0, :].drop('component')

15                                           @profile
16                                           def np_timing(num_comp, num_x):
17         1         1027   1027.0     50.2      da1 = np.random.random([num_comp, num_x]).astype(np.float32)
18         1          977    977.0     47.8      da2 = np.random.random([num_comp, num_x]).astype(np.float32)
19         1           41     41.0      2.0      da3 = da1[0, :] * da2[0, :]

最快的 xarray 乘法大约需要 numpy 版本的时间的 150 倍。这只是我的代码中的操作之一，但我发现它们中的大多数比 numpy 等价物慢很多倍，这是不幸的，因为 xarray 使代码更加清晰。我做错了吗？

更新：即使 da1[0, :].values * da2[0, :].values（失去了使用 xarray 的许多好处）也需要 2464 个时间单位。

我正在使用 xarray 0.9.6、pandas 0.21.0、numpy 1.13.3 和 Python 3.5.2。

更新 2：按照@Maximilian 的要求，这里重新运行 num_x=1000000：

Line #      Hits   Time    Per Hit   % Time  Line Contents
# xarray
 9         5       408596  81719.2     11.3      da3 = da1.isel(component=0).drop('component') * da2.isel(component=0).drop('component')
10         5       407003  81400.6     11.3      da4 = da1[dict(component=0)].drop('component') * da2[dict(component=0)].drop('component')
11         5       411248  82249.6     11.4      da5 = da1.sel(component='a').drop('component') * da2.sel(component='a').drop('component')
12         5       411730  82346.0     11.4      da6 = da1.loc[dict(component='a')].drop('component') * da2.loc[dict(component='a')].drop('component')
13         5       406757  81351.4     11.3      da7 = da1[0, :].drop('component') * da2[0, :].drop('component')
14         5        48800   9760.0      1.4      da8 = da1[0, :].values * da2[0, :].values

# numpy
20         5        37476   7495.2      2.9      da3 = da1[0, :] * da2[0, :]

正如预期的那样，性能差异已大幅减少（现在仅慢了大约 10 倍），但我仍然很高兴在下一个版本的文档中会提到这个问题，因为即使如此大的差异也可能会让一些人感到惊讶。

【问题讨论】：

你能重新比较大数组吗？性能差异可能会显着降低，因为计算部分本身应该是相似的速度

标签： performance python-xarray

【解决方案1】：

是的，这是 xarray 的一个已知限制。对于 xarray，使用小数组的性能敏感代码比 NumPy 慢得多。我在我们的文档中为下一个版本写了一个新的部分： http://xarray.pydata.org/en/stable/computation.html#wrapping-custom-computation

你基本上有两种选择：

在未包装的数组上编写对性能敏感的代码，然后将它们包装回 xarray 数据结构中。 Xarray v0.10 有一个新的辅助函数 (apply_ufunc)，它使这更容易一些。如果您对此感兴趣，请参阅上面的链接。
使用 xarray/Python 以外的其他工具进行计算。这也很有意义，因为 Python 本身会增加大量开销。 Julia 的 AxisArrays.jl 看起来很有趣，虽然我自己没有尝试过。

我认为选项 3 是用 C++ 重写 xarray 本身（例如，在 xtensor 之上），但这会涉及更多！

【讨论】：