python - 如何减少方位角计算标准差的执行时间

Question

我使用下面的代码来计算从图像转换的二维数组的方位角标准偏差。但是，该代码需要几分钟才能运行。谁能建议如何使其更快？

import numpy as np

def radial_profile(data, center):
    y, x = np.indices((data.shape))

    r = np.sqrt((x - center[0])**2 + (y - center[1])**2)
    r = r.astype(np.int64)

    radialprofile = np.zeros(np.max(r) + 1)
    for i in range (np.max(r) + 1):
        radialprofile[i] = np.std(data[r == i])

    return radialprofile

data = random.randint(10, size=(4096, 4096))
std_azi = radial_profile(data, [2048,2048])

CPU times: total: 1min 1s
Wall time: 1min 1s

正如您在 CPU 和 Wall time 中看到的，运行它至少需要一分钟。我有 10,000 张这样的天文图像需要处理。因此，任何有关如何使此代码更快的建议都将受到高度赞赏。

Answer 1

主要问题是 data[r == i] 已完成超过 2000 次。这是非常低效的，尤其是因为数据需求非常大。

可以使用分组策略来一次性完成所有这些操作。话虽这么说，Numpy 不支持这一点。 AFAIK，有一些包可以做到这一点，但我们可以根据此代码的特定属性设计更有效的实现(例如，组索引是相对较小的整数)。我们可以在 Numba 中实现这一点。为了准确起见，我们需要进行两遍实现。

这是生成的代码:

import numpy as np
import numba as nb

@nb.njit('(int32[:,::1], UniTuple(int32,2))')
def radial_profile(data, center):
    dh, dw = data.shape
    cx, cy = center
    rMax = np.int64(np.sqrt(max(cx**2 + cy**2, cx**2+(dh-cy)**2, (dw-cx)**2+cy**2, (dw-cx)**2+(dh-cy)**2)))

    # Step 1: computes the mean by group
    sumByGroup = np.zeros(rMax + 1)
    groupSize = np.zeros(rMax + 1, dtype=np.int64)
    for y in range(dh):
        for x in range(dw):
            r = np.int64(np.sqrt((x - cx)**2 + (y - cy)**2))
            sumByGroup[r] += data[y, x]
            groupSize[r] += 1
    meanByGroup = sumByGroup / groupSize

    # Step 2: computes the variance by group based on the mean by group
    sqSumByGroup = np.zeros(rMax + 1)
    for y in range(dh):
        for x in range(dw):
            r = np.int64(np.sqrt((x - cx)**2 + (y - cy)**2))
            sqSumByGroup[r] += (data[y, x] - meanByGroup[r])**2
    varianceByGroup = sqSumByGroup / groupSize

    return np.sqrt(varianceByGroup)

std_azi = radial_profile(data, (2048,2048))

---

性能结果

以下是我的配备 i5-9600KF 处理器的机器上的结果:

Initial code:  33154 ms
New code:         63 ms

新代码比初始代码快 526 倍。

请注意，此代码是连续的。每个图像都可以并行计算。生成的并行代码将很好地扩展，因为该代码显然受计算限制。