如何优化椭圆转图像的嵌套循环Python代码？

问题背景 #

输入是观测值列表，每个观测值包含固定数量的椭圆（每个椭圆由7个参数表示：中心x、中心y、R、G、B、长半轴a、短半轴b）。输出是图像列表，每个观测值对应一张纯白背景图像，椭圆重叠区域取RGB均值。现有代码逻辑正常，但因需反复执行，嵌套循环的性能瓶颈明显，需要优化。

原代码：

import numpy as np

n, m = 图像像素尺寸, 图像用(n, m, 3)的numpy数组表示（3对应RGB通道）
N = 每个观测值中的椭圆数量

xx, yy = np.mgrid[:n, :m]

def elipses_population_to_img_population(elipses_population):
        population_size = elipses_population.shape[0]
        img_population = np.empty((population_size, n, m, 3))
        for j in range(population_size):
            imarray = np.empty((N, n, m, 3))
            imarray.fill(np.nan)
            for i in range(N):
                x = elipses_population[j, i, 0]
                y = elipses_population[j, i, 1]
                R = elipses_population[j, i, 2]
                G = elipses_population[j, i, 3]
                B = elipses_population[j, i, 4]
                a = elipses_population[j, i, 5]
                b = elipses_population[j, i, 6]
                xx_centered = xx - x
                yy_centered = yy - y
                elipse = (xx_centered / a)**2 + (yy_centered / b)**2 < 1
                imarray[i, elipse, :] = np.array([R, G, B])
            means_img = np.nanmean(imarray, axis=0)
            means_img = np.nan_to_num(means_img, nan=255)
            img_population[j, :, :, :] = means_img
        return img_population

性能优化方案 #

1. 全向量化改造，消除嵌套循环 #

利用numpy的广播机制，把逐个处理椭圆、逐个处理观测值的循环全部替换为向量运算，彻底规避Python循环的开销。

2. 替换nanmean为加权均值计算 #

原代码用nan标记未覆盖区域，会占用大量内存且计算效率低。改为统计每个像素被椭圆覆盖的次数，再用RGB总和除以覆盖次数得到均值，未覆盖的直接设为255。

3. 预计算重复值，减少计算量 #

提前计算网格坐标的平方值，避免每次循环重复计算；一次性提取所有参数并扩展维度，匹配广播要求。

优化后的代码：

import numpy as np

n, m = 图像像素尺寸, 图像用(n, m, 3)的numpy数组表示（3对应RGB通道）
N = 每个观测值中的椭圆数量

# 预计算网格坐标的平方，避免重复计算
xx, yy = np.mgrid[:n, :m]
xx_sq = xx ** 2
yy_sq = yy ** 2

def elipses_population_to_img_population(elipses_population):
    population_size = elipses_population.shape[0]
    
    # 提取所有参数并扩展维度，适配广播
    x = elipses_population[..., 0][:, :, np.newaxis, np.newaxis]
    y = elipses_population[..., 1][:, :, np.newaxis, np.newaxis]
    rgb = elipses_population[..., 2:5]  # 形状: (观测数, 椭圆数, 3)
    a = elipses_population[..., 5][:, :, np.newaxis, np.newaxis]
    b = elipses_population[..., 6][:, :, np.newaxis, np.newaxis]
    
    # 向量化计算所有椭圆的覆盖掩码，形状: (观测数, 椭圆数, n, m)
    a_sq = a ** 2
    b_sq = b ** 2
    ellipse_mask = ((xx_sq - 2 * x * xx + x ** 2) / a_sq) + ((yy_sq - 2 * y * yy + y ** 2) / b_sq) < 1
    
    # 扩展RGB维度，匹配掩码的形状
    rgb_expanded = rgb[:, :, np.newaxis, np.newaxis, :]
    # 计算每个像素的RGB总和与被覆盖次数
    rgb_sum = np.where(ellipse_mask[..., np.newaxis], rgb_expanded, 0).sum(axis=1)
    count = ellipse_mask.sum(axis=1)[..., np.newaxis]
    
    # 计算均值，未覆盖的像素设为255
    img_population = np.where(count > 0, rgb_sum / count, 255)
    return img_population.astype(np.float32)

4. 备选方案：用Numba JIT编译原循环 #

如果不想大幅修改原代码结构，可以用Numba对循环进行即时编译，直接提升循环执行效率：

import numpy as np
from numba import jit

n, m = 图像像素尺寸, 图像用(n, m, 3)的numpy数组表示（3对应RGB通道）
N = 每个观测值中的椭圆数量

xx, yy = np.mgrid[:n, :m]

@jit(nopython=True)
def elipses_population_to_img_population_numba(elipses_population):
    population_size = elipses_population.shape[0]
    img_population = np.empty((population_size, n, m, 3), dtype=np.float32)
    
    for j in range(population_size):
        rgb_sum = np.zeros((n, m, 3), dtype=np.float32)
        count = np.zeros((n, m), dtype=np.int32)
        
        for i in range(N):
            x = elipses_population[j, i, 0]
            y = elipses_population[j, i, 1]
            R = elipses_population[j, i, 2]
            G = elipses_population[j, i, 3]
            B = elipses_population[j, i, 4]
            a = elipses_population[j, i, 5]
            b = elipses_population[j, i, 6]
            
            for xi in range(n):
                for yi in range(m):
                    xx_centered = xi - x
                    yy_centered = yi - y
                    if (xx_centered / a)**2 + (yy_centered / b)**2 < 1:
                        rgb_sum[xi, yi, 0] += R
                        rgb_sum[xi, yi, 1] += G
                        rgb_sum[xi, yi, 2] += B
                        count[xi, yi] += 1
        
        # 填充结果图像
        for xi in range(n):
            for yi in range(m):
                if count[xi, yi] > 0:
                    img_population[j, xi, yi] = rgb_sum[xi, yi] / count[xi, yi]
                else:
                    img_population[j, xi, yi] = 255
    
    return img_population

内容的提问来源于stack exchange，提问作者Vizik