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

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 -因为我没有安装那个包。我假设你知道你在做什么，当传递arr的3d形状给它，而不是更长的1d数组。看看标签号，numpy有5000个粉丝，xarray有23个粉丝。

xarray coords参数也可以从dims(带有一个额外的名称列表)构造。

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

Answer 2

我已经忘记了einsum!如果你能让你的函数适应，这将会更快(1.5ms的时间)

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