an example SIR model fitting

fit_SIR_model.py

from symfit import Parameter, variables, Fit, D, ODEModel
import numpy as np

# define number of susceptible population
n_susceptible = 195000

# Some data
data_I = np.array([1, 2, 11, 23, 36, 75, 104, 137, 166, 209, 313, 400, 496, 693, 914, 1635, 2391, 5213, 8054, 11293, 15157, 19938, 24336, 29010])
data_R = np.array([0, 0, 0, 0, 0, 1, 2, 5, 7, 11, 15, 21, 29, 39, 53, 71, 104, 152, 256, 417, 643, 946, 1345, 1831])
data_S = [n_susceptible - x - data_R[idx] for idx, x in enumerate(data_I)]

# define variables of the SIR model
S, I, R, t = variables('S, I, R, t')

# define parameters of the SIR model with sensible initial parameter values
dIdt = np.diff(data_I)
dRdt = np.diff(data_R)
gamma_0 = 0.02
beta_0 = np.mean([(x + gamma_0 * data_I[idx] ) * n_susceptible / (data_S[idx] * data_I[idx]) for idx, x in enumerate(dIdt)])
beta = Parameter('beta', beta_0)
gamma = Parameter("gamma", gamma_0)

# define ODE equations 
model_dict = { D(S, t): - beta * S * I / n_susceptible,
               D(I, t): beta * S* I / n_susceptible - gamma * I,
               D(R, t): gamma * I}

# set initial values
I0 = data_I[0]
S0 = n_susceptible - I0
R0 = data_R[0]

# define the model
model = ODEModel(model_dict, initial = { t : 0, S : S0, I : I0, R : R0 })

# fit model parameters
fit = Fit(model, t = np.array(range(0,data_I.shape[0])), I = data_I, S = None, R = data_R )
fit_result = fit.execute()

# get predictions for 100 days
tvec = np.linspace(0, 99, 100)
outcome = model(t=tvec, **fit_result.params)
I_pred, S_pred, R_pred = outcome.I, outcome.S, outcome.R