116 lines
4.5 KiB
Python
116 lines
4.5 KiB
Python
import pygame
|
|
import numpy as np
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
# Constants
|
|
WIDTH, HEIGHT = 800, 600
|
|
MAX_ITER = 256
|
|
ZOOM_FACTOR = 1.2
|
|
RENDER_WIDTH, RENDER_HEIGHT = 400, 300 # Lower resolution for rendering
|
|
|
|
# Initialize Pygame
|
|
pygame.init()
|
|
screen = pygame.display.set_mode((WIDTH, HEIGHT))
|
|
pygame.display.set_caption("Mandelbrot Set Zoom")
|
|
|
|
# Functions to calculate Mandelbrot set
|
|
def mandelbrot(c, max_iter):
|
|
z = c
|
|
for n in range(max_iter):
|
|
if abs(z) > 2:
|
|
return n
|
|
z = z * z + c
|
|
return max_iter
|
|
|
|
def mandelbrot_row(xmin, xmax, ymin, ymax, width, height, max_iter, y):
|
|
r1 = np.linspace(xmin, xmax, width)
|
|
r2 = np.linspace(ymin, ymax, height)
|
|
row = np.empty(width, dtype=np.int32)
|
|
for i in range(width):
|
|
row[i] = mandelbrot(r1[i] + 1j * r2[y], max_iter)
|
|
return row
|
|
|
|
def mandelbrot_set(xmin, xmax, ymin, ymax, width, height, max_iter):
|
|
with ThreadPoolExecutor() as executor:
|
|
rows = list(executor.map(lambda y: mandelbrot_row(xmin, xmax, ymin, ymax, width, height, max_iter, y), range(height)))
|
|
return np.array(rows)
|
|
|
|
def draw_mandelbrot(surface, mandelbrot_image, width, height):
|
|
for x in range(width):
|
|
for y in range(height):
|
|
color_value = mandelbrot_image[y, x] / MAX_ITER * 255
|
|
color = (int(color_value), int(color_value), int(color_value))
|
|
surface.set_at((x, y), color)
|
|
|
|
# Initial boundaries
|
|
xmin, xmax = -2.0, 1.0
|
|
ymin, ymax = -1.5, 1.5
|
|
|
|
# Create a surface for rendering
|
|
render_surface = pygame.Surface((RENDER_WIDTH, RENDER_HEIGHT))
|
|
|
|
# Main loop
|
|
running = True
|
|
needs_update = True
|
|
|
|
while running:
|
|
for event in pygame.event.get():
|
|
if event.type == pygame.QUIT:
|
|
running = False
|
|
elif event.type == pygame.MOUSEBUTTONDOWN:
|
|
if event.button == 1: # Left mouse button for zoom in
|
|
mouse_x, mouse_y = event.pos
|
|
mouse_x = mouse_x * RENDER_WIDTH // WIDTH
|
|
mouse_y = mouse_y * RENDER_HEIGHT // HEIGHT
|
|
r1 = np.linspace(xmin, xmax, RENDER_WIDTH)
|
|
r2 = np.linspace(ymin, ymax, RENDER_HEIGHT)
|
|
x_center = r1[mouse_x]
|
|
y_center = r2[mouse_y]
|
|
width = (xmax - xmin) / ZOOM_FACTOR
|
|
height = (ymax - ymin) / ZOOM_FACTOR
|
|
xmin, xmax = x_center - width / 2, x_center + width / 2
|
|
ymin, ymax = y_center - height / 2, y_center + height / 2
|
|
needs_update = True
|
|
elif event.button == 3: # Right mouse button for zoom out
|
|
mouse_x, mouse_y = event.pos
|
|
mouse_x = mouse_x * RENDER_WIDTH // WIDTH
|
|
mouse_y = mouse_y * RENDER_HEIGHT // HEIGHT
|
|
r1 = np.linspace(xmin, xmax, RENDER_WIDTH)
|
|
r2 = np.linspace(ymin, ymax, RENDER_HEIGHT)
|
|
x_center = r1[mouse_x]
|
|
y_center = r2[mouse_y]
|
|
width = (xmax - xmin) * ZOOM_FACTOR
|
|
height = (ymax - ymin) * ZOOM_FACTOR
|
|
xmin, xmax = x_center - width / 2, x_center + width / 2
|
|
ymin, ymax = y_center - height / 2, y_center + height / 2
|
|
needs_update = True
|
|
elif event.type == pygame.MOUSEWHEEL:
|
|
if event.y > 0: # Scroll up for zoom in
|
|
x_center = (xmin + xmax) / 2
|
|
y_center = (ymin + ymax) / 2
|
|
width = (xmax - xmin) / ZOOM_FACTOR
|
|
height = (ymax - ymin) / ZOOM_FACTOR
|
|
xmin, xmax = x_center - width / 2, x_center + width / 2
|
|
ymin, ymax = y_center - height / 2, y_center + height / 2
|
|
needs_update = True
|
|
elif event.y < 0: # Scroll down for zoom out
|
|
x_center = (xmin + xmax) / 2
|
|
y_center = (ymin + ymax) / 2
|
|
width = (xmax - xmin) * ZOOM_FACTOR
|
|
height = (ymax - ymin) * ZOOM_FACTOR
|
|
xmin, xmax = x_center - width / 2, x_center + width / 2
|
|
ymin, ymax = y_center - height / 2, y_center + height / 2
|
|
needs_update = True
|
|
|
|
if needs_update:
|
|
mandelbrot_image = mandelbrot_set(xmin, xmax, ymin, ymax, RENDER_WIDTH, RENDER_HEIGHT, MAX_ITER)
|
|
draw_mandelbrot(render_surface, mandelbrot_image, RENDER_WIDTH, RENDER_HEIGHT)
|
|
needs_update = False
|
|
|
|
# Scale the render surface to the screen size and blit it
|
|
scaled_surface = pygame.transform.scale(render_surface, (WIDTH, HEIGHT))
|
|
screen.blit(scaled_surface, (0, 0))
|
|
pygame.display.flip()
|
|
|
|
pygame.quit()
|