Numpy函数可以获得添加数组的形状

时间:2015-08-07 13:49:57

标签: python arrays numpy numpy-broadcasting

tl; dr:我如何预测numpy广播在几个数组中返回的形状,而不必实际添加数组?

我有很多脚本使用numpy(Python)广播规则,因此基本上1D输入会产生多维输出。对于一个基本的例子,理想气体定律(压力= rho * R_d *温度)可能看起来像

def rhoIdeal(pressure,temperature):
    rho = np.zeros_like(pressure + temperature)
    rho += pressure / (287.05 * temperature)
    return rho

这里没有必要,但在更复杂的功能中,初始化具有正确形状的阵列非常有用。如果压力和温度具有相同的形状,则rho也具有该形状。如果压力有形状(n,),温度有形状(m,),我可以打电话

rhoIdeal(pressure[:,np.newaxis], temperature[np.newaxis,:])

获得形状为rho(n,m)。这使得我可以制作具有多个温度值的绘图而无需循环rhoIdeal,同时仍然允许脚本接受相同形状的数组并逐个元素地计算结果。

我的问题是:是否有内置函数可以返回与多个输入兼容的形状?表现得像

的东西
def returnShape(list_of_arrays):
    return np.zeros_like(sum(list_of_arrays)).shape

没有实际需要对数组求和?如果没有内置函数,那么良好的实现会是什么样的呢?

1 个答案:

答案 0 :(得分:6)

您可以使用np.broadcast。此函数返回一个对象,该对象封装了将两个或多个数组一起广播的结果。没有执行实际操作(例如添加) - 对象只具有一些与通过其他操作生成的数组相同的属性(shape,ndim等)。

例如:

x = np.array([1,2,3])  # shape (3,)
y = x.reshape(3,1)     # shape (3, 1)
z = np.ones((5,1,1))   # shape (5, 1, 1)

然后,您可以通过检查shape属性来检查广播xyz返回的数组形状:

>>> np.broadcast(x, y, z).shape
(5, 3, 3)

这意味着您可以按如下方式实现您的功能:

def returnShape(*args):
    return np.broadcast(*args).shape