Predator Prey Simulation

Question

Below is a simple random walk predator prey simulation that is optimized to the best of my abilities. I would love to hear about any improvements that can made.

import numpy as np
import time
from matplotlib import pylab as plt
import random

def run():
    # Initialize grid
    size = 100
    dims = 2
    # Each point in the 2D grid can hold the counts of: [prey,predators]
    grid = np.zeros((size,) * dims, dtype=(int, 2))
    num_rows, num_cols, identifiers = grid.shape

    num_predators = 10
    num_prey = 500
    prey_countdown = 500
    grid[50, 50, 1] = num_predators  # Manually inserting a few predators/prey
    grid[0, 0, 0] = num_prey

    # Coordinates for all non-empty grid locations
    coords = np.transpose(np.nonzero(grid != 0))
    x_pts, y_pts, idents = zip(*coords)

    # Please do not consider matplotlib the choke point of this program,
    # It will be commented out, and is only used for testing and amusement.
    # (But if you do have a way to speed it up, I'm very curious!)

    # Initialize figure and axes
    fig, ax1 = plt.subplots(1)
    # Cosmetics
    ax1.set_aspect('equal')
    ax1.set_xlim(0, size)
    ax1.set_ylim(0, size)
    # Display
    ax1.hold(True)
    plt.show(False)
    plt.draw()
    # Background is not to be redrawn each loop
    background = fig.canvas.copy_from_bbox(ax1.bbox)

    # Plot all initial positions as blue circles
    points = ax1.plot(x_pts, y_pts, 'bo')[0]

    # I would like to have blue for prey and red for predators,
    # I'm not sure how to do so quickly. I think multiple calls to axes.plot are needed.
    #colors = ['ro' if (ident==2) else 'bo' for ident in idents]

    time_steps = 1000
    for idx in range(time_steps):
        for coord in coords:
            direction = random.sample(range(1, 5), 1)[0]

            x, y, ident = coord
            count = grid[x, y, ident]

            # Random walk
            # Prey first
            if ident == 0:
                if count:  # A predator may have eaten the prey by now
                    grid[x, y, ident] -= 1  # Remove old value
                    if direction == 1:  # Move right
                        grid[(x+1) % num_rows, y, ident] += 1
                    elif direction == 2:  # Move left
                        grid[(x-1) % num_rows, y, ident] += 1
                    elif direction == 3:  # Move up
                        grid[x, (y+1) % num_cols, ident] += 1
                    elif direction == 4:  # Move down
                        grid[x, (y-1) % num_cols, ident] += 1
            # Predators do not die
            else:  # Predators will consume prey if prey exists at new location
                grid[x, y, ident] -= 1  # Remove old value
                if direction == 1:  # Move right
                    xnew = (x+1) % num_rows
                    grid[xnew, y, ident] += 1  # Move predator to new grid location
                    if grid[xnew, y, 0]:  # If there is prey at the new location...
                        grid[xnew, y, 0] -= 1  # Remove prey
                        prey_countdown -= 1
                        print 'Crunch! Prey left:', prey_countdown
                elif direction == 2:  # Move left
                    xnew = (x-1) % num_rows
                    grid[xnew, y, ident] += 1
                    if grid[xnew, y, 0]:
                        grid[xnew, y, 0] -= 1
                        prey_countdown -= 1
                        print 'Crunch! Prey left:', prey_countdown
                elif direction == 3:  # Move up
                    ynew = (y+1) % num_cols
                    grid[x, ynew, ident] += 1
                    if grid[x, ynew, 0]:
                        grid[x, ynew, 0] -= 1
                        prey_countdown -= 1
                        print 'Crunch! Prey left:', prey_countdown
                elif direction == 4:  # Move down
                    ynew = (y-1) % num_cols
                    grid[x, ynew % num_cols, ident] += 1
                    if grid[x, ynew, 0]:
                        grid[x, ynew, 0] -= 1
                        prey_countdown -= 1
                        print 'Crunch! Prey left:', prey_countdown

        # Redraw...
        coords = np.transpose(np.nonzero(grid != 0))
        x_pts, y_pts, idents = zip(*coords)
        points.set_data(x_pts, y_pts)
        fig.canvas.restore_region(background)  # Restore background
        ax1.draw_artist(points)  # Redraw just the points
        fig.canvas.blit(ax1.bbox)  # Fill in the axes rectangle

        time.sleep(0.01)  # Prevents figure from freezing

    plt.close(fig)

    # Do we have as many left as we should?
    coords = np.transpose(np.nonzero(grid[:,:,0] != 0))
    print (len(coords) != prey_countdown)

if __name__ == '__main__':
    run()  # Start button for cProfile

I hope to continue expanding this simulation to include a "chance to escape" for the prey as well. Any comments about the program's structure that will help ease growth will be greatly appreciated.

Answer 1

The biggest issue I see is code organization. There are two large tasks:

1. Numerical Simulation
2. Data Visualization

They are mixed together in ways that create arbitrary dependencies and make extending and maintaining the code difficult.

Names

Assigning meaningful names to numbers will improve readability and reduce comments:

right = 1
left = 2
up = 3
down = 4

if direction == up
    ...
if direction == down
    ...
etc.

Decoupling

Predators and prey are a different level of abstraction than anything in matplotlib. Predators/prey respond to other predators/prey and move of their own volition. 2d grids are something else entirely. Predator/prey logical abstractions should work for whales and giant squid in the sea, leopards and chimps in the forest or lions and wildebeests on the savannah. The visualization methods should not leak into the simulation's abstractions.

The simulation should be able to be run independently of any visualization. This allows running a million iterations, followed by statistical analysis using the Monte Carlo or similar methods.

This means that the simulation portion of the program has its own methods and data structures. Likewise the display portion of the program should have its own. In between, the main loop reads the simulation data structures and translates them into a visualization data structure [i.e. the main loop draws a pretty picture].

  #pseudo code for visualization
  on_clock_tick()
      my_display.draw(translate(simulate(my_simulation)))

An alternative use of a decoupled simulation:

  #pseudo code for later statistical analysis
  on_clock_tick()
      my_file.write(translate(simulate(my_simulation)))

Note that this also makes it possible to reuse the visualization code for a different simulation:

  #pseudo code for visualization
  on_clock_tick()
      my_display(translate(simulate(my_other_simulation)))

Final thoughts

The real scientific work is in simulating the behaviors of predators and prey. Getting the display code out of the way will allow you to focus on the core problem clearly. It will make the code more readable and prioritize bugs - it doesn't matter how right the visualization is if the underlying simulation is bad.

Answer 2

1. Functions: There are many places in your code were you have reused the same code without using a function.

   Ex:

   if grid[x, ynew, 0]:
       grid[x, ynew, 0] -= 1
       prey_countdown -= 1
       print 'Crush! Prey left:', prey_coutdown
   

   As well as its xnew counterpart shows up with the exact same code twice in your program. It may be advisable to turn these two sets of code into two functions and then call them were ever they are needed.

2. Naming conventions: I would advise that you follow the python PEP-8 standard for code structure and naming conventions.

   xnew and ynew would become x_new and y_new respectively.

Other than that I think that your code is quite nice looking.