How do I use Python multiprocessing to get rid of nested while loops

Question

I made this ray tracer to simulate two lens. Each time a ray hits a screen, the position on that screen is voted for. However, this is running so slowly. It's nothing with the ray tracer, is just the sheer number of rays (128^4). I know multithreading won't work for this, but how do I use multiprocessing to vote for pixels on the screen. It just boggles my mind - how do I have each child vote for a pixel. Each ray is independent. Here is some sample code. There's also an issue with creating the input data. It will take a long time just to make an array of rays to send to a multiprocessing class.

Basically, I want the inner loop to spawn upto 8 child processes - have those child processes vote - then free up a slot...I think even that would run slow.

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt
from matplotlib.mlab import griddata
from matplotlib import cm
import random

A = 0
B = 0
C = 0
D = 0
E_initial = np.zeros((512,512))

while A <= 50:
    B = 0
    while B <= 50:
        C = 0
        while C <= 50:
            D = 0
            while D <= 50:
                #Ray tracing here             
                x_contact = random.randint(0,511)
                y_contact = random.randint(0,511)               
                E_initial[x_contact,y_contact] += 1
                D = D + 1
            C = C +1
        B = B + 1
    print A
    A = A + 1
fig = plt.figure()
ax = fig.add_subplot(111)
plt.imshow(E_initial)
cbar = plt.colorbar(orientation='vertical')
cbar.set_label('# of contacts', rotation=270, labelpad=10)
plt.show()

Answer 1

Before trying multiprocessing, note that you can eliminate the quadruple while-loops by taking advantage of some numpy functions.

indices = np.random.randint(0, 512**2, size=(50**4))

generates random index locations between 0 (inclusive) and 512**2 (exclusive). Each index location corresponds to a location in E_initial .

Then you could use np.bincount to count how many times each index location occurs. The array returned by np.bincount will be of size 512**2, and will contain an integer count at each location corresponding to the number of times that index location appears in indices . But ahah -- that's essentially what we want E_initial to equal.

import numpy as np
import matplotlib.pyplot as plt

indices = np.random.randint(0, 512**2, size=(50**4))
E_initial = np.bincount(indices).reshape((512, 512))

fig = plt.figure()
ax = fig.add_subplot(111)
plt.imshow(E_initial)
cbar = plt.colorbar(orientation='vertical')
cbar.set_label('# of contacts', rotation=270, labelpad=10)
plt.show()

PS. Having 8 processes modify one array is going to be slow because array must be locked during assignments to avoid race conditions which would lead to erroneous results. Having to lock with every assignments will make the multiprocessing code much slower than equivalent single-process code.

PPS. Although

E_initial[x_contact,y_contact] += 1

is standard practice in some languages such as C, it is dead slow in Python when using a NumPy array. It's slow because each assignment involves at least three Python function calls ( __getitem__ , __iadd__ , __setitem__ ). If you put that inside the quadruple loop you get (at least) 3*50**4 Python function calls. If you can reduce all those Python calls to one NumPy function call you get a big performance boost. So avoid assigning values to individual locations in a NumPy array one-by-one if at all possible! To take advantage of NumPy, you want to minimize the number of function calls and loops, and maximize the amount of work down by as few calls to NumPy functions as possible. That usually means forming a big array and shoving the whole thing into a NumPy function which does the calculation in C or Fortran.