Solving ODE with Discontinuous Forcing Function

I’m writing a toy Covid model in Python as a beginner and I’m struggling with some strange output. My model works for a standard implementation with constant values but when I try to intoduce a discontinuous change I get a period of negative output - something that’s impossible within the model.

Specifically, the issue arises when I change my covid quarantining policy \gamma_1 to a step function step. The time path for quarantined Q_path jumps to around negative 5000 from zero at the discontinuity. This shouldn’t be possible.

If anyone can provide a hint as to where I’ve gone wrong I would be very greatful. Any other comments on my code are welcome. I’m very much a beginner.

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp
import math

def step(t):
    if t<=275-1e-10: # so the discontinuity doesn't happen at a t_eval point. 
        return 0.0 
    else:
        return 0.95
    
def linear(t): 
    if t<200:
        return 0.0
    elif ((t>=200) and (t<400)):
        return -0.95 + (0.95/200) * t 
    else: 
        return 0.95

    
quarantine_policies = {1:step, 2:linear}

def F(t, x, β, ϵ_E, ϵ_Q, ϵ_J, μ, p, γ_1, γ_2, κ_1, κ_2, σ_1, σ_2, d_1, d_2, idx):
    
    S, E, Q, I, J, R, N = x
    Π = 136.0
    
    if (idx > 0):
        γ_1 = quarantine_policies[idx](t)
    else:
        γ_1=γ_1

    dS = Π - (S*(β*I + ϵ_E*β*E + ϵ_Q*β*Q + ϵ_J*β*J)) / N - μ*S
    dE = p + (S*(β*I + ϵ_E*β*E + ϵ_Q*β*Q + ϵ_J*β*J)) / N - (γ_1 + κ_1 + μ)*E
    dQ = γ_1 * E - (κ_2 + μ) * Q 
    dI = κ_1 * E - (γ_2 + σ_1 + d_1 + μ) * I 
    dJ = γ_2 * I + κ_2 * Q - (σ_2 + d_2 + μ) * J 
    dR = σ_1 * I + σ_2 * J - μ * R 
    dN = dS + dE + dQ + dI + dJ + dR  
    
    return [dS, dE, dQ, dI, dJ, dR, dN]
    
def covid_model(β=0.2, 
                E_E=0.3, 
                E_Q=0.2, 
                E_J=0.36, 
                μ=0.000034, 
                p=0.06, 
                γ_1=0.4, 
                γ_2=0.5, 
                κ_1=0.1, 
                κ_2=0.125, 
                σ_1=0.0337, 
                σ_2=0.0386, 
                d_1=0.0079, 
                d_2=0.0068,
                idx=0):
    
    S_0 = 4000000.0
    E_0 = 6.0
    Q_0 = 0.0 
    I_0 = 1.0 
    J_0 = 0.0
    R_0 = 0.0
    N_0 = np.sum([S_0, E_0, Q_0, I_0, J_0, R_0])
    
    t_length = 700
    grid_size = 701
    t_vec = np.linspace(0.0, t_length, grid_size)
    
    x_init = [S_0, E_0, Q_0, I_0, J_0, R_0, N_0]
    args = β, E_E, E_Q, E_J, μ, p, γ_1, γ_2, κ_1, κ_2, σ_2, σ_2, d_1, d_2, idx
    
    sol = solve_ivp(F, [0,700], x_init, args=args, t_eval=t_vec)
    
    S_path = sol.y[0]
    E_path = sol.y[1]
    Q_path = sol.y[2] 
    I_path = sol.y[3] 
    J_path = sol.y[4] 
    R_path = sol.y[5]
    N_path = sol.y[6]
    
    F_path = np.zeros(grid_size)
    F_path = ((N_path*0.61)**0.2 - ((S_path + R_path + E_path)*0.61)**0.2) / (N_path*0.61)**0.2
    D_path = d_1*I_path + d_2*J_path 
    
    max_infected = np.max(I_path)
    days_to_peak = np.argmax(I_path)
    cum_deaths = np.sum(D_path)
    max_output_lost = np.max(F_path)
    
    return S_path, E_path, Q_path, I_path, J_path, R_path, D_path, F_path, N_path

S_path, E_path, Q_path, I_path, J_path, R_path, D_path, F_path, N_path = covid_model(idx=1)
np.min(Q_path)

t_vec = np.linspace(0,700,701)

fig, ax = plt.subplots(3,3)
fig.set_size_inches(18.5, 10.5)

ax[0,0].plot(t_vec, S_path)
ax[0,0].title.set_text('Susceptible')

ax[0,1].plot(t_vec, E_path)
ax[0,1].title.set_text('Asymptomatic')

ax[0,2].plot(t_vec, Q_path)
ax[0,2].title.set_text('Quarantined')

ax[1,0].plot(t_vec, I_path)
ax[1,0].title.set_text('Symptomatic')

ax[1,1].plot(t_vec, J_path)
ax[1,1].title.set_text("Isolated")

ax[1,2].plot(t_vec, R_path)
ax[1,2].title.set_text("Recovered")

ax[2,0].plot(t_vec, D_path)
ax[2,0].title.set_text("Deaths")

ax[2,1].plot(t_vec, F_path)
ax[2,1].title.set_text("Output Lost")

ax[2,2].plot(t_vec, N_path)
ax[2,2].title.set_text("Population")



1 Like