aspire-inference 0.1.0a7__py3-none-any.whl

aspire/history.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass, field
+
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+
+from .utils import recursively_save_to_h5_file
+
+
+@dataclass
+class History:
+    """Base class for storing history of a sampler."""
+
+    def save(self, h5_file, path="history"):
+        """Save the history to an HDF5 file."""
+        dictionary = copy.deepcopy(self.__dict__)
+        recursively_save_to_h5_file(h5_file, path, dictionary)
+
+
+@dataclass
+class FlowHistory(History):
+    training_loss: list[float] = field(default_factory=list)
+    validation_loss: list[float] = field(default_factory=list)
+
+    def plot_loss(self) -> Figure:
+        """Plot the training and validation loss."""
+        fig = plt.figure()
+        plt.plot(self.training_loss, label="Training loss")
+        plt.plot(self.validation_loss, label="Validation loss")
+        plt.legend()
+        plt.xlabel("Epoch")
+        plt.ylabel("Loss")
+        return fig
+
+    def save(self, h5_file, path="flow_history"):
+        """Save the history to an HDF5 file."""
+        super().save(h5_file, path=path)
+
+
+@dataclass
+class SMCHistory(History):
+    log_norm_ratio: list[float] = field(default_factory=list)
+    log_norm_ratio_var: list[float] = field(default_factory=list)
+    beta: list[float] = field(default_factory=list)
+    ess: list[float] = field(default_factory=list)
+    ess_target: list[float] = field(default_factory=list)
+    eff_target: list[float] = field(default_factory=list)
+    mcmc_autocorr: list[float] = field(default_factory=list)
+    mcmc_acceptance: list[float] = field(default_factory=list)
+
+    def save(self, h5_file, path="smc_history"):
+        """Save the history to an HDF5 file."""
+        super().save(h5_file, path=path)
+
+    def plot_beta(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.beta)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel(r"$\beta$")
+        return fig
+
+    def plot_log_norm_ratio(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.log_norm_ratio)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("Log evidence ratio")
+        return fig
+
+    def plot_ess(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.ess)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("ESS")
+        return fig
+
+    def plot_ess_target(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.ess_target)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("ESS target")
+        return fig
+
+    def plot_eff_target(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.eff_target)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("Efficiency target")
+        return fig
+
+    def plot_mcmc_acceptance(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.mcmc_acceptance)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("MCMC Acceptance")
+        return fig
+
+    def plot_mcmc_autocorr(self, ax=None) -> Figure | None:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = None
+        ax.plot(self.mcmc_autocorr)
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("MCMC Autocorr")
+        return fig
+
+    def plot(self, fig: Figure | None = None) -> Figure:
+        methods = [
+            self.plot_beta,
+            self.plot_log_norm_ratio,
+            self.plot_ess,
+            self.plot_ess_target,
+            self.plot_eff_target,
+            self.plot_mcmc_acceptance,
+        ]
+
+        if fig is None:
+            fig, axs = plt.subplots(len(methods), 1, sharex=True)
+        else:
+            axs = fig.axes
+
+        for method, ax in zip(methods, axs):
+            method(ax)
+
+        for ax in axs[:-1]:
+            ax.set_xlabel("")
+
+        return fig

aspire/plot.py

@@ -0,0 +1,50 @@
+import copy
+
+
+def plot_comparison(
+    *samples, parameters=None, per_samples_kwargs=None, labels=None, **kwargs
+):
+    """
+    Plot a comparison of multiple samples.
+    """
+    default_kwargs = dict(
+        density=True,
+        bins=30,
+        color="C0",
+        smooth=1.0,
+        plot_datapoints=True,
+        plot_density=False,
+        hist_kwargs=dict(density=True, color="C0"),
+    )
+    default_kwargs.update(kwargs)
+
+    if per_samples_kwargs is None:
+        per_samples_kwargs = [{}] * len(samples)
+
+    fig = None
+    for i, sample in enumerate(samples):
+        kwds = copy.deepcopy(default_kwargs)
+        color = per_samples_kwargs[i].pop("color", f"C{i}")
+        kwds["color"] = color
+        kwds["hist_kwargs"]["color"] = color
+        kwds.update(per_samples_kwargs[i])
+        fig = sample.plot_corner(fig=fig, parameters=parameters, **kwds)
+
+    if labels:
+        fig.legend(
+            labels=labels,
+            loc="upper right",
+            bbox_to_anchor=(0.9, 0.9),
+            bbox_transform=fig.transFigure,
+        )
+    return fig
+
+
+def plot_history_comparison(*histories):
+    # Assert that all histories are of the same type
+    if not all(isinstance(h, histories[0].__class__) for h in histories):
+        raise ValueError("All histories must be of the same type")
+    fig = histories[0].plot()
+    for history in histories[1:]:
+        fig = history.plot(fig=fig)
+    return fig

