Mandelbrot Generator

gistfile1.py

import numpy as np
from scipy import misc
import math
import time

"""
MANDLEBROT GENERATOR
by Josiah Beverton

Many Thanks to:
http://jehiah.cz/a/creating-images-with-numpy
http://www.lfd.uci.edu/~gohlke/pythonlibs/
http://stackoverflow.com/questions/369438/smooth-spectrum-for-mandelbrot-set-rendering
"""

def m_member(complex,iterations=40):
    """
    If a complex value escapes from the boundary 2 after infinite iterations over
    Z**2+Z0 then it is not a member of the set.

    """
    Z = complex
    #Iterate over the complex number till it escapes the bounds.
    #Return an RGB triple of 0,0,0 or 255,0,0 dependent on whether or not it escapes
    for i in xrange(1,iterations+1):
        if np.absolute(Z)>2:
            r,g,b = 0,0,0
            return np.array([r,g,b], dtype=np.uint8)
        Z = Z**2 + complex
    return np.array([255,0,0],dtype=np.uint8)

def nsmooth(n,zn):
    #Incomplete implementation of the colour grading algorithm
    #Literally no idea how to implement this
    a = n+1 - math.log(math.log(np.absolute(zn)))/math.log(2)
    r= int((a/40)*255)
    
    return int((1-r/40)*255)

def generate_map(realrange, complexrange, realstep = 0.05, complexstep = 0.05):
    """generate_map
        - realrange: The Range of values over which the function should calculate.
             - Tuple of size 2
        - complexrange: The Range of values over which the function should calculate.
            - Tuple of size 2
    Returns a numpy array of complex numbers
"""
    width = np.arange(realrange[0],realrange[1],realstep)
    height = np.arange(complexrange[0],complexrange[1],complexstep)
    #Create an empty numpy array over which to calculate
    map = np.zeros((len(width),len(height)),dtype=complex)
    for index, value in np.ndenumerate(map):
        x,y = index
        #Propagate through with values of complex numbers over which to calculate
        map[index[::-1]] = np.complex(width[x],height[y])
    return map

def calculate_set(generated_map, func):
    #Takes a map of complex numbers
    #Generate a function to take this map and replace with an RGB triple based on a supplied function
    mapper = np.vectorize(func, otypes = [np.ndarray])
    img = mapper(generated_map)
    nimg = np.empty((img.shape[0],img.shape[1],3),dtype=np.uint8)
    #Transform into a proper NxMx3 Array instead of NxM filled with 1x3 arrays
    for key, triple in np.ndenumerate(img):
        nimg[key] = triple
    misc.imsave(r"C:\Users\Josiah\Downloads\a.png",nimg)
    
if __name__ == "__main__":
    calculate_set(generate_map((-2,0.6),(-1.3,1.3),0.0025,0.0025),m_member)