Computation of population patterns in simple agent-based model

Question

I've modeled the population changes for a group of agents that inhabit a 2D lattice (with wrapped boundaries) made up of M x M number of grids. At each time step:

1. A parent reproduces offsprings that disperse independently -- the number depends on the grid location.
2. The offsprings live if they land on a "patch" grid but are removed otherwise.
3. Neighboring offsprings compete based on local density.
4. Survived offsprings grows up, becoming parents the next time step.

As an inexperienced programmer, my current Python script is pretty inefficient, taking as long as 19 minutes(!!) for the given parameters. I suspect I'm relying too much on NumPy functions but cannot see clear ways to optimize it.

import numpy as np
import random
import math as mp
import matplotlib.pyplot as plt
import time

# Lattice parameters
M = 60
patches = np.random.choice(2, M * M).reshape(M, -1)

# Biological parameters
mu = 14.5
alpha = .1
sigma = 5.

def reproduction(a_):
    da = np.floor(a_)    # discretize parent locations
    z = np.random.normal(0, 0.1, M**2).reshape(M,-1) # generate noise in reproduction
    muloc = np.ones(len(a_))   # mean reproduction rate per location
    for i in xrange(len(a_)):
        muloc[i] = mu * da[i,0]**0.01 + z[da[i,0], da[i,1]]   # grid-dependency
        if muloc[i]<0:     # mean reproduction rate cannot be negative
            muloc[i]=0
    return muloc

def dispersal(self): 
    return np.random.multivariate_normal([self[0], self[1]], [[sigma**2*1, 0], [0, 1*sigma**2]])%M    # Gaussian dispersal

def landed(offspring_list): 
    Inlist = []   # landed offspring list
    og = np.zeros((M, M))    # lattice defined by the number of landed offsprings per grid

    dl = np.floor(offspring_list)    # discretize offspring locations
    for i in range(len(offspring_list)):
        if patches[dl[i,0],dl[i,1]] == 1:
            Inlist.append(offspring_list[i])
    return Inlist

# competition kernel
ker = np.zeros((M, M))
for i in range(M):
    for j in range(M):
        r2 = min((i-0), M-(i-0))**2 + min((j-0), M-(j-0))**2
        rbf = mp.sqrt(1/(2*sigma**2))
        ker[i,j] = np.exp(-(rbf*r2)**2)
ker = ker/np.sum(ker)

def competition(offspring_list):
    dl = np.floor(offspring_list)   # discretize offspring location
    og = np.zeros((M,M))      # MxM lattice defined by the number of seeds at each grid
    for i in range(len(offspring_list)):
        og[dl[i,0],dl[i,1]] += 1

    v1 = np.fft.fft2(ker)        # fast fourier transform of kernel
    v2 = np.fft.fft2(og)         # fast fourier transform of offspring densities across lattice
    v0 = np.fft.ifft2(v1*v2)     # convolution via inverse fast fourier transform 
    dd = np.around(np.abs(v0), decimals=3)   # rounding resulting values

    alive = []    # indicators for competition outcome: 1 for alive, 0 for dead
    density = np.ones((len(dl),1))
    for i in range(len(dl)):
        density[i] = dd[dl[i,0],dl[i,1]]    #FFT-approximated density at each grid 
        mortality = 1 - 1 / (1 + density[i]/1.)  # competition-based mortality rate
        if np.random.random() < mortality:
            alive.append(0)
        else:
            alive.append(1) 
    return alive, og, dd

# Simulation
sim_num = 5
terminal_time = 50

start_time = time.time()
for s in range(sim_num):
    print 'simulation number ', s+1, ' of ', sim_num

    # initial conditions
    n = 200 
    parent = np.random.random(size=(n, 2)) * M
    pop_density = [n/float(M**2)]

    for t in range(terminal_time):
        print 't = ', t+1
        muloc = reproduction(parent)

        offspring = []
        for parent_id in range(len(parent)):
            for m in range(np.random.poisson(muloc[parent_id])):   # offsprings reproduced
                offspring.append(dispersal(parent[parent_id]))     # offspring dispersed
        offspring = landed(offspring)       # offsprings landed
        if len(offspring) == 0:  # possible extinction
            parent = []
        else:
            offspring = np.vstack(offspring)
            n = len(offspring)

            alive = competition(offspring)[0]   # offspring in competition
            indexes = [i for i, x in enumerate(alive) if x==1]
            if len(indexes) == 0:  # possible extinction
                parent = []   
            else:
                parent = np.vstack([offspring[l] for l in indexes])

        n = len(parent)
        pop_density.append(n/float(M**2))

    plt.plot(range(terminal_time+1), pop_density, 'k--', alpha=0.4)

