Parallel Reaction-Diffusion: Strips and Halos in Python

Code along video

Learning objectives

Our serial reaction-diffusion simulator is finished, and in this code along, we will begin parallelizing it. We will write simulate_parallel(), which creates a pool of workers once and reuses it for every batch of steps, dividing the board into vertical strips and handing one strip to each worker with starmap(). Along the way, we will meet a new idea from parallel programming called a halo: the particles just across a strip’s boundary that a worker needs to see so that reactions crossing that boundary can still happen. We will then implement assign_strips_with_halos(), taking care to handle the halos that wrap around the edges of our toroidal board.

Code along summary

Setup

We will continue working in the reaction_diffusion directory from the previous code alongs, this time in parallel.py, which contains the signatures and docstrings of the parallel functions with empty bodies that we will fill in over this code along and the next.

Starting simulate_parallel()

Our parallel simulator, simulate_parallel(), takes the same inputs as simulate_serial(): the initial list of particles, the reaction table, and the simulation parameters. The only new parameter is an integer num_procs, the number of workers that we will use. Its output, a list of snapshots of the system, is exactly the same. This is a pattern that we have seen throughout the chapter: from the outside, a parallel function should look just like its serial counterpart, with an extra parameter saying how many workers to use.

def simulate_parallel(particles: list[Particle],
                      table: ReactionTable,
                      params: SimulationParameters,
                      num_procs: int) -> list[list[Particle]]:

The start of the function is the same as in the serial version, so we can borrow it from serial.py. We make a copy of the particles to serve as the current state of the system, start our list of snapshots with a copy of that state, and then determine how many batches of params.proximity_interval steps we will run, along with the number of leftover steps.

# first, let's make the current list of particles and the list of "snapshots" that we eventually want to return
current: list[Particle] = copy.deepcopy(particles)
frames: list[list[Particle]] = [copy.deepcopy(current)]

# question: how many snapshots will I eventually take?
# remember: proximity_interval = number of steps between snapshots
num_batches = params.num_steps // params.proximity_interval
remainder = params.num_steps % params.proximity_interval

return frames

A pool of workers

Here is where things start to change. When we parallelized diffusion, we created a new collection of workers at every step. For reaction-diffusion, which runs many batches, creating a fresh set of workers every time would add a huge amount of overhead, since every new process has to start up its own copy of the Python interpreter. Instead, we will create a worker pool one time, using multiprocessing.Pool(), and those workers will persist, ready to take on new work for as long as the simulation runs. All that Pool() needs is the number of workers that we want.

# let's make our worker pool
pool = multiprocessing.Pool(num_procs)

We then range over the batches, as we did in the serial simulator. At each batch, we will divide the board into vertical strips, one per worker, so that if we have 10 cores and 10 workers, each worker receives the particles that fall in its strip. The function that does this division needs the current particles, the number of strips (num_procs), and the width of the board. We do not need the height, because our strips are vertical; we could have made them horizontal instead, but vertical strips will do.

Halos

There is one more key point to keep in mind. Each worker needs to know the particles that fall in its strip, but that is not enough. A particle that sits right at the boundary of a strip could react with a particle just across the boundary, which belongs to the neighboring strip. If a worker can only see its own particles, then that reaction can never happen, and our parallel simulation would quietly behave differently from the serial one.

This is a common problem in parallel programming. If we divided a cellular automaton board into strips, then a cell on the boundary of a strip would need to know the states of the cells one step away on the other side of the boundary. The solution is to give each worker a halo: a region just outside its strip that it can see into, as if peering across the fence into its neighbor’s yard. For reaction-diffusion, a strip’s halo is the list of particles from neighboring strips that are close enough to the strip to possibly take part in a reaction with one of its particles.

How close is close enough? Two particles can only react if they are within params.activation_range of each other, so the width of the halo on each side of the strip should be the activation range. And because our board is a torus, the halos of the leftmost and rightmost strips wrap around the edges of the board, as shown in the figure below.

