用x计算x**k,k是任意维度的数组

Question

我有两个numpy数组:一个x有形状的(n, a0, a1, ...)数组和一个k有形状的数组(n, b0, b1, ...).我想计算和指数数组,使输出具有维度(a0, a1, ..., b0, b1, ...)和

out[i0, i1, ..., j0, j1, ...] == prod(x[:, i0, i1, ...] ** k[:, j0, j1, ...])

如果只有一个a_i和一个b_j,广播通过

import numpy

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None], axis=0)

如果x有更多维度,None则必须添加更多s:

x = numpy.random.rand(2, 31, 32, 33)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None, None, None], axis=0)

如果x有更多维度,None则必须在其他地方添加更多s:

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 51, 51))

out = numpy.prod(x[..., None, None]**k[:, None], axis=0)

如何out根据输入数组的维数计算泛型？

Answer 1

这是使用reshaping两个阵列的一个,这样它们可以相互播放,然后prod沿第一个轴执行那些操作和缩小 -

k0_shp = [k.shape[0]] + [1]*(x.ndim-1) + list(k.shape[1:])
x0_shp = list(x.shape) + [1]*(k.ndim-1)
out = (x.reshape(x0_shp) ** k.reshape(k0_shp)).prod(0)

这是另一种重塑两个输入的方法,3D允许每个输入允许一个单独调暗,并且它们可以相互播放,执行prod缩减以获得2D数组,然后重塑为多调光阵列 -

s = x.shape[1:] + k.shape[1:] # output shape
out = (x.reshape(x.shape[0],-1,1)**k.reshape(k.shape[0],1,-1)).prod(0).reshape(s)

必须注意的是,重新整形只是创建了一个输入数组的视图,因此在内存和性能方面几乎都是免费的.