aspire-inference 0.1.0a2__py3-none-any.whl

aspire/flows/torch/flows.py

@@ -0,0 +1,276 @@
+import copy
+import logging
+from typing import Callable
+
+import array_api_compat.torch as torch_api
+import torch
+import tqdm
+import zuko
+from array_api_compat import is_numpy_namespace, is_torch_array
+
+from ...history import FlowHistory
+from ..base import Flow
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTorchFlow(Flow):
+    _flow = None
+    xp = torch_api
+
+    def __init__(
+        self,
+        dims: int,
+        seed: int = 1234,
+        device: str = "cpu",
+        data_transform=None,
+    ):
+        super().__init__(
+            dims,
+            device=torch.device(device or "cpu"),
+            data_transform=data_transform,
+        )
+        torch.manual_seed(seed)
+        self.loc = None
+        self.scale = None
+
+    @property
+    def flow(self):
+        return self._flow
+
+    @flow.setter
+    def flow(self, flow):
+        self._flow = flow
+        self._flow.to(self.device)
+        self._flow.compile()
+
+    def fit(self, x) -> FlowHistory:
+        raise NotImplementedError()
+
+
+class ZukoFlow(BaseTorchFlow):
+    def __init__(
+        self,
+        dims,
+        flow_class: str | Callable = "MAF",
+        data_transform=None,
+        seed=1234,
+        device: str = "cpu",
+        **kwargs,
+    ):
+        super().__init__(
+            dims,
+            device=device,
+            data_transform=data_transform,
+            seed=seed,
+        )
+
+        if isinstance(flow_class, str):
+            FlowClass = getattr(zuko.flows, flow_class)
+        else:
+            FlowClass = flow_class
+
+        # Ints are some times passed as strings, so we convert them
+        if hidden_features := kwargs.pop("hidden_features", None):
+            kwargs["hidden_features"] = list(map(int, hidden_features))
+
+        self.flow = FlowClass(self.dims, 0, **kwargs)
+        logger.info(f"Initialized normalizing flow: \n {self.flow}\n")
+
+    def loss_fn(self, x):
+        return -self.flow().log_prob(x).mean()
+
+    def fit(
+        self,
+        x,
+        n_epochs: int = 100,
+        lr: float = 1e-3,
+        batch_size: int = 500,
+        validation_fraction: float = 0.2,
+        clip_grad: float | None = None,
+        lr_annealing: bool = False,
+    ):
+        from ...history import FlowHistory
+
+        if not is_torch_array(x):
+            x = torch.tensor(
+                x, dtype=torch.get_default_dtype(), device=self.device
+            )
+        else:
+            x = torch.clone(x)
+            x = x.type(torch.get_default_dtype())
+            x = x.to(self.device)
+        x_prime = self.fit_data_transform(x)
+        indices = torch.randperm(x_prime.shape[0])
+        x_prime = x_prime[indices, ...]
+
+        n = x_prime.shape[0]
+        x_train = torch.as_tensor(
+            x_prime[: -int(validation_fraction * n)],
+            dtype=torch.get_default_dtype(),
+            device=self.device,
+        )
+
+        logger.info(
+            f"Training on {x_train.shape[0]} samples, "
+            f"validating on {x_prime.shape[0] - x_train.shape[0]} samples."
+        )
+
+        if torch.isnan(x_train).any():
+            raise ValueError("Training data contains NaN values.")
+        if not torch.isfinite(x_train).all():
+            raise ValueError("Training data contains infinite values.")
+
+        x_val = torch.as_tensor(
+            x_prime[-int(validation_fraction * n) :],
+            dtype=torch.get_default_dtype(),
+            device=self.device,
+        )
+        if torch.isnan(x_val).any():
+            raise ValueError("Validation data contains infinite values.")
+
+        if not torch.isfinite(x_val).all():
+            raise ValueError("Validation data contains infinite values.")
+
+        dataset = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(x_train),
+            shuffle=True,
+            batch_size=batch_size,
+        )
+        val_dataset = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(x_val),
+            shuffle=False,
+            batch_size=batch_size,
+        )
+
+        # Train to maximize the log-likelihood
+        optimizer = torch.optim.Adam(self._flow.parameters(), lr=lr)
+        if lr_annealing:
+            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, n_epochs
+            )
+        history = FlowHistory()
+
+        best_val_loss = float("inf")
+        best_flow_state = None
+
+        with tqdm.tqdm(range(n_epochs), desc="Epochs") as pbar:
+            for _ in pbar:
+                self.flow.train()
+                loss_epoch = 0.0
+                for (x_batch,) in dataset:
+                    loss = self.loss_fn(x_batch)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    if clip_grad is not None:
+                        torch.nn.utils.clip_grad_norm_(
+                            self.flow.parameters(), clip_grad
+                        )
+                    optimizer.step()
+                    loss_epoch += loss.item()
+                if lr_annealing:
+                    scheduler.step()
+                avg_train_loss = loss_epoch / len(dataset)
+                history.training_loss.append(avg_train_loss)
+                self.flow.eval()
+                val_loss = 0.0
+                for (x_batch,) in val_dataset:
+                    with torch.no_grad():
+                        val_loss += self.loss_fn(x_batch).item()
+                avg_val_loss = val_loss / len(val_dataset)
+                if avg_val_loss < best_val_loss:
+                    best_val_loss = avg_val_loss
+                    best_flow_state = copy.deepcopy(self.flow.state_dict())
+
+                history.validation_loss.append(avg_val_loss)
+                pbar.set_postfix(
+                    train_loss=f"{avg_train_loss:.4f}",
+                    val_loss=f"{avg_val_loss:.4f}",
+                )
+        if best_flow_state is not None:
+            self.flow.load_state_dict(best_flow_state)
+            logger.info(f"Loaded best model with val loss {best_val_loss:.4f}")
+
+        self.flow.eval()
+        return history
+
+    def sample_and_log_prob(self, n_samples: int, xp=torch_api):
+        with torch.no_grad():
+            x_prime, log_prob = self.flow().rsample_and_log_prob((n_samples,))
+        x, log_abs_det_jacobian = self.inverse_rescale(x_prime)
+        return xp.asarray(x), xp.asarray(log_prob - log_abs_det_jacobian)
+
+    def sample(self, n_samples: int, xp=torch_api):
+        with torch.no_grad():
+            x_prime = self.flow().rsample((n_samples,))
+        x = self.inverse_rescale(x_prime)[0]
+        return xp.asarray(x)
+
+    def log_prob(self, x, xp=torch_api):
+        x = torch.as_tensor(
+            x, dtype=torch.get_default_dtype(), device=self.device
+        )
+        x_prime, log_abs_det_jacobian = self.rescale(x)
+        return xp.asarray(
+            self._flow().log_prob(x_prime) + log_abs_det_jacobian
+        )
+
+    def forward(self, x, xp=torch_api):
+        x = torch.as_tensor(
+            x, dtype=torch.get_default_dtype(), device=self.device
+        )
+        x_prime, log_j_rescale = self.rescale(x)
+        z, log_abs_det_jacobian = self._flow().transform.call_and_ladj(x_prime)
+        if is_numpy_namespace(xp):
+            # Convert to numpy namespace if needed
+            z = z.detach().numpy()
+            log_abs_det_jacobian = log_abs_det_jacobian.detach().numpy()
+            log_j_rescale = log_j_rescale.detach().numpy()
+        return xp.asarray(z), xp.asarray(log_abs_det_jacobian + log_j_rescale)
+
+    def inverse(self, z, xp=torch_api):
+        z = torch.as_tensor(
+            z, dtype=torch.get_default_dtype(), device=self.device
+        )
+        with torch.no_grad():
+            x_prime, log_abs_det_jacobian = (
+                self._flow().transform.inv.call_and_ladj(z)
+            )
+        x, log_j_rescale = self.inverse_rescale(x_prime)
+        if is_numpy_namespace(xp):
+            # Convert to numpy namespace if needed
+            x = x.detach().numpy()
+            log_abs_det_jacobian = log_abs_det_jacobian.detach().numpy()
+            log_j_rescale = log_j_rescale.detach().numpy()
+        return xp.asarray(x), xp.asarray(log_j_rescale + log_abs_det_jacobian)
+
+
+class ZukoFlowMatching(ZukoFlow):
+    def __init__(
+        self,
+        dims,
+        data_transform=None,
+        seed=1234,
+        device="cpu",
+        eta: float = 1e-3,
+        **kwargs,
+    ):
+        kwargs.setdefault("hidden_features", 4 * [100])
+        super().__init__(
+            dims,
+            seed=seed,
+            device=device,
+            data_transform=data_transform,
+            flow_class="CNF",
+        )
+        self.eta = eta
+
+    def loss_fn(self, theta: torch.Tensor):
+        t = torch.rand(
+            theta.shape[:-1], dtype=theta.dtype, device=theta.device
+        )
+        t_ = t[..., None]
+        eps = torch.randn_like(theta)
+        theta_prime = (1 - t_) * theta + (t_ + self.eta) * eps
+        v = eps - theta
+        return (self._flow.transform.f(t, theta_prime) - v).square().mean()

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,92 @@
+import logging
+from typing import 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,
+        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.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,18 @@
+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
+        )
+        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