We will therefore have each batch call a function assign_strips_with_halos() that returns two lists, each of length num_procs: the particles in each strip, and the particles in each strip’s halo. We pass params.activation_range as the width of the halo.

# again, we divide into batches
for _ in range(num_batches):
    # divide the space into vertical strips
    # key tricky point: each worker needs to know the collection of points that fall in its strip + the collection of particles that are close enough to its strip to maybe be involved in a reaction
    strips, halos = assign_strips_with_halos(current, num_procs, params.width, params.activation_range)

Handing out the work with starmap()

As we did with diffusion, we next build a list of tuples, worker_args, where the i-th tuple holds the inputs that the i-th worker needs: its strip strips[i], its halo halos[i], the reaction table, the parameters, and the number of steps that it should take, params.proximity_interval. (The worker could infer this last value from params, but passing it separately will come in handy.)

Then we call the critical function, the starmap() method of our pool, passing it the function that every worker will run, process_strip(), along with worker_args. Each worker runs its own instance of process_strip() on its own tuple of arguments, and this one line is where the parallelization happens. We have not written process_strip() yet; that will be the subject of the next code along. All that we need to know right now is that it returns the list of particles in a single strip after that strip has moved and reacted.

# next, create list of all the arguments that are passed as inputs to each worker
worker_args: list[tuple] = []
for i in range(num_procs):
    worker_args.append((strips[i], halos[i], table, params, params.proximity_interval))

# now we call the critical pool method (starmap)
# this is where the parallelization happens
results = pool.starmap(process_strip, worker_args)

Because each worker returns a list of particles, results is a two-dimensional list of particles, with one list per strip. To obtain the next snapshot of the whole system, we set current equal to an empty list and then extend it by each strip’s list in turn, before appending current to frames.

# now we compile the results, list[list[Particle]]
current = []
for strip in results:  # strip is list[Particle]
    current.extend(strip)

# now, current is set: all particles in our snapshot
frames.append(current)

Finally, once every batch is done, we no longer need our workers, so we fire them. Calling pool.close() followed by pool.join() shuts down the worker processes and makes sure that their memory is cleaned up. You will see these two lines at the end of any function that uses a pool.

Here is simulate_parallel() at this point.

def simulate_parallel(particles: list[Particle],
                      table: ReactionTable,
                      params: SimulationParameters,
                      num_procs: int) -> list[list[Particle]]:
    # first, let's make the current list of particles and the list of "snapshots" that we eventually want to return
    current: list[Particle] = copy.deepcopy(particles)
    frames: list[list[Particle]] = [copy.deepcopy(current)]

    # question: how many snapshots will I eventually take?
    # remember: proximity_interval = number of steps between snapshots
    num_batches = params.num_steps // params.proximity_interval
    remainder = params.num_steps % params.proximity_interval

    # let's make our worker pool
    pool = multiprocessing.Pool(num_procs)

    # again, we divide into batches
    for _ in range(num_batches):
        # divide the space into vertical strips
        # key tricky point: each worker needs to know the collection of points that fall in its strip + the collection of particles that are close enough to its strip to maybe be involved in a reaction
        strips, halos = assign_strips_with_halos(current, num_procs, params.width, params.activation_range)

        # next, create list of all the arguments that are passed as inputs to each worker
        worker_args: list[tuple] = []
        for i in range(num_procs):
            worker_args.append((strips[i], halos[i], table, params, params.proximity_interval))

        # now we call the critical pool method (starmap)
        # this is where the parallelization happens
        results = pool.starmap(process_strip, worker_args)

        # now we compile the results, list[list[Particle]]
        current = []
        for strip in results:  # strip is list[Particle]
            current.extend(strip)

        # now, current is set: all particles in our snapshot
        frames.append(current)

    # at the end of the process, fire all the workers
    pool.close()
    pool.join()

    return frames

This is a remarkable thing about our parallel program: everything that it needs in order to run in parallel is contained in this one function. Everything else that we write will be serial code, much of it similar to code that we have already written.

Assigning particles to strips

