logmap_mandel.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation
from numba import njit, prange

@njit
def linear_interp(x, in_min, in_max, out_min, out_max):
    return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

@njit
def pix2point(shape, ix, iy, ovs=1):
    cr = (ix / ovs) * 3.0 / shape[0] - 2.0  # [0, width] -> [-2, 1]
    ci = (iy / ovs) * 2.0 / shape[1] - 1.0  # [0, height] -> [-1, 1]
    return cr, ci

@njit(parallel=True)
def make_mandelbrot(width, height, max_iterations):
    result = np.zeros((width, height))
    for iy in prange(height):
        for ix in prange(width):
            x0, y0 = pix2point(result.shape, ix, iy)
            x, y = 0.0, 0.0
            iteration = 0
            while (x*x + y*y <= 4.0) and (iteration < max_iterations):
                x_new = x*x - y*y + x0
                y = 2*x*y + y0
                x = x_new
                iteration += 1
            result[ix, iy] = iteration / max_iterations
    return result

@njit(parallel=True)
def make_triplebrot(mask, width, height, depth, max_iterations, cutoff=500, oversample=1):
    result = np.zeros((width, height, depth), dtype=np.float64)
    for iy in prange(height * oversample):
        for ix in prange(width * oversample):
            if mask[ix // oversample, iy // oversample] == 0:
                continue
            cr, ci = pix2point(mask.shape, ix, iy, ovs=oversample)
            zr, zi = 0.0, 0.0
            for iteration in range(max_iterations):
                zr_new = zr*zr - zi*zi + cr
                zi = 2*zr*zi + ci
                zr = zr_new
                if zr*zr + zi*zi > 4.0:
                    break
                if iteration <= cutoff:
                    continue
                x = result.shape[0] / 3.0 * (cr + 2.0)
                y = result.shape[1] / 2.0 * (ci + 1.0)
                z = result.shape[2] / 4.0 * (zr + 2.0)
                ix_nmo, iy_nmo, iz_nmo = int(round(x)), int(round(y)), int(round(z))
                if 0 <= ix_nmo < result.shape[0] and 0 <= iy_nmo < result.shape[1] and 0 <= iz_nmo < result.shape[2]:
                    result[ix_nmo, iy_nmo, iz_nmo] += 1e-6
    return result

def generate_mandelbrot_3d(width, height, depth, max_iterations, cutoff=500, oversample=1):
    mask = make_mandelbrot(width, height, max_iterations)
    volume = make_triplebrot(mask, width, height, depth, max_iterations, cutoff, oversample)
    return volume

def plot_mandelbrot_3d(ax, volume, threshold=1e-5):
    mask = volume > threshold
    X, Y, Z = np.where(mask)
    ax.scatter(X, Y, Z, c=volume[mask], cmap='inferno', marker='.', alpha=0.5, s=1)

def animate_combined_view(volume):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.set_box_aspect(volume.shape)
    ax.set_xlim([0, volume.shape[0]])
    ax.set_ylim([0, volume.shape[1]])
    ax.set_zlim([0, volume.shape[2]])

    def update(frame):
        ax.cla()
        ax.set_box_aspect(volume.shape)
        plot_mandelbrot_3d(ax, volume)
        ax.view_init(elev=30, azim=frame)
        ax.set_xlabel('Real part')
        ax.set_ylabel('Imaginary part')
        ax.set_zlabel('Iterations')
        ax.set_title('3D Mandelbrot Set')

    ani = FuncAnimation(fig, update, frames=np.arange(0, 360, 2), interval=50)
    plt.show()

if __name__ == '__main__':
    D = 150  # dimensions of the volume cube
    I = 1500  # maximum iterations in the Mandelbrot calculation
    Omax = 1  # oversampling

    volume = generate_mandelbrot_3d(D, D, D, I, cutoff=500, oversample=Omax)
    volume[:, volume.shape[1] // 2, :] = np.sqrt(volume[:, volume.shape[1] // 2, :])
    animate_combined_view(volume)