时间戳

首页 > 解决方案 > 从具有多个数组的函数创建矩阵而不使用 for 循环

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

问题描述

我有以下五个数组：

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)

现在对于 的每个第 i 个元素x、 的第 j 个元素y和 的第 k 个元素z，我想确定f(t, x, y, z)取 的所有值的值t。我还想将这些结果中的每一个与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 时非常耗时，因此我正在寻找一种减少计算时间的方法。有没有有效的方法来做到这一点？

标签： pythonnumpyfor-loopoptimizationnumpy-ndarray

解决方案

你不需要循环。您可以通过重塑参数变量让 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 位附近产生微小差异

推荐阅读