binference 0.1.0__py3-none-any.whl

binference/inference.py

@@ -0,0 +1,368 @@
+"""The one-call interface: :func:`infer`."""
+
+from __future__ import annotations
+
+import inspect
+import warnings
+from math import ceil
+from typing import Any
+
+import numpy as np
+
+from ._utils import RNGLike, check_positive, warn
+from .bins import Bins
+from .model import PopulationModel
+from .priors import PopulationPrior
+from .results import PointEstimate, PopulationResult
+from .samplers import MAP, NUTS, Bootstrap, Gibbs, Laplace, MeanFieldVB
+from .samplers.base import Engine
+from .samplers.laplace import PROPOSAL_FACTOR
+from .warnings import ApproximationWarning, BinferenceWarning
+
+__all__ = ["METHODS", "Auto", "infer"]
+
+#: Registry of inference methods accepted by :func:`infer`.
+METHODS = {
+    "nuts": NUTS,
+    "gibbs": Gibbs,
+    "map": MAP,
+    "em": MAP,
+    "laplace": Laplace,
+    "vb": MeanFieldVB,
+    "bootstrap": Bootstrap,
+}
+
+#: An extension of a Gibbs run aims for this many times ``target_ess``.
+EXTEND_TARGET = 1.2
+#: One extension of a Gibbs run adds at most this many times the draws so far.
+EXTEND_MAX_FACTOR = 3
+#: One extension of a Gibbs run adds at least this many times the draws so far.
+EXTEND_MIN_FACTOR = 0.5
+#: Proposals of the first pass that can stop a Laplace trial early.
+PROBE_PROPOSALS = 1000
+#: Draws of the first pass that can stop a Laplace trial early.
+PROBE_DRAWS = 500
+#: A Laplace trial stops when its first pass projects an importance ESS below this share of
+#: ``target_ess``.
+PROBE_ESS_SHARE = 0.5
+
+
+def _options(engine: type) -> dict[str, inspect.Parameter]:
+    return dict(inspect.signature(engine.__init__).parameters)
+
+
+def _default_draws(engine: type) -> int:
+    """The engine's own default ``n_samples``."""
+    return inspect.signature(engine.run).parameters["n_samples"].default
+
+
+class Auto(Engine):
+    r"""Picks Gibbs, Laplace or NUTS from rules set by a benchmark of the engines.
+
+    It tries Gibbs where it mixes well, then importance-corrected Laplace where its
+    check passes, then NUTS, keeping only engines that take every option in
+    ``engine_kwargs``. :doc:`/how-it-works/auto` states the rules and the benchmark.
+    The result records the choice and the reason in ``info["auto_method"]`` and
+    ``info["auto_reason"]``.
+
+    Parameters
+    ----------
+    model : PopulationModel
+    target_ess : float
+        Smallest bulk ESS that Gibbs's run length aims for and that Laplace's
+        importance ESS must reach.
+    gibbs_max_missing : float
+        Largest missing-information fraction for which Gibbs is chosen (0.98: an
+        integrated autocorrelation time of 99 iterations).
+    nnz_limit : float
+        Above this many stored likelihood entries Laplace is kept even when its check
+        fails, rather than running NUTS.
+    max_gibbs_draws : int
+        Most draws per chain a Gibbs run may be extended to before it gives way to the
+        next candidate.
+    rng, progress
+        As for every engine.
+    **engine_kwargs
+        Options for the chosen engine.
+    """
+
+    name = "auto"
+    joint = True
+
+    def __init__(
+        self,
+        model: PopulationModel,
+        *,
+        target_ess: float = 400.0,
+        gibbs_max_missing: float = 0.98,
+        nnz_limit: float = 2e7,
+        max_gibbs_draws: int = 25_000,
+        rng: RNGLike = None,
+        progress: bool = False,
+        **engine_kwargs: Any,
+    ) -> None:
+        super().__init__(model, rng=rng, progress=progress)
+        self.target_ess = float(check_positive(target_ess, "target_ess"))
+        if not 0.0 <= gibbs_max_missing < 1.0:
+            raise ValueError(f"gibbs_max_missing must be in [0, 1); got {gibbs_max_missing}.")
+        self.gibbs_max_missing = float(gibbs_max_missing)
+        self.nnz_limit = float(check_positive(nnz_limit, "nnz_limit"))
+        self.max_gibbs_draws = int(check_positive(max_gibbs_draws, "max_gibbs_draws"))
+        self.engine_kwargs = engine_kwargs
+        self._rng_arg = rng
+        self._map: tuple[Any, float] | None = None  # the MAP and its missing information
+
+    def candidates(self, n_samples: int | None = None) -> list[tuple[str, str, int | None]]:
+        """(engine, reason, Gibbs draws per chain) in the order they are tried.
+
+        Parameters
+        ----------
+        n_samples : int, optional
+            Draws per chain of the runs; ``None`` uses the engines' own lengths.
+        """
+        m, kw = self.model, self.engine_kwargs
+        out: list[tuple[str, str, int | None]] = []
+        if Gibbs.unsupported(m) is None and set(kw) <= set(_options(Gibbs)):
+            if m.joint:  # exact and fast (DECISIONS #80, #81)
+                out.append(("gibbs", "a joint model with a conjugate prior", None))
+            else:
+                if self._map is None:
+                    with warnings.catch_warnings(record=True) as caught:
+                        warnings.simplefilter("always", BinferenceWarning)
+                        est = MAP(m).run()
+                    # kept as best only when it raised no warning (best would show them)
+                    keep = est if not caught else None
+                    self._map = keep, m.missing_information(est.value)  # component weights
+                f = self._map[1]
+                if f <= self.gibbs_max_missing:
+                    chains = kw.get("n_chains", _options(Gibbs)["n_chains"].default)
+                    n = max(
+                        _default_draws(Gibbs),
+                        ceil(self.target_ess * (1 + f) / (1 - f) / chains),
+                    )
+                    if n_samples is not None or n <= self.max_gibbs_draws:
+                        out.append(("gibbs", f"conjugate prior, missing information {f:.2f}", n))
+        out += [
+            (key, "", None)
+            for key, engine in (("laplace", Laplace), ("nuts", NUTS))
+            if set(kw) <= set(_options(engine))
+        ]
+        if not out:
+            raise ValueError(
+                f"None of the engines method='auto' can use for this model (Laplace, NUTS, "
+                f"and Gibbs where it applies) takes all of {sorted(kw)}; pass method= "
+                "explicitly."
+            )
+        return out
+
+    def run(self, n_samples: int | None = None) -> PopulationResult:
+        """Run the first candidate that applies (see the class docstring).
+
+        Parameters
+        ----------
+        n_samples : int, optional
+            Draws per chain; ``None`` lets auto set the run length. Laplace returns as
+            many in total as NUTS's default four chains would (4000 by default).
+        """
+        plan = self.candidates(n_samples)
+        rejected: list[str] = []  # why earlier candidates were not kept
+        for i, (key, why, n_gibbs) in enumerate(plan):
+            last = i == len(plan) - 1
+            engine = METHODS[key](
+                self.model, rng=self._rng_arg, progress=self.progress, **self.engine_kwargs
+            )
+            if key == "laplace":
+                n = _options(NUTS)["n_chains"].default * (n_samples or _default_draws(NUTS))
+                result, why = (engine.run(n), "") if last else self._try(engine, n)
+                if result is None:
+                    rejected.append(why)
+                    continue
+                if not (result.approximate or self._passed(result)):
+                    # k-hat passed (else Laplace flagged and warned) but the ESS fell short
+                    result.approximate = True
+                    short = f"its importance ESS is below target_ess={self.target_ess:g}"
+                    why = "; ".join(filter(None, [why, short]))
+                    warn(
+                        f"auto kept Laplace draws although {short}; results are approximate.",
+                        ApproximationWarning,
+                    )
+            else:
+                n = n_samples or n_gibbs
+                result = engine.run() if n is None else engine.run(n)
+                # stored labels grow with the run: their length stays as asked
+                extend = key == "gibbs" and n_samples is None and engine.label_index is None
+                if extend and not self._extend(engine, result):
+                    if not last:
+                        rejected.append("Gibbs mixed too slowly")
+                        continue
+                    why += f"; it did not mix within max_gibbs_draws={self.max_gibbs_draws}"
+                if extend:
+                    result = engine.result
+            reasons = [*rejected, why] if why else rejected
+            result.info["auto_method"] = key
+            result.info["auto_reason"] = (
+                "; ".join(reasons) or "the only engine that takes these options"
+            )
+            if self._map is not None and self._map[0] is not None:  # not solved again
+                result._cache["best"] = self._map[0]
+            return result
+        raise AssertionError("unreachable: the last candidate always returns")  # pragma: no cover
+
+    def _extend(self, engine: Gibbs, result: PopulationResult) -> bool:
+        """Extend the Gibbs run until its bulk ESS reaches ``target_ess`` and it passes the
+        limits of ``check()``; False once that is projected to need more than
+        ``max_gibbs_draws`` draws per chain."""
+        while True:
+            d = result.diagnostics
+            ess = np.concatenate([np.ravel(v) for k, v in d.items() if k.startswith("ess_bulk")])
+            ess = float(np.min(ess[np.isfinite(ess)], initial=np.inf))  # constant: no ESS
+            if ess >= self.target_ess and not result._mixing_problems():
+                return True
+            done = result.chains.shape[1]
+            need = ceil(
+                done
+                * max(EXTEND_TARGET * self.target_ess / max(ess, 1.0) - 1.0, EXTEND_MIN_FACTOR)
+            )
+            if done + need > self.max_gibbs_draws:
+                return False
+            # bounded, so a poor first projection costs less
+            result = engine.run(min(need, EXTEND_MAX_FACTOR * done))
+
+    def _passed(self, res: PopulationResult) -> bool:
+        """Whether Laplace draws pass auto's check: Pareto k-hat and importance ESS."""
+        return not res.approximate and res.info.get("ess_importance", 0.0) >= self.target_ess
+
+    def _try(self, engine: Laplace, n: int) -> tuple[PopulationResult | None, str]:
+        """Laplace's result when it is kept (its warnings re-emitted), else None.
+
+        Laplace is only tried with no engine options, so it has its default
+        ``PROPOSAL_FACTOR * n`` proposals. A first pass of ``PROBE_PROPOSALS`` of them
+        stops the trial early when the importance ESS it projects for the full run is
+        below ``PROBE_ESS_SHARE`` of ``target_ess``.
+        """
+        failed = None, "the Laplace importance check failed"
+        big = self.model.likelihoods.nnz > self.nnz_limit  # kept whatever its check
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            try:
+                proposals = PROPOSAL_FACTOR * n
+                if not big and proposals > PROBE_PROPOSALS:
+                    with warnings.catch_warnings():
+                        warnings.simplefilter("ignore")
+                        engine.n_proposals = PROBE_PROPOSALS
+                        probe = engine.run(PROBE_DRAWS)
+                        engine.n_proposals = None
+                    projected = probe.info["ess_importance"] * proposals / PROBE_PROPOSALS
+                    if projected < PROBE_ESS_SHARE * self.target_ess:
+                        return failed
+                res = engine.run(n)
+            except (ValueError, FloatingPointError, np.linalg.LinAlgError):
+                return None, "the Laplace approximation failed"
+        if self._passed(res):
+            reason = "the Laplace importance check passed"
+        elif big:
+            reason = f"more than {self.nnz_limit:.0e} likelihood entries"
+        else:
+            return failed
+        for w in caught:
+            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)
+        return res, reason
+
+
+METHODS["auto"] = Auto
+
+
+def infer(
+    likelihoods: Any,
+    prior: PopulationPrior | None = None,
+    *,
+    bins: Bins | None = None,
+    selection: Any = None,
+    selection_mode: str | None = None,
+    method: str = "auto",
+    n_samples: int | None = None,
+    rng: RNGLike = None,
+    progress: bool = False,
+    check: bool = True,
+    **engine_kwargs: Any,
+) -> PopulationResult | PointEstimate:
+    r"""Infer a population distribution from per-object likelihoods in one call.
+
+    Parameters
+    ----------
+    likelihoods : Likelihoods, PopulationModel, or array_like (N, K)
+        Per-object likelihoods over bins (see :class:`~binference.Likelihoods` for
+        converting posteriors, samples, Gaussian measurements, response
+        matrices or classifier outputs). A raw array is treated as
+        likelihoods over categorical bins (or ``bins``).
+    prior : PopulationPrior, optional
+        Population prior (default ``Dirichlet(1.0)``). Use
+        :class:`~binference.priors.GammaRates` to infer rates.
+    bins : Bins, optional
+        Bins for a raw array input.
+    selection : array_like (K,), callable or Selection, optional
+        Detection probability (fractions) or exposure (rates) per bin.
+    selection_mode : {"on_measurement", "on_truth"}, optional
+        Required with an array or callable ``selection`` (see
+        :class:`~binference.selection.Selection`).
+    method : {"auto", "nuts", "gibbs", "map", "em", "laplace", "vb", "bootstrap"}
+        Inference engine. ``"auto"`` (default) picks Gibbs, Laplace or NUTS from
+        rules set by a benchmark of the engines (:class:`Auto`); its
+        options ``target_ess``, ``gibbs_max_missing``, ``max_gibbs_draws`` and
+        ``nnz_limit`` go in ``engine_kwargs``, and the choice and its reason are in
+        ``result.info["auto_method"]`` and ``result.info["auto_reason"]``.
+        ``"nuts"`` samples the exact posterior with any prior; ``"em"`` is MAP by
+        EM on the simplex (``algorithm="em"``) rather than mix-SQP; it needs a
+        conjugate prior, and a jointly inferred selection or response runs L-BFGS.
+    n_samples : int, optional
+        Draws per chain (NUTS, Gibbs, and ``"auto"``) or in total (Laplace, VB,
+        bootstrap). Under ``"auto"`` it fixes the run length (see :class:`Auto`).
+    rng : int, numpy.random.Generator, or None
+        Random number generator or seed (SPEC 7).
+    progress : bool
+        Show progress bars.
+    check : bool
+        Warn about convergence problems after sampling (NUTS, Gibbs).
+    **engine_kwargs
+        Passed to the engine (e.g. ``n_chains``, ``n_warmup``, ``target_accept``
+        for NUTS; see :mod:`binference.samplers`).
+
+    Returns
+    -------
+    PopulationResult or PointEstimate
+        Posterior draws, or the MAP estimate for ``method="map"`` and ``"em"``.
+
+    Examples
+    --------
+    >>> import numpy as np, binference
+    >>> rng = np.random.default_rng(1)
+    >>> bins = binference.Bins.linear(0.0, 1.0, 5)
+    >>> z = rng.uniform(0, 1, 300)
+    >>> lk = binference.Likelihoods.from_gaussian(z + rng.normal(0, 0.05, 300), 0.05, bins)
+    >>> res = binference.infer(lk, rng=0)
+    >>> res.mean.shape
+    (5,)
+    """
+    if isinstance(likelihoods, PopulationModel):
+        if any(v is not None for v in (prior, selection, selection_mode, bins)):
+            raise ValueError("Pass prior/selection/bins when building the model, not to infer().")
+        model = likelihoods
+    else:
+        model = PopulationModel(
+            likelihoods, prior, selection=selection, selection_mode=selection_mode, bins=bins
+        )
+    key = str(method).lower()
+    if key not in METHODS:
+        raise ValueError(f"Unknown method {method!r}; choose from {sorted(METHODS)}.")
+    if key == "em" and not model.joint:  # MAP by EM itself, not mix-SQP (audit O-17)
+        engine_kwargs.setdefault("algorithm", "em")
+    engine = METHODS[key](model, rng=rng, progress=progress, **engine_kwargs)
+    if n_samples is None:
+        result = engine.run()
+    elif isinstance(engine, MAP):
+        raise ValueError(f"n_samples does not apply to method={key!r}.")
+    else:
+        result = engine.run(int(n_samples))
+    if check and METHODS[result.method].markov:
+        result.check()
+    return result

