Optimization of the Mandelbrot set

I made a program in which you can approximate the Mandelbrot set

My question is how can it be optimized?

import pygame
from pygame import gfxdraw

pygame.init()
win = pygame.display.set_mode((500, 500))

I = (-1) ** .5
maxiter = 64

def f(z, c):
    return z * z + c

def mandelbrot(x, y, scale):
    win.fill((0, 0, 0))
    for i in range(500):
        for j in range(500):
            i0, j0 = i / scale + x, j / scale + y
            
            i0 = i0 / 250 * 2 - 2
            j0 = j0 / 250 * 2 - 2
            
            c = (i0 + j0 *I)
            
            z = 0
            
            for depth in range(maxiter):
                z = f(z, c)
                if abs(z) > 100:
                    break
            
            
            c = depth * (255 / maxiter)
            gfxdraw.pixel(win, i, j, (c, c, c))

            
        pygame.display.update()

mandelbrot(0, 0, 1)
x, y, scale = 0, 0, 1

while True:
    for ev in pygame.event.get():
        if ev.type == pygame.MOUSEBUTTONUP:
            x0, y0 = ev.pos
            
            x0 -= 50
            y0 -= 50
            
            x0 /= scale
            y0 /= scale
            
            x += x0
            y += y0
            
            scale *= 5
            
            mandelbrot(x, y, scale)
Author: Danis, 2020-07-30

2 answers

I made the image calculation faster using the Numpy and Numba libraries. I draw the entire image once at once, which also speeds up the process.

import pygame
from pygame import gfxdraw,surfarray

pygame.init()
win = pygame.display.set_mode((500, 500))

import numpy as np
from numba import njit

I = (-1) ** .5
maxiter = 64

@njit(fastmath=True)
def f(z, c):
    return z * z + c

@njit(fastmath=True)
def mandelbrot(x, y, scale):

    pic = np.zeros((500,500,3), dtype=np.float64)
    for i in range(500):
        for j in range(500):
            i0, j0 = i / scale + x, j / scale + y
            
            i0 = i0 / 250 * 2 - 2
            j0 = j0 / 250 * 2 - 2
            
            c = (i0 + j0 *I)
            
            z = 0
            
            for depth in range(maxiter):
                z = f(z, c)
                if abs(z) > 100:
                    break

            c = depth * (255 / maxiter)
            pic[i, j] = np.array((c.real,c.real,c.real), dtype=np.float64)

    return pic

def mandelbrot_draw(x, y, scale):
    pic = mandelbrot(x, y, scale)
    surfarray.blit_array(win, pic)
    pygame.display.update()

mandelbrot_draw(0, 0, 1)

Update: Fixed the data type and added code for one-time rendering.

 3
Author: CrazyElf, 2020-07-31 09:03:10

I suggest small computational optimizations:

from time import time
import pygame
from pygame import gfxdraw

pygame.init()
win = pygame.display.set_mode((500, 500))

I = (-1) ** .5
maxiter = 64
maxiter_255 = 255 / maxiter

def f(z, c):
    return z**2 + c

def mandelbrot(x, y, scale):
    win.fill((0, 0, 0))
    j_list = [I*((j / scale + y) / 125 - 2) for j in range(500)]
    for i in range(500):
        i0 = (i / scale + x) / 125 - 2
        for j in range(500):
            c = i0 + j_list[j]
            z = f(0, c)
            
            for depth in range(maxiter):
                if abs(z) > 100:
                    break
                z = f(z, c)                
            c = depth * maxiter_255
            gfxdraw.pixel(win, i, j, (c, c, c))

        pygame.display.update()
        
x, y, scale = 0, 0, 1
mandelbrot(x, y, scale)

while True:
    for ev in pygame.event.get():
        if ev.type == pygame.MOUSEBUTTONUP:
            x0, y0 = ev.pos
            
            x += (x0 - 50) / scale
            y += (y0 - 50) / scale
            scale *= 5
            
            mandelbrot(x, y, scale)
 1
Author: n1tr0xs, 2020-07-31 04:58:07