Merge 1e596a7 into 3bcc67b

arviz/stats/stats.py

@@ -5,6 +5,7 @@
 import pandas as pd
 from scipy.special import logsumexp
 import scipy.stats as st
+from scipy.signal import savgol_filter
 from scipy.optimize import minimize
 import xarray as xr
 
@@ -221,7 +222,14 @@ def _ic_matrix(ics, ic_i):
     return rows, cols, ic_i_val
 
 
-def hpd(x, credible_interval=0.94, transform=lambda x: x, circular=False):
+def hpd(
+    x,
+    credible_interval=0.94,
+    smooth=False,
+    transform=lambda x: x,
+    circular=False,
+    smooth_kwargs=None,
+):
     """
     Calculate highest posterior density (HPD) of array for given credible_interval.
 
@@ -236,24 +244,42 @@ def hpd(x, credible_interval=0.94, transform=lambda x: x, circular=False):
         Credible interval to plot. Defaults to 0.94.
     transform : callable
         Function to transform data (defaults to identity)
+    smooth : boolean
+        If True the result will be smoothed using the Savitzky-Golay filter. Defaults to False.
     circular : bool, optional
         Whether to compute the error taking into account `x` is a circular variable
         (in the range [-np.pi, np.pi]) or not. Defaults to False (i.e non-circular variables).
+    smooth_kwargs : dict, optional
+        Additional keywords modifying the Savitzky-Golay filter. See Scipy's documentation for
+        details
 
     Returns
     -------
     np.ndarray
         lower and upper value of the interval.
     """
     if x.ndim > 1:
-        return np.array(
+        hpd_array = np.array(
             [
                 hpd(
                     row, credible_interval=credible_interval, transform=transform, circular=circular
                 )
                 for row in x.T
             ]
         )
+        if smooth:
+            if smooth_kwargs is None:
+                smooth_kwargs = {}
+            window = x.shape[1] // 3
+            order = min(3, window - 1)
+            if window % 2 == 0:
+                window += 1
+            smooth_kwargs.setdefault("window_length", window)
+            smooth_kwargs.setdefault("polyorder", order)
+            smooth_kwargs.setdefault("mode", "mirror")
+            smooth_kwargs.setdefault("axis", 0)
+            hpd_array = savgol_filter(x=hpd_array, **smooth_kwargs)
+        return hpd_array
     # Make a copy of trace
     x = transform(x.copy())
     len_x = len(x)

arviz/tests/test_stats.py

@@ -32,6 +32,13 @@ def test_hpd():
     assert_array_almost_equal(interval, [-1.88, 1.88], 2)
 
 
+def test_hpd_smooth():
+    normal_sample = np.random.randn(500, 50)
+    hpd_s = hpd(normal_sample, smooth=True).var(0)
+    hpd_ns = hpd(normal_sample, smooth=False).var(0)
+    np.testing.assert_array_less(hpd_s, hpd_ns)
+
+
 def test_r2_score():
     x = np.linspace(0, 1, 100)
     y = np.random.normal(x, 1)