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_mse += (exact_s - optimal_traj[idx][0])**2 + (
exact_i - optimal_traj[idx][1])**2 + (exact_r -
optimal_traj[idx][2])**2
exact_mse /= len(exact_sub_traj)
optimal_mse /= len(optimal_traj)
# Scipy solver solution
t = np.linspace(0, t_final, t_final)
s_p, i_p, r_p = SIR_solution(t, optimal_s_0.item(), 1 - optimal_s_0.item(),
0, optimal_beta.item(), optimal_gamma.item())
# Generate points between 0 and t_final
grid = torch.arange(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 = []
for i, t in enumerate(t_dl, 0):
# Network solutions
s, i, r = sir.parametric_solution(t,
optimal_initial_conditions,
beta=optimal_beta,
gamma=optimal_gamma,
mode='bundle_total')
s_hat.append(s.item())
# Test between 0 and t_final
grid = torch.arange(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 = []
# Scipy solver solution
beta_t = 0.8
gamma_t = 0.3
s_0 = 0.9988
i_0 = 1 - s_0
r_0 = 0.0
t = np.linspace(0, t_final, t_final)
s_p, i_p, r_p = SIR_solution(t, s_0, i_0, r_0, beta_t, gamma_t)
# Convert initial conditions, beta and gamma to tensor for prediction
beta_t = torch.Tensor([beta_t]).reshape(-1, 1)
gamma_t = torch.Tensor([gamma_t]).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]
for i, t in enumerate(t_dl, 0):
# Network solutions
s, i, r = sir.parametric_solution(t,
initial_conditions_set,
beta=beta_t,
gamma=gamma_t,