mandelbrot.py

import numpy as np
from matplotlib.image import imsave
from numba import njit, prange, vectorize
import os

# Parameters
width, height = 3840 / 2, 3840 / 2
zoom_speed = 0.93
cx, cy = -0.743643887037158704752191506114774, 0.131825904205311970493132056385139
zoom = 1.0
max_frames = 3 * 60 * 60 * 23

@vectorize(['float64(complex128, int64)'], target='parallel', cache=True)
def mandelbrot(c, max_iter):
    z = c
    for n in prange(max_iter):
        if abs(z) > 2:
            return n
        z = z * z + c
    return max_iter

def render_mandelbrot(center_x, center_y, zoom, width, height, current_frame, max_frames):
    base_max_iter = 1024
    increment_rate = 45
    max_iter = int(base_max_iter + (current_frame * increment_rate))

    x_coords = (center_x + (np.arange(width) - width / 2) * zoom).astype(np.float64)
    y_coords = (center_y + (np.arange(height) - height / 2) * zoom).astype(np.float64)
    x_coords, y_coords = np.meshgrid(x_coords, y_coords, indexing='ij')
    complex_grid = x_coords + 1j * y_coords

    mandelbrot_set = mandelbrot(complex_grid, max_iter)

    return mandelbrot_set, max_iter

output_directory = "output"
os.makedirs(output_directory, exist_ok=True)

for frame in range(max_frames):
    mandelbrot_set, max_iter = render_mandelbrot(cx, cy, zoom, width, height, frame, max_frames)
    mandelbrot_set_norm = mandelbrot_set / max_iter
    imsave(f"{output_directory}/frame_{frame:04d}.png", mandelbrot_set_norm.T, cmap='nipy_spectral')
    zoom *= zoom_speed
    print(f"Frame {frame} saved")