Evolution of a trait in an asexual population (performance tuning)

Question

I just coded a simple simulation module that looks at the evolution of a continuous trait in a haploid asexually reproducing population under density dependent competition in discrete time (i.e. non-overlapping generations, using recurrence equations). What I am interested in is finding out whether evolution would always favour selecting for increased intrinsic growth rates R, perhaps eventually pushing the population to go extinct (a scenario known as "evolutionary suicide"), or if instead there would be selection for restrained growth rates, e.g. due to the fact that more selfish lineages with higher intrinsic growth rates would more often move into the chaotic regime and go extinct faster.

The module I have takes as arguments the desired fitness function (e.g. the discrete logistic model), the initial trait values of your individuals in the 1st generation, the mutation rate and the standard deviation of the normal deviation that is applied to mutants and the nr of generations to run the simulation :

EvolveHapl[fitnessfunc_, initpop_, mutrate_, stdev_, generations_] := 
  Module[{ndist, traitvalues, currpopsize, fastPoisson, fitnessinds, 
    numberoffspring, nrmutants, rnoise, rndelem},
   ndist = NormalDistribution[0, stdev] ;
   traitvalues = Table[{}, {generations + 1}]; (* 
   list of lists containing ind trait values in each generation *)
   traitvalues[[1]] = initpop; 
   currpopsize = Length[traitvalues[[1]]]; 
   (* fast Poisson random number generator *)
   fastPoisson = 
    Compile[{{\[Lambda], _Real}}, 
     Module[{b = 1., i, a = Exp[-\[Lambda]]}, 
      For[i = 0, b >= a, i++, b *= RandomReal[]];
      i - 1], RuntimeAttributes -> {Listable}, 
     Parallelization -> True];
   Do[fitnessinds = 
     Table[fitnessfunc[traitvalues[[gen - 1]][[i]], currpopsize], {i, 
       1, currpopsize}];            (* 
    fitness of every individual in the population, 
    in mean number of offspring *)
    numberoffspring = fastPoisson[fitnessinds];            (* 
    absolute number of offspring that every individual produces *)   
    traitvalues[[gen]] = 
     Flatten[Table[
       Table[traitvalues[[gen - 1]][[i]], {j, 1, 
         numberoffspring[[i]]}], {i, 1, currpopsize}]];   (* 
    expected offspring trait values before mutation *)
    currpopsize = Length[traitvalues[[gen]]];  (* 
    new population size *)
    nrmutants = 
     RandomVariate[BinomialDistribution[currpopsize, mutrate]]; (* 
    nr of offspring that should mutate *)
    rnoise = RandomReal[ndist, nrmutants]; (* 
    noise to be added to the trait values of the mutants *)
    Do[rndelem = RandomInteger[{1, currpopsize}]; 
     traitvalues[[gen]][[rndelem]] = 
      Max[traitvalues[[gen]][[rndelem]] + rnoise[[i]], 0];, {i, 1, 
      nrmutants}];, (* mutate trait values *)
    {gen, 2, generations + 1}];
   Return[traitvalues ]];

And to plot the resulting list of individual trait values I use

PlotResult[traitvalues_] := Module[{},
  generations = Length[traitvalues];
  Print["Mean phenotype at the beginning : " <> 
    ToString[Mean[traitvalues[[1]]]]];
  Print["Maximum mean phenotype at any generation : " <> 
    ToString[
     Max[Table[
        Mean[traitvalues[[i]]], {i, 1, generations}] /. {Mean[{}] -> 
         0}]]];
  Print["Mean phenotype after " <> ToString[ngens] <> 
    " generations : " <> 
    ToString[
     Mean[traitvalues[[generations]]] /. {Mean[{}] -> 
        "- (population extinct)"}]];
  Print["Final population size : " <> 
    ToString[Length[traitvalues[[generations]]]]];

  maxscaleN = 
   Max[Table[{Length[traitvalues[[i]]]}, {i, 1, generations}]];
  minscaleN = 
   Min[Table[{Length[traitvalues[[i]]]}, {i, 1, generations}]];
  maxscaleP = Max[Flatten[traitvalues]];

  GraphicsRow[{Show[
     ArrayPlot[
      Table[BinCounts[
         traitvalues[[i]], {0, maxscaleP + 0.5, 0.05}], {i, 1, 
         generations}]/
       Table[Length[traitvalues[[i]]] + 0.00001, {i, 1, generations}],
       DataRange -> {{0, maxscaleP + 0.5}, {1, generations}}, 
      DataReversed -> True], Frame -> True, 
     FrameLabel -> {"Phenotype frequency", "Generation"}, 
     FrameTicks -> True, AspectRatio -> 2],
    ListPlot[
     Table[{Length[traitvalues[[i]]], i}, {i, 1, generations}], 
     Joined -> True, Frame -> True, 
     FrameLabel -> {"Population size", "Generation"}, 
     AspectRatio -> 2, 
     PlotRange -> {{Clip[minscaleN - 50, {0, \[Infinity]}], 
        maxscaleN + 50}, {0, generations}}]}]
   ]

