python - 有没有办法通过使用 numpy 或其他资源来优化 Python 中的三重循环？

Question

我找不到在 Python 中优化三重循环的方法。我将直接给出代码，以便更好、更简单地表示我必须计算的内容:

给定两个名为 samples (M x N) 和 D(N x N) 的二维数组以及输出结果 (NxN):

for sigma in range(M):
        for i in range(N):
            for j in range(N):
                results[i, j]  += (1/N) * (samples[sigma, i]*samples[sigma, j] 
                                             - samples[sigma, i]*D[j, i] 
                                             - samples[sigma, j]*D[i, j])  
return results

它可以完成工作，但在 python 中根本无效。我试图解开 for i.. for j.. 循环，但我无法使用 sigma 正确计算它。

有人知道如何优化那几行吗？欢迎任何建议，例如 numpy、numexpr 等...

Answer 1

我发现改进代码(即减少循环次数)的一种方法是使用 np.meshgrid。

这是我发现的改进。它花了一些时间，但它提供了与三重循环代码相同的输出。我保留了相同的代码结构，因此您可以看到哪些部分对应于哪些部分。希望对你有用!

for sigma in range(M):
    xx, yy = np.meshgrid(samples[sigma], samples[sigma])

    results += (1/N) * (xx * yy 
                         - yy * D.T
                         - xx * D)       

print(results) # or return results

.

编辑:这里有一个小脚本来验证结果是否符合预期:

import numpy as np
M, N = 3, 4
rng = np.random.default_rng(seed=42)


samples = rng.random((M, N))
D       = rng.random((N, N))
results = rng.random((N, N))

results_old = results.copy()
results_new = results.copy()

for sigma in range(M):
        for i in range(N):
            for j in range(N):
                results_old[i, j]  += (1/N) * (samples[sigma, i]*samples[sigma, j]
                                             - samples[sigma, i]*D[j, i]
                                             - samples[sigma, j]*D[i, j])

print('\n\nresults_old', results_old, sep='\n')

for sigma in range(M):
    xx, yy = np.meshgrid(samples[sigma], samples[sigma])

    results_new += (1/N) * (xx * yy
                         - yy * D.T
                         - xx * D)

print('\n\nresults_new', results_new, sep='\n')

编辑 2:完全摆脱循环:它有点复杂，但本质上做同样的事情。

M, N = samples.shape
xxx, yyy = np.meshgrid(samples, samples)
split_x = np.array(np.hsplit(np.vsplit(xxx, M)[0], M))
split_y = np.array(np.vsplit(np.hsplit(yyy, M)[0], M))

results += np.sum(
    (1/N) * (split_x*split_y 
                 - split_y*D.T 
                 - split_x*D), axis=0)

print(results) # or return results