优化Python中基于Pygame的Mandelbrot集绘制程序

我采用两层for循环判断点是否属于Mandelbrot集，但程序运行耗时过长，导致绘制细节不足。请问如何优化代码以提升运行速度，从而增加绘制细节？注：Pygame原点位于左上角，我已手动将各点相对于中心定位。现有代码如下：

import math
import pygame


class Point:
    def __init__(self, x, y):
        self.x = x - 250
        self.y = y - 250


class Complex:
    def __init__(self, real, imaginary):
        self.real = real
        self.imaginary = imaginary

    def __eq__(self, other):
        if self.real == other.real:
            if self.imaginary == other.imaginary:
                return True
        return False

    def __abs__(self):
        return math.sqrt(self.real ** 2 + self.imaginary ** 2)

    def __round__(self, n=None):
        return Complex(round(self.real, 4), round(self.imaginary, 4))

    def __add__(self, other):
        return Complex(self.real + other.real, self.imaginary + other.imaginary)

    def __sub__(self, other):
        return Complex(self.real - other.real, self.imaginary - other.imaginary)

    def __mul__(self, other):
        return Complex((self.real * other.real) - (self.imaginary * other.imaginary),
                       (self.real * other.imaginary) + (self.imaginary * other.real))


def zheta(z, c):
    return round((z * z) + c, 3)


def mandelbrotNumber(num):
    z0 = Complex(0, 0)
    c = num
    checkList = [Complex(0, 0)]
    for i in range(10):
        z = zheta(z0, c)
        if abs(z) > 2:
            return False
        if z in checkList:
            return True
        checkList.append(z)
        z0 = z
    return True


def plotPoint(x, y):
    color = (0, 0, 0)
    position = (x, y)
    pygame.draw.circle(screen, color, position, 2)

pygame.init()
screen = pygame.display.set_mode((500, 500))
running = True
startPosX = (0, 250)
endPosX = (500, 250)
startPosY = (250, 0)
endPosY = (250, 500)

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    screen.fill("white")
    for i in range(-200, 200):
        for j in range(-200, 200):
            number = Complex(i / 100, j / 100)
            if mandelbrotNumber(number):
                plotPoint((number.real*100) + 300, (number.imaginary*100) + 250)

    pygame.display.flip()
pygame.quit()

核心优化措施 #

1. 替换自定义复数类为Python内置complex类型 #

Python内置的complex是C底层实现，运算速度远快于自定义类，且原生支持加减乘、实部虚部访问等操作，完全无需手动实现复数逻辑。

2. 移除无效的周期检测逻辑 #

代码中checkList用于检测迭代周期，但曼德博集合的判断规则是只要迭代中|z|>2则必然发散，无需检测重复点。列表的in操作是O(n)时间复杂度，会大幅拖慢运算速度，直接删除该逻辑即可。

3. 避免开根号运算，用平方值替代 #

判断|z|>2等价于z.real² + z.imag² > 4，开根号是高耗时运算，直接比较平方值可节省大量计算时间。

4. 批量绘制替代单点绘制 #

pygame.draw.circle是高频耗时操作，先收集所有需要绘制的点坐标，再用pygame.draw.points批量绘制，能大幅减少绘制开销。

5. 预计算并缓存图像 #

曼德博集合是固定图形，无需在主循环中重复计算。提前计算所有点并绘制到缓存Surface，主循环仅需显示缓存结果即可。

优化后的完整代码 #

import pygame

def mandelbrot(c, max_iter):
    z = 0j
    for _ in range(max_iter):
        z = z * z + c
        # 用平方值替代开根号，减少耗时运算
        if z.real ** 2 + z.imag ** 2 > 4:
            return False
    return True

pygame.init()
screen = pygame.display.set_mode((500, 500))
# 创建缓存Surface，预绘制曼德博集合
mandel_surface = pygame.Surface((500, 500))
mandel_surface.fill("white")

# 参数控制：scale越小显示范围越大，越大细节越丰富；max_iter越高判断越准确
scale = 100
offset_x = 250
offset_y = 250
max_iter = 50

# 预计算所有像素点
points = []
for x in range(500):
    for y in range(500):
        # 转换为复数平面坐标，翻转y轴适配Pygame原点位置
        real = (x - offset_x) / scale
        imag = -(y - offset_y) / scale
        c = complex(real, imag)
        if mandelbrot(c, max_iter):
            points.append((x, y))

# 批量绘制所有点
pygame.draw.points(mandel_surface, (0, 0, 0), points, 1)

running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    # 直接显示缓存好的图像
    screen.blit(mandel_surface, (0, 0))
    pygame.display.flip()

pygame.quit()

进阶优化建议 #

• 使用numpy批量计算所有点，将循环逻辑交给C实现的numpy，速度可再提升数倍。
• 为不同迭代次数的点设置不同颜色，生成彩色曼德博集合，同时更清晰地展示边界细节。
• 根据需求调整scale和max_iter参数：增大scale可放大局部细节，提高max_iter可更准确判断边界点（会增加计算时间）。

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