simple sir

sir.py

initial_susceptible # defaults to 3,600,000 https://github.com/CodeForPhilly/chime/blob/2895a9c4ddcf42b3c96bcf7e03a7e2a15f4983de/src/penn_chime/presentation.py#L200-L206
initial_infected # https://github.com/CodeForPhilly/chime/blob/2895a9c4ddcf42b3c96bcf7e03a7e2a15f4983de/src/penn_chime/models.py#L25-L27
initial_recovered # https://github.com/CodeForPhilly/chime/blob/2895a9c4ddcf42b3c96bcf7e03a7e2a15f4983de/src/penn_chime/models.py#L34
beta = # https://github.com/CodeForPhilly/chime/blob/2895a9c4ddcf42b3c96bcf7e03a7e2a15f4983de/src/penn_chime/models.py#L42-L45
gamma = # https://github.com/CodeForPhilly/chime/blob/2895a9c4ddcf42b3c96bcf7e03a7e2a15f4983de/src/penn_chime/models.py#L39
n_days = # User input, default to 60 or something

def sir(s, i, r, beta, gama, n):
    """The SIR model, one time step."""
    s_n = (-beta * s * i) + s
    i_n = (beta * s * i - gamma * i) + i
    r_n = gamma * i + r
    if s_n < 0.0:
        s_n = 0.0
    if i_n < 0.0:
        i_n = 0.0
    if r_n < 0.0:
        r_n = 0.0

    scale = n / (s_n + i_n + r_n)
    return s_n * scale, i_n * scale, r_n * scale

def gen_sir(init_susceptible, init_infected, init_recovered, beta, gamma, n_days):
    """Simulate SIR model forward in time yielding tuples."""
    current_susceptible = init_susceptible
    current_infected = init_infected
    current_recovered = init_recovered
    days_data = []
    susceptible_data = []
    infected_data = []
    recovered_data = []
    n = current_susceptible + current_infected + current_recovered
    for day in range(n_days + 1):
        days_data.append(day)
        susceptible_data.append(current_susceptible)
        infected_data.append(current_infected)
        recovered_data.append(current_recovered)
        current_susceptible, current_infected, current_recovered = sir(
            current_susceptible, current_infected, current_recovered, beta, gamma, n)

    return days_data, susceptible_data, infected_data, recovered_data

days_data, susceptible_data, infected_data, recovered_data = gen_sir(
    initial_susceptible,
    initial_infected,
    initial_recovered,
    beta,
    gamma,
    n_days,
)