def init_loss_train_H(params):
beta, gamma, sigma, mu, alpha1, rho1, theta1, S0, E0, I0 = params # H0, D0known
dyn_params = [beta, gamma, sigma, mu, alpha1, rho1, theta1]
# init = [S,E,I,R,H,D,ICU]
init = [S0 * N, E0 * confirm[0], I0 * confirm[0], (1 - I0) * confirm[0], H0, D0]
_, _, _, _, pred_confirm, pred_death, pred_hospital = model(
size, dyn_params, init)
return loss(pred_confirm, confirm) + loss(pred_death, death) + loss(pred_hospital, hospital)
def init_loss_train(self, params):
beta, gamma, sigma, mu, alpha1, rho1, theta1, alpha2, rho2, theta2, S0, E0, I0 = params # H0, D0, ICU0 known
dyn_params = [beta, gamma, sigma, mu, alpha1, rho1, theta1, alpha2, rho2, theta2]
# init = [S,E,I,R,H,D,ICU]
init = [S0 * self.N, E0 * self.confirm[0], I0 * self.confirm[0], (1 - I0) * self.confirm[0], self.H0, self.D0, self.ICU0]
_, _, _, _, pred_confirm, pred_death, pred_hospital, pred_ICU = self.model(
self.size, dyn_params, init)
return loss(pred_confirm, self.confirm) + loss(pred_death, self.death) + loss(pred_hospital, self.hospital) + loss(pred_ICU, self.icu)
def loss_train_H(params):
_, _, _, _, pred_confirm, pred_death, pred_hospital= model(size, params, init)
return loss(pred_confirm, confirm) + loss(pred_death, death) + loss(pred_hospital, hospital)
def loss_train(self, params):
_, _, _, _, pred_confirm, pred_death, pred_hospital, pred_ICU = self.model(self.size, params, self.init)
return loss(pred_confirm, self.confirm) + loss(pred_death, self.death) + loss(pred_hospital, self.hospital) + loss(pred_ICU, self.icu)