Iterate logic for generating buddhabrot given a set of non mandelbrot complex numbers and image size

generate_buddhabrot.py

import numpy as np


def generate_buddhabrot(non_mandelbrot: np.ndarray, size: int) -> np.ndarray:
    """Generate buddhabrot image array

    Args:
        non_manderlbrot (np.ndarray): Array of non-mandelbrot points on the complex plane
        size (int): Size of the image array

    Returns:
        np.ndarray: Image array containing buddhabrot trajectories
    """
    img_array = np.zeros([size, size], int)
    z = np.copy(non_mandelbrot)

    while len(z):

        # Mapping the complex plane to the image array
        x = np.array((z.real + 2.0) / 4 * size, int)
        y = np.array((z.imag + 2.0) / 4 * size, int)

        img_array[x, y] += 1

        z = z ** 2 + non_mandelbrot

        mask = np.abs(z) < 2
        non_mandelbrot = non_mandelbrot[mask]
        z = z[mask]

    return img_array