"""
Posterior statistics and diagnostics.
"""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast
import arviz as az
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from numpy.typing import ArrayLike
from liesel.__version__ import __version__
from liesel.goose.engine import ErrorLog, SamplingResults
from liesel.goose.epoch import EpochType
from liesel.goose.pytree import slice_leaves, stack_leaves
from liesel.goose.types import Position, PositionInput, TransitionInfo
from liesel.option import Option
class ErrorSummaryForOneCode(NamedTuple):
error_code: int
error_msg: str
count_per_chain: np.ndarray
count_per_chain_posterior: None
ErrorSummary = dict[str, dict[int, ErrorSummaryForOneCode]]
"""
See docstring of ``_make_error_summary``.
"""
def _make_error_summary(
error_log: ErrorLog,
posterior_error_log: Option[ErrorLog],
) -> ErrorSummary:
"""
Creates an error summary from the error log.
The returned value looks like this::
{
kernel_identifier: {
error_code: (error_code, error_msg, count, count_in_posterior),
error_code: (error_code, error_msg, count, count_in_posterior),
...
},
...
}
The ``error_msg`` is the empty string if the kernel class is not supplied in the
``error_log``.
"""
error_summary = {}
for kel in error_log.values():
counter_dict: dict[int, np.ndarray] = {}
# calculate the overall counts
ec_unique = np.unique(kel.error_codes)
for ec in ec_unique:
if ec == 0:
continue
occurences_per_chain = np.sum(kel.error_codes == ec, axis=1)
counter_dict[ec] = occurences_per_chain
krnl_summary: dict[int, ErrorSummaryForOneCode] = {}
for key, count in counter_dict.items():
ec = key
error_msg = (
kel.kernel_cls.unwrap().error_book[ec]
if kel.kernel_cls.is_some()
else ""
)
krnl_summary[ec] = ErrorSummaryForOneCode(ec, error_msg, count, None)
# calculate the counts in the posterior
if posterior_error_log.is_some():
posterior_error_log_unwrapped = posterior_error_log.unwrap()
kel_post = posterior_error_log_unwrapped[kel.kernel_ident]
for ec in ec_unique:
if ec == 0:
continue
occurences_per_chain = np.sum(kel_post.error_codes == ec, axis=1)
krnl_summary[ec] = krnl_summary[ec]._replace(
count_per_chain_posterior=occurences_per_chain
)
error_summary[kel.kernel_ident] = krnl_summary
return error_summary
SummaryQuantities = Literal[
"mean",
"sd",
"var",
"quantiles",
"hdi",
"ess_bulk",
"ess_tail",
"rhat",
"mcse_mean",
"mcse_sd",
]
summary_quantities: Sequence[SummaryQuantities] = (
"mean",
"sd",
"var",
"quantiles",
"hdi",
"ess_bulk",
"ess_tail",
"rhat",
"mcse_mean",
"mcse_sd",
)
def _summarize_acceptance_probabilities(
transition_infos: Mapping[str, TransitionInfo], phase: str
) -> list[dict[str, Any]]:
data = []
for k, tinfo in transition_infos.items():
ap = jnp.asarray(tinfo.acceptance_prob)
pm = jnp.asarray(tinfo.position_moved)
chains, _ = ap.shape
for c in range(chains):
chain = {"kernel": k, "phase": phase, "chain": c}
chain["acceptance_probability"] = float(ap[c, ...].mean())
position_moved = jnp.mean(pm[c, ...])
chain["position_moved"] = position_moved
if float(position_moved) > 1.0 and int(position_moved.round()) == 99:
chain["position_moved"] = jnp.nan
data.append(chain)
return data
def summarize_acceptance_probabilities(
results: SamplingResults,
) -> list[dict[str, Any]]:
warmup = _summarize_acceptance_probabilities(
results.get_warmup_transition_infos(), "warmup"
)
posterior = _summarize_acceptance_probabilities(
results.get_posterior_transition_infos(), "posterior"
)
return warmup + posterior
[docs]
class Summary:
"""
Posterior summary and diagnostics for :class:`~liesel.goose.SamplingResults`.
Offers two main use cases:
1. View an overall summary by printing a summary instance, including a summary table
of the posterior samples and a summary of sampling errors.
2. Programmatically access summary statistics via
``quantities[quantity_name][var_name]``. Please refer to the documentation of the
attribute :attr:`~liesel.goose.Summary.quantities` for details.
Additionally, the summary object can be turned into a :class:`~pandas.DataFrame`
using :meth:`~liesel.goose.Summary.to_dataframe`.
If ``per_chain=False``, statistics are computed over all posterior chains and
draws. If ``per_chain=True``, each chain is summarized separately.
The low-level computations for HDIs, effective sample sizes, R-hat, and Monte
Carlo standard errors are delegated to `ArviZ <https://python.arviz.org/>`_.
By default, the summary contains the following statistics:
- ``mean``: Posterior mean.
- ``sd``: Posterior standard deviation.
- ``var``: Posterior variance.
- ``quantiles``: Posterior quantiles at the probabilities given by
``quantiles``. These are stored as ``"quantile"`` in
:attr:`~liesel.goose.Summary.quantities`
and become columns named ``q_<probability>`` in
:meth:`~liesel.goose.Summary.to_dataframe`.
- ``hdi``: Highest density interval with probability mass ``hdi_prob``. This is
the narrowest posterior interval reported by ArviZ at that probability level.
In :meth:`~liesel.goose.Summary.to_dataframe`, it becomes ``hdi_low`` and
``hdi_high``.
- ``ess_bulk``: Bulk effective sample size, a diagnostic for Monte Carlo
precision in the central part of the posterior distribution.
- ``ess_tail``: Tail effective sample size, a diagnostic for Monte Carlo
precision in the posterior tails.
- ``rhat``: Rank-normalized split R-hat, a between-chain convergence diagnostic.
Values close to 1 indicate better agreement between chains. This statistic is
only computed when more than one chain is summarized together.
- ``mcse_mean``: Monte Carlo standard error of the posterior mean.
- ``mcse_sd``: Monte Carlo standard error of the posterior standard deviation.
Use ``which`` to compute only a subset of these statistics.
Parameters
----------
results
The sampling results to summarize.
additional_chain
Can be supplied to add more parameters to the summary output. Must be a position
chain which matches chain and time dimension of the posterior chain as returned
by :meth:`~liesel.goose.SamplingResults.get_posterior_samples`.
quantiles
Posterior quantile probabilities to compute when ``"quantiles"`` is included
in ``which``.
hdi_prob
Posterior probability mass of the highest density interval to compute when
``"hdi"`` is included in ``which``.
selected, deselected
Allow to get a summary only for a subset of the position keys.
per_chain
If *True*, the summary is calculated on a per-chain basis. Certain measures like
``rhat`` are not available if ``per_chain`` is *True*.
which
Names of the summary statistics to compute. Supported values are ``"mean"``,
``"sd"``, ``"var"``, ``"quantiles"``, ``"hdi"``, ``"ess_bulk"``,
``"ess_tail"``, ``"rhat"``, ``"mcse_mean"``, and ``"mcse_sd"``.
Notes
-----
This class is still considered experimental. The API may still undergo larger
changes.
"""
per_chain: bool
"""
Whether results are summarized for individual chains (*True*), or aggregated
over chains (*False*).
"""
quantities: dict[str, dict[str, np.ndarray]]
"""
Dict of summarizing quantities.
Let ``summary`` be a :class:`~liesel.goose.Summary` instance. The hierarchy is::
q = summary.quantities["quantity_name"]["parameter_name"]
Available quantity names are ``"mean"``, ``"sd"``, ``"var"``, ``"quantile"``,
``"hdi"``, ``"ess_bulk"``, ``"ess_tail"``, ``"rhat"``, ``"mcse_mean"``, and
``"mcse_sd"``, depending on the ``which`` argument. Note that ``which`` uses
``"quantiles"`` to request quantiles, while :attr:`~liesel.goose.Summary.quantities`
stores the
result under ``"quantile"``.
The extracted object is an ``np.ndarray``. If ``per_chain=True``, the arrays for
the ``"quantile"`` and ``"hdi"`` quantities have the following dimensions:
1. First index refers to the chain
2. Second index refers to the quantile/interval
3. Third and subsequent indices refer to individual parameters.
If ``per_chain=True``, the arrays for the other quantities have the dimensions:
1. First index refers to the chain
2. Second and subsequent indices refer to individual parameters.
If ``per_chain=False``, the first index is removed for all quantities.
"""
config: dict
"""
A dictionary of config settings for this summary object. Should NOT be changed
after initialization; such changes have no effect on the computed summary values.
"""
sample_info: dict
"""
Dictionary of meta-information about the mcmc samples used to create this summary
object.
Contains ``num_chains``, ``sample_size_per_chain``, and ``warmup_size_per_chain``.
"""
error_summary: ErrorSummary
"""
Contains error information for each kernel.
"""
kernels_by_pos_key: dict[str, str]
"""
A dict, linking parameter names (the keys) to the kernel identifier (the values).
The identifier refers to the kernel that was used to sample the respective
parameter.
"""
liesel_version: str
"""
The specific version of Liesel used to produce the results.
"""
if TYPE_CHECKING:
pos_keys_by_kernels_df: pd.DataFrame
"""Table associating position keys with sampling kernels."""
def __init__(
self,
results: SamplingResults,
additional_chain: PositionInput | None = None,
quantiles: Sequence[float] = (0.05, 0.5, 0.95),
hdi_prob: float = 0.9,
selected: list[str] | None = None,
deselected: list[str] | None = None,
per_chain: bool = False,
which: Sequence[SummaryQuantities] = summary_quantities,
):
if not which:
raise ValueError(
f"Argument 'which' must not be empty. "
f"Supported keys are: {summary_quantities}"
)
for _which_key in which:
if _which_key not in summary_quantities:
raise ValueError(
f"Key {which} in 'which' is not supported. "
f"Supported keys are: {summary_quantities}"
)
posterior_chain = results.get_posterior_samples()
if additional_chain:
for k, v in additional_chain.items():
posterior_chain[k] = v
if selected:
posterior_chain = Position(
{
key: value
for key, value in posterior_chain.items()
if key in selected
}
)
if deselected is not None:
for key in deselected:
del posterior_chain[key]
# get some general infos on the sampling
param_chain = next(iter(posterior_chain.values()))
epochs = results.positions.get_epochs()
warmup_size = np.sum(
[
int(epoch.duration / epoch.thinning)
for epoch in epochs
if epoch.type.is_warmup(epoch.type)
]
)
thinning_warmup = np.unique(
[epoch.thinning for epoch in epochs if epoch.type.is_warmup(epoch.type)]
)
thinning_posterior = np.unique(
[epoch.thinning for epoch in epochs if epoch.type is EpochType.POSTERIOR]
)
sample_info = {
"num_chains": param_chain.shape[0],
"sample_size_per_chain": param_chain.shape[1],
"warmup_size_per_chain": warmup_size,
"thinning_warmup": thinning_warmup.squeeze(),
"thinning_posterior": thinning_posterior.squeeze(),
}
# convert everything to numpy array
for key in posterior_chain:
posterior_chain[key] = np.asarray(posterior_chain[key])
# calculate quantiles either per chain and merge the results or all at once
single_chain_summaries = []
if per_chain:
for chain_idx in range(sample_info["num_chains"]):
single_chain = slice_leaves(
posterior_chain, jnp.s_[None, chain_idx, ...]
)
qdict = _create_quantity_dict(single_chain, quantiles, hdi_prob, which)
single_chain_summaries.append(qdict)
quantities = stack_leaves(single_chain_summaries, axis=0)
else:
quantities = _create_quantity_dict(
posterior_chain, quantiles, hdi_prob, which
)
config = {
"quantiles": quantiles,
"hdi_prob": hdi_prob,
"chains_merged": not per_chain,
}
error_summary = _make_error_summary(
results.get_error_log(False).unwrap(), results.get_error_log(True)
)
pos_keys_by_kernels = []
for k, v in results.get_pos_keys_by_kernels().items():
pos_keys_by_kernels.append({"kernel": k, "positions": ", ".join(v)})
if len(pos_keys_by_kernels) == 1:
posdf = pd.DataFrame(pos_keys_by_kernels, index=pd.Index([0]))
else:
posdf = pd.DataFrame(pos_keys_by_kernels)
posdf = posdf.set_index(["kernel"])
self._which = which
self.per_chain = per_chain
self.quantities = quantities
self.config = config
self.sample_info = sample_info
self.error_summary = error_summary
self.pos_keys_by_kernels_df = posdf
self._acceptance_prob_summary = summarize_acceptance_probabilities(results)
self.kernels_by_pos_key = results.get_kernels_by_pos_key()
self.liesel_version = __version__
[docs]
def to_dataframe(self) -> pd.DataFrame:
"""Turns Summary object into a :class:`~pandas.DataFrame` object."""
# don't change the original data
quants = self.quantities.copy()
# make new entries for the quantiles
if self.per_chain:
if "quantiles" in self._which:
for i, q in enumerate(self.config["quantiles"]):
quants[f"q_{q}"] = {
k: v[:, i, ...] for k, v in quants["quantile"].items()
}
if "hdi" in self._which:
quants["hdi_low"] = {k: v[:, 0, ...] for k, v in quants["hdi"].items()}
quants["hdi_high"] = {k: v[:, 1, ...] for k, v in quants["hdi"].items()}
else:
if "quantiles" in self._which:
for i, q in enumerate(self.config["quantiles"]):
quants[f"q_{q}"] = {
k: v[i, ...] for k, v in quants["quantile"].items()
}
if "hdi" in self._which:
quants["hdi_low"] = {k: v[0, ...] for k, v in quants["hdi"].items()}
quants["hdi_high"] = {k: v[1, ...] for k, v in quants["hdi"].items()}
# remove the old entries
if "hdi" in self._which:
del quants["hdi"]
if "quantiles" in self._which:
del quants["quantile"]
# create one row per entry
df_dict = {}
first_quant = next(iter(quants.values()))
for var in first_quant:
it = np.nditer(first_quant[var], flags=["multi_index"])
for _ in it:
var_fqn = (
var if len(it.multi_index) == 0 else f"{var}{list(it.multi_index)}"
)
quant_per_elem: dict[str, Any] = {}
quant_per_elem["variable"] = var
quant_per_elem["kernel"] = self.kernels_by_pos_key.get(var, "-")
if self.config["chains_merged"]:
quant_per_elem["var_index"] = it.multi_index
quant_per_elem["sample_size"] = (
self.sample_info["sample_size_per_chain"]
* self.sample_info["num_chains"]
)
else:
quant_per_elem["chain_index"] = it.multi_index[0]
quant_per_elem["var_index"] = it.multi_index[1:]
quant_per_elem["sample_size"] = self.sample_info[
"sample_size_per_chain"
]
for quant_name, quant_dict in quants.items():
quant_per_elem[quant_name] = quant_dict[var][it.multi_index]
# convert jax.Arrays (scalar) to floats so that pandas treats them
# correctly
for key, val in quant_per_elem.items():
if isinstance(val, jax.Array):
# value should be a scalar
assert val.shape == ()
# replace dict element with value casted to float32
quant_per_elem[key] = float(val)
df_dict[var_fqn] = quant_per_elem
# convert to dataframe and use varname as index
df = pd.DataFrame.from_dict(df_dict, orient="index")
df = df.reset_index()
df = df.rename(columns={"index": "var_fqn"})
df = df.set_index("variable")
return df
[docs]
def aggregate_diagnostics(
self,
by: Literal[
"min/max", "mean", "median", "std", "var", "min", "max"
] = "min/max",
) -> pd.DataFrame:
"""
Aggregates effective sample sizes (ESS) and rhat.
Parameters
----------
by
How to aggregate. The three current options are:
- ``"min/max"``: Aggregate ESS by taking the minimum ESS per parameter
block, and the rhat by taking the maximum rhat per parameter block. This
corresponds to a worst-case summary.
- ``"mean"``: Aggregate ESS and rhat by averaging inside parameter blocks.
- ``"median"``: Aggregate ESS and rhat by taking the median inside parameter
blocks.
- ``"std"``: Aggregate ESS and rhat by taking the standard deviation inside
parameter blocks.
- ``"var"``: Aggregate ESS and rhat by taking the variance inside parameter
blocks.
- ``"min"``: Aggregate ESS and rhat by taking the minimum inside parameter
blocks.
- ``"max"``: Aggregate ESS and rhat by taking the maximum inside parameter
blocks.
Notes
------
If :attr:`~liesel.goose.Summary.per_chain` is ``True``, rhat cannot be computed
and there is not
present in the output dataframe.
"""
df = self.to_dataframe()
df.index.name = "parameter"
if by == "min/max":
by_ess = "min"
by_rhat = "max"
else:
by_ess = by
by_rhat = by
if self.per_chain:
diagnostics = (
df.loc[:, ["chain_index", "ess_bulk", "ess_tail"]]
.groupby(["parameter", "chain_index"])
.agg(
ess_bulk=("ess_bulk", by_ess),
ess_tail=("ess_tail", by_ess),
)
)
else:
diagnostics = (
df.loc[:, ["rhat", "ess_bulk", "ess_tail"]]
.groupby("parameter")
.agg(
ess_bulk=("ess_bulk", by_ess),
ess_tail=("ess_tail", by_ess),
rhat=("rhat", by_rhat),
)
)
if by == "min/max":
diagnostics["aggregated_by"] = "min (ess); max (rhat)"
else:
diagnostics["aggregated_by"] = by
return diagnostics
def _param_df(self):
df = self.to_dataframe()
df.index.name = "parameter"
df = df.rename(columns={"var_index": "index"})
df = df.set_index("index", append=True)
qtls = [f"q_{qtl}" for qtl in self.config["quantiles"]]
cols = (
["kernel", "mean", "sd"]
+ qtls
+ ["sample_size", "ess_bulk", "ess_tail", "rhat"]
)
cols = [col for col in cols if col in df.columns]
df = df[cols]
return df
[docs]
def acceptance_prob_df(self) -> pd.DataFrame:
"""Returns an overview of acceptance probabilities as a dataframe."""
apdf = pd.DataFrame(self._acceptance_prob_summary)
apdf = apdf.set_index(["kernel", "phase", "chain"])
apdf = apdf.join(self.pos_keys_by_kernels_df, on="kernel")
apdf = apdf.reset_index().set_index(["kernel", "positions", "phase", "chain"])
if not self.per_chain:
apdf = apdf.groupby(level=["kernel", "positions", "phase"]).agg(
{"acceptance_probability": "mean", "position_moved": "mean"}
)
return apdf
[docs]
def error_df(self, per_chain: bool = False) -> pd.DataFrame:
"""
Returns an overview of the errors recorded during sampling as a dataframe.
"""
return self._error_df(per_chain=per_chain)
def _error_df(self, per_chain: bool = False) -> pd.DataFrame:
# fmt: off
error_summaries = {k: v for k, v in self.error_summary.items() if v}
if error_summaries:
df = pd.concat({
kernel: pd.DataFrame.from_dict(code_summary, orient="index")
for kernel, code_summary in error_summaries.items()
})
else:
return pd.DataFrame()
# fmt: on
df = df.reset_index(level=1, drop=True)
df["error_code"] = df["error_code"].astype(int)
df = df.set_index(["error_code", "error_msg"], append=True)
df.index.names = ["kernel", "error_code", "error_msg"]
# fmt: off
df = df.rename(columns={
"count_per_chain": "total",
"count_per_chain_posterior": "posterior",
})
# fmt: on
df = df.explode(["total", "posterior"])
df["warmup"] = df["total"] - df["posterior"]
df = df.drop(columns="total")
df = df.melt(
value_vars=["warmup", "posterior"],
var_name="phase",
value_name="count",
ignore_index=False,
)
df["phase"] = pd.Categorical(df["phase"], categories=["warmup", "posterior"])
df = df.set_index("phase", append=True)
df["chain"] = df.groupby(level=[0, 1, 2, 3], observed=True).cumcount()
df = df.set_index("chain", append=True)
df = df.sort_index()
df["sample_size"] = None
warmup_size = self.sample_info["warmup_size_per_chain"]
posterior_size = self.sample_info["sample_size_per_chain"]
df.loc[pd.IndexSlice[:, :, :, "warmup"], "sample_size"] = warmup_size
df.loc[pd.IndexSlice[:, :, :, "posterior"], "sample_size"] = posterior_size
df["thinning"] = None
warmup_thinning = self.sample_info["thinning_warmup"]
posterior_thinning = self.sample_info["thinning_posterior"]
df.loc[pd.IndexSlice[:, :, :, "warmup"], "thinning"] = warmup_thinning
df.loc[pd.IndexSlice[:, :, :, "posterior"], "thinning"] = posterior_thinning
df["sample_size_total"] = df["sample_size"] * df["thinning"]
df["relative"] = df["count"] / df["sample_size_total"]
df = (
df.join(self.pos_keys_by_kernels_df, on="kernel")
.reset_index()
.set_index(["kernel", "positions", "error_code", "error_msg", "phase"])
)
# df = df.drop(columns="sample_size")
if not per_chain:
df = df.groupby(level=[0, 1, 2, 3, 4], observed=True)
df = df.aggregate(
{
"count": "sum",
"relative": "mean",
"sample_size": "sum",
"sample_size_total": "sum",
}
)
df = df.sort_index()
# re-order columns
cols = ["count", "sample_size", "sample_size_total", "relative"]
return df[cols]
def __repr__(self):
param_df = self._param_df()
error_df = self._error_df()
txt = "Parameter summary:\n\n" + repr(param_df)
if not error_df.empty:
txt += "\n\nError summary:\n\n" + repr(error_df)
return txt
def _repr_html_(self):
param_df = self._param_df()
apdf = self.acceptance_prob_df()
error_df = self._error_df()
html = "\n<p><strong>Parameter summary:</strong></p>\n" + param_df.to_html()
html += "\n<p><strong>Acceptance probabilities:</strong></p>\n" + apdf.to_html()
if not error_df.empty:
html += "\n<p><strong>Error summary:</strong></p>\n" + error_df.to_html()
html += "\n"
return html
def _repr_markdown_(self):
param_df = self._param_df()
apdf = self.acceptance_prob_df()
error_df = self._error_df()
try:
param_md = param_df.to_markdown()
apdf_md = apdf.to_markdown()
error_md = error_df.to_markdown()
except ImportError:
param_md = f"```\n{param_df!r}\n```"
apdf_md = f"```\n{apdf!r}\n```"
error_md = f"```\n{error_df!r}\n```"
md = "\n\n**Parameter summary:**\n\n" + param_md
md += "\n\n**Acceptance probabilities:**\n\n" + apdf_md
if not error_df.empty:
md += "\n\n**Error summary:**\n\n" + error_md
md += "\n\n"
return md
def __str__(self):
return str(self.to_dataframe())
def _create_quantity_dict(
chain: Position,
quantiles: Sequence[float],
hdi_prob: float,
which: Sequence[SummaryQuantities] = summary_quantities,
) -> dict[str, dict[str, np.ndarray]]:
azchain = az.from_dict({"posterior": chain})["posterior"].dataset
quantities = {}
# calculate quantities
if "mean" in which:
quantities["mean"] = azchain.mean(dim=["chain", "draw"])
if "var" in which:
quantities["var"] = azchain.var(dim=["chain", "draw"])
if "sd" in which:
quantities["sd"] = azchain.std(dim=["chain", "draw"])
if "quantiles" in which:
quantities["quantile"] = azchain.quantile(q=quantiles, dim=["chain", "draw"])
if "hdi" in which:
quantities["hdi"] = az.hdi(azchain, prob=hdi_prob)
if "ess_bulk" in which:
quantities["ess_bulk"] = az.ess(azchain, method="bulk")
if "ess_tail" in which:
quantities["ess_tail"] = az.ess(azchain, method="tail")
if "mcse_mean" in which:
quantities["mcse_mean"] = az.mcse(azchain, method="mean")
if "mcse_sd" in which:
quantities["mcse_sd"] = az.mcse(azchain, method="sd")
if "rhat" in which and azchain.sizes["chain"] > 1:
quantities["rhat"] = az.rhat(azchain)
# convert to simple dict[str, np.ndarray]
for key, val in quantities.items():
quantity = {}
for k, v in val.data_vars.items():
if key == "hdi":
remaining_dims = [dim for dim in v.dims if dim != "ci_bound"]
v = v.transpose("ci_bound", *remaining_dims)
quantity[k] = v.values
quantities[key] = quantity
return quantities
[docs]
class SamplesSummary:
"""
Posterior summary and diagnostics for a dictionary of sample arrays.
Accepts read-only mappings and does not modify the supplied mapping or its arrays.
See :class:`~liesel.goose.Summary` for the full description of the computed
statistics, their
interpretation, the ``quantities`` layout, and the behavior of ``quantiles``,
``hdi_prob``, ``per_chain``, and ``which``. This class computes the same
sample-based statistics as :class:`~liesel.goose.Summary`, but takes a plain
dictionary of
sample arrays instead of a :class:`~liesel.goose.SamplingResults` object and does
not include
sampling-error or acceptance-probability diagnostics.
Offers two main use cases:
1. The summary object can be turned into a :class:`~pandas.DataFrame`
using :meth:`~liesel.goose.SamplesSummary.to_dataframe`.
2. Programmatically access summary statistics via
``quantities[quantity_name][var_name]``. Please refer to the documentation of the
attribute :attr:`~liesel.goose.SamplesSummary.quantities` for details.
Parameters
----------
samples
The dictionary of samples to summarize. Each array is expected to have leading
dimensions ``(nchains, ndraws, ...)``.
quantiles
Posterior quantile probabilities to compute when ``"quantiles"`` is included
in ``which``.
hdi_prob
Posterior probability mass of the highest density interval to compute when
``"hdi"`` is included in ``which``.
selected, deselected
Allow to get a summary only for a subset of the position keys.
per_chain
If *True*, the summary is calculated on a per-chain basis. Certain measures like
``rhat`` are not available if ``per_chain`` is *True*.
which
Names of the summary statistics to compute. Supported values are the same as
for :class:`~liesel.goose.Summary`.
Notes
-----
This class is still considered experimental. The API may still undergo larger
changes.
"""
config: dict
"""Configuration used to compute and format the summaries."""
if TYPE_CHECKING:
per_chain: bool
"""Whether summaries are computed separately for each chain."""
quantities: dict
"""Computed summary quantities for the selected samples."""
sample_info: dict
"""Metadata describing the chains and posterior draws."""
def __init__(
self,
samples: Mapping[str, ArrayLike],
quantiles: Sequence[float] = (0.05, 0.5, 0.95),
hdi_prob: float = 0.9,
selected: list[str] | None = None,
deselected: list[str] | None = None,
per_chain: bool = False,
which: Sequence[SummaryQuantities] = summary_quantities,
):
if not which:
raise ValueError(
f"Argument 'which' must not be empty. "
f"Supported keys are: {summary_quantities}"
)
for _which_key in which:
if _which_key not in summary_quantities:
raise ValueError(
f"Key {which} in 'which' is not supported. "
f"Supported keys are: {summary_quantities}"
)
posterior_chain = Position(dict(samples))
if selected:
posterior_chain = Position(
{
key: value
for key, value in posterior_chain.items()
if key in selected
}
)
if deselected is not None:
for key in deselected:
del posterior_chain[key]
# get some general infos on the sampling
param_chain = np.asarray(next(iter(posterior_chain.values())))
sample_info = {
"num_chains": param_chain.shape[0],
"sample_size_per_chain": param_chain.shape[1],
}
# convert everything to numpy array
for key in posterior_chain:
posterior_chain[key] = np.asarray(posterior_chain[key])
# calculate quantiles either per chain and merge the results or all at once
single_chain_summaries = []
if per_chain:
for chain_idx in range(sample_info["num_chains"]):
single_chain = slice_leaves(
posterior_chain, jnp.s_[None, chain_idx, ...]
)
qdict = _create_quantity_dict(single_chain, quantiles, hdi_prob, which)
single_chain_summaries.append(qdict)
quantities = stack_leaves(single_chain_summaries, axis=0)
else:
quantities = _create_quantity_dict(
posterior_chain, quantiles, hdi_prob, which
)
config = {
"quantiles": quantiles,
"hdi_prob": hdi_prob,
"chains_merged": not per_chain,
}
self._which = which
self.per_chain = per_chain
self.quantities = quantities
self.config = config
self.sample_info = sample_info
[docs]
@classmethod
def from_array(
cls,
a: ArrayLike,
quantiles: Sequence[float] = (0.05, 0.5, 0.95),
hdi_prob: float = 0.9,
selected: list[str] | None = None,
deselected: list[str] | None = None,
per_chain: bool = False,
name: str = "v",
which: Sequence[SummaryQuantities] = summary_quantities,
) -> SamplesSummary:
"""
Initializes the summary from an array of samples.
Parameters
----------
a
The array of samples to summarize. Expected to have leading dimensions
``(nchains, ndraws, ...)``.
quantiles
Posterior quantile probabilities to compute when ``"quantiles"`` is
included in ``which``.
hdi_prob
Posterior probability mass of the highest density interval to compute when
``"hdi"`` is included in ``which``.
selected, deselected
Allow to get a summary only for a subset of the position keys.
per_chain
If *True*, the summary is calculated on a per-chain basis. Certain \
measures like ``rhat`` are not available if ``per_chain`` is *True*.
name
Variable name to use for labelling in
:meth:`~liesel.goose.SamplesSummary.to_dataframe`.
which
Names of the summary statistics to compute. Supported values are the same
as for :class:`~liesel.goose.Summary`.
"""
samples = {name: a}
return cls(samples, quantiles, hdi_prob, selected, deselected, per_chain, which)
[docs]
def to_dataframe(self) -> pd.DataFrame:
"""Turns SamplesSummary object into a :class:`~pandas.DataFrame` object."""
# don't change the original data
quants = self.quantities.copy()
# make new entries for the quantiles
if self.per_chain:
if "quantiles" in self._which:
for i, q in enumerate(self.config["quantiles"]):
quants[f"q_{q}"] = {
k: v[:, i, ...] for k, v in quants["quantile"].items()
}
if "hdi" in self._which:
quants["hdi_low"] = {k: v[:, 0, ...] for k, v in quants["hdi"].items()}
quants["hdi_high"] = {k: v[:, 1, ...] for k, v in quants["hdi"].items()}
else:
if "quantiles" in self._which:
for i, q in enumerate(self.config["quantiles"]):
quants[f"q_{q}"] = {
k: v[i, ...] for k, v in quants["quantile"].items()
}
if "hdi" in self._which:
quants["hdi_low"] = {k: v[0, ...] for k, v in quants["hdi"].items()}
quants["hdi_high"] = {k: v[1, ...] for k, v in quants["hdi"].items()}
# remove the old entries
if "hdi" in self._which:
del quants["hdi"]
if "quantiles" in self._which:
del quants["quantile"]
# create one row per entry
df_dict = {}
first_quant = next(iter(quants.values()))
for var in first_quant:
it = np.nditer(first_quant[var], flags=["multi_index"])
for _ in it:
var_fqn = (
var if len(it.multi_index) == 0 else f"{var}{list(it.multi_index)}"
)
quant_per_elem: dict[str, Any] = {}
quant_per_elem["variable"] = var
if self.config["chains_merged"]:
quant_per_elem["var_index"] = it.multi_index
quant_per_elem["sample_size"] = (
self.sample_info["sample_size_per_chain"]
* self.sample_info["num_chains"]
)
else:
quant_per_elem["chain_index"] = it.multi_index[0]
quant_per_elem["var_index"] = it.multi_index[1:]
quant_per_elem["sample_size"] = self.sample_info[
"sample_size_per_chain"
]
for quant_name, quant_dict in quants.items():
quant_per_elem[quant_name] = quant_dict[var][it.multi_index]
# convert jax.Arrays (scalar) to floats so that pandas treats them
# correctly
for key, val in quant_per_elem.items():
if isinstance(val, jax.Array):
# value should be a scalar
assert val.shape == ()
# replace dict element with value casted to float32
quant_per_elem[key] = float(val)
df_dict[var_fqn] = quant_per_elem
# convert to dataframe and use varname as index
df = pd.DataFrame.from_dict(df_dict, orient="index")
df = df.reset_index()
df = df.rename(columns={"index": "var_fqn"})
df = df.set_index("variable")
return df
def _param_df(self):
df = self.to_dataframe()
df.index.name = "parameter"
df = df.rename(columns={"var_index": "index"})
df = df.set_index("index", append=True)
qtls = [f"q_{qtl}" for qtl in self.config["quantiles"]]
cols = (
["kernel", "mean", "sd"]
+ qtls
+ ["sample_size", "ess_bulk", "ess_tail", "rhat"]
)
cols = [col for col in cols if col in df.columns]
df = df[cols]
return df
[docs]
def aggregate_diagnostics(
self,
by: Literal[
"min/max", "mean", "median", "std", "var", "min", "max"
] = "min/max",
) -> pd.DataFrame:
"""
Aggregates effective sample sizes (ESS) and rhat.
Parameters
----------
by
How to aggregate. The three current options are:
- ``"min/max"``: Aggregate ESS by taking the minimum ESS per parameter
block, and the rhat by taking the maximum rhat per parameter block. This
corresponds to a worst-case summary.
- ``"mean"``: Aggregate ESS and rhat by averaging inside parameter blocks.
- ``"median"``: Aggregate ESS and rhat by taking the median inside parameter
blocks.
- ``"std"``: Aggregate ESS and rhat by taking the standard deviation inside
parameter blocks.
- ``"var"``: Aggregate ESS and rhat by taking the variance inside parameter
blocks.
- ``"min"``: Aggregate ESS and rhat by taking the minimum inside parameter
blocks.
- ``"max"``: Aggregate ESS and rhat by taking the maximum inside parameter
blocks.
Notes
------
If :attr:`~liesel.goose.SamplesSummary.per_chain` is ``True``, rhat cannot be
computed and there is not
present in the output dataframe.
"""
df = self.to_dataframe()
df.index.name = "parameter"
if by == "min/max":
by_ess = "min"
by_rhat = "max"
else:
by_ess = by
by_rhat = by
if self.per_chain:
diagnostics = (
df.loc[:, ["chain_index", "ess_bulk", "ess_tail"]]
.groupby(["parameter", "chain_index"])
.agg(
ess_bulk=("ess_bulk", by_ess),
ess_tail=("ess_tail", by_ess),
)
)
else:
diagnostics = (
df.loc[:, ["rhat", "ess_bulk", "ess_tail"]]
.groupby("parameter")
.agg(
ess_bulk=("ess_bulk", by_ess),
ess_tail=("ess_tail", by_ess),
rhat=("rhat", by_rhat),
)
)
if by == "min/max":
diagnostics["aggregated_by"] = "min (ess); max (rhat)"
else:
diagnostics["aggregated_by"] = by
return diagnostics
def concatenate_arrays_in_dict(
x: Mapping[str, jax.typing.ArrayLike], n_leading_axes: int = 2
) -> jax.Array:
"""
Concatenates all arrays in the supplied dictionary into a single array.
Returns
-------
An array of dimension ``(leading1, leading1, nobs)``, where ``nobs`` is the total
number of elements in the non-leading axes of the dictionary values. As the default
case, we expect the output to have shape ``(nchains, nsamples, nobs)``.
"""
flat_arrays = []
for v in x.values():
# assumed to have shape (s, c, ...)
v = jnp.atleast_3d(jnp.asarray(v))
vshape = v.shape[:n_leading_axes] + (-1,)
flat_arrays.append(jnp.reshape(v, vshape))
out_array = jnp.concatenate(flat_arrays, axis=-1)
return out_array
class _LieselELPDData(az.ELPDData):
"""An ArviZ LOO result with Liesel's legacy attribute aliases."""
elpd_loo: float
p_loo: float
def _apply_loo_scale(
result: az.ELPDData,
scale: Literal["log", "negative_log", "deviance"],
) -> _LieselELPDData:
if scale != "log":
multiplier = -1 if scale == "negative_log" else -2
result.elpd = multiplier * result.elpd
result.se = abs(multiplier) * result.se
result.elpd_i = multiplier * result.elpd_i
result.scale = scale
liesel_result = cast(_LieselELPDData, result)
liesel_result.elpd_loo = result.elpd
liesel_result.p_loo = result.p
return liesel_result
[docs]
def loo(
lpp: Mapping[str, jax.typing.ArrayLike] | jax.typing.ArrayLike,
samples: Mapping[str, jax.typing.ArrayLike] | None = None,
reff: float | None = None,
scale: Literal["log", "negative_log", "deviance"] = "log",
) -> _LieselELPDData:
"""
Compute Pareto-smoothed importance sampling leave-one-out cross-validation
(PSIS-LOO-CV) statistic via ArviZ.
Parameters
----------
lpp
Dictionary or array of pointwise log probability evaluations.
If passed as a dictionary, each value is expected to have shape
``(nchains, ndraws, ...)``.
If passed as an array, it is assumed to have shape ``(nchains, ndraws, n)``.
samples
Retained for backwards compatibility and ignored.
reff
Relative MCMC efficiency, ess / n i.e. number of effective samples divided
by the number of actual samples. Computed from the likelihood values by
default. For independent draws, pass ``1.0``.
scale
Output scale. The options are:
- ``log``: (default) log probability scale.
- ``negative_log``: ``-1 * log``
- ``deviance``: ``-2 * log``
A higher log probability (or a lower deviance or negative log_score)
indicates a model with better predictive accuracy.
References
----------
- Computations are carried out via ArviZ: https://python.arviz.org/en/stable/
- Theoretical background: Vehtari, A., Gelman, A., & Gabry, J. (2017). Practical
Bayesian model evaluation using leave-one-out cross-validation and WAIC.
Statistics and Computing, 27(5), 1413–1432.
https://doi.org/10.1007/s11222-016-9696-4
"""
if isinstance(lpp, Mapping):
lpp_by_variable = cast(Mapping[str, jax.typing.ArrayLike], lpp)
lpp_array = concatenate_arrays_in_dict(lpp_by_variable)
else:
lpp_array = jnp.asarray(lpp)
lpp_array = np.asarray(lpp_array)
idat = az.from_dict({"log_likelihood": {"observed": lpp_array}})
if reff is None:
nsamples = lpp_array.shape[0] * lpp_array.shape[1]
shift = lpp_array.max(axis=(0, 1), keepdims=True)
likelihood = np.exp(lpp_array - shift)
reff = float(np.mean(az.ess(likelihood, method="mean") / nsamples))
# now we assume reff is not None
return _apply_loo_scale(az.loo(idat, reff=reff), scale)