aspire/samplers/base.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Any, Callable
+
+from ..flows.base import Flow
+from ..samples import Samples
+from ..transforms import IdentityTransform
+from ..utils import track_calls
+
+logger = logging.getLogger(__name__)
+
+
+class Sampler:
+    """Base class for all samplers.
+
+    Parameters
+    ----------
+    log_likelihood : Callable
+        The log likelihood function.
+    log_prior : Callable
+        The log prior function.
+    dims : int
+        The number of dimensions.
+    flow : Flow
+        The flow object.
+    xp : Callable
+        The array backend to use.
+    parameters : list[str] | None
+        The list of parameter names. If None, any samples objects will not
+        have the parameters names specified.
+    """
+
+    def __init__(
+        self,
+        log_likelihood: Callable,
+        log_prior: Callable,
+        dims: int,
+        prior_flow: Flow,
+        xp: Callable,
+        dtype: Any | str | None = None,
+        parameters: list[str] | None = None,
+        preconditioning_transform: Callable | None = None,
+    ):
+        self.prior_flow = prior_flow
+        self._log_likelihood = log_likelihood
+        self.log_prior = log_prior
+        self.dims = dims
+        self.xp = xp
+        self.dtype = dtype
+        self.parameters = parameters
+        self.history = None
+        self.n_likelihood_evaluations = 0
+        if preconditioning_transform is None:
+            self.preconditioning_transform = IdentityTransform(xp=self.xp)
+        else:
+            self.preconditioning_transform = preconditioning_transform
+
+    def fit_preconditioning_transform(self, x):
+        """Fit the data transform to the data."""
+        x = self.preconditioning_transform.xp.asarray(x)
+        return self.preconditioning_transform.fit(x)
+
+    @track_calls
+    def sample(self, n_samples: int) -> Samples:
+        raise NotImplementedError
+
+    def log_likelihood(self, samples: Samples) -> Samples:
+        """Computes the log likelihood of the samples.
+
+        Also tracks the number of likelihood evaluations.
+        """
+        self.n_likelihood_evaluations += len(samples)
+        return self._log_likelihood(samples)
+
+    def config_dict(self, include_sample_calls: bool = True) -> dict:
+        """
+        Returns a dictionary with the configuration of the sampler.
+
+        Parameters
+        ----------
+        include_sample_calls : bool
+            Whether to include the sample calls in the configuration.
+            Default is True.
+        """
+        config = {}
+        if include_sample_calls:
+            if hasattr(self, "sample") and hasattr(self.sample, "calls"):
+                config["sample_calls"] = self.sample.calls.to_dict(
+                    list_to_dict=True
+                )
+            else:
+                logger.warning(
+                    "Sampler does not have a sample method with calls attribute."
+                )
+        return config

aspire/samplers/importance.py

@@ -0,0 +1,22 @@
+from ..samples import Samples
+from ..utils import track_calls
+from .base import Sampler
+
+
+class ImportanceSampler(Sampler):
+    @track_calls
+    def sample(self, n_samples: int) -> Samples:
+        x, log_q = self.prior_flow.sample_and_log_prob(n_samples)
+        samples = Samples(
+            x,
+            log_q=log_q,
+            xp=self.xp,
+            parameters=self.parameters,
+            dtype=self.dtype,
+        )
+        samples.log_prior = samples.array_to_namespace(self.log_prior(samples))
+        samples.log_likelihood = samples.array_to_namespace(
+            self.log_likelihood(samples)
+        )
+        samples.compute_weights()
+        return samples

aspire/samplers/mcmc.py

