Question

这就是我获取N-D数据的方式（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

有经验的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的5k粉丝，xarray的23个粉丝。

xarray coords参数也可以从dims构建（附加一个名称列表）。

如果这个答案不是你喜欢的，我建议关闭这个问题，然后用xarray标签开始一个新问题。这样你就不会吸引大量的numpy苍蝇。

Answer 2

第二次编辑我忘记了einsum！如果你可以折磨你的功能以适应这将更快（下面的时间1.5ms）

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

您需要重塑并广播到相同形状的阵列。正如巴尔佐拉所说，如果每个方向的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计算为4ms，itertools方法为150ms。我认为更多维度的差异会更大。

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警告，因为这可能影响了时间比较。

Answer 3

我的回答：[删除因为问题意外更改]

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

3 个答案: