bfflearn 0.1.0.dev1__py3-none-any.whl

bff/bayes/priors.py

@@ -0,0 +1,179 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable, Optional, Sequence, Union
+
+import numpy as np
+import torch
+from torch.distributions import Distribution, Normal, Uniform
+
+ArrayLike = Union[np.ndarray, torch.Tensor]
+PathLike = Union[str, Path]
+
+
+@dataclass(frozen=True, slots=True)
+class Prior:
+    kind: str
+    a: float
+    b: float
+    name: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        kind = self.kind.lower()
+        if kind not in {"normal", "uniform"}:
+            raise ValueError(
+                f'Unknown prior type "{self.kind}". Options are "normal" or "uniform".'
+            )
+        if kind == "normal" and self.b <= 0:
+            raise ValueError("Normal prior scale must be positive.")
+        if kind == "uniform" and self.a >= self.b:
+            raise ValueError("Uniform prior requires lower < upper.")
+
+        object.__setattr__(self, "kind", kind)
+        object.__setattr__(self, "a", float(self.a))
+        object.__setattr__(self, "b", float(self.b))
+
+    @property
+    def distribution(self) -> Distribution:
+        if self.kind == "normal":
+            return Normal(self.a, self.b, validate_args=False)
+        return Uniform(self.a, self.b, validate_args=False)
+
+    @property
+    def mean(self) -> float:
+        return self.a if self.kind == "normal" else 0.5 * (self.a + self.b)
+
+    @property
+    def scale(self) -> float:
+        return self.b if self.kind == "normal" else (self.b - self.a) / np.sqrt(12)
+
+    def to_dict(self) -> dict[str, float | str]:
+        data = {"kind": self.kind, "a": self.a, "b": self.b}
+        if self.name is not None:
+            data["name"] = self.name
+        return data
+
+
+@dataclass(slots=True)
+class Priors:
+    items: list[Prior] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __iter__(self):
+        return iter(self.items)
+
+    def __getitem__(self, idx: int) -> Prior:
+        return self.items[idx]
+
+    @property
+    def names(self) -> list[str]:
+        return [prior.name or f"theta_{i}" for i, prior in enumerate(self.items)]
+
+    @property
+    def distributions(self) -> list[Distribution]:
+        return [prior.distribution for prior in self.items]
+
+    @property
+    def means(self) -> np.ndarray:
+        return np.array([prior.mean for prior in self.items], dtype=float)
+
+    @property
+    def scales(self) -> np.ndarray:
+        return np.array([prior.scale for prior in self.items], dtype=float)
+
+    def log_prob(self, theta: torch.Tensor) -> torch.Tensor:
+        if theta.dim() == 1:
+            theta = theta.unsqueeze(0)
+        log_probs = [
+            prior.distribution.log_prob(theta[:, i]).to(theta.device)
+            for i, prior in enumerate(self.items)
+        ]
+        return torch.stack(log_probs, dim=1).sum(dim=1)
+
+    def write(self, fn_out: PathLike) -> None:
+        fn_out = Path(fn_out)
+        data = {"priors": [prior.to_dict() for prior in self.items]}
+        torch.save(data, fn_out)
+
+    @classmethod
+    def load(cls, fn_in: PathLike) -> "Priors":
+        raw = torch.load(Path(fn_in), weights_only=False)
+        entries = raw["priors"] if isinstance(raw, dict) and "priors" in raw else raw
+        return cls([cls._coerce_prior(entry) for entry in entries])
+
+    @classmethod
+    def from_any(cls, priors: Union["Priors", PathLike, Iterable]) -> "Priors":
+        if isinstance(priors, cls):
+            return priors
+        if isinstance(priors, (str, Path)):
+            return cls.load(priors)
+        return cls([cls._coerce_prior(entry) for entry in priors])
+
+    @staticmethod
+    def _coerce_prior(entry) -> Prior:
+        if isinstance(entry, Prior):
+            return entry
+
+        if isinstance(entry, Distribution):
+            name = type(entry).__name__.lower()
+            if name == "normal":
+                return Prior("normal", float(entry.mean), float(entry.scale))
+            if name == "uniform":
+                return Prior("uniform", float(entry.low), float(entry.high))
+            raise ValueError(f"Unsupported distribution type: {name}")
+
+        if isinstance(entry, dict):
+            return Prior(
+                kind=entry["kind"],
+                a=entry["a"],
+                b=entry["b"],
+                name=entry.get("name"),
+            )
+
+        raise TypeError(f"Unsupported prior entry type: {type(entry)}")
+
+    @classmethod
+    def from_bounds(
+        cls,
+        bounds: ArrayLike,
+        dist_type: str = "normal",
+        n_nuisance: int = 0,
+        names: Optional[Sequence[str]] = None,
+        nuisance_names: Optional[Sequence[str]] = None,
+    ) -> "Priors":
+        bounds = np.asarray(bounds, dtype=float)
+        dist_type = dist_type.lower()
+
+        if names is None:
+            names = [None] * len(bounds)
+        elif len(names) != len(bounds):
+            raise ValueError("names must match the number of parameter priors.")
+
+        if dist_type == "normal":
+            centers = bounds.mean(axis=1)
+            widths = 0.2 * (bounds[:, 1] - bounds[:, 0])
+            items = [
+                Prior("normal", center, width, name=name)
+                for center, width, name in zip(centers, widths, names)
+            ]
+        elif dist_type == "uniform":
+            items = [
+                Prior("uniform", lower, upper, name=name)
+                for (lower, upper), name in zip(bounds, names)
+            ]
+        else:
+            raise ValueError(
+                f'Unknown prior type "{dist_type}". Options are "normal" or "uniform".'
+            )
+
+        nuisance_names = nuisance_names or [f"nuisance_{i}" for i in range(n_nuisance)]
+        items.extend(Prior("normal", -2.0, 2.0, name=name) for name in nuisance_names)
+        return cls(items)
+
+
+def log_prior(
+    theta: torch.Tensor,
+    priors: Union[Priors, Sequence[Distribution]],
+) -> torch.Tensor:
+    return Priors.from_any(priors).log_prob(theta)

