python - Mandelbrot Numba/Numpy vectorization?

Question

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

Answer 1

Waitting for answers