mandelbrot.py

# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt


def mandelbrot(*, x_range, y_range, lod, max_n):
    colours = np.zeros((lod, lod))
    x, y = np.meshgrid(np.linspace(*x_range, lod, dtype=np.float64), np.linspace(*y_range, lod, dtype=np.float64))
    c = x + y * 1j
    z = np.zeros(c.shape, np.complex64)

    for n in np.arange(1, max_n + 1):
        z = np.add(z ** 2, c)
        condition = np.where(np.abs(z) > 2)
        z[condition], c[condition], colours[condition] = 0j, 0j, n
    return colours


if __name__ == '__main__':
    x_range = (-2.0, 0.7)
    y_range = (-1.4, 1.4)
    output = mandelbrot(x_range=x_range, y_range=y_range, lod=1_500, max_n=40)

    plt.imshow(output, extent=[*x_range, *y_range], origin='lower')
    plt.show()