290 lines
13 KiB
Python
290 lines
13 KiB
Python
#!/usr/bin/python3
|
|
from datetime import datetime, timedelta
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pandas as pd
|
|
from scipy.integrate import solve_ivp
|
|
from scipy.optimize import minimize
|
|
import argparse
|
|
import os
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument('--mode', '-m', dest = 'mode', help = 'change the mode of the model (SIR, Linear, ESIR, SEIR); default: SIR', default = 'SIR', choices = ['SIR', 'Linear', 'ESIR', 'SEIR'])
|
|
parser.add_argument('--data', '-d', dest = 'include_data', help = 'change the type of data to present in the graph (Actual, S, I, R, E); default: Actual S I R', nargs = '*', default = ['Actual', 'S', 'I', 'R'], choices = ['Actual', 'S', 'I', 'R', 'E'])
|
|
parser.add_argument('--folder', '-f', dest = 'folder', default = 'data', help = 'the folder in which to find the data files; defaults to looking in the data folder')
|
|
parser.add_argument('--disease', '-D', dest = 'disease', default = 'COVID-19', help = 'the disease to model; defaults to COVID-19')
|
|
parser.add_argument('--out', '-o', dest = 'out', default = None, help = 'the name of the graph and csv files; defaults to the name of the disease')
|
|
parser.add_argument('--start', '-s', dest = 'start', default = '1/22/20', help = 'the date where the data starts (defaults to the start date of COVID-19 (1/22/20))')
|
|
parser.add_argument('--end', '-e', dest = 'end', default = None, help = 'the date where the data stops (defaults to whereever the input data ends)')
|
|
parser.add_argument('--incubation', '-i', dest = 'incubation_period', default = None, help = 'the incubation period of the disease (only applicable if using SEIR model; ignored otherwise); none by default')
|
|
parser.add_argument('--predict', '-p', dest = 'prediction_range', default = None, help = 'the number of days to predict the course of the disease (defaults to None, meaning the model will not predict beyond the given data)')
|
|
parser.add_argument('--country', '-c', dest = 'country', default = 'US', help = 'the country that is being modeled (defaults to US)')
|
|
parser.add_argument('--popcountry', '-pc', dest = 'popcountry', default = '3328200000', help = 'the population of the country (defaults to US population)')
|
|
parser.add_argument('--popmodel', '-pm', dest = 'popmodel', default = '10000', help = 'the population of the model (defaults to 10000)')
|
|
parser.add_argument('--initial', '-I', dest = 'initial', default = '1', help = 'initial infected people (defaults to 1)')
|
|
args = parser.parse_args()
|
|
|
|
# Running a model for a million population is quite hard, so here we've reduced the population and modified the actual stats to match
|
|
correction_factor = int(args.popmodel) / int(args.popcountry)
|
|
|
|
S_0 = (int(args.popcountry) - int(args.initial)) / int(args.popcountry)
|
|
I_0 = int(args.initial) / int(args.popcountry)
|
|
R_0 = 0
|
|
E_0 = 0
|
|
|
|
class Learner(object):
|
|
def __init__(self, country):
|
|
self.country = country
|
|
|
|
def load_confirmed(self, country):
|
|
"""
|
|
Load confirmed cases
|
|
"""
|
|
df = pd.read_csv(f'{args.folder}/{args.disease}-Confirmed.csv')
|
|
country_df = df[df['Country/Region'] == country]
|
|
|
|
if args.end != None:
|
|
confirmed_sums = np.sum([reg.loc[args.start:args.end].values for reg in country_df.iloc], axis = 0)
|
|
else:
|
|
confirmed_sums = np.sum([reg.loc[args.start:].values for reg in country_df.iloc], axis = 0)
|
|
|
|
if args.end != None:
|
|
new_data = pd.DataFrame(confirmed_sums, country_df.iloc[0].loc[args.start:args.end].index.tolist())
|
|
else:
|
|
new_data = pd.DataFrame(confirmed_sums, country_df.iloc[0].loc[args.start:].index.tolist())
|
|
|
|
return new_data
|
|
|
|
|
|
def load_recovered(self, country):
|
|
"""
|
|
Load recovered cases
|
|
"""
|
|
df = pd.read_csv(f'{args.folder}/{args.disease}-Recovered.csv')
|
|
country_df = df[df['Country/Region'] == country]
|
|
|
|
if args.end != None:
|
|
out = country_df.iloc[0].loc[args.start:args.end]
|
|
else:
|
|
out = country_df.iloc[0].loc[args.start:]
|
|
|
|
return out
|
|
|
|
|
|
def extend_index(self, index, new_size):
|
|
values = index.values
|
|
current = datetime.strptime(index[-1], '%m/%d/%y')
|
|
while len(values) < new_size:
|
|
current = current + timedelta(days=1)
|
|
values = np.append(values, datetime.strftime(current, '%m/%d/%y'))
|
|
return values
|
|
|
|
def predict(self, data, beta = None, gamma = None, mu = None, sigma = None):
|
|
"""
|
|
Predict how the number of people in each compartment can be changed through time toward the future.
|
|
The model is formulated with the given beta and gamma (or others).
|
|
Returns the "solved" system of initial value problems, to be "graded" by the loss function
|
|
"""
|
|
new_index = self.extend_index(data.index, args.prediction_range if args.prediction_range != None else len(data.index))
|
|
size = len(new_index)
|
|
def model(t, y):
|
|
S = y[0]
|
|
I = y[1]
|
|
R = y[2]
|
|
|
|
if args.mode == 'SEIR':
|
|
E = y[3]
|
|
|
|
if args.mode == 'Linear':
|
|
return [-beta * S, beta * S - gamma * I, gamma * I]
|
|
elif args.mode == 'SIR':
|
|
return [-beta * S * I, beta * S * I - gamma * I, gamma * I]
|
|
elif args.mode == 'ESIR':
|
|
if mu != None:
|
|
return [mu - beta * S * I - mu * S, beta * S * I - gamma * I - mu * I, gamma * I - mu * R]
|
|
else:
|
|
raise Exception('Expected mu for ESIR model')
|
|
elif args.mode == 'SEIR':
|
|
if mu != None and sigma != None:
|
|
return [mu - beta * S * I - mu * S, beta * S * I - sigma * E - mu * E, sigma * E - gamma * I - mu * I, gamma * I - mu * R]
|
|
elif mu == None:
|
|
raise Exception('Expected mu for SEIR model')
|
|
elif sigma == None:
|
|
raise Exception('Expected sigma for SEIR model')
|
|
|
|
extended_actual = np.concatenate((data.values.flatten(), [0] * (size - len(data.values))))
|
|
|
|
if args.mode == 'SEIR':
|
|
result = solve_ivp(model, [0, size], [S_0,I_0,R_0,E_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
else:
|
|
result = solve_ivp(model, [0, size], [S_0,I_0,R_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
|
|
return new_index, extended_actual, result
|
|
|
|
def train(self):
|
|
"""
|
|
Run the optimization to estimate the beta and gamma fitting the given confirmed cases.
|
|
"""
|
|
confirmed_data = self.load_confirmed(self.country)
|
|
recovered_data = self.load_recovered(self.country)
|
|
|
|
if not os.path.isdir('out'):
|
|
os.mkdir('out')
|
|
|
|
if args.mode == 'Linear':
|
|
optimal = minimize(
|
|
loss_linear,
|
|
[0.01, 0.01],
|
|
args=(confirmed_data, recovered_data),
|
|
method='L-BFGS-B',
|
|
bounds=[(0.001, 1.0), (0.001, 1.0)]
|
|
)
|
|
beta, gamma = optimal.x
|
|
print(f'Beta: {beta}, Gamma: {gamma}, R0: {beta/gamma}')
|
|
new_index, extended_actual, prediction = self.predict(confirmed_data, beta = beta, gamma = gamma)
|
|
print(f'Predicted I: {prediction.y[1][-1] * int(args.popmodel)}, Actual I: {extended_actual[-1] * correction_factor}')
|
|
df = compose_df(prediction, extended_actual, correction_factor, new_index)
|
|
with open(f'out/{args.disease}-{args.mode}-data.csv', 'w+') as file:
|
|
file.write(f'Beta, {beta}\nGamma, {gamma}\nR0, {beta/gamma}\n')
|
|
file.write(f'Predicted I, {prediction.y[1][-1] * int(args.popmodel)}\nActual I, {extended_actual[-1] * correction_factor}')
|
|
elif args.mode == 'SIR':
|
|
optimal = minimize(
|
|
loss_sir,
|
|
[0.01, 0.01],
|
|
args=(confirmed_data, recovered_data),
|
|
method='L-BFGS-B',
|
|
bounds=[(0.001, 1.0), (0.001, 1.0)]
|
|
)
|
|
beta, gamma = optimal.x
|
|
print(f'Beta: {beta}, Gamma: {gamma}, R0: {beta/gamma}')
|
|
new_index, extended_actual, prediction = self.predict(confirmed_data, beta = beta, gamma = gamma)
|
|
print(f'Predicted I: {prediction.y[1][-1] * int(args.popmodel)}, Actual I: {extended_actual[-1] * correction_factor}')
|
|
df = compose_df(prediction, extended_actual, correction_factor, new_index)
|
|
with open(f'out/{args.disease}-{args.mode}-data.csv', 'w+') as file:
|
|
file.write(f'Beta, {beta}\nGamma, {gamma}\nR0, {beta/gamma}\n')
|
|
file.write(f'Predicted I, {prediction.y[1][-1] * int(args.popmodel)}\nActual I, {extended_actual[-1] * correction_factor}')
|
|
elif args.mode == 'ESIR':
|
|
optimal = minimize(
|
|
loss_esir,
|
|
[0.01, 0.01, 0.01],
|
|
args=(confirmed_data, recovered_data),
|
|
method='L-BFGS-B',
|
|
bounds=[(0.001, 1.0), (0.001, 1.0), (0.001, 1.0)]
|
|
)
|
|
beta, gamma, mu = optimal.x
|
|
print(f'Beta: {beta}, Gamma: {gamma}, Mu: {mu} R0: {beta/(gamma + mu)}')
|
|
new_index, extended_actual, prediction = self.predict(confirmed_data, beta = beta, gamma = gamma, mu = mu)
|
|
print(f'Predicted I: {prediction.y[1][-1] * int(args.popmodel)}, Actual I: {extended_actual[-1] * correction_factor}')
|
|
df = compose_df(prediction, extended_actual, correction_factor, new_index)
|
|
with open(f'out/{args.disease}-{args.mode}-data.csv', 'w+') as file:
|
|
file.write(f'Beta, {beta}\nGamma, {gamma}\nMu, {mu}\nR0, {beta/(gamma + mu)}\n')
|
|
file.write(f'Predicted I, {prediction.y[1][-1] * int(args.popmodel)}\nActual I, {extended_actual[-1] * correction_factor}')
|
|
elif args.mode == 'SEIR':
|
|
# exposed_data = self.load_exposed(self.country)
|
|
|
|
optimal = minimize(
|
|
loss_seir,
|
|
[0.01, 0.01, 0.01, 0.01],
|
|
args=(confirmed_data, recovered_data),
|
|
method='L-BFGS-B',
|
|
bounds=[(0.001, 1.0), (0.001, 1.0), (0.001, 1.0), (0.001, 1.0)]
|
|
)
|
|
beta, gamma, mu, sigma = optimal.x
|
|
print(f'Beta: {beta}, Gamma: {gamma}, Mu: {mu}, Sigma: {sigma} R0: {(beta * sigma)/((mu + gamma) * (mu + sigma))}')
|
|
new_index, extended_actual, prediction = self.predict(confirmed_data, beta = beta, gamma = gamma, mu = mu, sigma = sigma)
|
|
print(f'Predicted I: {prediction.y[1][-1] * int(args.popmodel)}, Actual I: {extended_actual[-1] * correction_factor}')
|
|
df = compose_df(prediction, extended_actual, correction_factor, new_index)
|
|
with open(f'out/{args.disease}-{args.mode}-data.csv', 'w+') as file:
|
|
file.write(f'Beta, {beta}\nGamma, {gamma}\nMu, {mu}\nSigma, {sigma}\nR0, {(beta * sigma)/((mu + gamma) * (mu + sigma))}\n')
|
|
file.write(f'Predicted I, {prediction.y[1][-1] * int(args.popmodel)}\nActual I, {extended_actual[-1] * correction_factor}')
|
|
fig, ax = plt.subplots(figsize=(15, 10))
|
|
ax.set_title(f'{args.disease} cases over time ({args.mode} Model)')
|
|
df.plot(ax=ax)
|
|
fig.savefig(f"{args.out if args.out != None else args.disease}-{args.mode}.png")
|
|
df.to_csv(f'out/{args.disease}-{args.mode}-prediction.csv')
|
|
|
|
def filter_zeroes(arr):
|
|
out = np.array(arr)
|
|
for index in range(len(out)):
|
|
if out[index] == 0:
|
|
out[index] = None
|
|
return out
|
|
|
|
def compose_df(prediction, actual, correction_factor, index):
|
|
df_dict = {}
|
|
|
|
for data in args.include_data:
|
|
if data == 'Actual':
|
|
df_dict['Actual'] = filter_zeroes(actual * correction_factor)
|
|
elif data == 'S':
|
|
df_dict['S'] = prediction.y[0] * int(args.popmodel)
|
|
elif data == 'I':
|
|
df_dict['I'] = prediction.y[1] * int(args.popmodel)
|
|
elif data == 'R':
|
|
df_dict['R'] = prediction.y[2] * int(args.popmodel)
|
|
elif data == 'E':
|
|
df_dict['E'] = prediction.y[3] * int(args.popmodel)
|
|
|
|
return pd.DataFrame(df_dict, index=index)
|
|
|
|
# Loss Functions - used to "train" the model
|
|
|
|
def loss_linear(point, confirmed, recovered):
|
|
size = len(confirmed)
|
|
beta, gamma = point
|
|
def model(t, y):
|
|
S = y[0]
|
|
I = y[1]
|
|
R = y[2]
|
|
return [-beta * S, beta * S - gamma * I, gamma * I]
|
|
solution = solve_ivp(model, [0, size], [S_0,I_0,R_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
sol_inf = np.sqrt(np.mean((solution.y[1] - (confirmed.values.flatten() * correction_factor/int(args.popmodel)))**2))
|
|
sol_rec = np.sqrt(np.mean((solution.y[2] - (recovered.values * correction_factor/int(args.popmodel)))**2))
|
|
return sol_inf * 0.5 + sol_rec * 0.5
|
|
|
|
def loss_sir(point, confirmed, recovered):
|
|
size = len(confirmed)
|
|
beta, gamma = point
|
|
def model(t, y):
|
|
S = y[0]
|
|
I = y[1]
|
|
R = y[2]
|
|
return [-beta * S * I, beta * S * I - gamma * I, gamma * I]
|
|
solution = solve_ivp(model, [0, size], [S_0,I_0,R_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
sol_inf = np.sqrt(np.mean((solution.y[1] - (confirmed.values.flatten() * correction_factor/int(args.popmodel)))**2))
|
|
sol_rec = np.sqrt(np.mean((solution.y[2] - (recovered.values * correction_factor/int(args.popmodel)))**2))
|
|
return sol_inf * 0.5 + sol_rec * 0.5
|
|
|
|
def loss_esir(point, confirmed, recovered):
|
|
size = len(confirmed)
|
|
beta, gamma, mu = point
|
|
def model(t, y):
|
|
S = y[0]
|
|
I = y[1]
|
|
R = y[2]
|
|
return [mu - beta * S * I - mu * S, beta * S * I - gamma * I - mu * I, gamma * I - mu * R]
|
|
solution = solve_ivp(model, [0, size], [S_0,I_0,R_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
sol_inf = np.sqrt(np.mean((solution.y[1] - (confirmed.values.flatten() * correction_factor/int(args.popmodel)))**2))
|
|
sol_rec = np.sqrt(np.mean((solution.y[2] - (recovered.values * correction_factor/int(args.popmodel)))**2))
|
|
return sol_inf * 0.5 + sol_rec * 0.5
|
|
|
|
def loss_seir(point, confirmed, recovered):
|
|
size = len(confirmed)
|
|
beta, gamma, mu, sigma = point
|
|
def model(t, y):
|
|
S = y[0]
|
|
I = y[1]
|
|
R = y[2]
|
|
E = y[3]
|
|
return [mu - beta * S * I - mu * S, beta * S * I - sigma * E - mu * E, sigma * E * I - gamma * I - mu * I, gamma * I - mu * R]
|
|
solution = solve_ivp(model, [0, size], [S_0,I_0,R_0,E_0], t_eval=np.arange(0, size, 1), vectorized=True)
|
|
sol_inf = np.sqrt(np.mean((solution.y[1] - (confirmed.values.flatten() * correction_factor/int(args.popmodel)))**2))
|
|
sol_rec = np.sqrt(np.mean((solution.y[2] - (recovered.values * correction_factor/int(args.popmodel)))**2))
|
|
# sol_exp = np.sqrt(np.mean((solution.y[3] - (exposed.values * correction_factor/int(args.popmodel)))**2))
|
|
return sol_inf * 0.5 + sol_rec * 0.5
|
|
|
|
my_learner = Learner(args.country)
|
|
my_learner.train()
|