什么时候对中等多维数据使用 xarray 而不是 numpy？

Question

我有一些多维数据，想知道我是否应该使用 xarray 当速度是我关注的一个（尽管不是最高的）时。

我有一个 4D 数组，所以它没有大到妨碍我使用 numpy。 coordinates/indices 对一个维度至关重要，但对所有其他维度并非如此。我将不得不做一些簿记工作，但作为主要开发人员，这对我来说没问题。对于在我之后继续迭代代码的开发人员来说，使用整数索引可能比使用基于标签的 (xarray/pandas) 方法更令人困惑。无论如何，如果我很好地记录了这个过程，我仍然可以使用 numpy。但我想使用 xarray 来提高可读性。

实施解决方案后，我注意到下面的 operations/indexing 在我的机器上将在大约 5 秒内完成。

for isotope in isotopes:
    for height in heights:
        for assm in assemblies:
            da.loc[dict(power=['NW','NE','SW','SE'],
                        assembly=assm,
                        height=height,
                        isotope=isotope)] = [3,5,1,20]

如果我在 xarray 上使用基于整数的方法做同样的事情，大约需要 2 秒。

for k,isotope in enumerate(isotopes):
    for j,height in enumerate(heights):
        for i,assm in enumerate(assemblies):
            da[i,[-4,-3,-2,-1],j,k] = [3,5,1,20]

最后，我注意到如果我在 numpy 中执行相同的基于整数的索引，它只需要不到半秒

arr = np.zeros((44,10,22,13))
for k,isotope in enumerate(isotopes):
    for j,height in enumerate(heights):
        for i,assm in enumerate(assemblies):
            arr[i,[-4,-3,-2,-1],j,k] = [3,5,1,20]

速度不是我最关心的问题，但如果 xarray 中基于标签的方法比标准的基于整数的方法慢 8 倍以上，而 xarray 中基于整数的方法比标准的基于整数的方法慢 4 倍，那就令人望而却步了我不再深入挖掘 xarray 以获得中等多维数据。

有什么想法、建议等吗？

Answer 1

我们无法真正告诉您使用哪个包，在不了解更多关于您的数据和用例的情况下当然也不能。

就其价值而言，虽然 xarray 的性能总是落后于 numpy，但在执行像这样的小操作时差异最为明显。您在三重 for 循环中使用索引分配少量数据，这是 xarray 的氪石。如果您同时进行所有分配，您会看到惩罚显着减少，因为索引开销变得不那么重要 相对于底层 numpy 操作 。 xarray 中的性能就是了解如何最大限度地减少开销并尽可能利用后端的性能，同时仍然提供 labels-based 索引的便利性。

看这个简单的例子。我创建了一个包含 100 万个 float64 的 3-D DataArray，索引为 (x, y, z):

In [11]: da = xr.DataArray(
    ...:     np.random.random(size=(100, 100, 100)),
    ...:     dims=list('xyz'),
    ...:     coords=[pd.Index([f'{d}{i}' for i in range(100)], name=d) for d in 'xyz'],
    ...: )

遍历 x 和 y，然后沿着 z 的前四个元素赋值会导致巨大的损失，对于相同的操作，xarray 的运行时间超过 numpy 的 100 倍：

In [12]: %%time
    ...: for xi, x in enumerate(da.x.values):
    ...:     for yi, y in enumerate(da.y.values):
    ...:         da.loc[{'x': x, 'y': y, 'z': ['z0', 'z1', 'z2', 'z3']}] = [1, 2, 3, 4]
    ...:
CPU times: user 2.96 s, sys: 38.3 ms, total: 3 s
Wall time: 2.97 s

In [13]: %%time
    ...: for xi, x in enumerate(da.x.values):
    ...:     for yi, y in enumerate(da.y.values):
    ...:         da.values[xi, yi, :4] = [1, 2, 3, 4]
    ...:
CPU times: user 25.7 ms, sys: 508 µs, total: 26.3 ms
Wall time: 25.8 ms

如果将同一操作重组为一次分配所有元素，xarray 的性能损失将降低到大约 numpy 运行时间的 6 倍。

In [15]: %%time
    ...: da.loc[{'z': ['z0', 'z1', 'z2', 'z3']}] = np.tile([1, 2, 3, 4], (100, 100, 1))
    ...:
    ...:
CPU times: user 1.4 ms, sys: 675 µs, total: 2.07 ms
Wall time: 2.99 ms

In [16]: %%time
    ...: da.values[:, :, :4] = np.tile([1, 2, 3, 4], (100, 100, 1))
    ...:
    ...:
CPU times: user 488 µs, sys: 222 µs, total: 710 µs
Wall time: 428 µs

分配整个数组将 xarray 的开销减少到大约 2 倍：

In [19]: %%time
    ...: da.loc[{'z': da.z}] = np.tile(np.random.random(100), (100, 100, 1))
    ...:
    ...:
CPU times: user 11.2 ms, sys: 9.43 ms, total: 20.7 ms
Wall time: 20.9 ms

In [20]: %%time
    ...: da.values[:, :, :] = np.tile(np.random.random(100), (100, 100, 1))
    ...:
    ...:
CPU times: user 3.08 ms, sys: 4.61 ms, total: 7.7 ms
Wall time: 6.72 ms

这是否值得，由您决定。但是无论你选择哪个，都不要使用嵌套的 for 循环进行赋值:)

什么时候对中等多维数据使用 xarray 而不是 numpy？

When to use xarray over numpy for medium rank multidimensional data?

python

numpy

python-xarray

scientific-computing