Mandelbrot in parallel

# -*- coding: utf-8 -*-
"""
The following code, building on discussion in [1] demonstrates how parallel processing
on a multicore CPU could speed up computations in Python. However, it hardly scratches the surface
of the complexities that parallel processing can entail.

The code produces the mandelbrot fractal in a trivial way. See [2] for how to do it a lot faster.

Some notes on parallel processing in Python:
    * Some things that work on Linux won't work in Windows, because of windows' lack of
      some "forking" capability.
    * You can run multiple threads in Python, but as explained in [3],
      Python will not be running these in parallel, so this can't help you with using multicores
    * This code uses the multiprocessing pool to create different "subprocesses"
      that windows then delegate to different cores. An alternative is "Process" in Python, see comments in [1]


On my 4-core CPU, with 8 "logical cores" (so it's really 4) the parallel code runs twice as fast
as the single-core code. You can use the Windows Resource Monitor to verify that the first run
only uses a single core, while the second run utilizes all cores.

[1] https://stackoverflow.com/questions/44521931/what-is-some-example-code-for-demonstrating-multicore-speedup-in-python-on-windo
[2] https://www.ibm.com/developerworks/community/blogs/jfp/entry/How_To_Compute_Mandelbrodt_Set_Quickly?lang=en
[3] https://www.ploggingdev.com/2017/01/multiprocessing-and-multithreading-in-python-3/

"""

import numpy as np
import matplotlib.pyplot as plt
import time
from multiprocessing import Pool, cpu_count
from functools import partial
import sys

#The mandelbrot fractal is stored in an array of size height x width x 3,
#where dataarray[height,width,2] stores the actual value of a pixel and [..,0] and [..,1] stores the
#real and imaginary components of the pixel's coordinate.


#The calculatePixel function calculates the value of a single pixel in the picture of
#the mandelbrot fractal.The function both returns the value and writes it directly to dataarray.
#To have parallel processes share data like this is usually bad coding practice and
#could even hinder parallel processing if each subprocess tries to "lock" access to it.
#However, in this case it seems to run fine.

#Note: Python must be able to "pickle" the function you want to parallelize [4]
#[4] https://stackoverflow.com/questions/8804830/python-multiprocessing-pickling-error
def calculatePixel(l,k,dataarray):
    #This function
    #print('koord:',l,k)
    z =  0
    temp = 0
    c =  complex(dataarray[l,k,1], dataarray[l,k,0])
    for m in range(40):
            z = z*z + c
            if abs(z) > 2:
                dataarray[l,k,2] = m
                temp = m
                break

    return temp



def runforestrun():
# multi-platform precision clock
    get_timer = time.clock if sys.platform == "win32" else time.time

    runSequential = True
    runParallel = True

#Populates the data array with real and complex coordinates for each pixel.
#A more elegant Python "comprehension" probably exists :-)

    for j in range(height):
       #kolonner fylles nedover med verdiar frå 1 til -1.
       #Dette representerar Im(c)
       #kunne bytta ut med range() men får størrelse autoamtisk her
        data[j,:,0] = [1 - 2*j/height  for x in data[j,:,0]]

    for j in range(width):
        #rader representerar Re(c)
        data[:,j,1] = [-2 + 3*j/width  for x in data[:,j,1]]
        #data[:,j,1] = [j/width - (3/2) for x in data[:,j,0]]





    if runSequential:
        print("Starting sequential run...")
        tic = get_timer()
        #Here we just loop over all pixels one by one
        for k in range(width):
            for l in range(height):
                calculatePixel(l,k,data)
        print("Time spent sequentially",get_timer()-tic)
        #Lets print the result
        printmandel()

    if runParallel:
        print("Starting parallel run with Pool...")
        tic = get_timer()

        #We initialize a pool of workers, the number of which is determined by the number of cores
        #on the CPU. You may want to make multiple runs with different numbers, to demonstrate speedup.
        p = Pool(cpu_count())

        #The multiprocessing package of Python has a function called "starmap"
        #that does exactly what "map" does in Python, except it only accepts a single input
        #for the functions it "maps". We jump through hoops with "partial" to circumvent this.
        #See [5] for a discussion and alternative approaches.
        #[5] https://stackoverflow.com/questions/5442910/python-multiprocessing-pool-map-for-multiple-arguments

        partial_calculatePixel = partial(calculatePixel, dataarray=data)
        koord = []
        for j in range(height):
            for k in range(width):
                koord.append((j,k))

        #Runs the calls to calculatePixel in a pool. "hmm" collects the output
        hmm = p.starmap(partial_calculatePixel,koord)

        print("Time spent in parallel run:",get_timer()-tic)

        #Since calculatePixel edits "data" directly, we don't really need the starmap output
        #but for complection, this is how you would put the output back into the data array:
        data[:,:,2] = np.array(hmm).reshape((height,width))

        #Print the result
        printmandel()

def printmandel():
    #Also see: stackoverflow.com/questions/6916978/how-do-i-tell-matplotlib-to-create-a-second-new-plot-then-later-plot-on-the-o
    plt.figure()
    plt.imshow(data[:,:,2],cmap=plt.plasma())
    plt.colorbar()

#We only want to run calculatePixel in parallel, but Python loads the whole file.
#We therefore use the "if main" check as below, to avoid an infinite loop of new runs being spawned
# "if __name__ == '__main__':" ensures that it's only run when YOU start it.

if __name__ == '__main__':
   print('Loading multicorespeedup.py')
   #height determines the resolution of the mandelbrot fractal image.
   #large values may give an out of memory error. A hint that something in this code
   #is far from optimal :-)
   height = 2000
   width = int(height*1.5)
   #"data" is declared here, so that all functions in the module will have access to it.
   data = np.zeros( (height,width,3), dtype=np.float )
   #The main part of the code:
   runforestrun()


#Typical output from a run with height = 2000 on a 4-core (8 logical) Intel CPU (i7-3632QM @2.20GHz)
#Loading multicorespeedup.py
#Starting sequential run...
#Time spent sequentially 40.12519295999846
#Starting parallel run with Pool...
#Time spent in parallel run: 20.85525935829719