use biased autocorrelation estimate by default (resolves #1785)

numpyro/diagnostics.py

@@ -96,12 +96,13 @@ def _fft_next_fast_len(target):
         target += 1
 
 
-def autocorrelation(x, axis=0):
+def autocorrelation(x, axis=0, unbiased=False):
     """
     Computes the autocorrelation of samples at dimension ``axis``.
 
     :param numpy.ndarray x: the input array.
     :param int axis: the dimension to calculate autocorrelation.
+    :param unbiased: use an unbiased estimator of the autocorrelation.
     :return: autocorrelation of ``x``.
     :rtype: numpy.ndarray
     """
@@ -127,7 +128,13 @@ def autocorrelation(x, axis=0):
 
     # truncate and normalize the result, then transpose back to original shape
     autocorr = autocorr[..., :N]
-    autocorr = autocorr / np.arange(N, 0.0, -1)
+
+    # the unbiased estimator is known to have "wild" tails, due to few samples at longer lags.
+    # see Geyer (1992) and Priestley (1981) for a discussion. also note that it is only strictly
+    # unbiased when the mean is known, whereas we it estimate from samples here.
+    if unbiased:
+        autocorr = autocorr / np.arange(N, 0.0, -1)
+
     with np.errstate(invalid="ignore", divide="ignore"):
         autocorr = autocorr / autocorr[..., :1]
     return np.swapaxes(autocorr, axis, -1)

test/test_diagnostics.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import numpy as np
-from numpy.testing import assert_allclose
+from numpy.testing import assert_, assert_allclose
 import pytest
 from scipy.fftpack import next_fast_len
 
@@ -60,10 +60,22 @@ def test_hpdi():
 
 def test_autocorrelation():
     x = np.arange(10.0)
-    actual = autocorrelation(x)
+    actual = autocorrelation(x, unbiased=True)
     expected = np.array([1, 0.78, 0.52, 0.21, -0.13, -0.52, -0.94, -1.4, -1.91, -2.45])
     assert_allclose(actual, expected, atol=0.01)
 
+    actual = autocorrelation(x, unbiased=False)
+    expected = expected * np.arange(len(x), 0.0, -1) / len(x)
+    assert_allclose(actual, expected, atol=0.01)
+
+    # the unbiased estimator has variance O(1) at large lags
+    x = np.random.normal(size=20000)
+    ac = autocorrelation(x, unbiased=True)
+    assert_(np.any(np.abs(ac[-100:]) > 0.01))
+
+    ac = autocorrelation(x, unbiased=False)
+    assert_allclose(np.abs(ac[-100:]), 0.0, atol=0.01)
+
 
 def test_autocovariance():
     x = np.arange(10.0)