Mandelbrot Set generator and fractal viewer

mandelbrotGenerator.py

#/bin/python

# Mandelbrot pattern generator by Bob Webb, 2012.
# Generates a Mandelbrot pattern as a 2D array of numerical values (number of iterations) that's then pickled.
# I need to get round to pickling metadata as well but not now. Har har.
import sys
import numpy
import pickle
import math


def mandelbrot(c, iters,threshold):
    z = 0
    i = 0
    while i < iters and abs(z) < threshold:
        z = z*z + c
        i += 1
    
    return i

def _pow(z,d):
    if z == 0j:
        return 0
    else:
        return pow(z,d)

def multibrot(c, d, iters,threshold):
    #print "Creating multibrot......"
    z = 0
    i = 0
    while i < iters and abs(z) < threshold:
        z = _pow(z,d) + c
        
        i += 1
    
    return i

def generateMandelbrot(winWidth,winHeight,iterations,threshold):
    outputValues = []
    
    for x in numpy.arange(0,winWidth):
        outputValues.append([])
        for y in numpy.arange(0,winHeight):
            if x % 100 == 0  and y % 100 == 0:
                print x,y
                #print outputValues
            mX = x*(3.0/winWidth) - 2
            mY = y*(2.0/winHeight) - 1
            
            c = complex(mX,mY)
            
            m = mandelbrot(c,iterations,threshold)
            outputValues[x].append(m)
    
    return outputValues

def generateMultibrot(d,winWidth,winHeight,iterations,threshold):
    outputValues = []
    
    for x in numpy.arange(0,winWidth):
        outputValues.append([])
        for y in numpy.arange(0,winHeight):
            if x % 100 == 0  and y % 100 == 0:
                print x,y
                #print outputValues
            mX = x*(3.0/winWidth) - 2
            mY = y*(2.0/winHeight) - 1
            
            c = complex(mX,mY)
            
            m = multibrot(d,c,iterations,threshold)
            outputValues[x].append(m)
    
    return outputValues


def main():
	outputFileName = sys.argv[1]
	
	winWidth = int(sys.argv[2])
	winHeight = int(sys.argv[3])
	
	numIterations = int(sys.argv[4])
	
	threshold = float(sys.argv[5])
	
	try:
		print "trying multibrot..."
		powers = int(sys.argv[6])
		print "got #powers"
		print "Creating a Multibrot set with power:",powers
		values = generateMultibrot(powers,winWidth,winHeight,numIterations,threshold)
	except IndexError:
		values = generateMandelbrot(winWidth,winHeight,numIterations,threshold)
	
	
	#values = generateMandelbrot(winWidth,winHeight,numIterations,threshold)
	
	
	f = open(outputFileName,'w')
	
	
	pickle.dump(values,f)
	
	


if __name__== '__main__' : main()

mandelbrotViewer.py

#!/bin/python

# Mandelbrot Set Viewer! By Bob Webb, 2012

import sys
import functools
import time
import pygame
import pickle
import colorsys
from pygame.locals import *

class memoized(object):
   '''Decorator. Caches a function's return value each time it is called.
   If called later with the same arguments, the cached value is returned 
   (not reevaluated).
   '''
   def __init__(self, func):
      self.func = func
      self.cache = {}
   def __call__(self, *args):
      try:
         return self.cache[args]
      except KeyError:
         value = self.func(*args)
         self.cache[args] = value
         return value
      except TypeError:
         # uncachable -- for instance, passing a list as an argument.
         # Better to not cache than to blow up entirely.
         return self.func(*args)
   def __repr__(self):
      '''Return the function's docstring.'''
      return self.func.__doc__
   def __get__(self, obj, objtype):
      '''Support instance methods.'''
      return functools.partial(self.__call__, obj)

def drawData(surface,data,colourScheme):
    '''
    Draws a 2D grid of values ('data') to 'surface' using the given 'colourScheme'
    '''
    for x in range(len(data) - 1):
        for y in range(len(data[x]) - 1):
            surface.set_at((x,y),colourScheme(data[x][y]))
            
    return surface


@memoized
def itersToColour(iters):
    #print iters/255
    rgb = [int(255*x) for x in colorsys.hls_to_rgb(1.0/(iters+1),0.5,0.8)]
    #print rgb
    
    return Color(*rgb)

def main():
    inputFileName = sys.argv[1] # the pickle file to be read ^^
    
    dataFile = open(inputFileName,'r')
    
    data = pickle.load(dataFile)
    
    pygame.init()
    
    winWidth,winHeight = len(data),len(data[0])
    
    surf = pygame.display.set_mode((winWidth,winHeight))
    pygame.display.set_caption("Mandelbrot set!")
    
    
    drawData(surf,data,itersToColour)
    
    pygame.image.save(surf,"mandelbrot "+time.asctime(time.gmtime())+".png")
    print "Image displayed!"
    
    pygame.display.flip()
    
    while True:
        # loop shite
        # handle input
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                pygame.event.post(pygame.event.Event(QUIT))


if __name__== '__main__' : main()