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())