We could parallelize assign_strips_with_halos() too, but that would be overkill. Recall that a task should be substantial enough to justify the overhead of parallelizing it, and dividing a few thousand particles among a handful of strips is fast. So we will write this function as ordinary serial code.

The function takes the list of particles, the number of strips, the width of the board, and the width of the halo, halo_width. We give this last parameter its own name because it answers the question, “If I am on the boundary between two strips, how far across the fence can I see?” When we call the function, we pass params.activation_range, the maximum distance at which two particles can react. Be careful not to confuse this parameter with params.proximity_interval, which is the number of random steps that each particle takes between rounds of reactions.

def assign_strips_with_halos(particles: list[Particle],
                             num_strips: int,
                             width: float,
                             halo_width: float) -> tuple[list[list[Particle]],
                                                         list[list[Particle]]]:

Dividing the particles into strips is something that we already did when we parallelized diffusion, so we can borrow that code from diffusion/parallel.py. (We called them “stripes” there and call them “strips” here; either word works.) The width of each strip is the width of the board divided by the number of strips, and each strip starts out empty. Then, we range over the particles, and we place each particle into the strip at index int(p.position.x / strip_width).

We add one small safety check. If a particle’s index is somehow at least num_strips, which would be outside the board, we set it to num_strips - 1, and if its index is negative, we set it to 0. These two checks cost us almost nothing, but they could save us from a program crash if a particle ever lands exactly on the edge of the board.

# first, set strip width and declare them
strip_width = width/num_strips
strips: list[list[Particle]] = []

# every strip starts as empty
for _ in range(num_strips):
    strips.append([])

# we range over particles and assign each one to the strip it belongs to
for p in particles:
    idx = int(p.position.x / strip_width)

    # it could be that the point has walked off the board but has not been wrapped
    if idx >= num_strips:
        idx = num_strips-1
    if idx < 0:
        idx = 0

    # now, assign the current particle to its appropriate strip
    strips[idx].append(p)

Building each strip’s halo

Now for the part that is new. We declare halos in the same way as strips, as a two-dimensional list with one empty list per strip. Then we range over the strips, and for each strip i, we compute the range of x-coordinates that the strip covers, from x_lo = i*strip_width to x_hi = (i+1)*strip_width. This idea of multiplying i and i+1 by a width has appeared again and again in our work on parallel algorithms.

Most of the time, the halo of strip i is straightforward. We range over the particles in the strip to the left, strips[i-1], and add a particle to halos[i] if its x-coordinate is at least x_lo - halo_width, meaning that it is close enough to strip i to react across the boundary. Similarly, we range over the particles in the strip to the right, strips[i+1], and add a particle if its x-coordinate is at most x_hi + halo_width.

The trouble is that these two checks only work when we are not wrapping around the board. The first strip has no strip to its left, and the last strip has no strip to its right, so we only check the left neighbor when i > 0 and the right neighbor when i < num_strips-1.

# we now need to build each strip's halo (the particles that could possibly be used in a reaction in a neighboring strip)
halos: list[list[Particle]] = []

# every halo starts as empty
for _ in range(num_strips):
    halos.append([])

# let's range over all the strips and determine the particles belonging to its halo
for i in range(num_strips):
    # what range of x values are possible in this strip
    x_lo = i*strip_width
    x_hi = (i+1)*strip_width

    # pull in the particles from the left strip (if not in first strip)
    if i > 0:
        for p in strips[i-1]:
            if p.position.x >= x_lo - halo_width:
                halos[i].append(p)

    # pull in the particles from the right strip (if not in last strip)
    if i < num_strips-1:
        for p in strips[i+1]:
            if p.position.x <= x_hi + halo_width:
                halos[i].append(p)

The toroidal wrap

Because our board is a torus, the first and last strips are neighbors across the edge of the board. If i is 0, then we range over the particles in the last strip, strips[num_strips-1], and add any whose x-coordinate is at least width-halo_width to the halo. For example, if the width of the board is 1000 and the halo width is 2, then we want every particle in the last strip whose x-coordinate is at least 998, since such a particle is within distance 2 of a particle at x = 0 when we wrap around.

