Creates Julia set using smaller sections

julia.py

# copied from a ipython notebook
import numpy as np
from matplotlib import pyplot as plt
import time
import warnings
%matplotlib inline

def julia(c, extent, pixels):
    t, l, b, r = extent
    ny, nx = pixels
    ys, xs = np.ogrid[t:b:ny*1j, l:r:nx*1j]
    im = ys * 1j + xs
    counts = np.zeros(pixels)
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        while True:
            im = im * im + c
            still_alive = [abs(im) < 2]
            counts[still_alive] += 1
            if not np.any(still_alive):
                break
    return counts
    
# do we speed up using many smaller sections?
def do_section(args):
    c, yargs, xargs = args
    ystart, yend, nystart, nyend = yargs
    xstart, xend, nxstart, nxend = xargs
    sub_im = julia(c, (ystart, xstart, yend, xend), (nyend-nystart, nxend-nxstart))
    #print(iy, iye, ix, ixe)
    return sub_im, nystart, nxstart
    
def julia_sections(c, extent, pixels, section_shape, _map=map):
    out = np.zeros(pixels)
    t, l, b, r = extent
    ny, nx = pixels
    sy, sx = section_shape
    # lets create the xs,ys ogrid and use that to map pixels to starting pos
    def get_args(start, end, n, sample):
        # return a list of (start, end, nstart, nend)
        ret = []
        sp = np.linspace(start, end, n)
        for istart in range(0, n, sample):
            iend = min(istart + sample, len(sp))
            sub_start, sub_end = sp[istart], sp[iend-1]
            ret.append((sub_start, sub_end, istart, iend))
        return ret
        
    y_args = get_args(t, b, ny, sy)
    x_args = get_args(l, r, nx, sx)

            
    sections = _map(do_section, [(c, ya, xa) for ya in y_args for xa in x_args])
    for subim, ny, nx in sections:
        out[ny:ny+subim.shape[0], nx:nx+subim.shape[1]] = subim
        
    return out
    
N = 1000
t = time.time()
j = julia(-0.835 - 0.2321j, (-1, -2, 1, 2), (N, 2*N))
t = time.time() - t
print("Single section took %s s" % t)
plt.imshow(j)

t = time.time()
j = julia_sections(-0.835 - 0.2321j, (-1, -2, 1, 2), (N, 2*N), (100,100))
t = time.time() - t
print("100x100 sections took %s s" % t)
plt.imshow(j)

# with multiprocessing takes ~28s for a 10kx20k image. Not too shabby.
N = 10000
import multiprocessing
with multiprocessing.Pool() as p:
    t = time.time()
    j = julia_sections(-0.835 - 0.2321j, (-1, -2, 1, 2), (N, 2*N), (100,100), p.map)
    t = time.time() - t
    print("took %s s" % t)
plt.imshow(j)