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

Question

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

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

Answer 1

你可以用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)

然后你可以检查广播返回的数组的形状x,y并z检查shape属性:

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

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

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