<jupyter_start><jupyter_text>Example notebook illusrtating the usage of the compartmental.Model class <NAME> (IST Austria) February 9, 2021<jupyter_code>import numpy as np
from matplotlib import pyplot as plt
import compartmental
from scipy.integrate import odeint
from IPython.display import display, Math
from scipy.optimize import curve_fit
# Generate a model

compartments = [
    ('S'    , {"layer" : 1}),    # Susceptible
    ('I'    , {"layer" : 2}),    # Infectious
    ('R'    , {"layer" : 3}),    # Removed (recovered/dead)
]

parameters = {
    'beta' : 0.6,     # avg cont. / pers. / day * prob(trans.)
    'gamma': 0.2,    # (days)^-1 Inverse mean infectious period  
    'N'    : 1e6,        # population
}

rates = [
    ('S', 'I', {"label" : "beta * S * I / N"}), 
    ('I', 'R', {"label" : "gamma * I"})
]

model = compartmental.Model()
model.set_compartments(compartments)
model.set_rates(rates)
model.set_parameters(list(parameters.keys()))
fig = model.visualize(figsize=(6,3))
ode_latex, ode_symbolic, ode = model.generate_ode()
display(Math(ode_latex))
# Run an epidemic simulation with 
# N = 1e6, gamma = 1/5 (days^-1), beta = 0.6 
# i.e. R_0 = beta / gamma = 3 ~ covid-19

initial_population = {'S': 1e6 - 100, 
                      'I': 100, 
                      'R': 0,}

simulation_time = np.arange(0, 365)

population = odeint(ode, 
                    list(initial_population.values()), 
                    simulation_time, 
                    args = (list(parameters.values()),))
fig = plt.figure(figsize=(12,6))
axes = fig.subplots(1,2)

for ax in axes:
    model.plot_compartment(simulation_time, population, 'S', ax=ax)
    model.plot_compartment(simulation_time, population, 'I', ax=ax)
    model.plot_compartment(simulation_time, population, 'R', ax=ax)

    ax.set_xlim(1, 50)
    ax.set_ylim(1e-1, 1.1e6)
    ax.grid(True)
    ax.set_xlabel('Days', fontsize=16)
    ax.set_ylabel('Population', fontsize=16)

axes[0
].get_legend().remove()
axes[1].set_yscale('log')
fig.suptitle('Dynamics with an initial transient', fontsize=16)
plt.tight_layout()
initial_population_exp = model.initiate_exponential(
    initial_population, parameters, np.arange(0, 10), ['R'], 'S'
    )
# Run an epidemic simulation with 
# N = 1e6, gamma = 1/5 (days^-1), beta = 0.6 
# i.e. R_0 = beta / gamma = 3 ~ covid-19

simulation_time = np.arange(0, 365)

population = odeint(ode, 
                    list(initial_population_exp.values()), 
                    simulation_time, 
                    args = (list(parameters.values()),))
fig = plt.figure(figsize=(12,6))
axes = fig.subplots(1,2)

for ax in axes:
    model.plot_compartment(simulation_time, population, 'S', ax=ax)
    model.plot_compartment(simulation_time, population, 'I', ax=ax)
    model.plot_compartment(simulation_time, population, 'R', ax=ax)

    ax.set_xlim(1, 50)
    ax.set_ylim(10, 1.1e6)
    ax.grid(True)
    ax.set_xlabel('Days', fontsize=16)
    ax.set_ylabel('Population', fontsize=16)

axes[0
].get_legend().remove()
axes[1].set_yscale('log')
fig.suptitle('Dynamics without the initial transient', fontsize=16)

plt.tight_layout()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15180265654648956, 'fraction_numerical': 0.03194180898165718, 'mean_word_length': 3.34478021978022, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '461714', 'license': 'MIT', 'n_tokens_mistral': 1249, 'n_tokens_neox': 1135, 'n_words': 247, 'path': 'examples/SIR_init_exp.ipynb', 'repo_name': 'burakbudanur/compartmental', 'size': '3175'}