Running it, however, is quite slow, e.g.

psize = 300; ngens = 5000; mutrate = 0.01; stdev = 0.05; K = 2*psize; 
f[R_, popsize_] := Max[(1 + R *(1 - popsize/K)), 0.00001];
First@AbsoluteTiming[
  traitvalues = 
    EvolveHapl[ f, RandomReal[{2.5, 2.6}, psize], mutrate, stdev, 
     ngens];]
PlotResult[traitvalues]

204.02

I was just wondering if there might be any obvious ways to speed up this routine in Mathematica? E.g. could I get rid of all my Do[] and Table[] loops somehow? Or could this whole routine or parts of it be compiled to speed it up? I've already replaced the Poisson number generation with a faster compiled version and used Max[] rather than Clip[] in my original post, but this hasn't improved speed as much as I had hoped it would... Or better to go to C or C++ for this kind of simple simulation?

Answer 1

Here is a modified version of your code. On my PC it completes your example run in about 5 seconds.

I won't try to describe every change but will point out the major features. Some of the changes are stylistic rather than performance-based. This is not a criticism of your style but a reflection of the way I broke the original code down in order to understand it.

A quick note on Map

The most widespread change was to use Map in a lot of places where you had used Table. Generally in Mathematica whenever you find yourself using Table to iterate over the elements in a pre-existing list, you can usually find a more functional approach. As a simple example

Table[Length[traitvalues[[i]]], {i, 1, generations}]

can be written as

Length /@ traitvalues

Often this change will provide a performance boost, and it makes the code more readable too (once you are accustomed to the notation)

The altered code

The main Do loop in EvolveHapl is an iteration to compute the next generation of traitvalues from the previous. This can be done very neatly using NestList. My version of EvolveHapl just calls NestList to do all the work:

