I wrote an interactive mandelbrot renderer in Python using kivy where you can zoom with the mouse pointer and am trying to optimize it as best as I can. I currently use this implementation to render the set/zoom (this is a small snippet, just the two functions used to render it):
import numba as nb
import numpy as np
@nb.njit(cache= True, parallel = True)
def mandelbrot(c_r, c_i,maxIt): #mandelbrot function
z_r = 0
z_i = 0
z_r2 = 0
z_i2= 0
for x in nb.prange(maxIt):
z_i = 2 * z_r * z_i + c_i
z_r = z_r2 - z_i2 + c_r
z_r2 = z_r * z_r
z_i2 = z_i * z_i
if z_r2 + z_i2 > 4:
return x
return maxIt
@nb.njit(cache= True, parallel = True)
def DrawSet(W, H, xStart, xDist, yStart, yDist, maxIt):
array = np.zeros((H, W, 3), dtype=np.uint8) #array that holds 'hsv' tuple for every pixel
for x in nb.prange(0, W):
c_r = (x/W)* xDist + xStart #some math to calculate real part
for y in range (0, H):
c_i = -((y/H) * yDist + yStart) #some more math to calculate imaginary part
cIt = mandelbrot(c_r, c_i, maxIt)
color = int((255 * cIt) / maxIt)
array[y,x] = (color, 255, 255) #adds hue value
return array #returns hsv array, gets later displayed using PIL
I currently get pretty good performance. It can render a 500x500 area where every point is bounded (so basically a black picture, the worst-case) in about 0.08 - 0.09 seconds, with 300 iterations. I'm using Numba JIT along with the parallelized range function 'prange()', which helps quite a lot.
However, I've heard that vectorization is usually the fastest for rendering these kind of fractals. After a lot of research (I'm pretty new to vectorization), I managed to put this implementation together:
import numba as nb
import numpy as np
def DrawSet(W, H, xStart, xEnd, yStart, yEnd, maxIt):
array = np.zeros((H,W,3), dtype = np.uint8) # 3D array containing 'hsv' tuple (hue,saturation,value) of each pixel
x = np.linspace(xStart, xEnd, W).reshape((1, W)) #scaling horizontal pixels to x-axis
y = np.linspace(yStart, yEnd, H).reshape((H, 1)) #scaling vertical pixels to y-axis
c = x + 1j * y #creating complex plane out of x axis (real) and y axis (imaginary)
z = np.zeros(c.shape, dtype= np.complex128)
div_time = np.zeros(z.shape, dtype= int)
m = np.full(c.shape, True, dtype= bool)
div_time = loop(z, c, div_time, m, maxIt)
array[:,:,0] = (div_time/maxIt) * 255 -20 #adding 'hue' value
array[:,:,1] = 255 #adding 'saturation' value
array[:,:,2] = 255 #adding 'value'
return array
@nb.vectorize(nb.int64[:,:](nb.complex128[:,:], nb.complex128[:,:], nb.int64[:,:], nb.boolean[:,:], nb.int64))
def loop(z, c, div_time, m, maxIt):
for i in range(maxIt):
z[m] = z[m]**2 + c[m]
diverged = np.greater(np.abs(z), 2, out=np.full(c.shape, False), where=m)
div_time[diverged] = i
m[np.abs(z) > 2] = False
return div_time
Without, the @nb.vectorize decorator, it runs horribly slow. (4 seconds worst-case for 500x500, 300 It.). With the @nb.vectorize decorator, I'm getting this error:
Traceback (most recent call last):
File "MandelBrot.py", line 13, in <module>
from test import DrawSet
File "C:UsersUserDocumentsCodePythonMandelbrot-GUIest.py", line 25, in <module>
def loop(z, c, div_time, m, maxIt):
File "C:UsersUserAppDataLocalProgramsPythonPython38libsite-packages
umba
pufuncdecorators.py", line 119, in wrap
for sig in ftylist:
TypeError: 'Signature' object is not iterable
What am I doing wrong? Am I defining all the numba signatures the correct way?
Would this vectorization method even top my current implementation?
I would be thankful for every advice! Thank you in advance.
question from:
https://stackoverflow.com/questions/65852141/mandelbrot-numba-numpy-vectorization