Mandelbrot Set Basic Viewer

mb.py

# (c) copyright zuri bar yochay - [email protected]
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button

class MandelbrotViewer:
    def __init__(self):
        self.max_iter = 100
        self.xmin, self.xmax = -2, 1
        self.ymin, self.ymax = -1.5, 1.5
        self.history = []  # Store zoom history
        
        # Create the main figure and plot
        self.fig, self.ax = plt.subplots(figsize=(10, 8))
        plt.subplots_adjust(bottom=0.2)  # Make room for buttons
        
        # Add zoom out button
        self.bzoom_out_ax = plt.axes([0.7, 0.05, 0.1, 0.075])
        self.bzoom_out = Button(self.bzoom_out_ax, 'Zoom Out')
        self.bzoom_out.on_clicked(self.zoom_out)
        
        # Add reset button
        self.breset_ax = plt.axes([0.81, 0.05, 0.1, 0.075])
        self.breset = Button(self.breset_ax, 'Reset')
        self.breset.on_clicked(self.reset)
        
        # Connect mouse events
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        self.draw_mandelbrot()
        
    def mandelbrot(self, h, w):
        y, x = np.ogrid[self.ymin:self.ymax:h*1j, self.xmin:self.xmax:w*1j]
        c = x + y*1j
        z = c
        divtime = self.max_iter + np.zeros(z.shape, dtype=int)
        
        for i in range(self.max_iter):
            z = z**2 + c
            diverge = z*np.conj(z) > 2**2
            div_now = diverge & (divtime == self.max_iter)
            divtime[div_now] = i
            z[diverge] = 2
            
        return divtime
    
    def draw_mandelbrot(self):
        self.ax.clear()
        h, w = 1000, 1000
        self.ax.imshow(self.mandelbrot(h, w), cmap='hot', extent=[self.xmin, self.xmax, self.ymin, self.ymax])
        self.ax.set_title(f'Mandelbrot Set (click to zoom)\nCurrent view: [{self.xmin:.3f}, {self.xmax:.3f}] x [{self.ymin:.3f}, {self.ymax:.3f}]')
        plt.draw()
    
    def on_click(self, event):
        if event.inaxes != self.ax or event.button != 1:
            return
            
        # Store current view in history
        self.history.append((self.xmin, self.xmax, self.ymin, self.ymax))
        
        # Calculate new boundaries (zoom by factor of 4)
        zoom_factor = 4
        x, y = event.xdata, event.ydata
        dx = (self.xmax - self.xmin) / zoom_factor
        dy = (self.ymax - self.ymin) / zoom_factor
        
        self.xmin = x - dx/2
        self.xmax = x + dx/2
        self.ymin = y - dy/2
        self.ymax = y + dy/2
        
        # Example: set max_iter proportional to 2000 times the zoom level
        current_width = self.xmax - self.xmin
        self.max_iter = int(100 / current_width) 
        self.max_iter = (self.max_iter <= 2000 and self.max_iter) or 2000
        print (f"current_width {current_width} max_iter {self.max_iter}")
        self.draw_mandelbrot()
    
    def zoom_out(self, event):
        if self.history:
            self.xmin, self.xmax, self.ymin, self.ymax = self.history.pop()
            self.draw_mandelbrot()
    
    def reset(self, event):
        self.history = []
        self.xmin, self.xmax = -2, 1
        self.ymin, self.ymax = -1.5, 1.5
        self.draw_mandelbrot()

# Create and show the viewer
viewer = MandelbrotViewer()
plt.show()