Source

sketchbook / fractal / fractal.py

Full commit
 #! /usr/bin/env python
 # Plot random pixels on the screen.
from __future__ import division 
import pygame
import random
  
 # Window dimensions
screen_width = 600
screen_height = 600
max_iterations = 30

screen = pygame.display.set_mode((screen_width, screen_height))
clock = pygame.time.Clock()
zoom_factor = 10

spectrum_img = pygame.image.load('fire.bmp')
spectrum = []
for i in xrange(255):
    c = spectrum_img.get_at( (i, 1) )
    h = pygame.Color(c[0],c[1], c[2])
    spectrum.append(h)


class Brot(object):
    def __init__(self, re, im, width, height):
        self.set_range(re, im, width, height)
        self.mouse_pos = (0,0)

    def set_range(self, re, im, width, height):
        aspect =  height/width
        self.MinRe = re
        self.MaxRe = re + (height * (1/aspect))
        self.MinIm = im
        self.MaxIm = im + (width * aspect)
        self.Re_factor = (self.MaxRe-self.MinRe)/(screen_width-1);
        self.Im_factor = (self.MaxIm-self.MinIm)/(screen_height -1);
        print "Set Range (%s,%s) %s %s" % (re, im, width, height)

    def zoom(self, x, y):
        print x,y
        print self.MinRe, self.MinIm
        print self.MaxRe, self.MaxIm
        old_w = self.MaxRe - self.MinRe
        old_h = self.MaxIm - self.MinIm
        x_pos = self.MinRe + (old_w/screen_width) * x
        y_pos = self.MinIm + (old_h/screen_height) * (screen_height - y)

        new_w = old_w/zoom_factor
        new_h = old_h/zoom_factor

        new_l = x_pos - new_w/2
        new_t = y_pos - new_h/2

        self.set_range(new_l,new_t,new_w,new_h)

    def mandel(self, c):
        z=0
        inside = True
        for h in range(0,max_iterations):
            z = z**2 + c
            if abs(z) > 4:
                inside = False
                break
        if inside:
            return 0
        else:
            return int(h*8.5)


    def draw(self):
        for y in xrange(screen_height):
            c_im = self.MaxIm - y * self.Im_factor
            for x in xrange(screen_width):
                c_re = self.MinRe + x * self.Re_factor
                screen.set_at((x, y), spectrum[self.mandel(complex(c_re,c_im))])
            pygame.display.flip()


fractal = Brot(-2.0, -2.0, 4, 4)
running = True
first = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            break
        if event.type == pygame.MOUSEMOTION:
            fractal.mouse_pos = event.pos
        if event.type == pygame.MOUSEBUTTONUP:
            fractal.zoom(*fractal.mouse_pos)
            fractal.draw()

    if first:
        fractal.draw()
        first = False