def solve(self, i_0, r_0, beta, gamma, t_0=0, t_final=20):
s_0 = 1 - i_0 - r_0
# Test between 0 and t_final
grid = torch.arange(t_0, t_final,
out=torch.FloatTensor()).reshape(-1, 1)
t_dl = DataLoader(dataset=grid, batch_size=1, shuffle=False)
s_hat = []
i_hat = []
r_hat = []
# Convert initial conditions, beta and gamma to tensor for prediction
beta_t = torch.Tensor([beta]).reshape(-1, 1)
gamma_t = torch.Tensor([gamma]).reshape(-1, 1)
s_0_t = torch.Tensor([s_0]).reshape(-1, 1)
i_0_t = torch.Tensor([i_0]).reshape(-1, 1)
r_0_t = torch.Tensor([r_0]).reshape(-1, 1)
initial_conditions_set = [s_0_t, i_0_t, r_0_t]
de_loss = 0.
for i, t in enumerate(t_dl, 0):
t.requires_grad = True
# Network solutions
s, i, r = self.parametric_solution(t,
initial_conditions_set,
beta=beta_t,
gamma=gamma_t)
s_hat.append(s.item())
i_hat.append(i.item())
r_hat.append(r.item())
de_loss += sir_loss(t, s, i, r, beta_t, gamma_t)
return s_hat, i_hat, r_hat, de_loss
] # Solution of the network with the optimal found set of params
exact_de = 0.
optimal_de = 0.
for t, v in exact_points.items():
t_tensor = torch.Tensor([t]).reshape(-1, 1)
t_tensor.requires_grad = True
s_best, i_best, r_best = sir.parametric_solution(
t_tensor,
optimal_initial_conditions,
beta=optimal_beta,
gamma=optimal_gamma,
mode='bundle_total')
optimal_de += sir_loss(t_tensor, s_best, i_best, r_best, optimal_beta,
optimal_gamma)
optimal_traj.append([s_best.item(), i_best.item(), r_best.item()])
# Exact solution subset
exact_sub_traj = [exact_points[t] for t in exact_points.keys()]
exact_mse = 0.
optimal_mse = 0.
rnd_mse = 0.
for idx, p in enumerate(exact_sub_traj):
exact_s, exact_i, exact_r = p
exact_mse += (exact_s - exact_traj[idx][0])**2 + (
exact_i - exact_traj[idx][1])**2 + (exact_r -
exact_traj[idx][2])**2
optimal_initial_conditions = [optimal_s_0, optimal_i_0, optimal_r_0]
# Let's save the predicted trajectories in the known points
optimal_traj = [] # Solution of the network with the optimal found set of params
optimal_de = 0.
for t, v in data_prelock.items():
t_tensor = torch.Tensor([t]).reshape(-1, 1)
t_tensor.requires_grad = True
s_optimal, i_optimal, r_optimal = sir.parametric_solution(t_tensor, optimal_initial_conditions,
beta=optimal_beta,
gamma=optimal_gamma,
mode='bundle_total')
optimal_de += sir_loss(t_tensor, s_optimal, i_optimal, r_optimal, optimal_beta, optimal_gamma)
optimal_traj.append([s_optimal, i_optimal, r_optimal])
# Exact solution subset
exact_sub_traj = [data_prelock[t] for t in data_prelock.keys()]
optimal_mse = 0.
for idx, p in enumerate(exact_sub_traj):
exact_s, exact_i, exact_r = p
optimal_mse += (exact_s - optimal_traj[idx][0]) ** 2 + (exact_i - optimal_traj[idx][1]) ** 2 + (
exact_r - optimal_traj[idx][2]) ** 2
optimal_mse /= len(optimal_traj)
# Generate points between 0 and t_final
grid = torch.arange(0, t_final, out=torch.FloatTensor()).reshape(-1, 1)
r_0 = torch.Tensor([r_0]).reshape(-1, 1)
beta = torch.Tensor([beta]).reshape(-1, 1)
gamma = torch.Tensor([gamma]).reshape(-1, 1)
initial_conditions = [s_0, i_0, r_0]
for i, t in enumerate(t_dl, 0):
t.requires_grad = True
avg_loss = 0.
# Network solutions
for sir in models:
s, i, r = sir.parametric_solution(t,
initial_conditions,
beta=beta,
gamma=gamma,
mode='bundle_total')
avg_loss += sir_loss(t, s, i, r, beta, gamma)
avg_loss = avg_loss / max_model_index
current_loss += avg_loss
init_losses[(i_0.item(), r_0.item())] = current_loss
inits = list(init_losses.keys())
inits = [list(p) for p in inits]
i_0_sampled = [p[0] for p in inits]
r_0_sampled = [p[1] for p in inits]
inits_losses = list(init_losses.values())
inits_losses = [l.item() for l in inits_losses]
log_inits_losses = np.log(inits_losses)
params_losses = {}
plt.plot(range(len(r_hat)), r_hat, label='Recovered', color=green, linewidth=linewidth)
plt.plot(range(len(s_p)), s_p, label='Susceptible - Scipy', linestyle='--', color=blue, linewidth=linewidth)
plt.plot(range(len(i_p)), i_p, label='Infected - Scipy', linestyle='--', color=red, linewidth=linewidth)
plt.plot(range(len(r_p)), r_p, label='Recovered - Scipy', linestyle='--', color=green, linewidth=linewidth)
plt.title('Solving SIR model with Beta = {} | Gamma = {} \n'
'Starting conditions: S(0) = {:.2f} | I(0) = {:.2f} | R(0) = {:.2f} \n'.format(beta, gamma, s_0, i_0, r_0))
plt.legend(loc='lower right')
plt.xlabel('Time')
plt.ylabel('S(t), I(t), R(t)')
plt.savefig(
ROOT_DIR + '/plots/SIR_s0={:.2f}_i0={:.2f}_r0={:.2f}_beta={}_gamma={}.png'.format(s_0, i_0, r_0, beta, gamma))
plt.show()
# Compute loss as a function of the time
log_losses = []
for i, t in enumerate(t_dl, 0):
from losses import sir_loss
t.requires_grad = True
s, i, r = sir.parametric_solution(t, initial_conditions)
t_loss = sir_loss(t, s, i, r, beta, gamma)
log_losses.append(np.log(t_loss.item()))
plt.figure(figsize=(15, 5))
plt.plot(range(len(log_losses)), log_losses)
plt.xlabel('Time')
plt.ylabel('Logloss')
plt.title('Solving SIR model with Beta = {} | Gamma = {} \n'
'Starting conditions: S(0) = {:.2f} | I(0) = {:.2f} | R(0) = {:.2f} \n'.format(beta, gamma, s_0, i_0, r_0))
plt.show()
s_0 = torch.Tensor([s_0]).reshape(-1, 1)
i_0 = torch.Tensor([i_0]).reshape(-1, 1)
r_0 = torch.Tensor([r_0]).reshape(-1, 1)
beta = torch.Tensor([beta]).reshape(-1, 1)
gamma = torch.Tensor([gamma]).reshape(-1, 1)
initial_conditions = [s_0, i_0, r_0]
for i, t in enumerate(t_dl, 0):
t.requires_grad = True
# Network solutions
s, i, r = sir.parametric_solution(t,
initial_conditions,
beta=beta,
gamma=gamma,
mode='bundle_total')
current_loss += sir_loss(t, s, i, r, beta, gamma)
params_losses[(beta.item(), gamma.item())] = current_loss
init_losses[(i_0.item(), r_0.item())] = current_loss
beta_gamma = list(params_losses.keys())
beta_gamma = [list(p) for p in beta_gamma]
beta_sampled = [p[0] for p in beta_gamma]
gamma_sampled = [p[1] for p in beta_gamma]
inits = list(init_losses.keys())
inits = [list(p) for p in inits]
i_0_sampled = [p[0] for p in inits]
r_0_sampled = [p[1] for p in inits]
losses = list(params_losses.values())