Skip to content

Add death rate to the model #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
giulianobelinassi wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add death rate to the model #14

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 35 additions & 16 deletions solver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,14 +123,15 @@ def load_json(json_file_str):


class Learner(object):
def __init__(self, country, loss, start_date, predict_range,s_0, i_0, r_0):
def __init__(self, country, loss, start_date, predict_range,s_0, i_0, r_0, d_0):
self.country = country
self.loss = loss
self.start_date = start_date
self.predict_range = predict_range
self.s_0 = s_0
self.i_0 = i_0
self.r_0 = r_0
self.d_0 = d_0


def load_confirmed(self, country):
Expand Down Expand Up @@ -159,50 +160,68 @@ def extend_index(self, index, new_size):
values = np.append(values, datetime.strftime(current, '%m/%d/%y'))
return values

def predict(self, beta, gamma, data, recovered, death, country, s_0, i_0, r_0):
def predict(self, beta, a, b, data, recovered, death, country, s_0, i_0, r_0, d_0):
new_index = self.extend_index(data.index, self.predict_range)
size = len(new_index)
def SIR(t, y):
S = y[0]
I = y[1]
R = y[2]
return [-beta*S*I, beta*S*I-gamma*I, gamma*I]
D = y[3]
return [-beta*S*I, beta*S*I-(a+b)*I, a*I, b*I]
extended_actual = np.concatenate((data.values, [None] * (size - len(data.values))))
extended_recovered = np.concatenate((recovered.values, [None] * (size - len(recovered.values))))
extended_death = np.concatenate((death.values, [None] * (size - len(death.values))))
return new_index, extended_actual, extended_recovered, extended_death, solve_ivp(SIR, [0, size], [s_0,i_0,r_0], t_eval=np.arange(0, size, 1))
ivp = solve_ivp(SIR, [0, size], [s_0,i_0,r_0,d_0], t_eval=np.arange(0, size, 1))
return new_index, extended_actual, extended_recovered, extended_death, ivp


def train(self):
recovered = self.load_recovered(self.country)
death = self.load_dead(self.country)
data = (self.load_confirmed(self.country) - recovered - death)

optimal = minimize(loss, [0.001, 0.001], args=(data, recovered, self.s_0, self.i_0, self.r_0), method='L-BFGS-B', bounds=[(0.00000001, 0.4), (0.00000001, 0.4)])

optimal = minimize(loss,
[0.001, 0.001, 0.001],
args=(data, recovered, death, self.s_0, self.i_0, self.r_0, self.d_0),
method='L-BFGS-B',
bounds=[(0.00000001, 0.4), (0.00000001, 0.4), (0.00000001, 0.4)])

print(optimal)
beta, gamma = optimal.x
new_index, extended_actual, extended_recovered, extended_death, prediction = self.predict(beta, gamma, data, recovered, death, self.country, self.s_0, self.i_0, self.r_0)
df = pd.DataFrame({'Infected data': extended_actual, 'Recovered data': extended_recovered, 'Death data': extended_death, 'Susceptible': prediction.y[0], 'Infected': prediction.y[1], 'Recovered': prediction.y[2]}, index=new_index)
beta, a, b = optimal.x
new_index, extended_actual, extended_recovered, extended_death, prediction = self.predict(beta, a, b, data, recovered, death, self.country, self.s_0, self.i_0, self.r_0, self.d_0)

df = pd.DataFrame({
'Infected data': extended_actual,
'Recovered data': extended_recovered,
'Death data': extended_death,
'Susceptible': prediction.y[0],
'Infected': prediction.y[1],
'Recovered': prediction.y[2],
'Extimated Deaths': prediction.y[3]},
index=new_index)
fig, ax = plt.subplots(figsize=(15, 10))
ax.set_title(self.country)
df.plot(ax=ax)
print(f"country={self.country}, beta={beta:.8f}, gamma={gamma:.8f}, r_0:{(beta/gamma):.8f}")
print(f"country={self.country}, beta={beta:.8f}, a={a:.8f}, b={b:.8f}, gamma={(a+b):.8f}, r_0:{(beta/(a+b)):.8f}")
fig.savefig(f"{self.country}.png")


def loss(point, data, recovered, s_0, i_0, r_0):
def loss(point, data, recovered, death, s_0, i_0, r_0, d_0):
size = len(data)
beta, gamma = point
beta, a, b = point
def SIR(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(SIR, [0, size], [s_0,i_0,r_0], t_eval=np.arange(0, size, 1), vectorized=True)
D = y[3]
return [-beta*S*I, beta*S*I-(a+b)*I, a*I, b*I]
solution = solve_ivp(SIR, [0, size], [s_0,i_0,r_0,d_0], t_eval=np.arange(0, size, 1), vectorized=True)
l1 = np.sqrt(np.mean((solution.y[1] - data)**2))
l2 = np.sqrt(np.mean((solution.y[2] - recovered)**2))
l3 = np.sqrt(np.mean((solution.y[3] - death)**2))
alpha = 0.1
return alpha * l1 + (1 - alpha) * l2
return alpha * l1 + (1 - alpha) * l2 + 0.9*l3


def main():
Expand All @@ -218,7 +237,7 @@ def main():
remove_province('data/time_series_19-covid-Deaths.csv', 'data/time_series_19-covid-Deaths-country.csv')

for country in countries:
learner = Learner(country, loss, startdate, predict_range, s_0, i_0, r_0)
learner = Learner(country, loss, startdate, predict_range, s_0, i_0, r_0, 0)
#try:
learner.train()
#except BaseException:
Expand Down