print time.time() - start_time, 'seconds'
plt.xlabel('Time')
plt.ylabel('Density of Individuals')
plt.show()

Answer 1

I profiled your code using pycallgraph to find the bottlenecks. Two observations:

• Over 80% of the time is spent in a numpy function called ndarray.dispersal, so we need to find ways to cut down the use of numpy’s array functions.
• The most expensive user function is competition(), so we should start there first.

Python’s builtin structures are often a lot faster than numpy’s for simple data, so that’s always where I start when I’m trying to speed up some numpy code. The fact that ndarray is the main culprit strengthens the idea that overuse of numpy’s structures is the problem.

---

• In competition, rather than declaring go with np.zeros, I create it as a list of lists:

  og = [[0] * M] * M
  

  You need to make a few tweaks in competition() to modify the way indices are handled; you need to drop in this line instead:

  og[int(dl[i,0])][int(dl[i,1])] += 1
  

  This means that every time you access og, you’re using native Python lists rather than Numpy’s functions. This gave me a big speed boost.

• Running the profiler again, competition() is still the main culprit. What else can we change?

  I haven’t found anything that gets as good an improvement as fixing og, but here are some other small optimisations we can make:

  • Since the ker variable never changes, recalculating v1 every time we call competition() is wasteful. Pull the computation of v1 out of the function, and put it next to ker.

  • Build density as a list of lists, not a numpy array. Again, we make some small savings:

    density = [[1]] * len(dl)
    

• Running the profiler a third time, the reproduction() function is pretty close to the time of competition(). What can we do here?

  I haven’t looked into this in detail, but I think your random matrix is probably the cause of the problem. I think it’s slow to index, and that’s dragging you down. Consider converting it into a list of lists?

• After that, I’d probably start looking through landed(), which is the next slowest user function. You can drop the variable og, as it goes unused. Not sure what else, if anything, you could do here.

That’s about as much as I want to do for now, but hopefully it gives you a few pointers in the right direction. The code is certainly faster than I left it, anyway.

Summary: don’t over-rely on numpy’s structures. Sometimes the Python builtin structures are as good as, or better. Use a profiler to determine what’s more appropriate.

---

Some other general optimisations and/or code tweaks that I'd suggest:

• In your print statements, you often have extra spaces. When you print a series of strings separated by commas, Python always adds a space between them for you. Compare the output of these two lines:

  print 'simulation number ', s+1, ' of ', sim_num
  print 'simulation number', s+1, 'of', sim_num
  

  and you should see the extra spaces go away.

• The parameters seem somewhat scattered. Put all your variables right at the top of the file, so they’re easy to find and edit.

• Rather than running code in the top-level, move it into a main() function and put this at the bottom of the file:

  if __name__ == '__main__':
      main()
  

  This means that if the script is run directly, you get all the plotting as before, but you can also import functions from this file into another script. That’s useful for reusing code – it’s the best of both worlds.

• Don't use self as a variable name in dispersal; by convention, that's used for classes, and it's best not to use it to avoid confusion.

• In your dispersal function, the mean is the input of the function, and the covariance matrix is constant (as sigma is constant). I'd thus suggest rewriting the function this way:

  # Convention dictates that constants are in UPPERCASE_WITH_UNDERSCORES 
  M = 60
  SIGMA = 5
  COV = [[SIGMA ** 2, 0], [0, SIGMA ** 2]]
  
  def dispersal(mean):
      return np.random.multivariate_normal(mean=mean, cov=COV)
  

• Doing import math as mp just causes unnecessary confusion. It’s fine to do import math.

Answer 2

I ran the script in iPython Notebook using python 2.7. It took roughly eleven minutes... When timing the individual functions, i got the following per time step (i ran only 10 time steps to measure this, it might therefore reach larger numbers at the end of the simulation):

• reproduction: ~0.01 s
• landing: ~0.1 s
• competition: ~0.3 s

So it makes sense to look at competition first. I honestly don't fully understand what is happening here. But some hints that may or may not help:

• Did you try scipy.signal.fftconvolve instead of convoluting manually via fft?
• Could you calculate density as an array at once using numpy instead of using the for-loop? You could calculate mortality as well as the random numbers on this grid as well and then simply read out the `alive? values.

In general, you may be better off using one grid over the whole simulation (at least per time step). You create one MxM grid in landed and one in competition. You could use a bit flag array (np.dtype='uint8') and set flags for offspring, alive, dead, etc. Creating all these arrays consumes time and memory.