binference/io.py

@@ -0,0 +1,309 @@
+"""Saving, loading and converting results."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from ._version import __version__
+from .basis import Basis
+from .bins import Bins
+from .measure import Measure
+from .selection import Selection
+from .selection import _rebuild as _rebuild_selection
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .results import PopulationResult
+
+__all__ = ["load_result", "result_to_dict", "save_result", "to_datatree"]
+
+SCHEMA_VERSION = 1
+
+
+def _bins_to_parts(bins: Bins) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
+    if bins.is_categorical:
+        # JSON keeps integer, float and string labels as they are (audit O-10)
+        labels = [lab.item() if isinstance(lab, np.generic) else lab for lab in bins.labels]
+        if not all(isinstance(lab, int | float | str) for lab in labels):
+            raise ValueError("Category labels must be integers, floats or strings to be saved.")
+        return {"bins": "categories", "labels": labels}, {}
+    arrays = {f"edges_{d}": np.asarray(e) for d, e in enumerate(bins.edges)}
+    names = []
+    for d, (m, e) in enumerate(zip(bins.measures, bins.edges, strict=True)):
+        names.append(m.name)
+        if not m.is_named:  # other measures are saved as tables
+            arrays[f"measure_z_{d}"], arrays[f"measure_u_{d}"] = m._tabulate(e)
+    if bins.is_masked:
+        arrays["mask"] = bins.mask
+    return {"bins": "grid", "ndim": bins.ndim, "measures": names}, arrays
+
+
+def _bins_from_parts(meta: dict[str, Any], arrays: Any) -> Bins:
+    if meta["bins"] == "categories":
+        labels = meta["labels"]
+        if all(isinstance(lab, int) for lab in labels) and labels == list(range(len(labels))):
+            return Bins.categories(len(labels))
+        return Bins.categories(labels)
+    nd = int(meta["ndim"])
+    measures = []
+    for d, name in enumerate(meta.get("measures", ["linear"] * nd)):
+        if f"measure_z_{d}" not in arrays:
+            measures.append(name)
+        else:
+            z, u = arrays[f"measure_z_{d}"], arrays[f"measure_u_{d}"]
+            measures.append(Measure.tabulated(z, u, name=name))
+    mask = np.asarray(arrays["mask"]) if "mask" in arrays else None
+    edges = [np.asarray(arrays[f"edges_{d}"]) for d in range(nd)]
+    return Bins.grid(*edges, measure=measures, mask=mask)
+
+
+def _basis_meta(basis: Basis) -> dict[str, Any]:
+    return {
+        "order": basis.order,
+        "resolution": list(basis.resolution),
+        "oversample": basis.oversample,
+    }
+
+
+def _selection_parts(sel: Selection) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
+    values, row, scale, _, injections = sel._parts()
+    meta: dict[str, Any] = {
+        "scale": scale,
+        "row": row if row is None or isinstance(row, str) else "array",
+    }
+    arrays = {"selection_values": np.asarray(values)}
+    if meta["row"] == "array":
+        arrays["selection_row"] = np.asarray(row)
+    if injections is not None:
+        meta["injections"] = {"prior": list(injections.prior), "smooth": injections.smooth}
+        arrays["selection_injected"] = np.asarray(injections.injected)
+        arrays["selection_recovered"] = np.asarray(injections.recovered)
+    return meta, arrays
+
+
+def _selection_from_parts(meta: dict[str, Any], arrays: Any, basis: Basis | None) -> Selection:
+    row = np.asarray(arrays["selection_row"]) if meta["row"] == "array" else meta["row"]
+    inj = meta.get("injections")
+    if inj is not None:
+        inj = (
+            np.asarray(arrays["selection_injected"]),
+            np.asarray(arrays["selection_recovered"]),
+            tuple(inj["prior"]),
+            inj["smooth"],
+        )
+    return _rebuild_selection(
+        np.asarray(arrays["selection_values"]), row, meta["scale"], basis, inj
+    )
+
+
+def result_to_dict(result: PopulationResult) -> dict[str, Any]:
+    """Arrays and metadata of a result as a plain dictionary: what :func:`save_result`
+    writes, with ``"meta"`` the JSON-compatible metadata and every other entry an array.
+
+    Engine statistics that are neither numbers, strings nor numeric arrays (e.g. lists
+    of NUTS metrics) are left out and named in ``meta["info_dropped"]``.
+
+    Parameters
+    ----------
+    result : PopulationResult
+        The result.
+    """
+    meta, arrays = _bins_to_parts(result.bins)
+    meta.update(
+        schema=SCHEMA_VERSION,
+        binference_version=__version__,
+        method=result.method,
+        kind=result.kind,
+        approximate=bool(result.approximate),
+        prior=result.prior_repr,
+        prior_key=result.prior_key,
+    )
+    scalars, dropped = {}, []
+    for k, v in result.info.items():
+        if isinstance(v, int | float | bool | str | np.integer | np.floating | np.bool_):
+            scalars[k] = v.item() if hasattr(v, "item") else v
+        elif isinstance(v, np.ndarray) and v.dtype.kind in "biuf":
+            arrays[f"info_{k}"] = v
+        elif k == "hyperparameters" and isinstance(v, dict):
+            for name, draws in v.items():
+                arrays[f"hyper_{name}"] = np.asarray(draws, dtype=float)
+        else:
+            dropped.append(k)
+    meta["info"] = scalars
+    if dropped:
+        meta["info_dropped"] = dropped
+    out: dict[str, Any] = {"draws": result._draws}
+    if result.basis is not None:  # the draws per bin follow from the component draws
+        meta["basis"] = _basis_meta(result.basis)
+        out["component_draws"] = result.component_draws
+    if result.selection is not None:  # per-object methods need it again (audit O-1)
+        meta["selection"], sel_arrays = _selection_parts(result.selection)
+        arrays.update(sel_arrays)
+    total = result.total if result.kind == "fractions" else None
+    if total is not None:  # kept so a result loaded without its likelihoods has rates
+        out["total"] = total.reshape(result._draws.shape[:2])
+    if result.log_density_unconstrained is not None:
+        out["log_density_unconstrained"] = np.asarray(result.log_density_unconstrained)
+    if result.x is not None:  # where prior_information_ratio evaluates, as before saving
+        out["x"] = np.asarray(result.x)
+    if result.labels is not None:  # draws of the objects' bins (or basis components)
+        out["label_draws"] = np.asarray(result.labels)
+        if result.label_index is not None:
+            out["label_index"] = np.asarray(result.label_index)
+    out.update(arrays)
+    out["meta"] = meta
+    return out
+
+
+def save_result(result: PopulationResult, path: str) -> None:
+    """Save a result to a compressed ``.npz`` file (versioned schema;
+    :func:`result_to_dict`).
+
+    A :meth:`Measure.transform <binference.Measure.transform>` axis is saved as
+    a table of points, so it loads as a tabulated measure: the loaded bins
+    compare unequal to the original ones.
+
+    Parameters
+    ----------
+    result : PopulationResult
+        The result.
+    path : str
+        File to write; ``.npz`` is appended when missing.
+    """
+    payload = result_to_dict(result)
+    payload["meta"] = np.array(json.dumps(payload["meta"]))
+    np.savez_compressed(path, **payload)
+
+
+def load_result(path: str) -> PopulationResult:
+    """Load a result saved with :func:`save_result` (``.npz`` may be left off ``path``).
+
+    The returned result is detached from its likelihoods; call
+    ``result.attach(likelihoods)`` to use per-object methods.
+
+    Parameters
+    ----------
+    path : str
+        File written by :func:`save_result`.
+    """
+    from .results import PopulationResult
+
+    if isinstance(path, str | os.PathLike):  # np.savez_compressed adds .npz (audit O-10)
+        name = os.fspath(path)
+        if not os.path.exists(name) and os.path.exists(name + ".npz"):
+            path = name + ".npz"
+    with np.load(path, allow_pickle=False) as data:
+        meta = json.loads(str(data["meta"]))
+        if meta.get("schema", 0) > SCHEMA_VERSION:
+            raise ValueError("This file was written by a newer version of binference.")
+        bins = _bins_from_parts(meta, data)
+        info: dict[str, Any] = dict(meta.get("info", {}))
+        hyper = {}
+        for key in data.files:
+            if key.startswith("info_"):
+                info[key[5:]] = np.array(data[key])
+            elif key.startswith("hyper_"):
+                hyper[key[6:]] = np.array(data[key])
+        if hyper:
+            info["hyperparameters"] = hyper
+        basis = Basis(bins, **meta["basis"]) if "basis" in meta else None
+        selection = (
+            _selection_from_parts(meta["selection"], data, basis) if "selection" in meta else None
+        )
+        return PopulationResult(
+            None,
+            np.array(data["component_draws" if basis is not None else "draws"]),
+            method=meta["method"],
+            kind=meta["kind"],
+            bins=bins,
+            x=np.array(data["x"]) if "x" in data.files else None,
+            log_density_unconstrained=(
+                np.array(data["log_density_unconstrained"])
+                if "log_density_unconstrained" in data.files
+                else None
+            ),
+            approximate=meta["approximate"],
+            info=info,
+            labels=np.array(data["label_draws"]) if "label_draws" in data.files else None,
+            label_index=np.array(data["label_index"]) if "label_index" in data.files else None,
+            prior_repr=meta.get("prior"),
+            prior_key=meta.get("prior_key"),
+            basis=basis,
+            total=np.array(data["total"]) if "total" in data.files else None,
+            selection=selection,
+        )
+
+
+def to_datatree(result: PopulationResult):
+    """Convert to an :class:`xarray.DataTree` with ``posterior`` and ``sample_stats`` groups.
+
+    Follows the ArviZ (>= 1.0) layout: dimensions ``chain``, ``draw`` and
+    ``bin``, with the bins' centres as numeric coordinates (``bin_center``, or
+    ``bin_center_<d>`` per axis of a grid), hyperparameter draws in the posterior,
+    and the NUTS statistics under ArviZ's names (``acceptance_rate``, ``n_steps``,
+    boolean ``diverging``). Requires xarray >= 2024.10.
+
+    Parameters
+    ----------
+    result : PopulationResult
+        The result.
+    """
+    try:
+        import xarray as xr
+    except ImportError as err:  # pragma: no cover - depends on environment
+        raise ImportError("to_datatree needs xarray: pip install 'binference[xarray]'.") from err
+    C = result.chains
+    n_chains, n_draws, _ = C.shape
+    from .results import _bin_labels
+
+    coords: dict[str, Any] = {
+        "chain": np.arange(n_chains),
+        "draw": np.arange(n_draws),
+        "bin": np.asarray(_bin_labels(result.bins), dtype=str),
+    }
+    if not result.bins.is_categorical:
+        centers = np.asarray(result.bins.centers).reshape(result.bins.K, -1)
+        for d in range(centers.shape[1]):
+            name = "bin_center" if centers.shape[1] == 1 else f"bin_center_{d}"
+            coords[name] = ("bin", centers[:, d])
+    data_vars = {"rho": (("chain", "draw", "bin"), C)}
+    rates = result.rates
+    if rates is not None:
+        data_vars["rate"] = (("chain", "draw", "bin"), rates.reshape(C.shape))
+        data_vars["total"] = (("chain", "draw"), result.total.reshape(C.shape[:2]))
+    if "selection_draws" in result.info:  # an uncertain selection, per draw
+        data_vars["selection"] = (("chain", "draw", "bin"), result.info["selection_draws"])
+    if "response_draws" in result.info:  # an inferred response, per draw
+        dims = ("chain", "draw", "measured_bin", "bin")
+        data_vars["response"] = (dims, result.info["response_draws"])
+    for name, draws in result.info.get("hyperparameters", {}).items():
+        data_vars[name] = (("chain", "draw"), np.asarray(draws, dtype=float))
+    attrs = {
+        "method": result.method,
+        "kind": result.kind,
+        "approximate": int(result.approximate),
+        "prior": str(result.prior_repr),
+        "binference_version": __version__,
+    }
+    post = xr.Dataset(data_vars, coords=coords, attrs=attrs)
+    groups = {"posterior": post}
+    arviz_names = {"accept_stat": "acceptance_rate", "n_leapfrog": "n_steps"}
+    stats = {
+        arviz_names.get(k, k): (
+            ("chain", "draw"),
+            np.asarray(v, dtype=bool if k == "diverging" else None),
+        )
+        for k, v in result.info.items()
+        if isinstance(v, np.ndarray) and v.shape == (n_chains, n_draws)
+    }
+    lp = result.log_density_unconstrained
+    if lp is not None and np.shape(lp) == (n_chains, n_draws):
+        stats["lp"] = (("chain", "draw"), np.asarray(lp))
+    if stats:
+        groups["sample_stats"] = xr.Dataset(
+            stats, coords={k: coords[k] for k in ("chain", "draw")}
+        )
+    return xr.DataTree.from_dict(groups)