如何优化具有多个循环的代码？

Question

所以我尝试了这个，我发现它真的不可行。我不太了解执行以下操作的聪明方法。有人可以帮忙吗？列表的输入也很大。

此任务是根据我生成的值创建图像。 center_star 包含 [x,y] 对的列表，它们是各种点状对象的中心。

1800 值表示要生成的图像为 1800x1800 像素。

Sigma 变量默认值为 2。

final=[[0]*1800]*1800
for i in range(len(center_stars)):
    xi=center_stars[i][0]
    yi=center_stars[i][1]
    print(i)
    for j in range(1800):
        for k in range(1800):
            final[j][k]+=gauss_amplitude[i]*(math.e**((-1*((xi-j)**2+(yi-k)**2))/2*sigma*sigma))

有没有更聪明的方法来使用一些 numpy 操作来节省时间并在更短的时间内执行这段代码？

Answer 1

类似的东西：

import math
import numpy as np
N = 1800
final = np.empty((N, N))
final1 = np.empty((N, N))
j = np.arange(N)
k = np.arange(N)
jj, kk = np.meshgrid(j, k)
sigma = 2.
s = 0.5 / (sigma * sigma)
for i in range(len(center_stars)):
    xi = center_stars[i][0]
    yi = center_stars[i][1]
    final += gauss_amplitude[i] * np.exp(- ((xi - jj.T)**2 + (yi - kk.T)**2) * s)
# Code below is for comparison:
    for j in range(N):
        for k in range(N):
            final1[j][k]+=gauss_amplitude[i] * (math.e** (-((xi-j)**2+(yi-k)**2)/(2*sigma*sigma)))

此外，我假设您错过了 2*sigma*sigma 周围的括号

Answer 2

你可以尝试像这样压缩你的代码：

Gauss=lambda i,j,k,xi,yi:gauss_amplitude[i]*(math.e**((-((xi-j)**2+(yi-k)**2))/(2*sigma*sigma)))
final=[[Gauss(i,j,k,x[0],x[1]) for j in range(1800) for k in range(1800)] for i,x in enumerate(center_starts)]

Answer 3

如果你的 sigma 都是一样的，你可以通过使用没有任何循环来实现这一点 scipy.signal.convolve2d.

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve2d
from scipy.stats import multivariate_normal

sigma = 3
width = 200   # smaller than yours so you can see the single pixels
n_stars = 50

# draw some random stars
star_row = np.random.randint(0, width, n_stars)
star_col = np.random.randint(0, width, n_stars)
star_amplitude = np.random.normal(50, 10, n_stars)

# assign amplitudes to center pixel of stars
amplitudes = np.zeros((width, width))
amplitudes[star_row, star_col] = star_amplitude


# create 2d gaussian kernel
row = col = np.arange(-4 * sigma, 4 * sigma + 1)
grid = np.stack(np.meshgrid(row, col)).T
kernel = multivariate_normal(
    [0, 0],
    [[sigma**2, 0], [0, sigma**2]]
).pdf(grid)
kernel /= kernel.sum()


# convolve with 2d gaussian
final = convolve2d(amplitudes, kernel, mode='same')


fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(9, 3))

img = ax1.imshow(amplitudes)
fig.colorbar(img, ax=ax1)
ax1.set_label('Before Convolution')

img = ax2.imshow(kernel)
fig.colorbar(img, ax=ax2)
ax2.set_label('Convolution Kernel')

img = ax3.imshow(final)
fig.colorbar(img, ax=ax3)
ax3.set_label('After Convolution')

fig.tight_layout()
fig.savefig('conv2d.png', dpi=300)

结果：

如果 sigma 不同，您可以通过单个循环遍历可能的 sigma。