"""Compute moment matching for problematic observations in PSIS-LOO-CV."""
import warnings
from collections import namedtuple
from copy import deepcopy
import arviz_base as azb
import numpy as np
import xarray as xr
from arviz_base import dataset_to_dataarray, rcParams
from xarray_einstats.stats import logsumexp
from arviz_stats.loo.helper_loo import (
_get_log_likelihood_i,
_get_r_eff,
_get_r_eff_i,
_get_weights_and_k_i,
_prepare_loo_inputs,
_shift,
_shift_and_cov,
_shift_and_scale,
_warn_pareto_k,
)
from arviz_stats.sampling_diagnostics import ess
from arviz_stats.utils import ELPDData
SplitMomentMatch = namedtuple("SplitMomentMatch", ["lwi", "lwfi", "log_liki", "reff"])
UpdateQuantities = namedtuple("UpdateQuantities", ["lwi", "lwfi", "ki", "kfi", "log_liki"])
LooMomentMatchResult = namedtuple(
"LooMomentMatchResult",
["final_log_liki", "final_lwi", "final_ki", "kfs_i", "reff_i", "original_ki", "i"],
)
[docs]
def loo_moment_match(
data,
loo_orig,
log_prob_upars_fn,
log_lik_i_upars_fn,
upars=None,
var_name=None,
reff=None,
max_iters=30,
k_threshold=None,
split=True,
cov=False,
pointwise=None,
):
r"""Compute moment matching for problematic observations in PSIS-LOO-CV.
Adjusts the results of a previously computed Pareto smoothed importance sampling leave-one-out
cross-validation (PSIS-LOO-CV) object by applying a moment matching algorithm to
observations with high Pareto k diagnostic values. The moment matching algorithm iteratively
adjusts the posterior draws in the unconstrained parameter space to better approximate the
leave-one-out posterior.
The moment matching algorithm is described in [1]_ and the PSIS-LOO-CV method is described in
[2]_ and [3]_.
Parameters
----------
data : DataTree or InferenceData
Input data. It should contain the posterior and the log_likelihood groups.
loo_orig : ELPDData
An existing ELPDData object from a previous `loo` result. Must contain
pointwise Pareto k values (`pointwise=True` must have been used).
log_prob_upars_fn : callable
Function that computes the log probability density of the full posterior
distribution evaluated at unconstrained parameter draws.
The function signature is ``log_prob_upars_fn(upars)`` where ``upars``
is a :class:`~xarray.DataArray` of unconstrained parameter draws with dimensions
``chain``, ``draw``, and a parameter dimension. It should return a
:class:`~xarray.DataArray` with dimensions ``chain``, ``draw``.
log_lik_i_upars_fn : callable
Function that computes the log-likelihood of a single left-out observation
evaluated at unconstrained parameter draws.
The function signature is ``log_lik_i_upars_fn(upars, i)`` where ``upars``
is a :class:`~xarray.DataArray` of unconstrained parameter draws and ``i``
is the integer index of the left-out observation. It should return a
:class:`~xarray.DataArray` with dimensions ``chain``, ``draw``.
upars : DataArray, optional
Posterior draws transformed to the unconstrained parameter space. Must have
``chain`` and ``draw`` dimensions, plus one additional dimension containing all
parameters. Parameter names can be provided as coordinate values on this
dimension. If not provided, will attempt to use the ``unconstrained_posterior``
group from the input data if available.
var_name : str, optional
The name of the variable in log_likelihood group storing the pointwise log
likelihood data to use for loo computation.
reff: float, optional
Relative MCMC efficiency, ``ess / n`` i.e. number of effective samples divided by the number
of actual samples. Computed from trace by default.
max_iters : int, default 30
Maximum number of moment matching iterations for each problematic observation.
k_threshold : float, optional
Threshold value for Pareto k values above which moment matching is applied.
Defaults to :math:`\min(1 - 1/\log_{10}(S), 0.7)`, where S is the number of samples.
split : bool, default True
If True, only transform half of the draws and use multiple importance sampling to combine
them with untransformed draws.
cov : bool, default False
If True, match the covariance structure during the transformation, in addition
to the mean and marginal variances. Ignored if ``split=False``.
pointwise: bool, optional
If True, the pointwise predictive accuracy will be returned. Defaults to
``rcParams["stats.ic_pointwise"]``. Moment matching always requires
pointwise data from ``loo_orig``. This argument controls whether the returned
object includes pointwise data.
Returns
-------
ELPDData
Object 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.
Examples
--------
Moment matching can improve PSIS-LOO-CV estimates for observations with high Pareto k values
without having to refit the model for each problematic observation. We will use the non-centered
eight schools data which has 1 problematic observation. In practice, moment matching is useful
when you have a potentially large number of problematic observations:
.. ipython::
:okwarning:
In [1]: import arviz_base as azb
...: import numpy as np
...: import xarray as xr
...: from scipy import stats
...: from arviz_stats import loo
...:
...: idata = azb.load_arviz_data("non_centered_eight")
...: posterior = idata.posterior
...: schools = posterior.theta_t.coords["school"].values
...: y_obs = idata.observed_data.obs
...: obs_dim = y_obs.dims[0]
...:
...: loo_orig = loo(idata, pointwise=True, var_name="obs")
...: loo_orig
The moment matching algorithm applies affine transformations to posterior draws in
unconstrained parameter space. To enable this, we need to collect the posterior
parameters from their original space, transform them to unconstrained space if
needed, and stack them into a single :class:`xarray.DataArray` that matches the
expected ``(chain, draw, param)`` structure. Some parameters may already be in
unconstrained space, so we don't need to transform them. This will depend on the
model and the choice of parameterization:
.. ipython::
:okwarning:
In [2]: upars_ds = xr.Dataset(
...: {
...: **{
...: f"theta_t_{school}": posterior.theta_t.sel(school=school, drop=True)
...: for school in schools
...: },
...: "mu": posterior.mu,
...: "log_tau": xr.apply_ufunc(np.log, posterior.tau),
...: }
...: )
...: upars = azb.dataset_to_dataarray(
...: upars_ds, sample_dims=["chain", "draw"], new_dim="upars_dim"
...: )
Moment matching requires two functions: one for the joint log probability (likelihood + priors)
and another for the pointwise log-likelihood of a single observation. We first define functions
that accept the data they need as keyword-only arguments:
.. ipython::
:okwarning:
In [3]: sigmas = xr.DataArray(
...: [15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0],
...: dims=[obs_dim],
...: )
...:
...: def log_prob_upars(upars, *, sigmas, y, schools, obs_dim):
...: theta_t = xr.concat(
...: [upars.sel(upars_dim=f"theta_t_{school}") for school in schools],
...: dim=obs_dim,
...: )
...: mu = upars.sel(upars_dim="mu")
...: log_tau = upars.sel(upars_dim="log_tau")
...: tau = xr.apply_ufunc(np.exp, log_tau)
...: theta = mu + tau * theta_t
...:
...: log_prior = xr.apply_ufunc(stats.norm(0, 5).logpdf, mu)
...: log_prior = log_prior + xr.apply_ufunc(
...: stats.halfcauchy(0, 5).logpdf,
...: tau,
...: )
...: log_prior = log_prior + log_tau
...: log_prior = log_prior + xr.apply_ufunc(
...: stats.norm(0, 1).logpdf,
...: theta_t,
...: ).sum(obs_dim)
...:
...: const = -0.5 * np.log(2 * np.pi)
...: log_like = const - np.log(sigmas) - 0.5 * ((y - theta) / sigmas) ** 2
...: log_like = log_like.sum(obs_dim)
...: return log_prior + log_like
...:
...: def log_lik_i_upars(upars, i, *, sigmas, y, schools, obs_dim):
...: mu = upars.sel(upars_dim="mu")
...: log_tau = upars.sel(upars_dim="log_tau")
...: tau = xr.apply_ufunc(np.exp, log_tau)
...:
...: theta_t_i = upars.sel(upars_dim=f"theta_t_{schools[i]}")
...: theta_i = mu + tau * theta_t_i
...:
...: sigma_i = sigmas.isel({obs_dim: i})
...: y_i = y.isel({obs_dim: i})
...: const = -0.5 * np.log(2 * np.pi)
...: return const - np.log(sigma_i) - 0.5 * ((y_i - theta_i) / sigma_i) ** 2
Now, we can specialise these functions with :func:`functools.partial` so the resulting functions
match the signature expected by :func:`loo_moment_match()`:
.. ipython::
:okwarning:
In [4]: from functools import partial
...: log_prob_fn = partial(
...: log_prob_upars,
...: sigmas=sigmas,
...: y=y_obs,
...: schools=schools,
...: obs_dim=obs_dim,
...: )
...: log_lik_i_fn = partial(
...: log_lik_i_upars,
...: sigmas=sigmas,
...: y=y_obs,
...: schools=schools,
...: obs_dim=obs_dim,
...: )
Finally, we can run moment matching using the prepared inputs. Now, we
have no problematic observations anymore:
.. ipython::
:okwarning:
In [5]: from arviz_stats import loo_moment_match
...: loo_mm = loo_moment_match(
...: idata,
...: loo_orig,
...: upars=upars,
...: log_prob_upars_fn=log_prob_fn,
...: log_lik_i_upars_fn=log_lik_i_fn,
...: var_name="obs",
...: split=True,
...: )
...: loo_mm
Notes
-----
The moment matching algorithm considers three affine transformations of the posterior draws:
For a specific draw :math:`\theta^{(s)}`, a generic affine transformation includes a square
matrix :math:`\mathbf{A}` representing a linear map and a vector :math:`\mathbf{b}`
representing a translation such that
.. math::
T : \theta^{(s)} \mapsto \mathbf{A}\theta^{(s)} + \mathbf{b}
=: \theta^{*{(s)}}.
The first transformation, :math:`T_1`, is a translation that matches the mean of the sample
to its importance weighted mean given by
.. math::
\mathbf{\theta^{*{(s)}}} = T_1(\mathbf{\theta^{(s)}}) =
\mathbf{\theta^{(s)}} - \bar{\theta} + \bar{\theta}_w,
where :math:`\bar{\theta}` is the mean of the sample and :math:`\bar{\theta}_w` is the
importance weighted mean of the sample. The second transformation, :math:`T_2`, is a scaling
that matches the marginal variances in addition to the means given by
.. math::
\mathbf{\theta^{*{(s)}}} = T_2(\mathbf{\theta^{(s)}}) =
\mathbf{v}^{1/2}_w \circ \mathbf{v}^{-1/2} \circ (\mathbf{\theta^{(s)}} - \bar{\theta}) +
\bar{\theta}_w,
where :math:`\mathbf{v}` and :math:`\mathbf{v}_w` are the sample and weighted variances, and
:math:`\circ` denotes the pointwise product of the elements of two vectors. The third
transformation, :math:`T_3`, is a covariance transformation that matches the covariance matrix
of the sample to its importance weighted covariance matrix given by
.. math::
\mathbf{\theta^{*{(s)}}} = T_3(\mathbf{\theta^{(s)}}) =
\mathbf{L}_w \mathbf{L}^{-1} (\mathbf{\theta^{(s)}} - \bar{\theta}) + \bar{\theta}_w,
where :math:`\mathbf{L}` and :math:`\mathbf{L}_w` are the Cholesky decompositions of the
covariance matrix and the weighted covariance matrix, respectively, e.g.,
.. math::
\mathbf{LL}^T = \mathbf{\Sigma} = \frac{1}{S} \sum_{s=1}^S (\mathbf{\theta^{(s)}} -
\bar{\theta}) (\mathbf{\theta^{(s)}} - \bar{\theta})^T
and
.. math::
\mathbf{L}_w \mathbf{L}_w^T = \mathbf{\Sigma}_w = \frac{\frac{1}{S} \sum_{s=1}^S
w^{(s)} (\mathbf{\theta^{(s)}} - \bar{\theta}_w) (\mathbf{\theta^{(s)}} -
\bar{\theta}_w)^T}{\sum_{s=1}^S w^{(s)}}.
We iterate on :math:`T_1` repeatedly and move onto :math:`T_2` and :math:`T_3` only
if :math:`T_1` fails to yield a Pareto-k statistic below the threshold.
See Also
--------
loo : Standard PSIS-LOO-CV.
reloo : Exact re-fitting for problematic observations.
References
----------
.. [1] Paananen, T., Piironen, J., Buerkner, P.-C., Vehtari, A. (2021). Implicitly Adaptive
Importance Sampling. Statistics and Computing. 31(2) (2021)
https://doi.org/10.1007/s11222-020-09982-2
arXiv preprint https://arxiv.org/abs/1906.08850.
.. [2] 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.
.. [3] 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(loo_orig, ELPDData):
raise TypeError("loo_orig must be an ELPDData object.")
if loo_orig.pareto_k is None or loo_orig.elpd_i is None:
raise ValueError(
"Moment matching requires pointwise LOO results with Pareto k values. "
"Please compute the initial LOO with pointwise=True."
)
sample_dims = ["chain", "draw"]
if upars is None:
if hasattr(data, "unconstrained_posterior"):
upars_ds = azb.get_unconstrained_samples(data, return_dataset=True)
upars = dataset_to_dataarray(
upars_ds, sample_dims=sample_dims, new_dim="unconstrained_parameter"
)
else:
raise ValueError(
"upars must be provided or data must contain an 'unconstrained_posterior' group."
)
if not isinstance(upars, xr.DataArray):
raise TypeError("upars must be a DataArray.")
if not all(dim_name in upars.dims for dim_name in sample_dims):
raise ValueError(f"upars must have dimensions {sample_dims}.")
param_dim_list = [dim for dim in upars.dims if dim not in sample_dims]
if len(param_dim_list) == 0:
param_dim_name = "upars_dim"
upars = upars.expand_dims(dim={param_dim_name: 1})
elif len(param_dim_list) == 1:
param_dim_name = param_dim_list[0]
else:
raise ValueError("upars must have at most one dimension besides 'chain' and 'draw'.")
loo_data = deepcopy(loo_orig)
loo_data.method = "loo_moment_match"
pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise
loo_inputs = _prepare_loo_inputs(data, var_name)
log_likelihood = loo_inputs.log_likelihood
obs_dims = loo_inputs.obs_dims
n_samples = loo_inputs.n_samples
var_name = loo_inputs.var_name
n_params = upars.sizes[param_dim_name]
n_data_points = loo_orig.n_data_points
if reff is None:
reff = _get_r_eff(data, n_samples)
try:
orig_log_prob = log_prob_upars_fn(upars)
if not isinstance(orig_log_prob, xr.DataArray):
raise TypeError("log_prob_upars_fn must return a DataArray.")
if not all(dim in orig_log_prob.dims for dim in sample_dims):
raise ValueError(f"Original log probability must have dimensions {sample_dims}.")
if len(orig_log_prob.dims) != len(sample_dims):
raise ValueError(
f"Original log probability should only have dimensions {sample_dims}, "
f"found {orig_log_prob.dims}"
)
except Exception as e:
raise ValueError(f"Error executing log_prob_upars_fn: {e}") from e
if k_threshold is None:
k_threshold = min(1 - 1 / np.log10(n_samples), 0.7) if n_samples > 1 else 0.7
ks = loo_data.pareto_k.stack(__obs__=obs_dims).transpose("__obs__").values
bad_obs_indices = np.where(ks > k_threshold)[0]
if len(bad_obs_indices) == 0:
warnings.warn("No Pareto k values exceed the threshold. Returning original LOO data.")
if not pointwise:
loo_data.elpd_i = None
loo_data.pareto_k = None
if hasattr(loo_data, "p_loo_i"):
loo_data.p_loo_i = None
return loo_data
lpd = logsumexp(log_likelihood, dims=sample_dims, b=1 / n_samples)
loo_data.p_loo_i = lpd - loo_data.elpd_i
kfs = np.zeros(n_data_points)
log_weights = getattr(loo_data, "log_weights", None)
r_eff_data = getattr(loo_data, "r_eff", reff)
# Moment matching algorithm
for i in bad_obs_indices:
mm_result = _loo_moment_match_i(
i=i,
upars=upars,
log_likelihood=log_likelihood,
log_prob_upars_fn=log_prob_upars_fn,
log_lik_i_upars_fn=log_lik_i_upars_fn,
max_iters=max_iters,
k_threshold=k_threshold,
split=split,
cov=cov,
orig_log_prob=orig_log_prob,
ks=ks,
log_weights=log_weights,
pareto_k=loo_data.pareto_k,
r_eff=r_eff_data,
sample_dims=sample_dims,
obs_dims=obs_dims,
n_samples=n_samples,
n_params=n_params,
param_dim_name=param_dim_name,
var_name=var_name,
)
kfs[i] = mm_result.kfs_i
if mm_result.final_ki < mm_result.original_ki:
new_elpd_i = logsumexp(
mm_result.final_log_liki + mm_result.final_lwi, dims=sample_dims
).item()
original_log_liki = _get_log_likelihood_i(log_likelihood, i, obs_dims)
_update_loo_data_i(
loo_data,
i,
new_elpd_i,
mm_result.final_ki,
mm_result.final_log_liki,
sample_dims,
obs_dims,
n_samples,
original_log_liki,
suppress_warnings=True,
)
else:
warnings.warn(
f"Observation {i}: Moment matching did not improve k "
f"({mm_result.original_ki:.2f} -> {mm_result.final_ki:.2f}). Reverting.",
UserWarning,
stacklevel=2,
)
if hasattr(loo_orig, "p_loo_i") and loo_orig.p_loo_i is not None:
if len(obs_dims) == 1:
idx_dict = {obs_dims[0]: i}
else:
coords = np.unravel_index(i, tuple(loo_data.elpd_i.sizes[d] for d in obs_dims))
idx_dict = dict(zip(obs_dims, coords))
loo_data.p_loo_i[idx_dict] = loo_orig.p_loo_i[idx_dict]
final_ks = loo_data.pareto_k.stack(__obs__=obs_dims).transpose("__obs__").values
if np.any(final_ks[bad_obs_indices] > k_threshold):
warnings.warn(
f"After Moment Matching, {np.sum(final_ks > k_threshold)} observations still have "
f"Pareto k > {k_threshold:.2f}.",
UserWarning,
stacklevel=2,
)
if not split and np.any(kfs > k_threshold):
warnings.warn(
"The accuracy of self-normalized importance sampling may be bad. "
"Setting the argument 'split' to 'True' will likely improve accuracy.",
UserWarning,
stacklevel=2,
)
elpd_raw = logsumexp(log_likelihood, dims=sample_dims, b=1 / n_samples).sum().values
loo_data.p = elpd_raw - loo_data.elpd
if not pointwise:
loo_data.elpd_i = None
loo_data.pareto_k = None
if hasattr(loo_data, "p_loo_i"):
loo_data.p_loo_i = None
return loo_data
def _split_moment_match(
upars,
cov,
total_shift,
total_scaling,
total_mapping,
i,
reff,
log_prob_upars_fn,
log_lik_i_upars_fn,
):
r"""Split moment matching importance sampling for PSIS-LOO-CV.
Applies affine transformations based on the total moment matching transformation
to half of the posterior draws, leaving the other half unchanged. These approximations
to the leave-one-out posterior are then combined using multiple importance sampling.
Based on the implicit adaptive importance sampling algorithm of [1]_ and the
PSIS-LOO-CV method of [2]_ and [3]_.
Parameters
----------
upars : DataArray
A DataArray representing the posterior draws of the model parameters in the
unconstrained space. Must contain the dimensions `chain` and `draw` and a final
dimension representing the different unconstrained parameters.
cov : bool
Whether to match the full covariance matrix of the samples (True) or just the
marginal variances (False). Using the full covariance is more computationally
expensive.
total_shift : ndarray
Vector containing the total shift (translation) applied to the parameters. Shape should
match the parameter dimension of ``upars``.
total_scaling : ndarray
Vector containing the total scaling factors for the marginal variances. Shape should
match the parameter dimension of ``upars``.
total_mapping : ndarray
Square matrix representing the linear transformation applied to the covariance matrix.
Shape should be (d, d) where d is the parameter dimension.
i : int
Index of the specific observation to be left out for computing leave-one-out
likelihood.
reff : float
Relative MCMC efficiency, ``ess / n`` i.e. number of effective samples divided by the number
of actual samples.
log_prob_upars_fn : callable
Function that computes the log probability density of the *full posterior*
distribution evaluated at unconstrained parameter draws.
The function signature is ``log_prob_upars_fn(upars)`` where ``upars``
is a :class:`~xarray.DataArray` of unconstrained parameter draws.
It should return a :class:`~xarray.DataArray` with dimensions ``chain``, ``draw``.
log_lik_i_upars_fn : callable
Function that computes the log-likelihood of the *left-out observation* ``i``
evaluated at unconstrained parameter draws.
The function signature is ``log_lik_i_upars_fn(upars, i)`` where ``upars``
is a :class:`~xarray.DataArray` of unconstrained parameter draws and ``i``
is the integer index of the observation.
It should return a :class:`~xarray.DataArray` with dimensions ``chain``, ``draw``.
Returns
-------
SplitMomentMatch
A namedtuple containing:
- lwi: Updated log importance weights for each sample
- lwfi: Updated log importance weights for full distribution
- log_liki: Updated log likelihood values for the specific observation
- reff: Relative MCMC efficiency (updated based on the split samples)
References
----------
.. [1] Paananen, T., Piironen, J., Buerkner, P.-C., Vehtari, A. (2021). *Implicitly Adaptive
Importance Sampling*. Statistics and Computing. 31(2) (2021)
https://doi.org/10.1007/s11222-020-09982-2
arXiv preprint https://arxiv.org/abs/1906.08850.
.. [2] 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.
.. [3] 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
"""
sample_dims = ["chain", "draw"]
param_dim = next(dim for dim in upars.dims if dim not in sample_dims)
dim = upars.sizes[param_dim]
n_chains = upars.sizes["chain"]
n_draws = upars.sizes["draw"]
n_samples = n_chains * n_draws
n_samples_half = n_samples // 2
stack_dims = ["draw", "chain"]
upars_stacked = upars.stack(__sample__=stack_dims).transpose("__sample__", param_dim)
mean_original = upars_stacked.mean(dim="__sample__")
# Forward transformation
upars_trans = upars_stacked - mean_original
upars_trans = upars_trans * xr.DataArray(total_scaling, dims=param_dim)
if cov and dim > 0:
upars_trans = xr.DataArray(
upars_trans.data @ total_mapping.T,
coords=upars_trans.coords,
dims=upars_trans.dims,
)
# Inverse Transformation
upars_trans = upars_trans + (xr.DataArray(total_shift, dims=param_dim) + mean_original)
upars_trans_inv = upars_stacked - (xr.DataArray(total_shift, dims=param_dim) + mean_original)
if cov and dim > 0:
try:
inv_mapping_t = np.linalg.inv(total_mapping.T)
upars_trans_inv = xr.DataArray(
upars_trans_inv.data @ inv_mapping_t,
coords=upars_trans_inv.coords,
dims=upars_trans_inv.dims,
)
except np.linalg.LinAlgError:
warnings.warn("Could not invert mapping matrix. Using identity.", UserWarning)
upars_trans_inv = upars_trans_inv / xr.DataArray(total_scaling, dims=param_dim)
upars_trans_inv = upars_trans_inv + (mean_original - xr.DataArray(total_shift, dims=param_dim))
upars_trans_half_stacked = upars_stacked.copy(deep=True)
upars_trans_half_stacked.data[:n_samples_half, :] = upars_trans.data[:n_samples_half, :]
upars_trans_half = upars_trans_half_stacked.unstack("__sample__").transpose(
*reversed(stack_dims), param_dim
)
upars_trans_half_inv_stacked = upars_stacked.copy(deep=True)
upars_trans_half_inv_stacked.data[n_samples_half:, :] = upars_trans_inv.data[n_samples_half:, :]
upars_trans_half_inv = upars_trans_half_inv_stacked.unstack("__sample__").transpose(
*reversed(stack_dims), param_dim
)
try:
log_prob_half_trans = log_prob_upars_fn(upars_trans_half)
if not isinstance(log_prob_half_trans, xr.DataArray):
raise TypeError("log_prob_upars_fn must return a DataArray.")
if not all(dim in log_prob_half_trans.dims for dim in sample_dims) or len(
log_prob_half_trans.dims
) != len(sample_dims):
raise ValueError(
f"log_prob_upars_fn must return a DataArray with dimensions {sample_dims}, "
f"but got {log_prob_half_trans.dims}"
)
log_prob_half_trans_inv = log_prob_upars_fn(upars_trans_half_inv)
if not isinstance(log_prob_half_trans_inv, xr.DataArray):
raise TypeError("log_prob_upars_fn must return a DataArray.")
if not all(dim in log_prob_half_trans_inv.dims for dim in sample_dims) or len(
log_prob_half_trans_inv.dims
) != len(sample_dims):
raise ValueError(
f"log_prob_upars_fn must return a DataArray with dimensions {sample_dims}, "
f"but got {log_prob_half_trans_inv.dims}"
)
except Exception as e:
raise ValueError(
f"Could not compute log probabilities for transformed parameters: {e}"
) from e
try:
log_liki_half = log_lik_i_upars_fn(upars_trans_half, i)
if not all(dim in log_liki_half.dims for dim in sample_dims) or len(
log_liki_half.dims
) != len(sample_dims):
raise ValueError(
f"log_lik_i_upars_fn must return a DataArray with dimensions {sample_dims}"
)
if (
log_liki_half.sizes["chain"] != upars.sizes["chain"]
or log_liki_half.sizes["draw"] != upars.sizes["draw"]
):
raise ValueError(
"log_lik_i_upars_fn output shape does not match input sample dimensions"
)
except Exception as e:
raise ValueError(f"Could not compute log likelihood for observation {i}: {e}") from e
log_jacobian_det = 0.0
if dim > 0:
log_jacobian_det = -np.sum(np.log(total_scaling))
try:
log_jacobian_det -= np.log(np.linalg.det(total_mapping))
except np.linalg.LinAlgError:
log_jacobian_det -= np.inf
log_prob_half_trans_inv_adj = log_prob_half_trans_inv + log_jacobian_det
# Multiple importance sampling
use_forward_log_prob = log_prob_half_trans > log_prob_half_trans_inv_adj
raw_log_weights_half = -log_liki_half + log_prob_half_trans
log_sum_terms = xr.where(
use_forward_log_prob,
log_prob_half_trans
+ xr.ufuncs.log1p(np.exp(log_prob_half_trans_inv_adj - log_prob_half_trans)),
log_prob_half_trans_inv_adj
+ xr.ufuncs.log1p(np.exp(log_prob_half_trans - log_prob_half_trans_inv_adj)),
)
raw_log_weights_half -= log_sum_terms
raw_log_weights_half = xr.where(np.isnan(raw_log_weights_half), -np.inf, raw_log_weights_half)
raw_log_weights_half = xr.where(
np.isposinf(raw_log_weights_half), -np.inf, raw_log_weights_half
)
# PSIS smoothing for half posterior
lwi_psis_da, _ = _wrap__psislw(raw_log_weights_half, sample_dims, reff)
lr_full = lwi_psis_da + log_liki_half
lr_full = xr.where(np.isnan(lr_full) | (np.isinf(lr_full) & (lr_full > 0)), -np.inf, lr_full)
# PSIS smoothing for full posterior
lwfi_psis_da, _ = _wrap__psislw(lr_full, sample_dims, reff)
n_chains = upars.sizes["chain"]
if n_chains == 1:
reff_updated = reff
else:
log_liki_half_1 = log_liki_half.isel(
chain=slice(None), draw=slice(0, n_samples_half // n_chains)
)
log_liki_half_2 = log_liki_half.isel(
chain=slice(None), draw=slice(n_samples_half // n_chains, None)
)
liki_half_1 = np.exp(log_liki_half_1)
liki_half_2 = np.exp(log_liki_half_2)
ess_1 = liki_half_1.azstats.ess(method="mean")
ess_2 = liki_half_2.azstats.ess(method="mean")
ess_1_value = ess_1.values if hasattr(ess_1, "values") else ess_1
ess_2_value = ess_2.values if hasattr(ess_2, "values") else ess_2
n_samples_1 = log_liki_half_1.size
n_samples_2 = log_liki_half_2.size
r_eff_1 = ess_1_value / n_samples_1
r_eff_2 = ess_2_value / n_samples_2
reff_updated = min(r_eff_1, r_eff_2)
return SplitMomentMatch(
lwi=lwi_psis_da,
lwfi=lwfi_psis_da,
log_liki=log_liki_half,
reff=reff_updated,
)
def _loo_moment_match_i(
i,
upars,
log_likelihood,
log_prob_upars_fn,
log_lik_i_upars_fn,
max_iters,
k_threshold,
split,
cov,
orig_log_prob,
ks,
log_weights,
pareto_k,
r_eff,
sample_dims,
obs_dims,
n_samples,
n_params,
param_dim_name,
var_name,
):
"""Compute moment matching for a single observation."""
n_chains = upars.sizes["chain"]
n_draws = upars.sizes["draw"]
log_liki = _get_log_likelihood_i(log_likelihood, i, obs_dims).squeeze(drop=True)
if isinstance(r_eff, xr.DataArray):
reff_i = _get_r_eff_i(r_eff, i, obs_dims)
elif r_eff is not None:
reff_i = r_eff
else:
liki = np.exp(log_liki)
liki_reshaped = liki.values.reshape(n_chains, n_draws).T
ess_val = ess(liki_reshaped, method="mean").item()
reff_i = ess_val / n_samples if n_samples > 0 else 1.0
original_ki = ks[i]
if log_weights is not None:
log_weights_i, ki = _get_weights_and_k_i(
log_weights=log_weights,
pareto_k=pareto_k,
i=i,
obs_dims=obs_dims,
sample_dims=sample_dims,
data=log_likelihood,
n_samples=n_samples,
reff=reff_i,
log_lik_i=log_liki,
var_name=var_name,
)
lwi = log_weights_i.squeeze(drop=True).transpose(*sample_dims).astype(np.float64)
else:
log_ratio_i_init = -log_liki
lwi, ki = _wrap__psislw(log_ratio_i_init, sample_dims, reff_i)
upars_i = upars.copy(deep=True)
total_shift = np.zeros(upars_i.sizes[param_dim_name])
total_scaling = np.ones(upars_i.sizes[param_dim_name])
total_mapping = np.eye(upars_i.sizes[param_dim_name])
iterind = 1
transformations_applied = False
kfs_i = 0
while iterind <= max_iters and ki > k_threshold:
if iterind == max_iters:
warnings.warn(
f"Maximum number of moment matching iterations ({max_iters}) reached "
f"for observation {i}. Final Pareto k is {ki:.2f}.",
UserWarning,
stacklevel=2,
)
break
# Try Mean Shift
try:
shift_res = _shift(upars_i, lwi)
quantities_i = _update_quantities_i(
shift_res.upars,
i,
orig_log_prob,
log_prob_upars_fn,
log_lik_i_upars_fn,
reff_i,
sample_dims,
)
if quantities_i.ki < ki:
ki = quantities_i.ki
lwi = quantities_i.lwi
log_liki = quantities_i.log_liki
kfs_i = quantities_i.kfi
upars_i = shift_res.upars
total_shift = total_shift + shift_res.shift
transformations_applied = True
iterind += 1
continue # Restart, try mean shift again
except RuntimeError as e:
warnings.warn(
f"Error during mean shift calculation for observation {i}: {e}. "
"Stopping moment matching for this observation.",
UserWarning,
stacklevel=2,
)
break
# Try Scale Shift
try:
scale_res = _shift_and_scale(upars_i, lwi)
quantities_i = _update_quantities_i(
scale_res.upars,
i,
orig_log_prob,
log_prob_upars_fn,
log_lik_i_upars_fn,
reff_i,
sample_dims,
)
if quantities_i.ki < ki:
ki = quantities_i.ki
lwi = quantities_i.lwi
log_liki = quantities_i.log_liki
kfs_i = quantities_i.kfi
upars_i = scale_res.upars
total_shift = total_shift + scale_res.shift
total_scaling = total_scaling * scale_res.scaling
transformations_applied = True
iterind += 1
continue # Restart, try mean shift again
except RuntimeError as e:
warnings.warn(
f"Error during scale shift calculation for observation {i}: {e}. "
"Stopping moment matching for this observation.",
UserWarning,
stacklevel=2,
)
break
# Try Covariance Shift
if cov and n_samples >= 10 * n_params:
try:
cov_res = _shift_and_cov(upars_i, lwi)
quantities_i = _update_quantities_i(
cov_res.upars,
i,
orig_log_prob,
log_prob_upars_fn,
log_lik_i_upars_fn,
reff_i,
sample_dims,
)
if quantities_i.ki < ki:
ki = quantities_i.ki
lwi = quantities_i.lwi
log_liki = quantities_i.log_liki
kfs_i = quantities_i.kfi
upars_i = cov_res.upars
total_shift = total_shift + cov_res.shift
total_mapping = cov_res.mapping @ total_mapping
transformations_applied = True
iterind += 1
continue # Restart, try mean shift again
except RuntimeError as e:
warnings.warn(
f"Error during covariance shift calculation for observation {i}: {e}. "
"Stopping moment matching for this observation.",
UserWarning,
stacklevel=2,
)
break
break
if split and transformations_applied:
try:
split_res = _split_moment_match(
upars=upars,
cov=cov,
total_shift=total_shift,
total_scaling=total_scaling,
total_mapping=total_mapping,
i=i,
reff=reff_i,
log_prob_upars_fn=log_prob_upars_fn,
log_lik_i_upars_fn=log_lik_i_upars_fn,
)
final_log_liki = split_res.log_liki
final_lwi = split_res.lwi
final_ki = ki
reff_i = split_res.reff
except RuntimeError as e:
warnings.warn(
f"Error during split moment matching for observation {i}: {e}. "
"Using non-split transformation result.",
UserWarning,
stacklevel=2,
)
final_log_liki = log_liki
final_lwi = lwi
final_ki = ki
else:
final_log_liki = log_liki
final_lwi = lwi
final_ki = ki
if transformations_applied:
liki_final = np.exp(final_log_liki)
liki_final_reshaped = liki_final.values.reshape(n_chains, n_draws).T
ess_val_final = ess(liki_final_reshaped, method="mean").item()
reff_i = ess_val_final / n_samples if n_samples > 0 else 1.0
return LooMomentMatchResult(
final_log_liki=final_log_liki,
final_lwi=final_lwi,
final_ki=final_ki,
kfs_i=kfs_i,
reff_i=reff_i,
original_ki=original_ki,
i=i,
)
def _update_loo_data_i(
loo_data,
i,
new_elpd_i,
new_pareto_k,
log_liki,
sample_dims,
obs_dims,
n_samples,
original_log_liki=None,
suppress_warnings=False,
):
"""Update the ELPDData object for a single observation."""
if loo_data.elpd_i is None or loo_data.pareto_k is None:
raise ValueError("loo_data must contain pointwise elpd_i and pareto_k values.")
lpd_i_log_lik = original_log_liki if original_log_liki is not None else log_liki
lpd_i = logsumexp(lpd_i_log_lik, dims=sample_dims, b=1 / n_samples).item()
p_loo_i = lpd_i - new_elpd_i
if len(obs_dims) == 1:
idx_dict = {obs_dims[0]: i}
else:
coords = np.unravel_index(i, tuple(loo_data.elpd_i.sizes[d] for d in obs_dims))
idx_dict = dict(zip(obs_dims, coords))
loo_data.elpd_i[idx_dict] = new_elpd_i
loo_data.pareto_k[idx_dict] = new_pareto_k
if not hasattr(loo_data, "p_loo_i") or loo_data.p_loo_i is None:
loo_data.p_loo_i = xr.full_like(loo_data.elpd_i, np.nan)
loo_data.p_loo_i[idx_dict] = p_loo_i
loo_data.elpd = np.nansum(loo_data.elpd_i.values)
loo_data.se = np.sqrt(loo_data.n_data_points * np.nanvar(loo_data.elpd_i.values, ddof=1))
loo_data.warning, loo_data.good_k = _warn_pareto_k(
loo_data.pareto_k.values[~np.isnan(loo_data.pareto_k.values)],
loo_data.n_samples,
suppress=suppress_warnings,
)
def _update_quantities_i(
upars,
i,
orig_log_prob,
log_prob_upars_fn,
log_lik_i_upars_fn,
reff_i,
sample_dims,
):
"""Update the moment matching quantities for a single observation."""
log_prob_new = log_prob_upars_fn(upars)
log_liki_new = log_lik_i_upars_fn(upars, i)
log_ratio_i = -log_liki_new + log_prob_new - orig_log_prob
lwi_new, ki_new = _wrap__psislw(log_ratio_i, sample_dims, reff_i)
log_ratio_full = log_prob_new - orig_log_prob
lwfi_new, kfi_new = _wrap__psislw(log_ratio_full, sample_dims, reff_i)
return UpdateQuantities(
lwi=lwi_new,
lwfi=lwfi_new,
ki=ki_new,
kfi=kfi_new,
log_liki=log_liki_new,
)
def _wrap__psislw(log_weights, sample_dims, r_eff):
"""Apply PSIS smoothing over sample dimensions."""
if not isinstance(log_weights, xr.DataArray):
raise TypeError("log_weights must be an xarray.DataArray")
missing_dims = [dim for dim in sample_dims if dim not in log_weights.dims]
if missing_dims:
raise ValueError(
f"All sample dimensions must be present in the input; missing {missing_dims}."
)
other_dims = [dim for dim in log_weights.dims if dim not in sample_dims]
if other_dims:
raise ValueError(
"_wrap__psislw expects `log_weights` to include only sample dimensions; "
f"found extra dims {other_dims}."
)
stacked = log_weights.stack(__sample__=sample_dims)
stacked_for_psis = -stacked
try:
lw_stacked, k = stacked_for_psis.azstats.psislw(dim="__sample__", r_eff=r_eff)
except ValueError as err:
err_message = str(err)
fallback_errors = ("All tail values are the same", "n_draws_tail must be at least 5")
if not any(msg in err_message for msg in fallback_errors):
raise
log_norm = logsumexp(stacked, dims="__sample__")
lw_stacked = stacked - log_norm
k = np.inf
lw = lw_stacked.unstack("__sample__").transpose(*log_weights.dims)
if isinstance(k, xr.DataArray):
if k.dims:
raise ValueError("Unexpected dimensions on Pareto k output; expected scalar result.")
k_val = k.item()
elif isinstance(k, np.ndarray):
if k.ndim != 0:
raise ValueError("Unexpected array shape for Pareto k; expected scalar result.")
k_val = k.item()
else:
try:
k_val = k
except (TypeError, ValueError) as exc:
raise TypeError("Unable to convert PSIS tail index to float") from exc
return lw, k_val
