plague.py

'''
We'll need three variables to solve this Eigenvalue Problem (based on a Lesile model):
A = the "Leslie matrix"
p_i = the initial population breakdown of healthy, sick, and deceased individuals
percentage = stop value (percentage of healthy + sick people)
'''

import numpy as np

# Our initial Leslie matrix for the system
A = np.matrix([[0.3, 0.2, 0.0],
               [0.6, 0.2, 0.0],
               [0.1, 0.6, 1.0]])

# Initial population breakdown of healthy, sick, and deceased individuals
p_i = np.matrix([[200],
                [0],
                [0]])

# Percentage that will trigger a "stop value" for number of people still alive
percentage = 0.10


def calculatePopulation(A, p_i, percentage):
    '''
    Create function to record the population of healthy/sick/deceased at the end
    of the year. A dictonary with values will be returned once the percent of people
    alive falls below our defined value of 0.10.
    
    Returns both the final year, and the numbers containing change in population by year.
    '''
    values_by_end_of_year = {'year_0': {'healthy': p_i.item(0),
                                        'sick': p_i.item(1),
                                        'deceased': p_i.item(2)
                                        }}
    year = 0
    percent_alive = 1.0
    while percentage < percent_alive:
        year +=1
        new = A * p_i
        percent_alive = (new.item(0) + new.item(1))/(sum(new).item(0))
        values_by_end_of_year[f'year_{year}'] = {'healthy': new.item(0),
                                        'sick': new.item(1),
                                        'deceased': new.item(2)
                                        }
        p_i = new

    return year, values_by_end_of_year
    
final_year, all_values = calculatePopulation(A, p_i, percentage)