diagnostics.py 文件源码

def effective_sample_size_1d(samples):
    """
    Compute the effective sample size of a chain of scalar samples.

    :param samples: A 1-D numpy array. The chain of samples.
    :return: A float. The effective sample size.
    """
    n = samples.shape[0]
    mu_hat = np.mean(samples)
    var = np.var(samples) * n / (n - 1)
    var_plus = var * (n - 1) / n

    def auto_covariance(lag):
        return np.mean((samples[:n - lag] - mu_hat) * (samples[lag:] - mu_hat))

    sum_rho = 0
    for t in range(0, n):
        rho = 1 - (var - auto_covariance(t)) / var_plus
        if rho < 0:
            break
        sum_rho += rho

    ess = n / (1 + 2 * sum_rho)
    return ess