60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
import pygame, colorsys
|
|
import numpy as np
|
|
from numba import jit
|
|
|
|
# Set the width and height of the window
|
|
width, height = 1100, 800
|
|
|
|
# Define the region of the complex plane to visualize
|
|
xmin, xmax = -2, 1
|
|
ymin, ymax = -1, 1
|
|
|
|
# Set the maximum number of iterations for the Mandelbrot calculation
|
|
max_iterations = 100
|
|
|
|
# Initialize Pygame
|
|
pygame.init()
|
|
|
|
# Create a window
|
|
screen = pygame.display.set_mode((width, height))
|
|
pygame.display.set_caption("Mandelbrot Set")
|
|
|
|
@jit(nopython=True)
|
|
def mandelbrot(c):
|
|
z = 0
|
|
for i in range(max_iterations):
|
|
z = z**2 + c
|
|
if abs(z) > 2:
|
|
return i
|
|
return max_iterations
|
|
|
|
# Function to convert iterations to a smooth gradient in the HSV color space
|
|
def smooth_color(iterations):
|
|
if iterations == max_iterations:
|
|
return 0, 0, 0 # Black for points in the set
|
|
hue = 0.7 + iterations / max_iterations
|
|
saturation = 1.0
|
|
value = iterations / max_iterations
|
|
return [int(c * 255) for c in colorsys.hsv_to_rgb(hue, saturation, value)]
|
|
|
|
# Function to update the display
|
|
def update_display():
|
|
for x in range(width):
|
|
for y in range(height):
|
|
c = complex(xmin + (x / width) * (xmax - xmin), ymin + (y / height) * (ymax - ymin))
|
|
iterations = mandelbrot(c)
|
|
color = smooth_color(iterations)
|
|
screen.set_at((x, y), color)
|
|
pygame.display.flip()
|
|
|
|
# Run the Pygame event loop
|
|
running = True
|
|
while running:
|
|
for event in pygame.event.get():
|
|
if event.type == pygame.QUIT:
|
|
running = False
|
|
|
|
update_display()
|
|
|
|
# Quit Pygame
|
|
pygame.quit()
|