mandelbrot.py

1import mrlypy.two as m22from mrlypy.core.colors import blue, gradient, green, red3import numpy as np4from config import DATA_DIR, IMAGE_SIZE56RESOLUTION = IMAGE_SIZE7MAX_ITER = 10089def mandelbrot_set(width, height, max_iter=100, x_min=-2.0, x_max=1.0, y_min=-1.5, y_max=1.5):10    x = np.linspace(x_min, x_max, width)11    y = np.linspace(y_min, y_max, height)12    X, Y = np.meshgrid(x, y)13    C = X + 1j * Y14    Z = np.zeros_like(C)15    div_time = np.zeros(Z.shape, dtype=np.int16)16    for i in range(max_iter):17        mask = np.abs(Z) <= 218        Z[mask] = Z[mask] * Z[mask] + C[mask]19        div_time[mask] = i20    return div_time2122def main():23    width, height = RESOLUTION24    grid = mandelbrot_set(width, height, MAX_ITER)25    cell = m2.Cell2d.from_array(grid)26    start = red27    mid = green28    end = blue29    colors = gradient([start, mid, end], MAX_ITER)30    palette = np.array([c.to_rgba() for c in colors], dtype=np.uint8)31    grid_clipped = np.clip(grid, 0, MAX_ITER - 1)32    cell.colors = palette[grid_clipped]33    fp = f"{DATA_DIR}/mandelbrot.png"34    image = cell.to_image(scale=10)35    image.save(fp)36    print(f"Saved: {fp}")3738if __name__ == "__main__":39    main()