# Define variables to be traced trace_func = utils.construct_trace_func(generate_params, data, dim_u, dim_v=dim_y) # Run experiment final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="robust_gp_regression", dir_prefix=f"{args.dataset}_data_subsampled_by_{args.data_subsample}", var_names=list(prior_specifications.keys()), var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=GeneralGaussianProcessModelSystem, constrained_system_kwargs={ "covar_func": covar_func, "noise_scale_func": noise_scale_func, "noise_transform_func": noise_transform_func, "data": data, "dim_u": dim_u, }, sample_initial_states=sample_initial_states, )
# Define variables to be traced jitted_generate_params = api.jit(api.partial(generate_params, data=data)) def trace_func(state): u = state.pos[:dim_u] params = jitted_generate_params(u) return {**params, "u": u} # Run experiment final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="fitzhugh_nagumo", dir_prefix=f"σ_{args.obs_noise_std:.0e}", var_names=list(prior_specifications.keys()), var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=IndependentAdditiveNoiseModelSystem, constrained_system_kwargs={ "generate_y": generate_y, "data": data, "dim_u": dim_u, }, sample_initial_states=sample_initial_states, )
data["y_obs"] = np.concatenate(data["obs_vals_g_Na"]) dim_u = compute_dim_u(data) # Set up seeded random number generator rng = onp.random.default_rng(args.seed) # Define variables to be traced trace_func = utils.construct_trace_func(generate_params, data, dim_u) # Run experiment final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="hh_voltage_clamp_sodium", var_names=list(prior_specifications.keys()), var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=IndependentAdditiveNoiseModelSystem, constrained_system_kwargs={ "generate_y": generate_y, "data": data, "dim_u": dim_u, }, sample_initial_states=sample_initial_states, )
# Define variables to be traced jitted_generate_from_model = api.jit( api.partial(generate_from_model, data=data)) def trace_func(state): u, v = state.pos[:dim_u], state.pos[dim_u:dim_u + dim_y] params, x = jitted_generate_from_model(u, v) return {**params, "x": x, "u": u, "v": v} # Run experiment final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="eight_schools", var_names=list(prior_specifications.keys()) + ["x"], var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=HierarchicalLatentVariableModelSystem, constrained_system_kwargs={ "generate_y": generate_y, "data": data, "dim_u": dim_u, }, sample_initial_states=sample_initial_states, )
generate_σ=generate_σ, grad_generate_σ=grad_generate_σ, hessian_generate_σ=hessian_generate_σ, ) euclidean_system_kwargs = pde.construct_euclidean_system_kwargs( y_obs=data["y_obs"], forward_func=forward_func, vjp_forward_func=vjp_forward_func, dim_y=dim_y, dim_z=dim_z, generate_σ=generate_σ, grad_generate_σ=grad_generate_σ, ) final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="poisson", dir_prefix=f"σ_{args.obs_noise_std:.0e}", var_names=["σ", "z_mean", "z_std"], var_trace_func=trace_func, constrained_system_class=mici.systems.DenseConstrainedEuclideanMetricSystem, constrained_system_kwargs=constrained_system_kwargs, euclidean_system_class=mici.systems.EuclideanMetricSystem, euclidean_system_kwargs=euclidean_system_kwargs, sample_initial_states=sample_initial_states, precompile_jax_functions=False )
# Run experiment ( constrained_system_class, constrained_system_kwargs, ) = utils.get_ssm_constrained_system_class_and_kwargs( args.use_manual_constraint_and_jacobian, generate_params, generate_x_0, forward_func, inverse_observation_func, constr_split, jacob_constr_split_blocks, ) constrained_system_kwargs.update(data=data, dim_u=dim_u) final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="stochastic_volatility", var_names=list(prior_specifications.keys()), var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=constrained_system_class, constrained_system_kwargs=constrained_system_kwargs, sample_initial_states=sample_initial_states, )
# Run experiment ( constrained_system_class, constrained_system_kwargs, ) = utils.get_ssm_constrained_system_class_and_kwargs( args.use_manual_constraint_and_jacobian, generate_params, generate_x_0, forward_func, inverse_observation_func, constr_split, jacob_constr_split_blocks, ) constrained_system_kwargs.update(data=data, dim_u=dim_u) final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="nonlinear_ssm", dir_prefix=f"σ_{args.obs_noise_std:.0e}", var_names=list(prior_specifications.keys()), var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=constrained_system_class, constrained_system_kwargs=constrained_system_kwargs, sample_initial_states=sample_initial_states, )
# Define variables to be traced jitted_generate_from_model = api.jit( api.partial(generate_from_model, data=data)) def trace_func(state): u = state.pos[:dim_u] params, x = jitted_generate_from_model(u) return {**params, "x": x, "u": u} # Run experiment final_states, traces, stats, summary_dict, sampler = utils.run_experiment( args=args, data=data, rng=rng, experiment_name="garch", var_names=["μ", "α_0", "α_1", "β_1"], var_trace_func=trace_func, posterior_neg_log_dens=posterior_neg_log_dens, extended_prior_neg_log_dens=extended_prior_neg_log_dens, constrained_system_class=IndependentAdditiveNoiseModelSystem, constrained_system_kwargs={ "generate_y": generate_y, "data": data, "dim_u": dim_u, }, sample_initial_states=sample_initial_states, )