sigma = 0.01
real_params = [0.07, 0.6, 0.05, 0.3, 0.017, 0.3]
sol = pr_ode_model.solve(real_params)
np.random.seed(121)
Y = sol + np.random.randn(len(times), 5) * sigma
### Run inference ###
param_names = [
r"$p_1$", r"$p_2$", r"$p_3$", r"$p_4$", r"$p_5$", r"$p_6$", r"$\sigma$"
]
real_params.append(0.01)
if not (args.adjoint):
print('Using VJP by Forward Sensitivity')
method = 'NUTS'
NUTS_samples = run_inference(Y, ProteinGenModel, pr_ode_model, method, \
iterations = args.num_samples, warmup_steps = args.warmup_steps)
mc_params = np.concatenate(
(NUTS_samples['ode_params'], NUTS_samples['scale'][:, None]),
axis=1)
method = 'VI'
lr = 0.5
vb_samples = run_inference(Y, ProteinGenModel, pr_ode_model, method, \
iterations = args.iterations, num_samples = args.num_qsamples, \
lr = lr, num_particles = 1, return_sites = ("ode_params","scale","_RETURN"))
vb_params = \
np.concatenate((vb_samples['ode_params'].detach().numpy().reshape((args.num_qsamples,6)), \
vb_samples['scale'].detach().numpy().reshape((args.num_qsamples,1))),axis=1)
plot_marginals(vb_params,
mc_params,
times = np.array([0, 5])
Y = np.array([[100, 100], [343, 213], [ 28, 354], [ 41, 137], [255, 101], \
[ 71, 416], [ 29, 136], [186, 76], [141, 449], [ 27, 189]])
### Run inference ###
param_names = [r"$c_1$", r"$c_2 \times 100$", r"$c_3$"]
real_params = np.array([0.53, 0.25, 0.3])
print('Using VJP by Forward Sensitivity')
lna_ode_model = ForwardSensManualJacobians(rhs_f, jac_x_f, jac_p_f, 6, 3, \
times, 1e-5, 1e-6, [100, 100, 0, 0 ,0 ,0])
method = 'VI'
lr = 0.5
vb_samples = run_inference(Y, LNAGenModel, lna_ode_model, method, iterations=args.iterations, \
lr = lr, num_particles = 1, num_samples = args.num_qsamples, \
return_sites = ("ode_params1","ode_params2","ode_params3"))
vb_params_for = np.concatenate(
(vb_samples['ode_params1'][:, None].detach().numpy(),
vb_samples['ode_params2'][:, None].detach().numpy(),
vb_samples['ode_params3'][:, None].detach().numpy()),
axis=1)
print('Using VJP by Adjoint Sensitivity')
lna_ode_model = AdjointSensManualJacobians(rhs_f, jac_x_f, jac_p_f, 6, 3, \
times, 1e-5, 1e-6, [100, 100, 0, 0 ,0 ,0])
vb_samples = run_inference(Y, LNAGenModel, lna_ode_model, method, iterations=args.iterations, \
lr = lr, num_particles = 1, num_samples = args.num_qsamples, \
return_sites = ("ode_params1","ode_params2","ode_params3"))
vb_params_adj = np.concatenate(
print('Using VJP by Forward Sensitivity')
sir_ode_model = ForwardSensManualJacobians(rhs_f, jac_x_f, jac_p_f, 3, 5, \
times, 1e-5, 1e-6, [0.9, 0.1, 0.0])
sir_ode_model.set_unknown_y0()
# method = 'NUTS'
# NUTS_samples = run_inference(Y, SIRGenModel, sir_ode_model, method, \
# iterations = args.num_samples, warmup_steps = args.warmup_steps)
# mc_params=np.concatenate((NUTS_samples['ode_params1'][:,None],
# NUTS_samples['ode_params2'][:,None],
# NUTS_samples['ode_params3'][:,None]
# ),axis=1)
method = 'VI'
lr = 0.5
vb_samples = run_inference(Y, SIRGenModel, sir_ode_model, method, \
iterations=args.iterations, num_samples=args.num_qsamples, \
lr=lr, num_particles=1, return_sites=("ode_params1", "ode_params2", "ode_params3"))
vb_params = np.concatenate(
(vb_samples['ode_params1'][:, None].detach().numpy(),
vb_samples['ode_params2'][:, None].detach().numpy(),
vb_samples['ode_params3'][:, None].detach().numpy()),
axis=1)
# plot_marginals(vb_params, mc_params, param_names, rows=2)
plot_marginals(vb_params, vb_params, param_names, rows=2)
else:
print('Using VJP by Adjoint Sensitivity')
sir_ode_model = AdjointSensManualJacobians(rhs_f, jac_x_f, jac_p_f, 3, 5, \
times, 1e-5, 1e-6, [0.9, 0.1, 0.0])
sir_ode_model.set_unknown_y0()
# method = 'NUTS'
# method = 'NUTS'
# NUTS_samples = run_inference(Y, SIRGenModel, sir_ode_model, method, \
# iterations = args.num_samples, warmup_steps = args.warmup_steps)
# mc_params=np.concatenate((NUTS_samples['ode_params1'][:,None],
# NUTS_samples['ode_params2'][:,None],
# NUTS_samples['ode_params3'][:,None]
# ),axis=1)
method = 'VI'
lr = 0.5
num_particles = 1
vb_samples = run_inference(Y,
PlantModel,
plant_ode_model,
method,
iterations=args.iterations,
num_samples=args.num_qsamples,
lr=lr,
num_particles=num_particles,
return_sites=("ode_params1", "ode_params2"))
vb_params = np.concatenate(
(vb_samples['ode_params1'][:, None].detach().numpy(),
vb_samples['ode_params2'][:, None].detach().numpy()),
axis=1)
# plot_marginals(vb_params, mc_params, param_names, rows=2)
plot_marginals(vb_params, vb_params, param_names, rows=2)
else:
print('Using VJP by Adjoint Sensitivity')
plant_ode_model = AdjointSensManualJacobians(rhs_f, jac_x_f, jac_p_f,
2, 2, times, 1e-5, 1e-6,