bff/bayes/results.py

@@ -0,0 +1,327 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Union
+
+import numpy as np
+import torch
+from scipy.stats import gaussian_kde
+
+from ..domain.specs import Specs
+from ..io.utils import save_yaml
+from ..mcmc.convergence import integrated_autocorr_time
+from .priors import Priors
+
+PathLike = Union[str, Path]
+
+
+@dataclass
+class InferenceResults:
+    posterior: np.ndarray
+    priors: Optional[Priors] = None
+    sample_labels: Optional[List[str]] = None
+    sample_transform: Optional[Callable[[np.ndarray], np.ndarray]] = None
+    specs: Optional[Specs] = None
+    include_implicit_charge: bool = False
+
+    _prepared_samples: Optional[np.ndarray] = field(
+        default=None,
+        init=False,
+        repr=False,
+    )
+    _tau: Optional[np.ndarray] = field(default=None, init=False, repr=False)
+    _prepared_labels: Optional[list[str]] = field(default=None, init=False, repr=False)
+
+    @classmethod
+    def load(
+        cls,
+        posterior: PathLike | np.ndarray | torch.Tensor,
+        priors: Optional[Union[Priors, PathLike, list]] = None,
+        sample_labels: Optional[List[str]] = None,
+        sample_transform: Optional[Callable[[np.ndarray], np.ndarray]] = None,
+        specs: Optional[Specs | PathLike] = None,
+        include_implicit_charge: bool = False,
+    ) -> "InferenceResults":
+        return cls(
+            posterior=posterior,
+            priors=priors,
+            sample_labels=sample_labels,
+            sample_transform=sample_transform,
+            specs=specs,
+            include_implicit_charge=include_implicit_charge,
+        )
+
+    def __post_init__(self) -> None:
+        self.posterior = self._coerce_posterior(self.posterior)
+        if self.priors is not None:
+            self.priors = self._coerce_priors(self.priors)
+        if self.specs is not None and not isinstance(self.specs, Specs):
+            self.specs = Specs(self.specs)
+        if self.sample_labels is None and self.priors is not None:
+            self.sample_labels = self._default_sample_labels()
+        if self.sample_transform is None and self._nuisance_indices:
+            self.sample_transform = self._default_transform
+
+    @property
+    def _nuisance_indices(self) -> list[int]:
+        if self.priors is None:
+            return []
+        return [
+            i for i, name in enumerate(self.priors.names)
+            if name.startswith("log_sigma_")
+        ]
+
+    def _default_sample_labels(self) -> list[str]:
+        if self.priors is None:
+            return [f"theta_{i}" for i in range(self.n_dim)]
+        labels = []
+        for i, name in enumerate(self.priors.names):
+            if name.startswith("log_sigma_"):
+                qoi = name.removeprefix("log_sigma_")
+                labels.append(f"$\\sigma_{{\\mathrm{{{qoi}}}}}$")
+            else:
+                labels.append(name or f"theta_{i}")
+        return labels
+
+    def _default_transform(self, samples: np.ndarray) -> np.ndarray:
+        transformed = np.asarray(samples, dtype=float).copy()
+        if self._nuisance_indices:
+            transformed[:, self._nuisance_indices] = np.exp(
+                transformed[:, self._nuisance_indices]
+            )
+        return transformed
+
+    @staticmethod
+    def _to_numpy(x: np.ndarray | torch.Tensor) -> np.ndarray:
+        if isinstance(x, torch.Tensor):
+            return x.detach().cpu().numpy()
+        return np.asarray(x, dtype=float)
+
+    @staticmethod
+    def _coerce_priors(priors: Priors | PathLike | list) -> Priors:
+        if isinstance(priors, Priors):
+            return priors
+        return Priors.from_any(priors)
+
+    @staticmethod
+    def _load_posterior_payload(fn_in: PathLike) -> dict[str, Any]:
+        fn_in = Path(fn_in)
+        if fn_in.suffix != ".pt":
+            raise ValueError(
+                f"Unsupported posterior file {fn_in!r}. "
+                "InferenceResults accepts only .pt files."
+            )
+        return torch.load(fn_in, weights_only=False)
+
+    @classmethod
+    def _coerce_posterior(
+        cls,
+        posterior: PathLike | np.ndarray | torch.Tensor,
+    ) -> np.ndarray:
+        if isinstance(posterior, (np.ndarray, torch.Tensor)):
+            posterior_array = cls._to_numpy(posterior)
+        elif isinstance(posterior, (str, Path)):
+            payload = cls._load_posterior_payload(posterior)
+            if "posterior" not in payload:
+                raise KeyError(
+                    f"Posterior file {posterior!r} does not contain posterior samples."
+                )
+            posterior_array = cls._to_numpy(payload["posterior"])
+        else:
+            raise TypeError(f"Unsupported posterior source: {type(posterior)}")
+
+        if posterior_array.ndim != 3:
+            raise ValueError(
+                "Expected raw posterior with shape "
+                f"(n_saved, n_walkers, n_dim), got {posterior_array.shape}."
+            )
+        return posterior_array
+
+    @property
+    def n_dim(self) -> int:
+        return self.posterior.shape[-1]
+
+    @property
+    def autocorr_time(self) -> np.ndarray:
+        if self._tau is None:
+            tau = integrated_autocorr_time(
+                torch.as_tensor(self.posterior, dtype=torch.float32)
+            ).mean(dim=0)
+            self._tau = np.atleast_1d(tau.detach().cpu().numpy().astype(float))
+        return self._tau
+
+    @property
+    def prepared_samples(self) -> np.ndarray:
+        if self._prepared_samples is None:
+            raise ValueError(
+                "No prepared samples available. Call prepare_samples() first."
+            )
+        return self._prepared_samples
+
+    @property
+    def has_prepared_samples(self) -> bool:
+        return self._prepared_samples is not None
+
+    @property
+    def labels(self) -> List[str]:
+        if self._prepared_samples is not None and self._prepared_labels is not None:
+            return self._prepared_labels
+        if self.sample_labels is not None:
+            if len(self.sample_labels) != self.n_dim:
+                raise ValueError(
+                    "sample_labels length does not match the posterior dimension."
+                )
+            return self.sample_labels
+        return [f"theta_{i}" for i in range(self.n_dim)]
+
+    def _labels_with_implicit_charge(self) -> list[str]:
+        if self.specs is None:
+            raise ValueError(
+                "Implicit-charge expansion requires Specs to be attached to "
+                "InferenceResults."
+            )
+        raw_labels = (
+            list(self.sample_labels)
+            if self.sample_labels is not None
+            else [f"theta_{i}" for i in range(self.n_dim)]
+        )
+        n_explicit = self.specs.explicit_bounds.n_params
+        if len(raw_labels) < n_explicit:
+            raise ValueError(
+                "Posterior labels do not cover all explicit parameters required "
+                "by the specs."
+            )
+        insert_at = self.specs.implicit_param_index
+        implicit_param = self.specs.implicit_param
+        return (
+            raw_labels[:insert_at]
+            + [implicit_param]
+            + raw_labels[insert_at:]
+        )
+
+    def prepare_samples(
+        self,
+        discard: Optional[int] = None,
+        thin: Optional[int] = None,
+        sample_transform: Optional[Callable[[np.ndarray], np.ndarray]] = None,
+        strip_outliers: bool = True
+    ) -> None:
+        tau = self.autocorr_time
+        discard = discard if discard is not None else int(2 * np.max(tau))
+        thin = thin if thin is not None else max(1, int(0.5 * np.min(tau)))
+
+        samples = self.posterior[discard::thin]
+        if samples.size == 0:
+            raise ValueError(
+                "No posterior samples remain after applying discard/thin. "
+                "Try smaller values."
+            )
+
+        prepared = samples.reshape(-1, samples.shape[-1]).copy()
+        transform_fn = sample_transform or self.sample_transform
+        if transform_fn is not None:
+            prepared = np.asarray(transform_fn(prepared), dtype=float)
+        if self.include_implicit_charge:
+            if self.specs is None:
+                raise ValueError(
+                    "include_implicit_charge=True requires Specs to be provided."
+                )
+            prepared = self.specs.with_implicit_charge(prepared)
+            self._prepared_labels = self._labels_with_implicit_charge()
+        else:
+            self._prepared_labels = (
+                list(self.sample_labels)
+                if self.sample_labels is not None
+                else [f"theta_{i}" for i in range(prepared.shape[1])]
+            )
+        self._prepared_samples = prepared
+
+        if strip_outliers:
+            # Remove samples outside the 99% central interval to avoid skewing plots.
+            q_low = 0.01 / 2
+            q_high = 1 - q_low
+            confint = np.quantile(prepared, [q_low, q_high], axis=0)
+            mask = np.all((prepared >= confint[0]) & (prepared <= confint[1]), axis=1)
+            self._prepared_samples = prepared[mask]
+
+    @property
+    def map_estimates(self) -> dict[str, float]:
+        modes: list[float] = []
+        for col in self.prepared_samples.T:
+            grid = np.linspace(col.min(), col.max(), 1000)
+            density = gaussian_kde(col)
+            modes.append(float(np.round(grid[np.argmax(density(grid))], 3)))
+        return dict(zip(self.labels, modes))
+
+    def quantiles(self, confidence: float = 0.95) -> np.ndarray:
+        q_low = (1 - confidence) / 2
+        q_high = 1 - q_low
+        return np.quantile(self.prepared_samples, [q_low, 0.5, q_high], axis=0)
+
+    def _export_parameter_indices_and_labels(self) -> tuple[list[int], list[str]]:
+        if self.specs is not None:
+            explicit_names = list(self.specs.parameter_names(explicit_only=True))
+            label_to_index = {label: i for i, label in enumerate(self.labels)}
+            missing = [name for name in explicit_names if name not in label_to_index]
+            if missing:
+                raise ValueError(
+                    "Prepared samples do not contain all explicit parameter "
+                    "labels required by the specs: "
+                    + ", ".join(repr(name) for name in missing)
+                )
+            return [label_to_index[name] for name in explicit_names], explicit_names
+
+        if self.priors is not None:
+            indices = [
+                i for i, name in enumerate(self.priors.names)
+                if not name.startswith("log_sigma_")
+            ]
+            labels = [self.labels[i] for i in indices]
+            return indices, labels
+
+        return list(range(self.prepared_samples.shape[1])), list(self.labels)
+
+    def _parameter_draw_source(self) -> tuple[np.ndarray, list[str]]:
+        indices, labels = self._export_parameter_indices_and_labels()
+        return np.asarray(self.prepared_samples[:, indices], dtype=float), labels
+
+    def sample_parameters(
+        self,
+        n_samples: int = 10,
+        distribution: str = "normal",
+        confidence: float = 0.9,
+        fn_out: Optional[str] = None,
+        overwrite: bool = False,
+    ) -> np.ndarray:
+        param_samples, export_labels = self._parameter_draw_source()
+        q_low = (1 - confidence) / 2
+        q_high = 1 - q_low
+        confint = np.quantile(param_samples, [q_low, q_high], axis=0)
+
+        if distribution == "normal":
+            mean = np.mean(param_samples, axis=0)
+            cov = np.cov(param_samples, rowvar=False)
+            if np.ndim(cov) == 0:
+                draws = np.random.normal(mean, np.sqrt(cov), size=n_samples)[:, None]
+            else:
+                draws = np.random.multivariate_normal(mean, cov, size=n_samples)
+        elif distribution == "uniform":
+            widths = np.diff(confint, axis=0).ravel()
+            draws = np.random.rand(n_samples, widths.size) * widths + confint[0]
+        else:
+            raise ValueError('distribution must be "uniform" or "normal".')
+
+        if fn_out:
+            fn_out = Path(fn_out).resolve()
+            if fn_out.exists() and not overwrite:
+                raise FileExistsError(f"File '{fn_out}' already exists.")
+            if fn_out.suffix != ".yaml":
+                raise ValueError("fn_out must end with .yaml")
+            save_yaml(
+                {
+                    label: draws[:, i].tolist()
+                    for i, label in enumerate(export_labels)
+                },
+                fn_out,
+            )
+
+        return np.asarray(draws, dtype=float)