From 2baf7c2bcf2004bcf6801e004ee4e79f4d737579 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 1 Dec 2025 10:02:53 +0000 Subject: [PATCH 1/6] ENH: add MCMC proposal with minipcn --- src/nessai/proposal/flowproposal/mcmc.py | 185 +++++++++++++++++++++++ src/nessai/proposal/utils.py | 2 + 2 files changed, 187 insertions(+) create mode 100644 src/nessai/proposal/flowproposal/mcmc.py diff --git a/src/nessai/proposal/flowproposal/mcmc.py b/src/nessai/proposal/flowproposal/mcmc.py new file mode 100644 index 00000000..33766951 --- /dev/null +++ b/src/nessai/proposal/flowproposal/mcmc.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import datetime +import os +from typing import Callable + +import numpy as np + +from ...livepoint import live_points_to_array +from .base import BaseFlowProposal + + +class MiniPCNFlowProposal(BaseFlowProposal): + """Version of FlowProposal that uses MiniPCN instead of rejection sampling. + + Parameters + ---------- + model: Model + Model with likelihood and prior + n_steps : int + Number of MCMC steps to take + step_fn: str + Step function to use. See the minipcn documentation for details. + minipcn_kwargs : dict + Dictionary of keyword arguments passed to :code:`minipcn.Sampler` and + `Sampler.sample`. + enforce_likelihood_threshold: bool + Check the likelihood constraint when sampling. If false, samples are drawn + from the prior using the flow. + plot_history: bool + Toggle plotting the MCMC history. + """ + + def __init__( + self, + model, + n_steps: int, + step_fn: str = "pcn", + minipcn_kwargs: dict | None = None, + enforce_likelihood_threshold: bool = True, + plot_history: bool = False, + **kwargs, + ): + super().__init__(model, **kwargs) + self.n_steps = n_steps + self.step_fn = step_fn + self.minipcn_kwargs = minipcn_kwargs or {} + self.enforce_likelihood_threshold = enforce_likelihood_threshold + self.mcmc_history = { + "acceptance": [], + "n_steps": [], + } + self._plot_history = plot_history + + def _get_log_prob( + self, logl_threshold: float + ) -> Callable[[np.ndarray], float]: + """Get the log-probability function for MiniPCN.""" + + def _log_prob(z): + """Log-probability for minipcn""" + self.backward_pass(z, rescale=True, return_unit_hypercube=False) + x, log_j_flow = self.backward_pass( + z, + rescale=True, + return_unit_hypercube=self.map_to_unit_hypercube, + ) + if self.map_to_unit_hypercube: + log_p = self.unit_hypercube_log_prior(x) + else: + log_p = self.log_prior(x) + finite_prior = np.isfinite(log_p) + + if self.enforce_likelihood_threshold: + x["logL"][finite_prior] = ( + self.model.batch_evaluate_log_likelihood( + x[finite_prior], + unit_hypercube=self.map_to_unit_hypercube, + ) + ) + above_threshold = x["logL"] > logl_threshold + log_p[~above_threshold] = -np.inf + + val = log_p + log_j_flow + return val + + return _log_prob + + def plot_history(self): + """Plot the history of MCMC acceptance and number of steps. + + This is useful for diagnosing the performance of the MCMC proposal over + the course of the run. + """ + import matplotlib.pyplot as plt + + fig, axs = plt.subplots(2, 1, sharex=True) + axs[0].plot(self.mcmc_history["acceptance"]) + axs[0].set_ylabel("Acceptance") + axs[1].plot(self.mcmc_history["n_steps"]) + axs[1].set_ylabel("Number of steps") + axs[-1].set_xlabel("Iteration") + plt.tight_layout() + fig.savefig(os.path.join(self.output, "mcmc_history.png")) + plt.close(fig) + + def populate( + self, + worst_point: np.ndarray, + n_samples: int | None = 10000, + plot: bool = True, + ) -> None: + """Populate the proposal pool using MiniPCN MCMC. + + Parameters + ---------- + worst_point : np.ndarray + The current worst point in the nested sampling run. Used to set the + likelihood threshold. + n_samples : int + Number of samples to generate. If None, uses the poolsize. + plot : bool + Whether to plot diagnostic plots. Each plot can be toggled on/off + use the corresponding keyword argument when initialising the proposal. + """ + from minipcn import Sampler + + st = datetime.datetime.now() + + log_prob_fn = self._get_log_prob(worst_point["logL"]) + + kwargs = self.minipcn_kwargs.copy() + verbose = kwargs.pop("verbose", False) + + sampler = Sampler( + log_prob_fn=log_prob_fn, + step_fn=self.step_fn, + rng=self.rng, + dims=self.rescaled_dims, + **kwargs, + ) + + x_prime_array = live_points_to_array( + self.training_data_prime, + self.prime_parameters, + copy=True, + ) + z_init, _ = self.flow.forward_and_log_prob(x_prime_array) + + n_walkers = n_samples if n_samples is not None else self.poolsize + + z_init = self.rng.choice(z_init, n_walkers, axis=0) + chain, history = sampler.sample( + z_init, + n_steps=self.n_steps, + verbose=verbose, + ) + + if self._plot_history and plot: + self.plot_history() + if self._plot_pool and plot: + self.plot_pool(self.samples) + + z_pool = chain[-1] + x_pool, _ = self.backward_pass( + z_pool, + rescale=True, + return_unit_hypercube=False, + ) + # This could be made more efficient + x_pool["logL"] = self.model.batch_evaluate_log_likelihood(x_pool) + self.samples = self.convert_to_samples(x_pool) + + self.mcmc_history["acceptance"].append( + np.mean(history.acceptance_rate) + ) + self.mcmc_history["n_steps"].append(self.n_steps) + + self.population_time += datetime.datetime.now() - st + + self.population_acceptance = np.mean(history.acceptance_rate) + self.indices = self.rng.permutation(self.samples.size).tolist() + self.populated_count += 1 + self.populated = True + self._checked_population = False diff --git a/src/nessai/proposal/utils.py b/src/nessai/proposal/utils.py index 02e48318..c864067d 100644 --- a/src/nessai/proposal/utils.py +++ b/src/nessai/proposal/utils.py @@ -85,9 +85,11 @@ def available_base_flow_proposal_classes(): from ..experimental.proposal.mcmc import MCMCFlowProposal from .augmented import AugmentedFlowProposal from .flowproposal import FlowProposal + from .flowproposal.mcmc import MiniPCNFlowProposal base_proposals = { "mcmcflowproposal": MCMCFlowProposal, + "minipcnflowproposal": MiniPCNFlowProposal, "clusteringflowproposal": ClusteringFlowProposal, "augmentedflowproposal": AugmentedFlowProposal, "flowproposal": FlowProposal, From 4bcd4083d558e8ed048dc36d51ff0eb8d460c4a3 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 1 Dec 2025 10:03:09 +0000 Subject: [PATCH 2/6] BLD: add optional dependency for MCMC --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index f0287e93..287fd002 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,9 @@ docs = [ nflows = [ "nflows", ] +mcmc = [ + "minipcn", +] [tool.setuptools_scm] version_scheme = "release-branch-semver" From ebd45e516dab41994b90f322a5ccb30f52c6ce85 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 1 Dec 2025 10:04:57 +0000 Subject: [PATCH 3/6] MAINT: add MiniPCNFlowProposal to init --- src/nessai/proposal/flowproposal/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/nessai/proposal/flowproposal/__init__.py b/src/nessai/proposal/flowproposal/__init__.py index 365e38fe..a984de9f 100644 --- a/src/nessai/proposal/flowproposal/__init__.py +++ b/src/nessai/proposal/flowproposal/__init__.py @@ -1,7 +1,9 @@ """Proposals that use normalising flows.""" from .flowproposal import FlowProposal +from .mcmc import MiniPCNFlowProposal __all__ = [ "FlowProposal", + "MiniPCNFlowProposal", ] From 052baabb468220d3eefb06e6ec201852d1d1c061 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Wed, 17 Dec 2025 11:47:10 +0000 Subject: [PATCH 4/6] TST: update proposal test --- tests/test_proposal/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_proposal/test_utils.py b/tests/test_proposal/test_utils.py index 01a482d6..f798190f 100644 --- a/tests/test_proposal/test_utils.py +++ b/tests/test_proposal/test_utils.py @@ -148,7 +148,7 @@ class FakeProposal: def test_available_base_flow_proposal_classes(): avail = available_base_flow_proposal_classes() - assert len(avail) == 4 + assert len(avail) == 5 @pytest.mark.parametrize("load", [True, False]) From fada2299c32900306ac718114f0cdb603d5731a2 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 6 Jan 2026 08:48:06 +0100 Subject: [PATCH 5/6] MAINT: change default to tpcn --- src/nessai/proposal/flowproposal/mcmc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/nessai/proposal/flowproposal/mcmc.py b/src/nessai/proposal/flowproposal/mcmc.py index 33766951..9ebb4428 100644 --- a/src/nessai/proposal/flowproposal/mcmc.py +++ b/src/nessai/proposal/flowproposal/mcmc.py @@ -35,7 +35,7 @@ def __init__( self, model, n_steps: int, - step_fn: str = "pcn", + step_fn: str = "tpcn", minipcn_kwargs: dict | None = None, enforce_likelihood_threshold: bool = True, plot_history: bool = False, From 95f53fdfe18f862ccd4a814965eefb059722f6d8 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 6 Jan 2026 08:49:24 +0100 Subject: [PATCH 6/6] EXAMP: update mcmc example to use minipcn --- examples/mcmc_example.py | 66 +++++++++++++++++++++------------------- 1 file changed, 35 insertions(+), 31 deletions(-) diff --git a/examples/mcmc_example.py b/examples/mcmc_example.py index e1348dad..0a8c5d47 100644 --- a/examples/mcmc_example.py +++ b/examples/mcmc_example.py @@ -1,6 +1,8 @@ #!/usr/bin/env python +"""Example of using the MiniPCNFlowProposal with the Rosenbrock likelihood. -# Example of using the experimental MCMCFlowProposal +Note that this example requires the `minipcn` package to be installed. +""" import numpy as np @@ -8,7 +10,7 @@ from nessai.model import Model from nessai.utils import configure_logger -output = "./outdir/mcmc_example_cvm_8d_flow_act/" +output = "./outdir/minipcn_mcmc_example/" logger = configure_logger(output=output) @@ -41,33 +43,35 @@ def log_likelihood(self, x): ) -# The MCMCFlowProposal class shares most of the configuration options with -# FlowProposal but also has some specific options -# Here we also set -# - n_accept : the number of accepted jumps required to stop -# - reset_flow : this is recommended with the MCMC proposal -fs = FlowSampler( - RosenbrockModel(8), - output=output, - resume=False, - seed=1234, - flow_proposal_class="mcmcflowproposal", - plot_history=True, - plot_chain=False, - step_type="diff", - enforce_likelihood_threshold=True, - # constant_volume_fraction=0.98, - n_steps=5000, - reset_flow=16, - flow_config=dict( - n_neurons=32, - n_blocks=8, - n_layers=2, - batch_norm_between_layers=True, - linear_transform=None, - net="mlp", - ), -) +def main(): + # The MCMCFlowProposal class shares most of the configuration options with + # FlowProposal but also has some specific options + # Here we also set + # - reset_flow : this is recommended with the MCMC proposal + fs = FlowSampler( + RosenbrockModel(8), + output=output, + nlive=1000, + resume=False, + seed=1234, + flow_proposal_class="minipcnflowproposal", + step_fn="tpcn", + n_steps=70, + reset_flow=10, + plot_history=True, + flow_config=dict( + n_neurons=32, + n_blocks=6, + n_layers=2, + batch_norm_between_layers=True, + linear_transform=None, + net="mlp", + ), + ) -# And go! -fs.run() + # And go! + fs.run() + + +if __name__ == "__main__": + main()