@@ -0,0 +1,160 @@
+import numpy as np
+
+from ..samples import Samples, to_numpy
+from ..utils import track_calls
+from .base import Sampler
+
+
+class MCMCSampler(Sampler):
+    def draw_initial_samples(self, n_samples: int) -> Samples:
+        """Draw initial samples from the prior flow."""
+        # Flow may propose samples outside prior bounds, so we may need
+        # to try multiple times to get enough valid samples.
+        n_samples_drawn = 0
+        samples = None
+        while n_samples_drawn < n_samples:
+            n_to_draw = n_samples - n_samples_drawn
+            x, log_q = self.prior_flow.sample_and_log_prob(n_to_draw)
+            new_samples = Samples(x, xp=self.xp, log_q=log_q, dtype=self.dtype)
+            new_samples.log_prior = new_samples.array_to_namespace(
+                self.log_prior(new_samples)
+            )
+            valid = self.xp.isfinite(new_samples.log_prior)
+            n_valid = self.xp.sum(valid)
+            if n_valid > 0:
+                if samples is None:
+                    samples = new_samples[valid]
+                else:
+                    samples = Samples.concatenate(
+                        [samples, new_samples[valid]]
+                    )
+                n_samples_drawn += n_valid
+
+        if n_samples_drawn > n_samples:
+            samples = samples[:n_samples]
+
+        samples.log_likelihood = samples.array_to_namespace(
+            self.log_likelihood(samples)
+        )
+        return samples
+
+    def log_prob(self, z):
+        """Compute the log probability of the samples.
+
+        Input samples are in the transformed space.
+        """
+        x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
+        samples = Samples(x, xp=self.xp, dtype=self.dtype)
+        samples.log_prior = self.log_prior(samples)
+        samples.log_likelihood = self.log_likelihood(samples)
+        log_prob = (
+            samples.log_likelihood
+            + samples.log_prior
+            + samples.array_to_namespace(log_abs_det_jacobian)
+        )
+        return to_numpy(log_prob).flatten()
+
+
+class Emcee(MCMCSampler):
+    @track_calls
+    def sample(
+        self,
+        n_samples: int,
+        nwalkers: int = None,
+        nsteps: int = 500,
+        rng=None,
+        discard=0,
+        **kwargs,
+    ) -> Samples:
+        from emcee import EnsembleSampler
+
+        nwalkers = nwalkers or n_samples
+        self.sampler = EnsembleSampler(
+            nwalkers,
+            self.dims,
+            log_prob_fn=self.log_prob,
+            vectorize=True,
+        )
+
+        rng = rng or np.random.default_rng()
+
+        samples = self.draw_initial_samples(nwalkers)
+        p0 = samples.x
+
+        z0 = to_numpy(self.preconditioning_transform.fit(p0))
+
+        self.sampler.run_mcmc(z0, nsteps, **kwargs)
+
+        z = self.sampler.get_chain(flat=True, discard=discard)
+        x = self.preconditioning_transform.inverse(z)[0]
+
+        x_evidence, log_q = self.prior_flow.sample_and_log_prob(n_samples)
+        samples_evidence = Samples(x_evidence, log_q=log_q, xp=self.xp)
+        samples_evidence.log_prior = self.log_prior(samples_evidence)
+        samples_evidence.log_likelihood = self.log_likelihood(samples_evidence)
+        samples_evidence.compute_weights()
+
+        samples_mcmc = Samples(
+            x, xp=self.xp, parameters=self.parameters, dtype=self.dtype
+        )
+        samples_mcmc.log_prior = samples_mcmc.array_to_namespace(
+            self.log_prior(samples_mcmc)
+        )
+        samples_mcmc.log_likelihood = samples_mcmc.array_to_namespace(
+            self.log_likelihood(samples_mcmc)
+        )
+        samples_mcmc.log_evidence = samples_mcmc.array_to_namespace(
+            samples_evidence.log_evidence
+        )
+        samples_mcmc.log_evidence_error = samples_mcmc.array_to_namespace(
+            samples_evidence.log_evidence_error
+        )
+
+        return samples_mcmc
+
+
+class MiniPCN(MCMCSampler):
+    @track_calls
+    def sample(
+        self,
+        n_samples,
+        rng=None,
+        target_acceptance_rate=0.234,
+        n_steps=100,
+        thin=1,
+        burnin=0,
+        last_step_only=False,
+        step_fn="tpcn",
+    ):
+        from minipcn import Sampler
+
+        rng = rng or np.random.default_rng()
+        p0 = self.draw_initial_samples(n_samples).x
+
+        z0 = to_numpy(self.preconditioning_transform.fit(p0))
+
+        self.sampler = Sampler(
+            log_prob_fn=self.log_prob,
+            step_fn=step_fn,
+            rng=rng,
+            dims=self.dims,
+            target_acceptance_rate=target_acceptance_rate,
+        )
+
+        chain, history = self.sampler.sample(z0, n_steps=n_steps)
+
+        if last_step_only:
+            z = chain[-1]
+        else:
+            z = chain[burnin::thin].reshape(-1, self.dims)
+
+        x = self.preconditioning_transform.inverse(z)[0]
+
+        samples_mcmc = Samples(x, xp=self.xp, parameters=self.parameters)
+        samples_mcmc.log_prior = samples_mcmc.array_to_namespace(
+            self.log_prior(samples_mcmc)
+        )
+        samples_mcmc.log_likelihood = samples_mcmc.array_to_namespace(
+            self.log_likelihood(samples_mcmc)
+        )
+        return samples_mcmc