Pareto-smoothed importance sampling leave-one-out cross-validation.

    Calculates leave-one-out (LOO) cross-validation for out of sample predictive model fit,
    following Vehtari et al. (2017). Cross-validation is computed using Pareto-smoothed
    importance sampling (PSIS).

    Parameters
    ----------
    data : result of MCMC run
    pointwise: bool, optional
        if True the pointwise predictive accuracy will be returned. Defaults to False
    reff : float, optional
        Relative MCMC efficiency, `effective_n / n` i.e. number of effective samples divided by
        the number of actual samples. Computed from trace by default.
    scale : str
        Output scale for loo. Available options are:

        - `deviance` : (default) -2 * (log-score)
        - `log` : 1 * log-score (after Vehtari et al. (2017))
        - `negative_log` : -1 * (log-score)

    Returns
    -------
    pandas.Series with the following columns:
    loo: approximated Leave-one-out cross-validation
    loo_se: standard error of loo
    p_loo: effective number of parameters
    shape_warn: 1 if the estimated shape parameter of
        Pareto distribution is greater than 0.7 for one or more samples
    loo_i: array of pointwise predictive accuracy, only if pointwise True
    pareto_k: array of Pareto shape values, only if pointwise True
    loo_scale: scale of the loo results

def loo(data, pointwise=False, reff=None, scale="deviance"):
    """Pareto-smoothed importance sampling leave-one-out cross-validation.

    Calculates leave-one-out (LOO) cross-validation for out of sample predictive model fit,
    following Vehtari et al. (2017). Cross-validation is computed using Pareto-smoothed
    importance sampling (PSIS).

    Parameters
    ----------
    data : result of MCMC run
    pointwise: bool, optional
        if True the pointwise predictive accuracy will be returned. Defaults to False
    reff : float, optional
        Relative MCMC efficiency, `effective_n / n` i.e. number of effective samples divided by
        the number of actual samples. Computed from trace by default.
    scale : str
        Output scale for loo. Available options are:

        - `deviance` : (default) -2 * (log-score)
        - `log` : 1 * log-score (after Vehtari et al. (2017))
        - `negative_log` : -1 * (log-score)

    Returns
    -------
    pandas.Series with the following columns:
    loo: approximated Leave-one-out cross-validation
    loo_se: standard error of loo
    p_loo: effective number of parameters
    shape_warn: 1 if the estimated shape parameter of
        Pareto distribution is greater than 0.7 for one or more samples
    loo_i: array of pointwise predictive accuracy, only if pointwise True
    pareto_k: array of Pareto shape values, only if pointwise True
    loo_scale: scale of the loo results
    """
    inference_data = convert_to_inference_data(data)
    for group in ("posterior", "sample_stats"):
        if not hasattr(inference_data, group):
            raise TypeError(
                "Must be able to extract a {group}" "group from data!".format(group=group)
            )
    if "log_likelihood" not in inference_data.sample_stats:
        raise TypeError("Data must include log_likelihood in sample_stats")
    posterior = inference_data.posterior
    log_likelihood = inference_data.sample_stats.log_likelihood
    n_samples = log_likelihood.chain.size * log_likelihood.draw.size
    new_shape = (n_samples,) + log_likelihood.shape[2:]
    log_likelihood = log_likelihood.values.reshape(*new_shape)

    if scale.lower() == "deviance":
        scale_value = -2
    elif scale.lower() == "log":
        scale_value = 1
    elif scale.lower() == "negative_log":
        scale_value = -1
    else:
        raise TypeError('Valid scale values are "deviance", "log", "negative_log"')

    if reff is None:
        n_chains = len(posterior.chain)
        if n_chains == 1:
            reff = 1.0
        else:
            ess = effective_sample_size(posterior)
            # this mean is over all data variables
            reff = np.hstack([ess[v].values.flatten() for v in ess.data_vars]).mean() / n_samples

    log_weights, pareto_shape = psislw(-log_likelihood, reff)
    log_weights += log_likelihood

    warn_mg = 0
    if np.any(pareto_shape > 0.7):
        warnings.warn(
            """Estimated shape parameter of Pareto distribution is greater than 0.7 for
        one or more samples. You should consider using a more robust model, this is because
        importance sampling is less likely to work well if the marginal posterior and LOO posterior
        are very different. This is more likely to happen with a non-robust model and highly
        influential observations."""
        )
        warn_mg = 1

    loo_lppd_i = scale_value * _logsumexp(log_weights, axis=0)
    loo_lppd = loo_lppd_i.sum()
    loo_lppd_se = (len(loo_lppd_i) * np.var(loo_lppd_i)) ** 0.5

    lppd = np.sum(_logsumexp(log_likelihood, axis=0, b_inv=log_likelihood.shape[0]))
    p_loo = lppd - loo_lppd / scale_value

    if pointwise:
        if np.equal(loo_lppd, loo_lppd_i).all():  # pylint: disable=no-member
            warnings.warn(
                """The point-wise LOO is the same with the sum LOO, please double check
                          the Observed RV in your model to make sure it returns element-wise logp.
                          """
            )
        return pd.Series(
            data=[loo_lppd, loo_lppd_se, p_loo, warn_mg, loo_lppd_i, pareto_shape, scale],
            index=["loo", "loo_se", "p_loo", "warning", "loo_i", "pareto_k", "loo_scale"],
        )

    else:
        return pd.Series(
            data=[loo_lppd, loo_lppd_se, p_loo, warn_mg, scale],
            index=["loo", "loo_se", "p_loo", "warning", "loo_scale"],
        )