Código Python

Programa que gera a animação do modelo SIR

← Voltar para o site

Como rodar

Baixe os 3 arquivos abaixo, coloque na mesma pasta, e rode no terminal:

python vis_matplotlib.py

O programa vai pedir os parâmetros e mostrar a animação do modelo SIR.

⬇ vis_matplotlib.py ⬇ sir_core.py ⬇ utils.py

vis_matplotlib.py

Arquivo principal - gera a animação interativa:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from sir_core import resolver_sir
from utils import obter_parametros_usuario


def executar_animacao_matplotlib():
    # pega os parâmetros do usuário no terminal
    (populacao_total, infectados_iniciais, recuperados_iniciais, 
     taxa_transmissao, taxa_recuperacao, dias_simulacao, taxa_vital) = obter_parametros_usuario()

    # resolve o modelo SIR
    tempo, suscetiveis, infectados, recuperados, r0 = resolver_sir(
        populacao_total, infectados_iniciais, recuperados_iniciais,
        taxa_transmissao, taxa_recuperacao, dias_simulacao, taxa_vital
    )

    # configuração do gráfico
    figura, eixo = plt.subplots(figsize=(10, 6))
    eixo.set_xlim(0, dias_simulacao)
    eixo.set_ylim(0, populacao_total)
    eixo.set_xlabel('Tempo (dias)')
    eixo.set_ylabel('População')
    
    sufixo_titulo = " (Pop. Aberta)" if taxa_vital > 0 else ""
    eixo.set_title(f'Modelo SIR{sufixo_titulo} (N={populacao_total})')
    
    # linhas do gráfico
    linha_s, = eixo.plot([], [], 'b-', label='Suscetíveis', linewidth=2)
    linha_i, = eixo.plot([], [], 'r-', label='Infectados', linewidth=2)
    linha_r, = eixo.plot([], [], 'g-', label='Recuperados', linewidth=2)
    
    eixo.legend()
    eixo.grid(True, alpha=0.3)

    def inicializar():
        linha_s.set_data([], [])
        linha_i.set_data([], [])
        linha_r.set_data([], [])
        return linha_s, linha_i, linha_r

    def atualizar(frame):
        tempo_atual = tempo[:frame]
        suscetiveis_atual = suscetiveis[:frame]
        infectados_atual = infectados[:frame]
        recuperados_atual = recuperados[:frame]
        
        linha_s.set_data(tempo_atual, suscetiveis_atual)
        linha_i.set_data(tempo_atual, infectados_atual)
        linha_r.set_data(tempo_atual, recuperados_atual)
        
        return linha_s, linha_i, linha_r

    # cria a animação
    animacao = FuncAnimation(
        figura, atualizar, 
        frames=range(1, len(tempo)+1, 2),
        init_func=inicializar, 
        blit=True, 
        interval=30
    )
    
    print("Gerando animação Matplotlib...")
    plt.show()


if __name__ == "__main__":
    executar_animacao_matplotlib()

sir_core.py

Módulo com as equações do modelo SIR e o solver:

import numpy as np
from scipy.integrate import solve_ivp


def derivadas_sir(tempo, estado, populacao_total, taxa_transmissao, taxa_recuperacao, taxa_vital):
    """
    Calcula as derivadas do sistema SIR.
    Retorna [dS/dt, dI/dt, dR/dt] dado o estado atual.
    """
    suscetiveis, infectados, recuperados = estado
    
    # força de infecção (λ = β * I / N)
    forca_infeccao = taxa_transmissao * infectados / populacao_total
    
    # as derivadas do sistema SIR
    dS_dt = taxa_vital * populacao_total - forca_infeccao * suscetiveis - taxa_vital * suscetiveis
    dI_dt = forca_infeccao * suscetiveis - taxa_recuperacao * infectados - taxa_vital * infectados
    dR_dt = taxa_recuperacao * infectados - taxa_vital * recuperados
    
    return [dS_dt, dI_dt, dR_dt]


def resolver_sir(populacao_total, infectados_iniciais, recuperados_iniciais, 
                 taxa_transmissao, taxa_recuperacao, dias, taxa_vital=0.0):
    """
    Resolve o modelo SIR usando Runge-Kutta do SciPy.
    """
    # condições iniciais
    suscetiveis_inicial = populacao_total - infectados_iniciais - recuperados_iniciais
    estado_inicial = [suscetiveis_inicial, infectados_iniciais, recuperados_iniciais]
    
    # intervalo de tempo
    intervalo = (0, dias)
    pontos_tempo = np.linspace(0, dias, dias + 1)
    
    # resolver usando RK do SciPy
    solucao = solve_ivp(
        derivadas_sir,
        intervalo,
        estado_inicial,
        method='RK45',
        t_eval=pontos_tempo,
        args=(populacao_total, taxa_transmissao, taxa_recuperacao, taxa_vital)
    )
    
    # extrair resultados
    tempo = solucao.t
    suscetiveis = solucao.y[0]
    infectados = solucao.y[1]
    recuperados = solucao.y[2]
    
    # calcular R₀
    r0 = taxa_transmissao / (taxa_recuperacao + taxa_vital)
    
    return tempo, suscetiveis, infectados, recuperados, r0

utils.py

Função que pega os parâmetros do usuário no terminal:

def obter_parametros_usuario():
    print("Apertar Enter para usar o valor padrão.\n")

    def obter_entrada(prompt, padrao, tipo=float):
        try:
            valor = input(f"{prompt} [{padrao}]: ").strip()
            if not valor:
                return padrao
            return tipo(valor)
        except ValueError:
            print(f"Valor inválido. usaremos então o padrão: {padrao}")
            return padrao

    populacao_total = obter_entrada("População Total (N)", 1000, int)
    infectados_iniciais = obter_entrada("Infectados Iniciais (I0)", 1, int)
    recuperados_iniciais = obter_entrada("Recuperados Iniciais (R0)", 0, int)
    dias_simulacao = obter_entrada("Dias de Simulação", 365, int)
    
    print("\nParâmetros da Doença")
    taxa_transmissao = obter_entrada("Taxa de Transmissão (0.3)", 0.3, float)
    taxa_recuperacao = obter_entrada("Taxa de Recuperação (0.1)", 0.1, float)
    
    print("\n--- Parâmetros Demográficos (Opcional) ---")
    print("Use 0 para população fechada (sem nascimentos/mortes).")
    taxa_vital = obter_entrada("Taxa de natalidade/mortalidade (mu) (ex: 0.01)", 0.0, float)

    return populacao_total, infectados_iniciais, recuperados_iniciais, taxa_transmissao, taxa_recuperacao, dias_simulacao, taxa_vital