从具有多个数组的函数创建矩阵而不使用 for 循环

Question

我有以下五个数组：

t = np.linspace(0,100,100)
Q = t**2

x = np.linspace(10,20,20)
y = np.linspace(5,8,5)
z = np.linspace(100,125,10)

我还定义了一个函数如下：

def f(t, x, y, z):
    return (1+x)**3.2 * y**x * np.cos(y**x)*(t-z)

现在对于x 的每个第 i 个元素、y 的第 j 个元素和 z 的第 k 个元素，我想确定 f(t, x, y, z) 在所有 @ 的值上的值987654327@。 我还想用Q 数组对这些结果中的每一个进行“内积”。 所以我最终想要的是一个 3 维矩阵，它存储大小为len(t) 的数组，该矩阵的每个元素 (i,j,k) 等于np.sum( Q * f(t, x[i], y[j], z[k]))。 所有这些都可以通过如下的 for 循环来实现：

result = np.zeros(len(x)*len(y)*len(z)).reshape((len(x),len(y), len(z)))
for i in range(len(x)):
    for j in range(len(y)):
        for k in range(len(z)):
            intermediate_result = f(t, x[i], y[j], z[k]) #This is an array of length len(t)
            result[i][j][k] = (np.nansum(intermediate_result * Q))**2

问题是这些 for 循环在计算超出这个简化示例的更大数组 x、y 和 z 时非常耗时，因此我正在寻找一种减少计算时间的方法。 有没有有效的方法来做到这一点？

Answer 1

您不需要循环。您可以通过重塑参数变量让 numpy 进行广播：

result = f(t,x[:,None,None,None],y[:,None,None],z[:,None])

result = np.sum(Q*result,axis=3)

请注意，您可以使用 np.ix_ 在函数内部执行此操作

def f(t, x, y, z):
    x,y,z,t = np.ix_(x,y,z,t)
    return (1+x)**3.2 * y**x * np.cos(y**x)*(t-z)

该函数将返回一个 4D 矩阵，然后您可以使用 np.sum 进行处理：

result = np.sum(Q*f(t,x,y,z),axis=3)

这将明显快于循环（对于这个 100 x 20x5x10 的示例，速度快 10 倍以上）

[编辑] 鉴于 4D 矩阵会消耗大量内存，随着维度变大，您最终会遇到另一种瓶颈。您可以通过将矩阵限制为 3 维并在循环中手动添加 t 维来规避此问题。这可确保您使用的 3D 矩阵的大小永远不会超过两倍：

def f2(t, x, y, z):
    x,y,z = np.ix_(x,y,z)
    return (1+x)**3.2 * y**x * np.cos(y**x)*(t-z)

result = np.zeros((x.size,y.size,z.size))
for tn,Qn in zip(t,Q):
    result += f2(tn,x,y,z)*Qn

通过这种方法，我能够在 19 秒内获得 t.size=300 和 x,y,z 为 256x256x256 的结果。 4D 矩阵方法使 IDLE shell 崩溃（可能是由于内存溢出），我没有耐心等待原始 for 循环完成。

请注意，由于 np.sum() 添加值的顺序，与原始解决方案相比，最后一种方法可能会在尾数的第 16 位附近产生微小差异