"""Compute exact Leave-One-Out cross validation refitting for problematic observations."""
from copy import deepcopy
import numpy as np
import xarray as xr
from arviz_base import rcParams
from xarray_einstats.stats import logsumexp
from arviz_stats.loo.loo import loo
from arviz_stats.loo.loo_helper import _prepare_loo_inputs
from arviz_stats.loo.wrapper import SamplingWrapper
from arviz_stats.utils import ELPDData, ELPDDataLOO
__all__ = ["reloo"]
[docs]
def reloo(
wrapper,
loo_orig=None,
var_name=None,
log_weights=None,
k_threshold=None,
pointwise=None,
):
r"""Recalculate exact Leave-One-Out cross validation refitting where the approximation fails.
:func:`arviz_stats.loo` estimates the values of Leave-One-Out (LOO) cross validation using
Pareto Smoothed Importance Sampling (PSIS) to approximate its value. PSIS works well when
the posterior and the posterior_i (excluding observation i from the data used to fit)
are similar. In some cases, there are highly influential observations for which PSIS
cannot approximate the LOO-CV, and a warning of a large Pareto shape is sent by ArviZ.
These cases typically have a handful of bad or very bad Pareto shapes, and a majority of
good or ok shapes.
Therefore, this may not indicate that the model is not robust enough nor that these
observations are inherently bad, only that PSIS cannot approximate LOO-CV correctly.
Thus, we can use PSIS for all observations where the Pareto shape is below a threshold
and refit the model to perform exact cross validation for the handful of observations
where PSIS cannot be used. This approach allows us to properly approximate LOO-CV with
only a handful of refits, which in most cases is still much less computationally expensive
than exact LOO-CV, which needs one refit per observation.
Parameters
----------
wrapper : SamplingWrapper
An instance of a SamplingWrapper subclass that implements the necessary
methods for model refitting. This wrapper allows ``reloo`` to work with
any modeling framework.
loo_orig : ELPDDataLOO, optional
Existing LOO results with pointwise data. If None, will compute
PSIS-LOO-CV first using the data from ``wrapper``.
var_name : str, optional
The name of the variable in log_likelihood groups storing the pointwise log
likelihood data to use for loo computation. Defaults to None, which will
use all log likelihood data.
log_weights : DataArray or ELPDDataLOO, optional
Smoothed log weights. Can be either:
- A DataArray with the same shape as the log likelihood data
- An ELPDDataLOO object from a previous :func:`arviz_stats.loo` call.
Defaults to None. If not provided, it will be computed using the PSIS-LOO method.
k_threshold : float, optional
Pareto shape threshold. Observations with k values above this
threshold will trigger a refit. Defaults to :math:`\min(1 - 1/\log_{10}(S), 0.7)`,
where S is the number of samples.
pointwise : bool, optional
If True, return pointwise LOO data. Defaults to
``rcParams["stats.ic_pointwise"]``.
Returns
-------
ELPDDataLOO
Updated LOO results with the following attributes:
- **kind**: "loo"
- **elpd**: expected log pointwise predictive density
- **se**: standard error of the elpd
- **p**: effective number of parameters
- **n_samples**: number of samples
- **n_data_points**: number of data points
- **scale**: "log"
- **warning**: True if the estimated shape parameter of Pareto distribution is greater
than ``good_k``.
- **good_k**: For a sample size S, the threshold is computed as
``min(1 - 1/log10(S), 0.7)``
- **elpd_i**: :class:`~xarray.DataArray` with the pointwise predictive accuracy, only if
``pointwise=True``
- **pareto_k**: :class:`~xarray.DataArray` with Pareto shape values,
only if ``pointwise=True``
- **approx_posterior**: False (not used for standard LOO)
- **log_weights**: Smoothed log weights. Set to NaN for refitted observations
when ``pointwise=True``.
- **p_loo_i**: :class:`~xarray.DataArray` with the pointwise effective number of
parameters, only if ``pointwise=True``.
Notes
-----
It is strongly recommended to first compute :func:`arviz_stats.loo` on the inference results to
confirm that the number of values above the threshold is small enough. Otherwise,
prohibitive computation time may be needed to perform all required refits.
As an extreme case, setting ``k_threshold=-np.inf`` will flag every observation and
make ``reloo`` perform the whole exact LOO-CV, requiring one refit per observation.
This is not generally recommended nor intended, however, if needed, this function can
be used to achieve the result.
The term `exact cross-validation` is used here to refer to the process of refitting the model
for each observation, which is also commonly called `brute-force cross-validation`. While
this eliminates the PSIS approximation, other approximations remain in the process, such as
those inherent in the fitting procedure itself and the use of cross-validation as a proxy
for expected log predictive density (ELPD).
Warnings
--------
Refitting can be computationally expensive. Check the number of high Pareto k
values before using ``reloo`` to ensure the computation time is acceptable.
See Also
--------
loo : Pareto smoothed importance sampling LOO-CV
loo_moment_match : Moment matching for problematic observations
References
----------
.. [1] Vehtari et al. *Practical Bayesian model evaluation using leave-one-out cross-validation
and WAIC*. Statistics and Computing. 27(5) (2017) https://doi.org/10.1007/s11222-016-9696-4
arXiv preprint https://arxiv.org/abs/1507.04544.
.. [2] Vehtari et al. *Pareto Smoothed Importance Sampling*.
Journal of Machine Learning Research, 25(72) (2024) https://jmlr.org/papers/v25/19-556.html
arXiv preprint https://arxiv.org/abs/1507.02646
"""
if not isinstance(wrapper, SamplingWrapper):
raise TypeError(
"wrapper must be an instance of SamplingWrapper or a subclass. "
"See the SamplingWrapper documentation for implementation details."
)
required_methods = ["sel_observations", "sample", "get_inference_data", "log_likelihood__i"]
not_implemented = wrapper.check_implemented_methods(required_methods)
if not_implemented:
raise TypeError(
"Passed wrapper instance does not implement all methods required for reloo "
f"to work. Check the documentation of SamplingWrapper. {not_implemented} must be "
"implemented and were not found."
)
if wrapper.idata_orig is None:
raise ValueError(
"wrapper.idata_orig must contain the original inference data. "
"Pass idata_orig when creating the SamplingWrapper."
)
pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise
if loo_orig is None:
pareto_k = None
if isinstance(log_weights, ELPDData):
if not isinstance(log_weights, ELPDDataLOO) or log_weights.log_weights is None:
raise ValueError("ELPDData object does not contain log_weights")
pareto_k = log_weights.pareto_k
log_weights = log_weights.log_weights
# Dataset case comes from loo_subsample
if isinstance(log_weights, xr.Dataset):
if var_name and var_name in log_weights:
log_weights = log_weights[var_name]
else:
log_weights = log_weights[list(log_weights.data_vars)[0]]
loo_orig = loo(
wrapper.idata_orig,
var_name=var_name,
pointwise=True,
log_weights=log_weights,
pareto_k=pareto_k,
)
if not isinstance(loo_orig, ELPDDataLOO):
raise TypeError("loo_orig must be an ELPDDataLOO object.")
if loo_orig.pareto_k is None or loo_orig.elpd_i is None:
raise ValueError(
"reloo requires pointwise LOO results with Pareto k values. "
"Please compute the initial LOO with pointwise=True."
)
sample_dims = ["chain", "draw"]
loo_refitted = deepcopy(loo_orig)
loo_inputs = _prepare_loo_inputs(wrapper.idata_orig, var_name)
obs_dims = loo_inputs.obs_dims
n_data_points = loo_inputs.n_data_points
n_samples = loo_inputs.n_samples
log_likelihood = loo_inputs.log_likelihood
obs_sizes = {dim: log_likelihood.sizes[dim] for dim in obs_dims}
loo_obs_sizes = {dim: loo_orig.pareto_k.sizes.get(dim) for dim in obs_dims}
if loo_obs_sizes != obs_sizes:
raise ValueError(
"Observation dimensions in loo_orig do not match the log likelihood in "
f"wrapper.idata_orig. loo_orig has {loo_obs_sizes}, idata_orig has {obs_sizes}."
)
if k_threshold is None:
k_threshold = min(1 - 1 / np.log10(n_samples), 0.7)
pareto_k_stacked = loo_orig.pareto_k.stack(__obs__=obs_dims).transpose("__obs__")
bad_k_mask = pareto_k_stacked > k_threshold
bad_obs_flat_indices = np.where(bad_k_mask.values)[0]
if len(bad_obs_flat_indices) == 0:
if not pointwise:
loo_refitted.elpd_i = None
loo_refitted.pareto_k = None
return loo_refitted
if loo_refitted.p_loo_i is None:
loo_refitted.p_loo_i = xr.full_like(loo_refitted.elpd_i, np.nan)
lpd_i = logsumexp(log_likelihood, b=1 / n_samples, dims=sample_dims)
loo_refitted.p_loo_i = lpd_i - loo_refitted.elpd_i
if len(obs_dims) == 1:
bad_obs_indices = bad_obs_flat_indices
else:
obs_shape = [loo_orig.pareto_k.sizes[dim] for dim in obs_dims]
bad_obs_indices = np.array(np.unravel_index(bad_obs_flat_indices, obs_shape)).T
refitted_obs = []
for i, _ in enumerate(bad_obs_flat_indices):
if len(obs_dims) == 1:
obs_idx = bad_obs_indices[i]
coord_value = loo_orig.pareto_k.coords[obs_dims[0]].values[obs_idx]
obs_idx_dict = {obs_dims[0]: coord_value}
sel_idx = obs_idx
else:
obs_idx = bad_obs_indices[i]
obs_idx_dict = {}
for j, dim in enumerate(obs_dims):
coord_value = loo_orig.pareto_k.coords[dim].values[obs_idx[j]]
obs_idx_dict[dim] = coord_value
sel_idx = obs_idx_dict
log_lik_i = log_likelihood.sel(obs_idx_dict)
hat_lpd_i = logsumexp(log_lik_i, dims=sample_dims, b=1 / n_samples).item()
new_obs, excluded_obs = wrapper.sel_observations(sel_idx)
fit = wrapper.sample(new_obs)
idata_idx = wrapper.get_inference_data(fit)
log_lik_idx = wrapper.log_likelihood__i(excluded_obs, idata_idx)
if not all(dim in log_lik_idx.dims for dim in sample_dims):
raise ValueError(
f"log_likelihood__i must return a DataArray with dimensions {sample_dims}, "
f"found {log_lik_idx.dims}."
)
other_dims = [dim for dim in log_lik_idx.dims if dim not in sample_dims]
if any(log_lik_idx.sizes[dim] != 1 for dim in other_dims):
raise ValueError(
"log_likelihood__i must return a single observation per call, found "
f"dimensions {other_dims} with sizes "
f"{[log_lik_idx.sizes[dim] for dim in other_dims]}."
)
elpd_loo_i = logsumexp(log_lik_idx, dims=sample_dims, b=1 / log_lik_idx.size).item()
loo_refitted.elpd_i.loc[obs_idx_dict] = elpd_loo_i
loo_refitted.pareto_k.loc[obs_idx_dict] = 0.0
loo_refitted.p_loo_i.loc[obs_idx_dict] = hat_lpd_i - elpd_loo_i
refitted_obs.append(obs_idx_dict)
loo_refitted.elpd = np.sum(loo_refitted.elpd_i.values)
loo_refitted.se = np.sqrt(n_data_points * np.var(loo_refitted.elpd_i.values, ddof=1))
loo_refitted.p = np.sum(loo_refitted.p_loo_i.values)
loo_refitted.warning = np.any(loo_refitted.pareto_k.values > loo_refitted.good_k)
if not pointwise:
loo_refitted.elpd_i = None
loo_refitted.pareto_k = None
loo_refitted.p_loo_i = None
if loo_orig.log_weights is not None:
loo_refitted.log_weights = loo_orig.log_weights
else:
if loo_orig.log_weights is not None:
log_weights_refitted = loo_orig.log_weights
if isinstance(log_weights_refitted, xr.Dataset):
log_weights_refitted = log_weights_refitted[loo_inputs.var_name]
log_weights_refitted = log_weights_refitted.copy()
for obs_idx_dict in refitted_obs:
log_weights_refitted.loc[obs_idx_dict] = np.nan
loo_refitted.log_weights = log_weights_refitted
return loo_refitted
