SIR_diffeq.py

GAMMA =1.0/10 # (1/D)

### Model ###
def sir_diff_eq(y, t, r0, gamma, N, t_y_interpolated):
    S, I, R = y

    def beta(t):
        try:
            return t_y_interpolated[int(t)] * r0 * gamma
        except:
            return t_y_interpolated[-1] * r0 * gamma

    dSdt = -beta(t) * I * S / N
    dIdt = beta(t) * I * S/N - gamma * I
    dRdt = gamma * I

    return dSdt, dIdt, dRdt