python - 如何避免在这个特定的 python 代码片段中使用循环？

Question

我有一个特定的Python问题，迫切需要通过避免使用循环来加速，但是，我不知道如何做到这一点。我需要读取适合的图像，将其转换为 numpy 数组(大小大约为 2000 x 2000 个元素)，然后为每个元素计算其周围元素环的统计数据。正如我现在的代码所示，元素周围环的统计数据是通过使用掩码的函数计算的。这很快，但是，当然，我调用这个函数 2000x2000 次(慢的部分)。我对 python 比较陌生。我认为使用 mask 函数很聪明，但我找不到单独寻址每个元素的方法。非常感谢您提供的任何帮助。

 # First, the function computing the statistics within a ring 
 around the central pixel:<br/>
 # flux       = image intensity at pixel (i,j)<br/>
 # rad1, rad2 = inner and outer radii<br/>
 # array      = image array<br/>_

def snr(flux, i, j, rad1, rad2, array):

     a, b   = i, j
     nx, ny = array.shape
     y, x   = np.ogrid[-a:nx-a, -b:ny-b]

     mask   = (x*x + y*y >= rad1*rad1) & (x*x + y*y <= rad2*rad2)

     Nmask  = np.count_nonzero(mask)

     noise  = 0.6052697 * abs(Nmask * flux - sum(array[mask]))

     return noise



 # Now, the call to snr for each pixel in the array data1:<br/>_


     frame1  = fits.open(in_frame, mode='readonly') # read in fits file            
     data1   = frame1[ext].data                     # convert to np array
     ny, nx  = data1.shape                          # array dimensions                                  


     noise1  = zeros((ny, nx), float)               # empty array                       

     r1 = 5                                     # inner radius (pixels)
     r2 = 7                                     # outer radius (pixels)


     # The function is fast, but calling it 2k x 2k times is not:

     for j in range(ny):
          for i in range(nx):

                  noise1[i,j] = der_snr(data1[i,j], i, j, r1, r2, data1)

Answer 1

您尝试执行的操作可以表示为 image convolution 。尝试这样的事情:

import numpy as np
import scipy.ndimage
from astropy.io import fits

def make_kernel(inner_radius, outer_radius):
    if inner_radius > outer_radius:
        raise ValueError
    x, y = np.ogrid[-outer_radius:outer_radius + 1, -outer_radius:outer_radius + 1]
    r2 = x * x + y * y
    kernel = (r2 >= inner_radius * inner_radius) & (r2 <= outer_radius * outer_radius)
    return kernel

in_frame = '<file path>'
ext = '...'
frame1 = fits.open(in_frame, mode='readonly')
data1 = frame1[ext].data
inner_radius = 5
outer_radius = 7
kernel = make_kernel(inner_radius, outer_radius)
n_kernel = np.count_nonzero(kernel)
conv = scipy.ndimage.convolve(data1, kernel, mode='constant')
noise1 = 0.6052697 * np.abs(n_kernel * data1 - conv)