seir_model.py

import numpy as np

def p_exposed(params, substep, state_history, previous_state):
  signal = previous_state["S"]*params["infection_rate"]*previous_state["I"]/(previous_state["S"]+previous_state["E"]+previous_state["I"]+previous_state["R"])
  return {"delta_s":np.ceil(signal)}
  
def p_infected(params, substep, state_history, previous_state):
  signal = params["exposure_rate"]*previous_state["E"]
  return {"delta_i":np.ceil(signal)}
  
def p_recovered(params, substep, state_history, previous_state):
  signal = params["recovery_rate"]*previous_state["I"]
  return {"delta_r":np.ceil(signal)}

def s_susceptible(params, substep, state_history, previous_state, policy_input):
  new_S = previous_state["S"] - policy_input["delta_s"]
  return ("S", max(new_S,0))
  
def s_exposed(params, substep, state_history, previous_state, policy_input):
  new_E = previous_state["E"] + policy_input["delta_s"] - policy_input["delta_i"]
  return ("E", max(new_E,0))

def s_infected(params, substep, state_history, previous_state, policy_input):
  new_I = previous_state["I"] + policy_input["delta_i"] - policy_input["delta_r"]
  return ("I", max(new_I,0))
  
def s_recovered(params, substep, state_history, previous_state, policy_input):
  new_R = previous_state["R"] + policy_input["delta_r"]
  return ("R", max(new_R, 0))