EvolveHapl[fitnessfunc_, initpop_, mutrate_, stdev_, generations_] := 
 NestList[genStep[#, fitnessfunc, mutrate, stdev] &, initpop, generations]

The function genStep computes the new population from the current one:

genStep[tv_, fitnessfunc_, mutrate_, stdev_] :=
 Module[{n, fitnessinds, numberoffspring, newtv},
  n = Length @ tv;
  fitnessinds = fitnessfunc[tv, n];
  numberoffspring = fastPoisson[fitnessinds];
  newtv = Join @@ MapThread[ConstantArray, {tv, numberoffspring}];
  addNoise[newtv, mutrate, stdev]]

The most significant speed-up is in the calculation of fitnessinds - I pass the whole population to the fitness function. Obviously this requires that the fitness function is able to operate on a list. In your case, it can (provided we use the original version with Clip rather than Max)

genStep calls another function addNoise to do the mutations:

addNoise[traitvalues_, mutrate_, stdev_] :=
 Module[{tv, n, nrmutants, rnoise, rndelem},
  tv = traitvalues;
  n = Length[tv];
  nrmutants = RandomVariate[BinomialDistribution[n, mutrate]]; 
  rnoise = RandomReal[NormalDistribution[0, stdev], nrmutants];
  rndelem = RandomChoice[Range[n], nrmutants];
  tv[[rndelem]] += rnoise;
  Clip[tv, {0, ∞}]]

One thing to point out about addNoise is that I used RandomChoice to identify which elements to mutate. This avoids picking the same one more than once (which I assume is desirable).

I pulled fastPoisson out as a global function, no changes to it but I'll copy it here for completeness:

fastPoisson = Compile[{{λ, _Real}},
 Module[{b = 1., i, a = Exp[-λ]},
  For[i = 0, b >= a, i++, b *= RandomReal[]]; i - 1],
 RuntimeAttributes -> {Listable}, Parallelization -> True];

I made some minor changes to the PlotResults function, mostly converting to use Map where possible.

PlotResult[traitvalues_] :=
 Module[{generations, pop, maxscaleN, minscaleN, maxscaleP, frequencydata},
  generations = Length[traitvalues];
  Print["Mean phenotype at the beginning : " <>
    ToString[Mean[traitvalues[[1]]]]];
  Print["Maximum mean phenotype at any generation : " <>
    ToString[Max[Mean /@ traitvalues /. Mean[{}] -> 0]]];
  Print["Mean phenotype after " <> ToString[generations] <> " generations : " <>
    ToString[Mean[traitvalues[[generations]]] /. Mean[{}] -> "- (population extinct)"]];
  Print["Final population size : " <>
    ToString[Length[traitvalues[[generations]]]]];
  pop = Length /@ traitvalues;
  maxscaleN = Max[pop];
  minscaleN = Min[pop];
  maxscaleP = Max[traitvalues];
  frequencydata = (BinCounts[#, {0, maxscaleP + 0.5, 0.05}] & /@ traitvalues)/(pop + 0.00001);
  GraphicsRow[{
    Show[ArrayPlot[frequencydata, 
      DataRange -> {{0, maxscaleP + 0.5}, {1, generations}}, 
      DataReversed -> True], Frame -> True, FrameTicks -> True, AspectRatio -> 2,
      FrameLabel -> {"Phenotype frequency", "Generation"}],
    ListPlot[Transpose[{pop, Range[generations]}],
     Joined -> True, Frame -> True, 
     FrameLabel -> {"Population size", "Generation"}, AspectRatio -> 2,
     PlotRange -> {{Clip[minscaleN - 50, {0, ∞}], maxscaleN + 50}, {0, generations}}]}]]

Test

As I mentioned I defined the fitness function using Clip so that it could work on a whole population at once. So the full test code is:

psize = 300; ngens = 5000; mutrate = 0.01; stdev = 0.05; K = 2*psize;

f[R_, popsize_] := Clip[(1 + R (1 - popsize/K)), {0.000001, ∞}];

First@AbsoluteTiming[
  traitvalues = EvolveHapl[f, RandomReal[{2.5, 2.6}, psize], mutrate, stdev, ngens];]

Answer 2

When building a simulation like yours you should test the performance of the individual components before incorporating them into the simulation. That is, you should know the cost of the components as well as the values they return. Here is a simple example based on your code.

You use Clip in a couple of places to limit values from below. That is suspect, because Clip is designed to confine value to an interval, i.e., from both below and above. It is likely that Max will be better. Let's test this theory with your fitness function.

initsize = 500;
With[{K = 2*initsize}, 
 fit1[R_, psize_] := Clip[(1 + R (1 - psize/K)), {0.000001, ∞}]];
 fit2[R_, psize_] := Max[(1 + R (1 - psize/K)), 0.000001]]

With[{R = 2.5, p = 1000, n = 100000},
  {First @ AbsoluteTiming[Do[fit1[R, p], n]], 
   First @ AbsoluteTiming[Do[fit2[R, p], n]]}]

{0.415949, 0.371045}

So there is a real advantage in using Max rather than Clip. Since your simulation evaluates its fitness function thousands of times, it would pay to use Max.

Thorough component testing will likely reveal many more places where your code can be improved. I personally think that generating so many tables over and over again is performance area which you should investigate.

When developing a simulation that will be run many thousands of iterations, testing each piece of component code is well worth the time and effort it takes. One of Mathematica's advantages for building simulations is that it makes the sort of incremental coding and unit testing for both correctness and performance much easier than traditional lower level languages such as C or C++.