Mosquito population model

mosquito.py

import numpy as np


class Mosquito:
    """Contains the details of each Mosquito"""
    def __init__(self, mother_gene_infected, father_gene_infected, sex):
        self.genes = [mother_gene_infected, father_gene_infected]
        self.sex = sex


class Simulation:
    """Regulates the population and simulation"""
    def __init__(self, n_healthy, n_infected):
        """
        Create a population (list of individuals)
        with a number of healthy and infected individuals
        """
        # The construct ['Male', 'Female'][x % 2] will make sure that
        # we produce mosquitos of alternating sex, resulting in a 0.5 sex ratio
        self.population = [Mosquito(mother_gene_infected=False,
                                    father_gene_infected=False,
                                    sex=['Male', 'Female'][x % 2]) for x in range(n_healthy)]
        self.population.extend([Mosquito(mother_gene_infected=True,
                                         father_gene_infected=False,
                                         sex=['Male', 'Female'][x % 2]) for x in range(n_infected)])

    def reproduction(self):
        """
        Replace the old population by a new one
        """
        # Divide the population into females and males.
        # Females with 2 copies of the infected gene are infertile and not selected
        females = [m for m in self.population if (m.sex == 'Female') and (sum(m.genes) < 2)]
        males = [m for m in self.population if m.sex == 'Male']

        # The old population is overwritten by an empty list and will be filled with new individuals
        self.population = []
        # If there is at least one fertile female..
        if females:
            # select a random male mate for each female
            lucky_males = np.random.choice(males, len(females))
            for i, f in enumerate(females):
                # Pass on genes to about 100 offspring (Poisson distribution).
                # Since gene drive is in place, an infected gene will always be passed on.
                # The construct ['Male', 'Female'][x % 2] will make sure that
                # we produce mosquitos of alternating sex, resulting in a 0.5 sex ratio
                self.population.extend([Mosquito(mother_gene_infected=sum(f.genes) > 0,
                                                 father_gene_infected=sum(lucky_males[i].genes) > 0,
                                                 sex=['Male', 'Female'][n % 2]) for n in range(np.random.poisson(100))])

    def density_dependent_survival(self, a=0.25, beta=0.5):
        """
        Reduce the current population in a density-dependent manner.
        When the population is big the survival rate is lower
        than when the population is low.
        Surviving individuals will go on to reproduce
        """
        # Single species density regulation based on Hassell & Comins
        survival_rate = (1 + a * len(self.population)) ** -beta
        n_survivors = int(len(self.population) * survival_rate)

        # if there are survivors we select them randomly from the current population
        if n_survivors:
            np.random.shuffle(self.population)
            self.population = self.population[0:(n_survivors + 1)]
        else:
            self.population = []

    def population_statistics(self):
        """
        Calculate the population statistics of interest
        """
        # Count the number of infected individuals
        infected = len([m for m in self.population if (sum(m.genes) > 0)])
        # Count the number of double infected (homozygote) individuals
        double_infected = len([m for m in self.population if (sum(m.genes) == 2)])
        return [len(self.population), infected, double_infected]

    def main_loop(self, n_generations):
        """
        The main loop function of a single simulation
        Each iteration is a generation
        """
        # create an array to store per-generation results
        results_array = np.zeros(shape=(n_generations + 1, 3))
        results_array[0] = self.population_statistics()

        for n in range(n_generations):
            self.density_dependent_survival()
            self.reproduction()
            results_array[n + 1] = self.population_statistics()

        # return the per-generation evolution of the infection
        # and whether the population went extinct
        return results_array, len(self.population) == 0


def plot_meta_results(results):
    """Results plotter for results over multiple simulations"""
    import matplotlib.pyplot as plt
    # switch from proportions to percentages
    plt.plot(100 * results[:, 0], 100 * results[:, 1])
    plt.xlabel('Population percentage of infected individuals introduced')
    plt.ylabel('Percentage of simulations with extinction')
    plt.show()


def plot_simulation_results(results):
    """Results plotter for results of single simulation"""
    import pylab
    x = np.arange(0, len(results))
    popsize = results[:, 0]
    all_infected = results[:, 1]
    double_infected = results[:, 2]

    pylab.plot(x, popsize, '-b', label='Population')
    pylab.plot(x, all_infected, '-r', label='All infected')
    pylab.plot(x, double_infected, '-g', label='Double infected')
    pylab.legend(loc='lower right')
    pylab.xlabel('Generation')
    pylab.ylabel('Number of individuals')
    pylab.ylim(0, max(popsize) * 1.1)
    pylab.show()


if __name__ == '__main__':
    # -------------------
    # simulation settings
    # -------------------
    # While there is randomness in our simulation,
    # for a particular random seed the outcome will be the same.
    np.random.seed(1)
    # The number of uninfected mosquitos present in the 1st generation
    n_healthy_mosquitos = 10000
    n_infected_mosquitos = 100
    dyn = Simulation(n_healthy=n_healthy_mosquitos,
                     n_infected=n_infected_mosquitos)

    # a single run of the simulation
    results = dyn.main_loop(n_generations=20)[0]
    print(results)
    plot_simulation_results(results)

# if __name__ == '__main__':
#     # -------------------
#     # simulation settings
#     # -------------------
#     # While there is randomness in our simulation,
#     # for a particular random seed the outcome will be the same.
#     np.random.seed(1)
#     # The number of infected mosquitos to introduce in the 1st generation
#     infections_to_test = [1, 10, 25, 50, 75, 100, 150, 200, 300]
#     # The number of simulations to run for each of the values in the list above
#     n_simulations = 100
#     # The number of uninfected mosquitos present in the 1st generation
#     n_healthy_mosquitos = 10000
#
#     # create an array to store per-simulation setting results
#     results = np.zeros(shape=(len(infections_to_test), 2))
#     # loop over the settings and run simulations
#     for i, n_infected_mosquitos in enumerate(infections_to_test):
#         n_extinct = 0
#         print(n_infected_mosquitos)
#         for n in range(n_simulations):
#             dyn = Simulation(n_healthy=n_healthy_mosquitos,
#                              n_infected=n_infected_mosquitos)
#             n_extinct += dyn.main_loop(n_generations=20)[1]
#         # write to the results matrix what the proportion of infected vs healthy
#         # mosquitos was + the proportion of simulations where extinction took place
#         results[i] = [n_infected_mosquitos / n_healthy_mosquitos,
#                       n_extinct / n_simulations]
#
#     print(results)
#     plot_meta_results(results)