Repository: arviz-devs/arviz

Issue:
arviz's psislw differs from loo's psis
ArviZ's implementation of PSIS sometimes agrees with the loo package's implementation and sometimes does not. Here's an example where `arviz.psislw` and `loo.psis` disagree:

```python
import arviz, rpy2
import numpy as np
import scipy as sp
from rpy2.robjects import numpy2ri, packages
numpy2ri.activate()

# install and load loo
utils = packages.importr('utils')
utils.install_packages('loo', repos="https://cloud.r-project.org")
loo = packages.importr('loo')

# random vector
rng = np.random.default_rng(44)
x = rng.normal(size=100)

# compute PSIS
x_smoothed_arviz, k_arviz = arviz.psislw(x.copy())
result_loo = loo.psis(x.reshape(1, 100, 1).copy())
x_smoothed_loo = result_loo[0].ravel()
k_loo = result_loo[1][0]

# check the shapes are the same
np.isclose(k_arviz, k_loo) # True

# check that weights are not
# loo doesn't normalize the log weights; arviz does
np.all(np.isclose(x_smoothed_arviz, x_smoothed_loo - sp.special.logsumexp(x_smoothed_loo)))  # False
np.all(np.isclose(x_smoothed_arviz, x - sp.special.logsumexp(x))) # True
```

So ArviZ and loo agree on the Pareto shape diagnostic, but ArviZ doesn't smooth the weights at all. If we change the seed to 42 or 43, then ArviZ does smooth the weights. The issue seems to be that `k` here is approximately 0.22, which falls below the `k_min=1/3` threshold below which ArviZ does not smooth the weights:
https://github.com/arviz-devs/arviz/blob/afd895cb0fde11bbcaa633bb363dcb0a0caaa22f/arviz/stats/stats.py#L1001-L1012

This check is absent from loo's implementation (https://github.com/stan-dev/loo/blob/409de413f56a453d6d4e5a096efceb8388b2d39b/R/psis.R#L254) and the PSIS paper but present in the reference Python and Matlab implementations of PSIS (https://github.com/avehtari/PSIS/blob/cac66c0d688e06c0e69d324b7c31db23fb138b51/py/psis.py#L187).

I propose we replace this check with `np.isfinite(k)`.


Discussion:

Resolving patch:
diff --git a/arviz/stats/stats.py b/arviz/stats/stats.py
--- a/arviz/stats/stats.py
+++ b/arviz/stats/stats.py
@@ -933,13 +933,12 @@ def psislw(log_weights, reff=1.0):
     # precalculate constants
     cutoff_ind = -int(np.ceil(min(n_samples / 5.0, 3 * (n_samples / reff) ** 0.5))) - 1
     cutoffmin = np.log(np.finfo(float).tiny)  # pylint: disable=no-member, assignment-from-no-return
-    k_min = 1.0 / 3
 
     # create output array with proper dimensions
     out = tuple([np.empty_like(log_weights), np.empty(shape)])
 
     # define kwargs
-    func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "k_min": k_min, "out": out}
+    func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "out": out}
     ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
     kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
     log_weights, pareto_shape = _wrap_xarray_ufunc(
@@ -956,7 +955,7 @@ def psislw(log_weights, reff=1.0):
     return log_weights, pareto_shape
 
 
-def _psislw(log_weights, cutoff_ind, cutoffmin, k_min=1.0 / 3):
+def _psislw(log_weights, cutoff_ind, cutoffmin):
     """
     Pareto smoothed importance sampling (PSIS) for a 1D vector.
 
@@ -998,8 +997,8 @@ def _psislw(log_weights, cutoff_ind, cutoffmin, k_min=1.0 / 3):
         x_tail = np.exp(x_tail) - expxcutoff
         k, sigma = _gpdfit(x_tail[x_tail_si])
 
-        if k >= k_min:
-            # no smoothing if short tail or GPD fit failed
+        if np.isfinite(k):
+            # no smoothing if GPD fit failed
             # compute ordered statistic for the fit
             sti = np.arange(0.5, tail_len) / tail_len
             smoothed_tail = _gpinv(sti, k, sigma)


Test patch:
diff --git a/arviz/tests/base_tests/test_stats.py b/arviz/tests/base_tests/test_stats.py
--- a/arviz/tests/base_tests/test_stats.py
+++ b/arviz/tests/base_tests/test_stats.py
@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_allclose, assert_array_almost_equal, assert_array_equal
+from scipy.special import logsumexp
 from scipy.stats import linregress
 from xarray import DataArray, Dataset
 
@@ -534,6 +535,16 @@ def test_psislw(centered_eight):
     assert_allclose(pareto_k, psislw(-log_likelihood, 0.7)[1])
 
 
+def test_psislw_smooths_for_low_k():
+    # check that log-weights are smoothed even when k < 1/3
+    # https://github.com/arviz-devs/arviz/issues/2010
+    rng = np.random.default_rng(44)
+    x = rng.normal(size=100)
+    x_smoothed, k = psislw(x.copy())
+    assert k < 1 / 3
+    assert not np.allclose(x - logsumexp(x), x_smoothed)
+
+
 @pytest.mark.parametrize("probs", [True, False])
 @pytest.mark.parametrize("kappa", [-1, -0.5, 1e-30, 0.5, 1])
 @pytest.mark.parametrize("sigma", [0, 2])