Optimise Mandelbrot set plotter in python using pygame

Question

I have used 2 for loops to check if the points are within the set, but apparently it takes too much time, giving less details on the plot itself. How can I optimise the code to run faster, and thus add more detail to the plot ? The pygame origin is at the top left corner, so I manually set each point relative to the centre.

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()

Answer 1

How can I optimise the code to run faster

1. Use the complex type built into Python.
2. Drop the rounding, that's just unnecessary extra steps.
3. Drop the checkList since exact equality is very unlikely and not worth the effort.
4. Instead of drawing a circle for each pixel, construct a raster image in memory (pygame docs suggest this is called a Surface), set pixels on that then draw that whole image to your application windowin a single drawing operation.

Answer 2

1. optimizing calculations

Python has a built-in class complex which should be faster, as it's internally implemented in C, but I guess that you implemented the class as an exercise.

The algorithm itself can be improved by itself - see what @MvG listed, but there is more that you can do to improve speed - calculating a square root of a number is a quite slow operation, x**2 + y**2 > 4 will work as good as sqrt(x**2 + y**2) > 2, but will avoid unnecessary square root. There are more optimizations available that let to reduce the number of multiplications (less costly operation than sqrt(), but still it takes a bit). Wikipedia article has a section dedicated on optimizing mandelbrot set drawing algorithm, see this.

Too slow? There are advanced libraries that help to speedup operations, numpy lets you to do operations on multiple values "at once", but it won't be very useful here (there is no trivial way of avoiding unnecessary operations on points that already escaped). More useful library would be numba - it is a jit compiler for a subset of python, it lets you to write code that would have a comparable speed to a compiled language.

Using third-party library to boost performance should be used if you really want to give everything up for speed, but this is definitely not a beginner friendly solution.

2. optimizing drawing

Drawing circles - that's a complicated operation, better use something like surface.set_at() to set color of a selected pixel.

Calculating Mandelbrot fractal each frame is not a good idea, you can have a separate surface, that you can draw Mandelbrot set once (when you draw on screen, you really just draw on surface that is linked to the window, you can imagine that surfaces are some virtual windows, that you can't see, but you can draw on them as you would do normally on window)

If you really want to use numba, using any functions related to pygame surfaces to draw on them in the function optimized by numba` will reduce the speed gained for using numba, to avoid this, edit values from an object received from pygame.surfarray``.

3. limitations

There are always some limits, with the propositions I've listed above (even without that one about using numba), you should get a decent program for drawing mandelbrot set, but the fact is that if you zoom more and more, the more heavy it becomes to calculate whether selected point might belong to mandelbrot set. Just be aware of that.

Also the question is how much time do you want to spend on optimizing this code.