获取一维结果列表并将其转换为 N-D xarray.DataArray

Question

这就是我获取 N 维数据的方式（func IRL 不可向量化）：

import numpy
import xarray
import itertools

xs = numpy.linspace(0, 10, 100)
ys = numpy.linspace(0, 0.1, 20)
zs = numpy.linspace(0, 5, 200)

def func(x, y, z):
    return x * y / z

vals = list(itertools.product(xs, ys, zs))
result = [func(x, y, z) for x, y, z in vals]

我觉得我所做的事情可以简化。我想把它放在 xarray.DataArray 中而不重塑数据。但是，这就是我现在的做法：

arr = np.array(result).reshape(len(xs), len(ys), len(zs))
da = xarray.DataArray(arr, coords=[('x', xs), ('y', ys), ('z', zs)])

这是一个简单的例子，但我通常使用通过映射 itertools.product（并行）获得的 ~10D 数据。

我的问题：如何在不重塑我的数据并使用 vals 并且不使用 xs、ys 和 zs 的长度的情况下做到这一点？

与您的处理方式类似：

index = pandas.MultiIndex.from_tuples(vals, names=['x', 'y', 'z'])
df = pandas.DataFrame(result, columns=['result'], index=index)

编辑： 这就是我解决它的方法，灵感来自@hpaulj 的答案，谢谢！

import numpy
import xarray
import itertools

coords = dict(x=numpy.linspace(0, 10, 100),
              y=numpy.linspace(0, 0.1, 20),
              z=numpy.linspace(0, 5, 200))

def func(x, y, z):
    return x * y / z

result = [func(x, y, z) for x, y, z in itertools.product(*coords.values())]

xarray.DataArray(numpy.reshape(result, [len(i) for i in coords.values()]), coords=coords)

编辑 2 看到这个问题：https://github.com/pydata/xarray/issues/1914

Answer 1

第二次编辑我忘记了 einsum！如果你可以折磨你的函数来适应它会更快（下面的时间为 1.5 毫秒）

result = np.einsum('i,j,k', xs, ys, 1.0 / zs)

您需要重塑和广播到相同形状的数组。正如 Balzola 所说，如果每个方向都是 10D 和 100（10**20 个元素），这将非常大。正如 hpaulj 所说，重塑 numpy 数组通常是微不足道的，在这种情况下也是如此，尽管广播确实需要一些工作。但是比 itertools.product() 方法少得多。举个例子

import numpy as np

xs = np.linspace(0, 10, 100)
ys = np.linspace(0, 0.1, 20)
zs = np.linspace(0.1, 5, 200)

xn, yn, zn = len(xs), len(ys), len(zs)

xs_b = np.broadcast_to(xs.reshape(xn, 1, 1), (xn, yn, zn))
ys_b = np.broadcast_to(ys.reshape(1, yn, 1), (xn, yn, zn))
zs_b = np.broadcast_to(zs.reshape(1, 1, zn), (xn, yn, zn))

result = xs_b * ys_b / zs_b

使用 timeit 如下所示，我得到的 numpy 计算为 4 毫秒，而 itertools 方法为 150 毫秒。我认为对于更多维度，差异会更大。

import timeit

init = '''
import itertools
import numpy as np

def func(x, y, z):
    return x * y / z

xs = np.linspace(0, 10, 100)
ys = np.linspace(0, 0.1, 20)
zs = np.linspace(0.1, 5, 200)

xn, yn, zn = len(xs), len(ys), len(zs)
'''
funcs = ['''
xs_b = np.broadcast_to(xs.reshape(xn, 1, 1), (xn, yn, zn))
ys_b = np.broadcast_to(ys.reshape(1, yn, 1), (xn, yn, zn))
zs_b = np.broadcast_to(zs.reshape(1, 1, zn), (xn, yn, zn))

result = xs_b * ys_b / zs_b
''','''
vals = list(itertools.product(xs, ys, zs))
result = [func(x, y, z) for x, y, z in vals]
''']

for f in funcs:
  print(timeit.timeit(f, setup=init, number=100))

编辑PS。我更改了你的 zs 以通过除以零来防止 numpy 警告，因为这可能会影响 timeit 比较。

Answer 2

有经验的 numpy 用户倾向于专注于删除迭代步骤。因此，我们放大了您的 result 计算，并将 reshape 视为微不足道的事情。因此，到目前为止的答案都集中在广播和计算您的函数上。

但我开始怀疑真正困扰你的是

reshape(len(xs), len(ys), len(zs))

如果您有 10 个这样的维度，而不仅仅是 3 个，那么

可能会变得笨拙。这与其说是计算速度，不如说是输入 len(..) 10 次所需的努力。或者可能是代码看起来很难看。

无论如何，这里有一种绕过所有输入的方法。关键是将维度数组收集在列表中

In [495]: dims = [np.linspace(0,10,4), np.linspace(0,.1,3), np.linspace(0,5,5)]
In [496]: from itertools import product
In [497]: vals = list(product(*dims))
In [498]: len(vals)
Out[498]: 60
In [499]: result = [sum(ijk) for ijk in vals] # a simple func

现在只需通过简单的列表理解即可获得 len's：

In [501]: arr=np.array(result).reshape([len(i) for i in dims])
In [502]: arr.shape
Out[502]: (4, 3, 5)

另一种可能性是将 linspace 参数放在列表的开头。

In [504]: ldims=[4,3,5]
In [505]: ends=[10,.1,5]
In [506]: dims=[np.linspace(0,e,l) for e,l in zip(ends, ldims)]
In [507]: vals = list(product(*dims))
In [508]: result=[sum(ijk) for ijk in vals]
In [509]: arr=np.array(result).reshape(ldims)

reshape 本身并不是一个昂贵的操作。通常它会创建一个视图，这是您可以使用数组做的最快的事情之一。

@Divakar 在他删除的答案中暗示了这种解决方案，用 *np.meshgrid(*A) 代替你的 product(xs,ys).

顺便说一下，我的回答也不涉及 xarray - 因为我没有安装那个包。我假设您知道将 3d 形状的 arr 传递给它时您在做什么，而不是较长的 1d 数组。查看标签编号，numpy 有 5000 个关注者，xarray 有 23 个关注者。

xarray coords 参数也可以从 dims 构造（带有额外的名称列表）。

如果您不喜欢这个答案，我建议您关闭问题，并开始一个仅包含 xarray 标签的新问题。这样你就不会吸引无数的 numpy 苍蝇。

获取一维结果列表并将其转换为 N-D xarray.DataArray

Take a 1D list of results and convert it to a N-D xarray.DataArray

python

pandas

numpy

python-xarray