pytransc.utils.autocorr

  1"""Autocorrelation time estimation utilities.
  2
  3This module implements autocorrelation time calculation routines following
  4Goodman & Weare (2010) and improved methods suggested in the emcee documentation.
  5These functions are used for assessing MCMC chain convergence and determining
  6appropriate thinning intervals.
  7"""
  8
  9import numpy as np
 10import numpy.typing as npt
 11
 12from .types import FloatArray
 13
 14
 15def autocorr_gw2010(y: FloatArray, c: float = 5.0) -> float:
 16    """Calculate autocorrelation time using Goodman & Weare (2010) method.
 17
 18    Parameters
 19    ----------
 20    y : FloatArray
 21        MCMC chain data with shape (n_walkers, n_steps) or (n_steps,).
 22    c : float, optional
 23        Window size factor for automatic windowing. Default is 5.0.
 24
 25    Returns
 26    -------
 27    float
 28        Autocorrelation time estimate.
 29
 30    References
 31    ----------
 32    Goodman, J. & Weare, J. (2010). Ensemble samplers with affine invariance.
 33    Communications in Applied Mathematics and Computational Science, 5(1), 65-80.
 34    """
 35    f = autocorr_func_1d(np.mean(y, axis=0))
 36    taus = 2.0 * np.cumsum(f) - 1.0
 37    window = auto_window(taus, c)
 38    return taus[window]
 39
 40
 41def autocorr_fardal(y: FloatArray, c: float = 5.0) -> float:
 42    """Calculate autocorrelation time using improved estimator.
 43
 44    This function implements an improved autocorrelation time estimator
 45    that averages the autocorrelation function across all walkers before
 46    calculating the integrated time, as recommended in the emcee documentation.
 47
 48    Parameters
 49    ----------
 50    y : FloatArray
 51        MCMC chain data.
 52    c : float, optional
 53        Window size factor for automatic windowing. Default is 5.0.
 54
 55    Returns
 56    -------
 57    float
 58        Autocorrelation time estimate.
 59
 60    References
 61    ----------
 62    https://emcee.readthedocs.io/en/stable/tutorials/autocorr/
 63    """
 64    f = np.zeros(y.shape[1])
 65    for yy in y:
 66        f += autocorr_func_1d(yy)
 67    f /= len(y)
 68    taus = 2.0 * np.cumsum(f) - 1.0
 69    window = auto_window(taus, c)
 70    return taus[window]
 71
 72
 73def next_pow_two(n: int) -> int:
 74    """Get the next power of two greater than or equal to n.
 75
 76    Parameters
 77    ----------
 78    n : int
 79        Input number.
 80
 81    Returns
 82    -------
 83    int
 84        Next power of two >= n.
 85    """
 86    i = 1
 87    while i < n:
 88        i = i << 1
 89    return i
 90
 91
 92def autocorr_func_1d(x: npt.ArrayLike, norm: bool = True) -> FloatArray:
 93    """Calculate 1D autocorrelation function using FFT.
 94
 95    Parameters
 96    ----------
 97    x : array_like
 98        1D input sequence.
 99    norm : bool, optional
100        Whether to normalize by the zero-lag value. Default is True.
101
102    Returns
103    -------
104    FloatArray
105        Autocorrelation function values.
106
107    Raises
108    ------
109    ValueError
110        If input is not 1-dimensional.
111    """
112    x = np.atleast_1d(x)
113    if len(x.shape) != 1:
114        raise ValueError("invalid dimensions for 1D autocorrelation function")
115    n = next_pow_two(len(x))
116
117    # Compute the FFT and then (from that) the auto-correlation function
118    f = np.fft.fft(x - np.mean(x), n=2 * n)
119    acf = np.fft.ifft(f * np.conjugate(f))[: len(x)].real
120    acf /= 4 * n
121
122    # Optionally normalize
123    if norm:
124        acf /= acf[0]
125
126    return acf
127
128
129def auto_window(taus: FloatArray, c: float) -> int:
130    """Determine optimal windowing for autocorrelation time calculation.
131
132    This function implements the automatic windowing procedure following
133    Sokal (1989) to determine the appropriate cutoff for integrated
134    autocorrelation time calculations.
135
136    Parameters
137    ----------
138    taus : FloatArray
139        Array of cumulative autocorrelation times.
140    c : float
141        Window size factor. The window is determined as the first index
142        where c * tau[i] < i.
143
144    Returns
145    -------
146    int
147        Optimal window index for autocorrelation time calculation.
148
149    References
150    ----------
151    Sokal, A. (1989). Monte Carlo methods in statistical mechanics: foundations
152    and new algorithms. NATO Advanced Science Institutes Series E, 188, 131-192.
153    """
154    m = np.arange(len(taus)) < c * taus
155    if np.any(m):
156        return int(np.argmin(m))
157    return len(taus) - 1