Symmetrically, if i is num_strips - 1, then we range over the particles in the first strip, strips[0], and add any whose x-coordinate is at most halo_width. Once we have built every halo, we return both lists.

    # finally, apply "toroidal wrap" to the first and last strip
    if num_strips > 2:
        if i == 0:
            for p in strips[num_strips-1]:
                if p.position.x >= width-halo_width:
                    halos[i].append(p)
        elif i == num_strips - 1:
            # look in the first strip for anything close to left edge
            for p in strips[0]:
                if p.position.x <= halo_width:
                    halos[i].append(p)

return strips, halos

When there is only one strip, it is both the first and the last strip, and its “neighbor” across the edge of the board is itself. Its particles already belong to the strip, so they must not also appear in its halo, which is why we need the condition. But as Phillip realizes in the video, > 2 is not quite right. When there are two strips, strip 0 sees strip 1 across their shared boundary in the middle of the board, but strip 0 also touches strip 1 across the edge of the board at x = 0. Without the wrap, a particle at x = 995 would never be added to the halo of strip 0, and so it could never react with a particle at x = 3 in strip 0, even though the two are only 8 units apart on the torus. So we should apply the wrap whenever there are at least two strips.

# finally, apply "toroidal wrap" to the first and last strip
if num_strips >= 2:

Here is the finished function, with the docstring from the starter code omitted.

def assign_strips_with_halos(particles: list[Particle],
                             num_strips: int,
                             width: float,
                             halo_width: float) -> tuple[list[list[Particle]],
                                                         list[list[Particle]]]:
    # first, set strip width and declare them
    strip_width = width/num_strips
    strips: list[list[Particle]] = []

    # every strip starts as empty
    for _ in range(num_strips):
        strips.append([])

    # we range over particles and assign each one to the strip it belongs to
    for p in particles:
        idx = int(p.position.x / strip_width)

        # it could be that the point has walked off the board but has not been wrapped
        if idx >= num_strips:
            idx = num_strips-1
        if idx < 0:
            idx = 0

        # now, assign the current particle to its appropriate strip
        strips[idx].append(p)

    # we now need to build each strip's halo (the particles that could possibly be used in a reaction in a neighboring strip)
    halos: list[list[Particle]] = []

    # every halo starts as empty
    for _ in range(num_strips):
        halos.append([])

    # let's range over all the strips and determine the particles belonging to its halo
    for i in range(num_strips):
        # what range of x values are possible in this strip
        x_lo = i*strip_width
        x_hi = (i+1)*strip_width

        # pull in the particles from the left strip (if not in first strip)
        if i > 0:
            for p in strips[i-1]:
                if p.position.x >= x_lo - halo_width:
                    halos[i].append(p)

        # pull in the particles from the right strip (if not in last strip)
        if i < num_strips-1:
            for p in strips[i+1]:
                if p.position.x <= x_hi + halo_width:
                    halos[i].append(p)

        # finally, apply "toroidal wrap" to the first and last strip
        if num_strips >= 2:
            if i == 0:
                for p in strips[num_strips-1]:
                    if p.position.x >= width-halo_width:
                        halos[i].append(p)
            elif i == num_strips - 1:
                # look in the first strip for anything close to left edge
                for p in strips[0]:
                    if p.position.x <= halo_width:
                        halos[i].append(p)

    return strips, halos

This function is a bit of a nightmare: it is long, and it is not clear how we could hand any of its work to a subroutine. Making our simulation work on a torus was mostly painless until now, but setting up strips and halos is where the torus really makes us pay. Fortunately, it is done.

Looking ahead

We now have the engine of our parallelization: simulate_parallel() sets up the pool and hands out the work, and assign_strips_with_halos() decides which particles each worker owns and which ones it can see. What remains is the function that each worker actually runs inside starmap(). In the next code along, we will write process_strip(), which moves the particles in a strip and its halo and then applies reactions, and we will take care of the leftover steps that we skipped over in simulate_parallel().