<jupyter_start><jupyter_text>Numba performance enhancements Contents1&nbsp;&nbsp;Introduction2&nbsp;&nbsp;Numpy2.1&nbsp;&nbsp;The bad way2.2&nbsp;&nbsp;A better way2.3&nbsp;&nbsp;Viewing the image3&nbsp;&nbsp;Numba @jit, serial4&nbsp;&nbsp;Numba @jit, parallel5&nbsp;&nbsp;Numba @vectorize6&nbsp;&nbsp;CUDA7&nbsp;&nbsp;Results  IntroductionThe Julia set example is based on Gorelick, "High Performance Python: Practical Performant Programming for Humans", as it gives us something simple and CPU-intensive to work with.<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

%matplotlib inline<jupyter_output><empty_output><jupyter_text>Area of complex space to investigate, as global variables:<jupyter_code>x1, x2, y1, y2 = -1.8, 1.8, -1.8, 1.8
c_real, c_imag = -0.62772, -.42193
c = complex(c_real, c_imag)<jupyter_output><empty_output><jupyter_text>Set reasonable defaults for a laptop<jupyter_code>desired_width = 1000
max_iterations = 300
def set_coords():
    x_step = (x2 - x1) / desired_width
    y_step = (y2 - y1) / desired_width
    x = np.arange(x1, x2, x_step)
    y = np.arange(y1, y2, y_step)

    # set width and height to the generated pixel counts, rather than the
    # pre-rounding desired width and height
    width = len(x)
    height = len(y)
    
    # build a list of co-ordinates 
    zs = [complex(xcoord, ycoord) for ycoord in y for xcoord in x]
    return np.array(zs), width, height
zs, width, height = set_coords()
zlen = len(zs)<jupyter_output><empty_output><jupyter_text>Numpy  The bad way<jupyter_code>def calculate_z_serial_numpy():
    """Calculate output list using Julia update rule"""
    output = np.zeros(zlen)
    for i in range(zlen):
        n = 0
        z = zs[i]
        while abs(z) < 2 and n < max_iterations:
            z = z*z + c
            n += 1
        output[i] = n
    return output
def calc_numpy():
    output = calculate_z_serial_numpy()

    # this sum is expected for 1000^2 grid with 300 iterations
    assert sum(output) == 33219980

    return output
%time output = calc_numpy()<jupyter_output>CPU times: user 23.9 s, sys: 0 ns, total: 23.9 s
Wall time: 23.9 s<jupyter_text>This is really slow!  A better waySplit out the per-pixel calculation and let Numpy vectorize it.<jupyter_code>def calc_iter(z):
    n = 0
    while abs(z) < 2 and n < max_iterations:
        z = z*z + c
        n += 1
    return n
calc_iter_vec = np.vectorize(calc_iter)
%time output = calc_iter_vec(zs)<jupyter_output>CPU times: user 7.6 s, sys: 31.8 ms, total: 7.63 s
Wall time: 7.63 s<jupyter_text>About a 3-fold improvement.  Viewing the image<jupyter_code>def show_image(output_raw):
    
    # sanity check our 1D array and desired 2D form
    assert width * height == len(output_raw)
    
    output_raw = np.array(output_raw)
    plt.imshow(output_raw.reshape((width, height)), origin='lower') #, cmap='Greys')
show_image(output)<jupyter_output><empty_output><jupyter_text>Numba @jit, serial<jupyter_code>from numba.decorators import jit

@jit(nopython=True)
def calc_iter_jit(z):
    n = 0
    while abs(z) < 2 and n < max_iterations:
        z = z*z + c
        n += 1
    return n

@jit(nopython=True)
def calculate_z_serial_jit():
    output = np.zeros(zlen)
    for i in range(len(zs)):
        output[i] = calc_iter_jit(zs[i])
    return output
def calc_jit():
    output = calculate_z_serial_jit()

    # this sum is expected for 1000^2 grid with 300 iterations
    assert sum(output) == 33219980

    return output
%timeit output = calc_jit()<jupyter_output>548 ms ± 6.67 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Numba @jit, parallel<jupyter_code>@jit(nopython=True, parallel=True)
def calculate_z_parallel_jit():
    output = np.zeros(zlen)
    for i in range(len(zs)):
        output[i] = calc_iter_jit(zs[i])
    return output
def calc_jit_parallel():
    output = calculate_z_parallel_jit()

    # this sum is expected for 1000^2 grid with 300 iterations
    assert sum(output) == 33219980

    return output
%timeit output = calc_jit_parallel()<jupyter_output>636 ms ± 29.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>It got slightly worse!  Numba @vectorize<jupyter_code>from numba import vectorize, complex128

@vectorize([complex128(complex128)])
def calc_iter_numba_vec(z):
    n = 0
    while abs(z) < 2 and n < max_iterations:
        z = z*z + c
        n += 1
    return n

# no @jit here: it wouldn't help and also conflicts with assert()
def calculate_z_numba_vec():
    output = calc_iter_numba_vec(zs)
    
    # this sum is expected for 1000^2 grid with 300 iterations
    assert sum(output) == 33219980
    
    return output
%timeit output = calculate_z_numba_vec()<jupyter_output>555 ms ± 6.62 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>That's about the same as @jit serial.  CUDAThis depends on specific hardware and drivers. Will it work here?<jupyter_code>from numba import cuda

if cuda.is_available():
    cuda.detect()
else:
    print("CUDA is not available on this machine")<jupyter_output>CUDA is not available on this machine<jupyter_text>The simplest usage is in the @vectorize decorator. This will throw CUDA_ERROR_UNKNOWN on machines lacking a CUDA-capable graphics system.<jupyter_code>@vectorize([complex128(complex128)], target='cuda')
def calc_iter_numba_cuda(z):
    n = 0
    while abs(z) < 2 and n < max_iterations:
        z = z*z + c
        n += 1
    return n

# no @jit here: it wouldn't help and also conflicts with assert()
def calculate_z_numba_cuda():
    output = calc_iter_numba_cuda(zs)
    
    # this sum is expected for 1000^2 grid with 300 iterations
    assert sum(output) == 33219980
    
    return output
%timeit output = calculate_z_numba_cuda()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10153157804164516, 'fraction_numerical': 0.035966270865599725, 'mean_word_length': 3.9087837837837838, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '188777', 'license': 'MIT', 'n_tokens_mistral': 2221, 'n_tokens_neox': 2065, 'n_words': 608, 'path': 'genpython/performance/Numba.ipynb', 'repo_name': 'colinleach/tinkerings', 'size': '5811'}