diff --git a/README.md b/README.md index 7aff6a8..dabd7fc 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ The following model calibration methods have been evaluated. * [Bayesian Optimisation for Likelihood-Free Inference](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/bolfi/run.py) * [Differential Evolution Adaptive Metropolis](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/dream/run.py) * [Experimental Design via Gaussian Process Emulation](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/experimental_design/run.py) +* [Flow Matching Posterior Estimation](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/fmpe/run.py) * [Tree-structured Parzen Estimator](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/optimisation/run.py) * [Polynomial Chaos Expansion](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/poly_chaos/run.py) * [Polynomial Chaos Kriging](https://github.com/JBris/model-calibration-evaluation/tree/main/pipelines/poly_chaos_kriging/run.py) diff --git a/data/fmpe/config.yaml b/data/fmpe/config.yaml new file mode 100644 index 0000000..adbbbf8 --- /dev/null +++ b/data/fmpe/config.yaml @@ -0,0 +1,8 @@ +mu: 52.5 +mu_lb: 45 +mu_ub: 60 +n: 100 +seed: 100 +sigma: 13.5 +sigma_lb: 5 +sigma_ub: 20 diff --git a/data/fmpe/corner.png b/data/fmpe/corner.png new file mode 100644 index 0000000..fb2fa5c Binary files /dev/null and b/data/fmpe/corner.png differ diff --git a/data/fmpe/coverage.png b/data/fmpe/coverage.png new file mode 100644 index 0000000..67007b7 Binary files /dev/null and b/data/fmpe/coverage.png differ diff --git a/pipelines/fmpe/run.py b/pipelines/fmpe/run.py new file mode 100644 index 0000000..cafd10f --- /dev/null +++ b/pipelines/fmpe/run.py @@ -0,0 +1,94 @@ +from lampe.data import JointLoader, JointDataset +from lampe.diagnostics import expected_coverage_mc, expected_coverage_ni +from lampe.inference import FMPE, FMPELoss +from lampe.plots import nice_rc, corner, mark_point, coverage_plot +from lampe.utils import GDStep +import torch +import torch.nn as nn +import torch.optim as optim +import zuko +from itertools import islice +from matplotlib import pyplot as plt +import numpy as np +import torch +from os.path import join as path_join + +from model_calibration_lib.gaussian_sim import ( + read_sim_config, simulator_func, load_data, DATA_DIR, INPUT_FILE, + CONFIG_FILE, make_dir, write_sim_config +) + +def main(): + infile = path_join(DATA_DIR, INPUT_FILE) + obs_data = load_data(infile) + config = read_sim_config() + + n, seed = config["n"], config["seed"] + mu_lb, mu_ub, sigma_lb, sigma_ub = config["mu_lb"], config["mu_ub"], config["sigma_lb"], config["sigma_ub"] + + def simulator(theta: torch.Tensor) -> torch.Tensor: + mu, sigma = theta.cpu().detach().numpy().T + y = simulator_func(mu, sigma, n, seed) + return torch.Tensor(y).float() + + lower = torch.Tensor([mu_lb, sigma_lb]) + upper = torch.Tensor([mu_ub, sigma_ub]) + prior = zuko.distributions.BoxUniform(lower, upper) + + loader = JointLoader( + prior, simulator, batch_size=1, vectorized=True + ) + + estimator = FMPE( + 2, n, hidden_features=[64] * 5, activation=nn.ELU + ) + + loss = FMPELoss(estimator) + optimizer = optim.AdamW(estimator.parameters(), lr=1e-3) + scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, 128) + step = GDStep(optimizer, clip=1.0) # gradient descent step with gradient clipping + + estimator.train() + + for epoch in range(120): + losses = torch.stack([ + step(loss(theta, x)) + for theta, x in islice(loader, 128) + ]) + scheduler.step() + print(f"Epoch {epoch}; Loss {losses.mean().item()}") + + theta_star = torch.Tensor( + np.array([config["mu"], config["sigma"]]) + ) + x_star = torch.Tensor(obs_data).float() + estimator.eval() + with torch.no_grad(): + samples = estimator.flow(x_star).sample((2**14,)) + + output_dir = path_join(DATA_DIR, "fmpe") + make_dir(output_dir) + labels = ["mu", "sigma"] + + fig = corner( + samples, + smooth=2, + domain=(lower, upper), + labels=labels, + legend=r'$p_\phi(\theta | x^*)$', + figsize=(4.8, 4.8), + ) + + mark_point(fig, theta_star) + plt.savefig(path_join(output_dir, "corner.png")) + + joint_dataset = JointDataset(theta_star.reshape((1, 1, 2)), x_star.reshape(1, n)) + fmpe_levels, fmpe_coverages = expected_coverage_mc(estimator.flow, joint_dataset) + fig = coverage_plot(fmpe_levels, fmpe_coverages, legend='FMPE') + plt.savefig(path_join(output_dir, "coverage.png")) + + outfile = path_join(output_dir, CONFIG_FILE) + write_sim_config(outfile, config) + +if __name__ == "__main__": + main() \ No newline at end of file