Update diagnostics (#1366)

* update mcse indexing * remove rint func * add median to mcse * update tests * fix rank transformation * Update backtransformation docstring * Fix reference * update reference tests * fix whitespace in docstring * update changelog * fix mcse median * fix lint * pandas read with high accuracy * add float accuracy to rcparams * use 1e-8 as a limit for diagnostics Co-authored-by: ahartikainen <ahartikainen@github.com>
arviz-devs · Sep 2, 2020 · 11068d5 · 11068d5
1 parent b7221d7
commit 11068d5
Show file tree

Hide file tree

Showing 10 changed files with 172 additions and 135 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@
 * Replaced `_fast_kde()` with `kde()` which now also supports circular variables via the argument `circular` ([1284](https://github.com/arviz-devs/arviz/pull/1284)).
 * Increased `from_pystan` attrs information content ([1353](https://github.com/arviz-devs/arviz/pull/1353))
 * Allow `plot_trace` to return and accept axes ([#1361](https://github.com/arviz-devs/arviz/pull/1361))
+* Update diagnostics to be on par with posterior package ([#1366](https://github.com/arviz-devs/arviz/pull/1366))
 
 ### Deprecation
 

diff --git a/arviz/data/io_cmdstan.py b/arviz/data/io_cmdstan.py
@@ -626,7 +626,9 @@ def _read_output_file(path):
                 comments.append(line.decode("utf-8").strip())
 
     with open(path, "rb") as f_obj:
-        data = pd.read_csv(f_obj, comment="#")
+        data = pd.read_csv(
+            f_obj, comment="#", float_precision=rcParams["data.pandas_float_precision"]
+        )
 
     return data, comments
 

diff --git a/arviz/rcparams.py b/arviz/rcparams.py
@@ -228,6 +228,10 @@ def validate_iterable(value):
     "data.metagroups": (METAGROUPS, _validate_dict_of_lists),
     "data.index_origin": (0, _make_validate_choice({0, 1}, typeof=int)),
     "data.save_warmup": (False, _validate_boolean),
+    "data.pandas_float_precision": (
+        "high",
+        _make_validate_choice({"high", "round_trip"}, allow_none=True),
+    ),
     "plot.backend": ("matplotlib", _make_validate_choice({"matplotlib", "bokeh"})),
     "plot.max_subplots": (40, _validate_positive_int_or_none),
     "plot.point_estimate": (

diff --git a/arviz/stats/diagnostics.py b/arviz/stats/diagnostics.py
@@ -14,7 +14,6 @@
 from .stats_utils import autocov as _autocov
 from .stats_utils import not_valid as _not_valid
 from .stats_utils import quantile as _quantile
-from .stats_utils import rint as _rint
 from .stats_utils import stats_variance_2d as svar
 from .stats_utils import wrap_xarray_ufunc as _wrap_xarray_ufunc
 
@@ -321,6 +320,7 @@ def mcse(data, *, var_names=None, method="mean", prob=None):
         Select mcse method. Valid methods are:
         - "mean"
         - "sd"
+        - "median"
         - "quantile"
 
     prob : float
@@ -345,10 +345,15 @@ def mcse(data, *, var_names=None, method="mean", prob=None):
 
     .. ipython::
 
-        In [1]: az.mcse(data, method="quantile", prob=.7)
+        In [1]: az.mcse(data, method="quantile", prob=0.7)
 
     """
-    methods = {"mean": _mcse_mean, "sd": _mcse_sd, "quantile": _mcse_quantile}
+    methods = {
+        "mean": _mcse_mean,
+        "sd": _mcse_sd,
+        "median": _mcse_median,
+        "quantile": _mcse_quantile,
+    }
     if method not in methods:
         raise TypeError(
             "mcse method {} not found. Valid methods are:\n{}".format(
@@ -530,6 +535,29 @@ def _bfmi(energy):
     return num / den
 
 
+def _backtransform_ranks(arr, c=3 / 8):  # pylint: disable=invalid-name
+    """Backtransformation of ranks.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Ranks array
+    c : float
+        Fractional offset. Defaults to c = 3/8 as recommended by Blom (1958).
+
+    Returns
+    -------
+    np.ndarray
+
+    References
+    ----------
+    Blom, G. (1958). Statistical Estimates and Transformed Beta-Variables. Wiley; New York.
+    """
+    arr = np.asarray(arr)
+    size = arr.size
+    return (arr - c) / (size - 2 * c + 1)
+
+
 def _z_scale(ary):
     """Calculate z_scale.
 
@@ -542,9 +570,9 @@ def _z_scale(ary):
     np.ndarray
     """
     ary = np.asarray(ary)
-    size = ary.size
     rank = stats.rankdata(ary, method="average")
-    z = stats.norm.ppf((rank - 0.5) / size)
+    rank = _backtransform_ranks(rank)
+    z = stats.norm.ppf(rank)
     z = z.reshape(ary.shape)
     return z
 
@@ -811,26 +839,6 @@ def _ess_identity(ary, relative=False):
     return _ess(ary, relative=relative)
 
 
-def _conv_quantile(ary, prob):
-    """Return mcse, Q05, Q95, Seff."""
-    ary = np.asarray(ary)
-    if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
-        return np.nan, np.nan, np.nan, np.nan
-    ess = _ess_quantile(ary, prob)
-    probability = [0.1586553, 0.8413447, 0.05, 0.95]
-    with np.errstate(invalid="ignore"):
-        ppf = stats.beta.ppf(probability, ess * prob + 1, ess * (1 - prob) + 1)
-    sorted_ary = np.sort(ary.ravel())
-    size = sorted_ary.size
-    ppf_size = ppf * size - 1
-    th1 = sorted_ary[_rint(np.nanmax((ppf_size[0], 0)))]
-    th2 = sorted_ary[_rint(np.nanmin((ppf_size[1], size - 1)))]
-    mcse_quantile = (th2 - th1) / 2
-    th1 = sorted_ary[_rint(np.nanmax((ppf_size[2], 0)))]
-    th2 = sorted_ary[_rint(np.nanmin((ppf_size[3], size - 1)))]
-    return mcse_quantile, th1, th2, ess
-
-
 def _mcse_mean(ary):
     """Compute the Markov Chain mean error."""
     _numba_flag = Numba.numba_flag
@@ -862,13 +870,26 @@ def _mcse_sd(ary):
     return mcse_sd_value
 
 
+def _mcse_median(ary):
+    """Compute the Markov Chain median error."""
+    return _mcse_quantile(ary, 0.5)
+
+
 def _mcse_quantile(ary, prob):
     """Compute the Markov Chain quantile error at quantile=prob."""
     ary = np.asarray(ary)
     if _not_valid(ary, shape_kwargs=dict(min_draws=4, min_chains=1)):
         return np.nan
-    mcse_q, *_ = _conv_quantile(ary, prob)
-    return mcse_q
+    ess = _ess_quantile(ary, prob)
+    probability = [0.1586553, 0.8413447]
+    with np.errstate(invalid="ignore"):
+        ppf = stats.beta.ppf(probability, ess * prob + 1, ess * (1 - prob) + 1)
+    sorted_ary = np.sort(ary.ravel())
+    size = sorted_ary.size
+    ppf_size = ppf * size - 1
+    th1 = sorted_ary[int(np.floor(np.nanmax((ppf_size[0], 0))))]
+    th2 = sorted_ary[int(np.ceil(np.nanmin((ppf_size[1], size - 1))))]
+    return (th2 - th1) / 2
 
 
 def _mc_error(ary, batches=5, circular=False):

diff --git a/arviz/stats/stats_utils.py b/arviz/stats/stats_utils.py
@@ -318,12 +318,6 @@ def logsumexp(ary, *, b=None, b_inv=None, axis=None, keepdims=False, out=None, c
     return out if out.shape else dtype(out)
 
 
-def rint(num):
-    """Round and change to ingeter."""
-    rnum = np.rint(num)  # pylint: disable=assignment-from-no-return
-    return int(rnum)
-
-
 def quantile(ary, q, axis=None, limit=None):
     """Use same quantile function as R (Type 7)."""
     if limit is None:

diff --git a/arviz/tests/base_tests/test_diagnostics.py b/arviz/tests/base_tests/test_diagnostics.py
@@ -5,17 +5,17 @@
 import numpy as np
 import pandas as pd
 import pytest
-from numpy.testing import assert_almost_equal, assert_array_almost_equal
+from numpy.testing import assert_almost_equal
 
 from ...data import from_cmdstan, load_arviz_data
 from ...plots.plot_utils import xarray_var_iter
-from ...rcparams import rcParams
+from ...rcparams import rc_context, rcParams
 from ...stats import bfmi, ess, geweke, mcse, rhat
 from ...stats.diagnostics import (
-    _conv_quantile,
     _ess,
     _ess_quantile,
     _mc_error,
+    _mcse_quantile,
     _multichain_statistics,
     _rhat,
     _rhat_rank,
@@ -54,64 +54,79 @@ def test_bfmi_dataset_bad(self):
 
     def test_deterministic(self):
         """
-        Test algorithm against RStan monitor.R functions.
-        monitor.R :
-        https://github.com/stan-dev/rstan/blob/425d195565c4d9bcbcb8cccf513e140e6908ca62/rstan/rstan/R/monitor.R
+        Test algorithm against posterior (R) convergence functions.
+
+        posterior: https://github.com/stan-dev/posterior
         R code:
         ```
-        source('~/monitor.R')
+        library("posterior")
         data2 <- read.csv("blocker.2.csv", comment.char = "#")
         data1 <- read.csv("blocker.1.csv", comment.char = "#")
-        output <- matrix(ncol=15, nrow=length(names(data1))-4)
+        output <- matrix(ncol=17, nrow=length(names(data1))-4)
         j = 0
         for (i in 1:length(names(data1))) {
-          name = names(data1)[i]
-          ary = matrix(c(data1[,name], data2[,name]), 1000, 2)
-          if (!endsWith(name, "__"))
-              j <- j + 1
-              output[j,] <- c(
-                rhat(ary),
-                rhat_rfun(ary),
-                ess_bulk(ary),
-                ess_tail(ary),
-                ess_mean(ary),
-                ess_sd(ary),
-                ess_rfun(ary),
-                ess_quantile(ary, 0.01),
-                ess_quantile(ary, 0.1),
-                ess_quantile(ary, 0.3),
-                mcse_mean(ary),
-                mcse_sd(ary),
-                mcse_quantile(ary, prob=0.01),
-                mcse_quantile(ary, prob=0.1),
-                mcse_quantile(ary, prob=0.3))
-            }
+            name = names(data1)[i]
+            ary = matrix(c(data1[,name], data2[,name]), 1000, 2)
+            if (!endsWith(name, "__"))
+                j <- j + 1
+                output[j,] <- c(
+                    posterior::rhat(ary),
+                    posterior::rhat_basic(ary, FALSE),
+                    posterior::ess_bulk(ary),
+                    posterior::ess_tail(ary),
+                    posterior::ess_mean(ary),
+                    posterior::ess_sd(ary),
+                    posterior::ess_median(ary),
+                    posterior::ess_basic(ary, FALSE),
+                    posterior::ess_quantile(ary, 0.01),
+                    posterior::ess_quantile(ary, 0.1),
+                    posterior::ess_quantile(ary, 0.3),
+                    posterior::mcse_mean(ary),
+                    posterior::mcse_sd(ary),
+                    posterior::mcse_median(ary),
+                    posterior::mcse_quantile(ary, prob=0.01),
+                    posterior::mcse_quantile(ary, prob=0.1),
+                    posterior::mcse_quantile(ary, prob=0.3))
+        }
         df = data.frame(output, row.names = names(data1)[5:ncol(data1)])
         colnames(df) <- c("rhat_rank",
                           "rhat_raw",
                           "ess_bulk",
                           "ess_tail",
                           "ess_mean",
                           "ess_sd",
+                          "ess_median",
                           "ess_raw",
                           "ess_quantile01",
                           "ess_quantile10",
                           "ess_quantile30",
                           "mcse_mean",
                           "mcse_sd",
+                          "mcse_median",
                           "mcse_quantile01",
                           "mcse_quantile10",
                           "mcse_quantile30")
-        write.csv(df, "reference_values.csv")
+        write.csv(df, "reference_posterior.csv")
         ```
+        Reference file:
+
+        Created: 2020-08-31
+        System: Ubuntu 18.04.5 LTS
+        R version 4.0.2 (2020-06-22)
+        posterior 0.1.2
         """
         # download input files
         here = os.path.dirname(os.path.abspath(__file__))
         data_directory = os.path.join(here, "..", "saved_models")
         path = os.path.join(data_directory, "stan_diagnostics", "blocker.[0-9].csv")
-        posterior = from_cmdstan(path)
-        reference_path = os.path.join(data_directory, "stan_diagnostics", "reference_values.csv")
-        reference = pd.read_csv(reference_path, index_col=0).sort_index(axis=1).sort_index(axis=0)
+        with rc_context(rc={"data.pandas_float_precision": "high"}):
+            posterior = from_cmdstan(path)
+        reference_path = os.path.join(data_directory, "stan_diagnostics", "reference_posterior.csv")
+        reference = (
+            pd.read_csv(reference_path, index_col=0, float_precision="high")
+            .sort_index(axis=1)
+            .sort_index(axis=0)
+        )
         # test arviz functions
         funcs = {
             "rhat_rank": lambda x: rhat(x, method="rank"),
@@ -120,12 +135,14 @@ def test_deterministic(self):
             "ess_tail": lambda x: ess(x, method="tail"),
             "ess_mean": lambda x: ess(x, method="mean"),
             "ess_sd": lambda x: ess(x, method="sd"),
+            "ess_median": lambda x: ess(x, method="median"),
             "ess_raw": lambda x: ess(x, method="identity"),
             "ess_quantile01": lambda x: ess(x, method="quantile", prob=0.01),
             "ess_quantile10": lambda x: ess(x, method="quantile", prob=0.1),
             "ess_quantile30": lambda x: ess(x, method="quantile", prob=0.3),
             "mcse_mean": lambda x: mcse(x, method="mean"),
             "mcse_sd": lambda x: mcse(x, method="sd"),
+            "mcse_median": lambda x: mcse(x, method="median"),
             "mcse_quantile01": lambda x: mcse(x, method="quantile", prob=0.01),
             "mcse_quantile10": lambda x: mcse(x, method="quantile", prob=0.1),
             "mcse_quantile30": lambda x: mcse(x, method="quantile", prob=0.3),
@@ -136,21 +153,19 @@ def test_deterministic(self):
                 key = key + ".{}".format(list(coord_dict.values())[0] + 1)
             results[key] = {func_name: func(vals) for func_name, func in funcs.items()}
         arviz_data = pd.DataFrame.from_dict(results).T.sort_index(axis=1).sort_index(axis=0)
+
         # check column names
         assert set(arviz_data.columns) == set(reference.columns)
+
         # check parameter names
         assert set(arviz_data.index) == set(reference.index)
-        # check equality (rhat_rank has accuracy < 6e-5, atleast with this data, R vs Py)
-        # this is due to numerical accuracy in calculation leading to rankdata
-        # function, which scales minimal difference to larger scale
-        # test first with numpy
-        assert_array_almost_equal(reference, arviz_data, decimal=4)
-        # then test manually (more strict)
-        assert (abs(reference["rhat_rank"] - arviz_data["rhat_rank"]) < 6e-5).all(None)
-        assert abs(np.median(reference["rhat_rank"] - arviz_data["rhat_rank"]) < 1e-14).all(None)
-        not_rhat = [col for col in reference.columns if col != "rhat_rank"]
-        assert (abs((reference[not_rhat] - arviz_data[not_rhat])).values < 1e-8).all(None)
-        assert abs(np.median(reference[not_rhat] - arviz_data[not_rhat]) < 1e-14).all(None)
+
+        # show print with pytests '-s' tag
+        np.set_printoptions(16)
+        print(abs(reference - arviz_data).max())
+
+        # test absolute accuracy
+        assert (abs(reference - arviz_data).values < 1e-8).all(None)
 
     @pytest.mark.parametrize("method", ("rank", "split", "folded", "z_scale", "identity"))
     @pytest.mark.parametrize("var_names", (None, "mu", ["mu", "tau"]))
@@ -350,7 +365,7 @@ def test_effective_sample_size_dataset(self, data, method, var_names, relative):
             ess_hat = ess(data, var_names=var_names, method=method, relative=relative)
         assert np.all(ess_hat.mu.values > n_low)  # This might break if the data is regenerated
 
-    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "quantile"))
+    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "median", "quantile"))
     def test_mcse_array(self, mcse_method):
         if mcse_method == "quantile":
             mcse_hat = mcse(np.random.randn(4, 100), method=mcse_method, prob=0.34)
@@ -362,7 +377,7 @@ def test_mcse_ndarray(self):
         with pytest.raises(TypeError):
             mcse(np.random.randn(2, 300, 10))
 
-    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "quantile"))
+    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "median", "quantile"))
     @pytest.mark.parametrize("var_names", (None, "mu", ["mu", "tau"]))
     def test_mcse_dataset(self, data, mcse_method, var_names):
         if mcse_method == "quantile":
@@ -371,7 +386,7 @@ def test_mcse_dataset(self, data, mcse_method, var_names):
             mcse_hat = mcse(data, var_names=var_names, method=mcse_method)
         assert mcse_hat  # This might break if the data is regenerated
 
-    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "quantile"))
+    @pytest.mark.parametrize("mcse_method", ("mean", "sd", "median", "quantile"))
     @pytest.mark.parametrize("chain", (None, 1, 2))
     @pytest.mark.parametrize("draw", (1, 2, 3, 4))
     @pytest.mark.parametrize("use_nan", (True, False))
@@ -515,12 +530,12 @@ def test_mc_error_nan(self, size, ndim):
         else:
             assert np.isnan(_mc_error(x))
 
-    @pytest.mark.parametrize("func", ("_conv_quantile", "_z_scale"))
+    @pytest.mark.parametrize("func", ("_mcse_quantile", "_z_scale"))
     def test_nan_behaviour(self, func):
         data = np.random.randn(100, 4)
         data[0, 0] = np.nan  #  pylint: disable=unsupported-assignment-operation
-        if func == "_conv_quantile":
-            assert np.isnan(_conv_quantile(data, 0.5)).all(None)
+        if func == "_mcse_quantile":
+            assert np.isnan(_mcse_quantile(data, 0.5)).all(None)
         else:
             assert not np.isnan(_z_scale(data)).all(None)
             assert not np.isnan(_z_scale(data)).any(None)