tir.py

import matplotlib.pyplot as plt

def SIR(S, I, beta, gamma, tot_time):
    '''
    S: susceptibles
    I: infecteds
    beta: contagion rate
    gamma: recovery rate
    tot_time: total time
    '''
    R = 0       # initial recovery
    N = S+I+R   # total population

    txS = []    # ratio of susceptible population
    txI = []    # ratio of infected population
    txR = []    # ratio of recovered population

    # For each time step
    for t in range(tot_time):
        qtdSI = (beta*S*I/N)    # how many susceptible will be infected
        qtdIR = (gamma*I)       # how many infected will recover
        
        # Update current values
        S = S - qtdSI
        I = I + qtdSI - qtdIR
        R = R + qtdIR

        # store ratios
        txS.append(S/N)
        txI.append(I/N)
        txR.append(R/N)
    return txS, txI, txR

def plotInfecteds(tx):
    fig, ax = plt.subplots(figsize=(10,10))
    
    for (i,txI) in tx:
        ax.plot(txI, label=f'{i}')  
    ax.set_xlabel('Dias')           
    ax.set_ylabel('Taxa da População') 
    ax.set_title('Modelo SIR')         
    plt.legend()                       
    plt.grid()                         
    plt.show()                         

def plotSIR(txS, txI, txR):    
    fig, ax = plt.subplots(figsize=(10,10))
    
    ax.plot(txS, label='Susceptibles', color='blue')
    ax.plot(txI, label='Infected', color='red')
    ax.plot(txR, label='Recovered', color='green')
    
    ax.set_xlabel('Days')
    ax.set_ylabel('Population Ratio')
    ax.set_title('SIR Model')
    plt.legend()
    plt.grid()
    
    plt.show()


tx = []
for i in [0.12, 0.15, 0.18, 0.2]:
    txS, txI, txR = SIR(200000000, 100, i, 0.1, 365)
    tx.append((i,txI))
plotInfecteds(tx)

txS, txI, txR = SIR(200000000, 100, 0.18, 0.1, 365)
plotSIR(txS, txI, txR)
txS, txI, txR = SIR(200000000, 100, 0.2, 0.1, 365)
plotSIR(txS, txI, txR)