Python中如何优化逐像素绘制行星分层环形的运行效率

原实现的核心性能瓶颈来自三个方面：

• Python双层for循环逐像素遍历全图，解释器层循环开销极高
• 逐像素调用np.random.choice做随机采样、逐分支判断颜色，重复函数调用开销大
• 逐像素调用PIL的putpixel方法，单像素IO操作效率远低于批量内存操作

你提到的维护已绘制像素列表的方案依然没有摆脱Python级别的元素遍历、删除操作，性能提升有限，更高效的实现完全基于Numpy向量化操作，不需要任何显式的像素遍历循环，性能可以提升两个数量级以上。

核心优化思路 #

1. 一次性生成全图坐标网格：用np.mgrid直接生成整张图所有像素的x、y坐标矩阵，批量计算所有像素到行星中心的距离平方，不需要逐像素做距离判断
2. 布尔掩码快速筛选圈层像素：对每个圈层，直接用向量化的大小比较生成布尔掩码，一次性圈出属于该圈层的所有像素位置，整个过程没有Python层循环
3. 批量生成随机颜色：统计当前圈层的像素总数后，一次性按组分占比生成所有像素的颜色索引，直接映射为RGB值，避免逐像素采样和分支判断
4. 数组级批量赋值：直接在Numpy数组中完成所有像素的颜色填充，所有操作都在C层执行，最后再一次性转换为PIL图片对象，完全抛弃低效的putpixel逐像素写入
5. 按顺序绘制圈层无需额外去重：圈层本身是互不重叠的环形结构，只要配置好每个圈层的内外半径，掩码筛选天然不会出现像素重复覆盖的问题，不需要额外维护已绘制像素列表。

优化后参考代码 #

from PIL import Image
import numpy as np

width = 1200
height = 800
planet_radius = 1737
center_x = width // 2
center_y = height // 2

# 圈层配置统一管理，新增圈层直接追加即可：(圈层内半径(实际物理值), 圈层外半径(实际物理值), 组分占比, 组分对应RGB值)
shell_configs = [
    (0, 300, (0.8, 0.15, 0.05), [(105, 105, 105), (128, 128, 128), (192, 192, 192)]),  # 核心层
    (300, 600, (0.3, 0.6, 0.1), [(0, 100, 0), (154, 205, 50), (255, 140, 0)]),  # 第一壳层
    # 后续圈层按顺序追加即可
]

# 初始化RGB数组，背景为白色
img_arr = np.full((height, width, 3), 255, dtype=np.uint8)
# 一次性生成全图坐标网格，预计算所有像素到中心的距离平方
y_grid, x_grid = np.mgrid[0:height, 0:width]
dist_sq = (x_grid - center_x) ** 2 + (y_grid - center_y) ** 2

for r_in_phys, r_out_phys, props, colors in shell_configs:
    # 物理半径转像素半径
    r_in = round(r_in_phys * center_x / planet_radius)
    r_out = round(r_out_phys * center_x / planet_radius)
    # 生成当前圈层的像素掩码
    shell_mask = (dist_sq >= r_in ** 2) & (dist_sq < r_out ** 2)
    pixel_num = shell_mask.sum()
    # 批量生成所有像素的组分索引，映射为RGB颜色
    comp_idx = np.random.choice(len(colors), size=pixel_num, p=props)
    shell_colors = np.array(colors, dtype=np.uint8)[comp_idx]
    # 批量赋值到对应像素位置
    img_arr[shell_mask] = shell_colors

# 数组转PIL图像输出
im = Image.fromarray(img_arr)
im.save("planet_internal_structure.png")

额外优化提示 #

• 如果需要更真实的岩层边界效果，可以在距离计算时叠加小尺度的噪声扰动圈层边界，该操作同样可以用向量化方式实现，几乎不增加额外耗时
• 当分辨率超过4K时，可以只裁剪行星外接正方形区域做计算，进一步减少计算量，不过当前实现哪怕处理8K分辨率图像，耗时也在百毫秒级别，完全满足使用需求
• 不要在循环中混用Python原生随机数模块，所有随机采样尽量用Numpy的批量接口完成，减少函数调用开销

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