Python生成六边形密排圆盘mask的代码优化与循环向量化方案

你当前代码里的for循环确实是主要的性能瓶颈：一方面每次循环都要对整个1000*1000的网格做全量运算，另一方面循环中反复更新画布数据的操作完全冗余，占用了大量不必要的时间。

核心优化：完全向量化消除for循环 #

利用numpy的广播机制，可以一次性完成所有圆盘的判断逻辑，不需要循环遍历每个圆盘中心，优化后代码如下：

import numpy as np
import matplotlib.pyplot as plt

r = 0.01
r_sq = r ** 2  # 提前计算r平方，避免重复运算
X, Y = np.mgrid[0:1:1000j, 0:1:1000j]

# 生成六边形排列的圆盘中心
px, py = np.mgrid[r : 1 : 2 * r * np.sqrt(3), r : 1 + r + np.finfo(float).eps: 2 * r]
px = np.vstack((px, px +  r * np.sqrt(3))) 
py = np.vstack((py, py - r))

# 展平中心坐标为一维数组，新增维度适配广播规则
px_flat = px.ravel()[:, np.newaxis, np.newaxis]
py_flat = py.ravel()[:, np.newaxis, np.newaxis]

# 一次性计算所有网格点到所有圆盘中心的距离平方，判断是否在任意圆盘内
mask = ((X - px_flat)**2 + (Y - py_flat)**2 < r_sq).any(axis=0)

# 一次性渲染结果，不需要循环更新画布
fig, ax = plt.subplots(figsize= (12, 12), dpi=50)
ax.imshow(mask * 1, cmap = 'gray', vmin = 0, vmax = 1, extent = [0, 1, 0, 1])
ax.set_aspect(1)
plt.show()

其他可选优化思路 #

• 若网格尺寸极大导致广播时内存占用过高，可以用numba.jit装饰器加速原有循环，仅需少量代码修改即可获得接近C语言的运算速度
• 可以进一步通过裁剪每个圆盘的判断区域来减少运算量：仅对每个圆盘中心周围2r*2r范围内的网格点做距离判断，避免全量网格运算

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