u = sample_from_prior(rng, data)
if algorithm == "chmc":
params, x = generate_from_model(u, data)
n = (data["y_obs"] - x) / params["σ"]
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Hodgkin-Huxley model voltage-clamp data experiment")
args = parser.parse_args()
# Load data
with open(os.path.join(args.data_dir, "hodgkin-huxley-data.pkl"),
"r+b") as f:
data = pickle.load(f)
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)
y_mean = onp.linalg.cholesky(covar_func(u, data)) @ v
n = inverse_noise_transform_func(
(data["y_obs"] - y_mean) / noise_scale_func(u, data))
q = onp.concatenate((u, v, onp.asarray(n)))
else:
q = onp.concatenate((u, v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Gaussian process regression (Student's T likelihood, ν = 4) experiment"
)
parser.add_argument(
"--dataset",
default="yacht",
choices=("yacht", "slump", "marthe"),
help="Which data set to perform inference with",
)
parser.add_argument(
"--data-subsample",
type=int,
default=1,
help="Factor to subsample data by (1 corresponds to no subsampling)",
)
parser.add_argument(
"--covar-jitter",
v = rng.standard_normal(dim_y)
if algorithm == "chmc":
_, x = generate_from_model(u, v, data)
n = (data["y_obs"] - x) / data["σ"]
q = onp.concatenate((u, v, onp.asarray(n)))
else:
q = onp.concatenate((u, v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run eight-school hierarchical model experiment")
args = parser.parse_args()
# Load data
data = dict(np.load(os.path.join(args.data_dir, "eight-schools-data.npz")))
dim_u = compute_dim_u(data)
dim_y = data["y_obs"].shape[0]
# Set up seeded random number generator
rng = onp.random.default_rng(args.seed)
# Define variables to be traced
jitted_generate_from_model = api.jit(
else:
q = onp.asarray(u)
init_states.append(q)
num_tries += 1
if len(init_states) != num_chain:
raise RuntimeError(
f"Failed to find {num_chain} acceptable initial states in "
f"{max_init_tries} tries.")
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run FitzHuhgh-Nagumo model simulated data experiment")
parser.add_argument(
"--obs-noise-std",
type=float,
default=0.1,
help="Standard deviation of observation noise to use in simulated data",
)
args = parser.parse_args()
# Load data
data = dict(
onp.load(
os.path.join(args.data_dir, "fitzhugh-nagumo-simulated-data.npz")))
data["y_obs"] = data["y_mean_obs"] + args.obs_noise_std * data["n_obs"]
dim_u = compute_dim_u(data)
y_mean = forward_func(v)
σ = generate_σ(u)
n = (data["y_obs"] - y_mean) / σ
q = np.concatenate((np.array([u]), v, n))
else:
q = np.concatenate((np.array([u]), v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Poisson PDE model simulated data experiment"
)
parser.add_argument(
"--obs-noise-std",
type=float,
default=1e-2,
help="Standard deviation of observation noise to use in simulated data",
)
args = parser.parse_args()
# Load data
data = dict(np.load(os.path.join(args.data_dir, "poisson-simulated-data.npz")))
data["y_obs"] = data["y_mean"] + args.obs_noise_std * data["n_obs"]
# Set up seeded random number generator
u = sample_from_prior(rng, data)
if algorithm == "chmc":
chol_covar = onp.linalg.cholesky(covar_func(u, data))
n = sla.solve_triangular(chol_covar, data["y_obs"], lower=True)
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Gaussian process regression (Gaussian likelihood) experiment")
parser.add_argument(
"--dataset",
default="yacht",
choices=("yacht", "slump", "marthe"),
help="Which data set to perform inference with",
)
parser.add_argument(
"--data-subsample",
type=int,
default=1,
help="Factor to subsample data by (1 corresponds to no subsampling)",
)
args = parser.parse_args()
# Set up seeded random number generator
u = sample_from_prior(rng, data)
if algorithm == "chmc":
params, x = generate_from_model(u, data)
n = (data["y_obs"] - x) / params["σ"]
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run autoregressive order-K (AR-K) benchmark model experiment ")
args = parser.parse_args()
# Load data
data = dict(np.load(os.path.join(args.data_dir,
"ar-k-benchmark-data.npz")))
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)
init_states.append(q)
num_tries += 1
if len(init_states) != num_chain:
raise RuntimeError(
f"Failed to find {num_chain} acceptable initial states in "
f"{max_init_tries} tries."
)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Lotka-Volterra model Hudson-Lynx data experiment"
)
parser.add_argument(
"--obs-noise-std",
type=float,
default=1.0,
help="Standard deviation of observation noise to use in simulated data",
)
args = parser.parse_args()
# Load data
data = dict(onp.load(os.path.join(args.data_dir, "lotka-volterra-hudson-data.npz")))
dim_u = compute_dim_u(data)
# Set up seeded random number generator
if algorithm == "chmc":
_, x = generate_from_model(u, v, data)
n = data["y_obs"] / onp.exp(x / 2)
q = onp.concatenate((u, v, onp.asarray(n)))
else:
q = onp.concatenate((u, v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run stochastic-volatility model simulated data experiment"
)
utils.add_ssm_specific_args(parser)
args = parser.parse_args()
# Load data
data = dict(np.load(os.path.join(args.data_dir, "sv-simulated-data.npz")))
dim_u = compute_dim_u(data)
dim_y = data["y_obs"].shape[0]
# Set up seeded random number generator
rng = onp.random.default_rng(args.seed)
# Define variables to be traced
u = sample_from_prior(rng, data)
if algorithm == "chmc":
params, x = generate_from_model(u, data)
n = (data["y_obs"] - x) / params["σ"]
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run autoregressive moving average (ARMA) benchmark model experiment")
args = parser.parse_args()
# Load data
data = dict(np.load(os.path.join(args.data_dir,
"arma-benchmark-data.npz")))
dim_u = compute_dim_u(data)
dim_y = data["y_obs"].shape[0]
# Set up seeded random number generator
rng = onp.random.default_rng(args.seed)
# Define variables to be traced
if (n**2).mean() > mean_residual_sq_threshold:
# Skip initialisations with large residuals as can cause numerical issues
continue
if algorithm == "chmc":
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run soil-incubation model experiment")
args = parser.parse_args()
# Load data
data = dict(
onp.load(os.path.join(args.data_dir, "soil-incubation-data.npz")))
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)
init_states.append(q)
num_tries += 1
if len(init_states) != num_chain:
raise RuntimeError(
f"Failed to find {num_chain} acceptable initial states in "
f"{max_init_tries} tries."
)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run Hodgkin-Huxley model current stimulus simulated data experiment"
)
parser.add_argument(
"--obs-noise-std",
type=float,
default=1.0,
help="Standard deviation of observation noise to use in simulated data",
)
args = parser.parse_args()
# Load data
data = dict(
onp.load(
os.path.join(args.data_dir, "hodgkin-huxley-current-stimulus-data.npz")
)
v = rng.standard_normal(dim_y)
if algorithm == "chmc":
params, x = generate_from_model(u, v, data)
n = (data["y_obs"] - onp.exp(x)) / params["σ"]
q = onp.concatenate((u, v, onp.asarray(n)))
else:
q = onp.concatenate((u, v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run nonlinear state space model simulated data experiment")
parser.add_argument(
"--obs-noise-std",
type=float,
default=1.0,
help="Standard deviation of observation noise to use in simulated data",
)
utils.add_ssm_specific_args(parser)
args = parser.parse_args()
# Load data
data = dict(
onp.load(
os.path.join(args.data_dir, "nonlinear-ssm-simulated-data.npz")))
data["y_obs"] = onp.exp(data["x_obs"]) + args.obs_noise_std * data["n_obs"]
v = rng.standard_normal(dim_y)
if algorithm == "chmc":
params, x = generate_from_model(u, v, data)
n = (data["y_obs"] - x) / params["σ"]
q = onp.concatenate((u, v, onp.asarray(n)))
else:
q = onp.concatenate((u, v))
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run latent GARCH simulated data experiment")
parser.add_argument(
"--obs-noise-std",
type=float,
default=1.0,
help="Standard deviation of observation noise to use in simulated data",
)
utils.add_ssm_specific_args(parser)
args = parser.parse_args()
# Load data
data = dict(
onp.load(os.path.join(args.data_dir,
"latent-garch-simulated-data.npz")))
data["y_obs"] = data["x_obs"] + args.obs_noise_std * data["n_obs"]
if algorithm == "chmc":
params, x = generate_from_model(u, data)
n = (data["y_obs"] - params["μ"]) / onp.sqrt(x)
q = onp.concatenate((u, onp.asarray(n)))
else:
q = onp.asarray(u)
init_states.append(q)
return init_states
if __name__ == "__main__":
# Process command line arguments defining experiment parameters
parser = utils.set_up_argparser_with_standard_arguments(
"Run generalised autoregressive conditional heteroscedasticity (GARCH) "
"benchmark model experiment")
args = parser.parse_args()
# Load data
data = dict(
np.load(os.path.join(args.data_dir, "garch-benchmark-data.npz")))
dim_u = compute_dim_u(data)
# Set up seeded random number generator
rng = onp.random.default_rng(args.seed)
# Define variables to be traced