<jupyter_start><jupyter_text>Effective sample size in PyStanReference implementation in PyStan: [pystan/_chains.pyx](https://github.com/stan-dev/pystan/blob/develop/pystan/_chains.pyx)(related PR: https://github.com/stan-dev/pystan/pull/415)<jupyter_code>f1 = '/usr/local/lib/python3.5/dist-packages/pystan/tests/data/blocker.1.csv'
f2 = '/usr/local/lib/python3.5/dist-packages/pystan/tests/data/blocker.2.csv'

# f1 = '/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/pystan/tests/data/blocker.1.csv'
# f2 = '/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-packages/pystan/tests/data/blocker.2.csv'

# read csv using numpy
c1 = np.loadtxt(f1, skiprows=41, delimiter=',')[:, 4:]
c1_colnames = open(f1, 'r').readlines()[36].strip().split(',')[4:]
np.testing.assert_equal(c1_colnames[0], 'd')
c2 = np.loadtxt(f2, skiprows=41, delimiter=',')[:, 4:]
c2_colnames = open(f2, 'r').readlines()[36].strip().split(',')[4:]
np.testing.assert_equal(c1_colnames, c2_colnames)
np.testing.assert_equal(len(c1_colnames), c1.shape[1])

n_samples = len(c1)
np.testing.assert_equal(n_samples, 1000)

c1 = OrderedDict((k, v) for k, v in zip(c1_colnames, c1.T))
c2 = OrderedDict((k, v) for k, v in zip(c2_colnames, c2.T))

lst = dict(fnames_oi=c1_colnames, samples=[{'chains': c1}, {'chains': c2}],
           n_save=np.repeat(n_samples, 2), permutation=None,
           warmup=0, warmup2=[0, 0], chains=2, n_flatnames=len(c1))
n_eff = [
    466.099, 136.953, 1170.390, 541.256,
    518.051, 589.244, 764.813, 688.294,
    323.777, 502.892, 353.823, 588.142,
    654.336, 480.914, 176.978, 182.649,
    642.389, 470.949, 561.947, 581.187,
    446.389, 397.641, 338.511, 678.772,
    1442.250, 837.956, 869.865, 951.124,
    619.336, 875.805, 233.260, 786.568,
    910.144, 231.582, 907.666, 747.347,
    720.660, 195.195, 944.547, 767.271,
    723.665, 1077.030, 470.903, 954.924,
    497.338, 583.539, 697.204, 98.421
]

ess = []
for i in range(len(n_eff)):
    ess.append(pystan.chains.ess(lst, i))
    np.testing.assert_almost_equal(ess[i], n_eff[i], 2)
df_neff = pd.DataFrame(data=dict(Target=n_eff, PyStan=ess),
                       columns=['Target', 'PyStan'])
df_neff
lst.keys()
slst = lst['samples'][0]['chains']
slst.keys()
param_names = list(slst.keys())
slst['d'].shape
from copy import deepcopy
sim = deepcopy(lst)
m = sim['chains']

ns_save = sim['n_save']
ns_warmup2 = sim['warmup2']
ns_kept = [s - w for s, w in zip(sim['n_save'], sim['warmup2'])]

n_samples = min(ns_kept)
for im in range(m):
    plt.plot(sim['samples'][im]['chains']['d']);<jupyter_output><empty_output><jupyter_text>Validate PyStan autocov implementation<jupyter_code># test autocovariance function
chain = 0
var_id = 5
stan_acov = pystan._chains._test_autocovariance(sim, chain, var_id)
len(stan_acov)
print(pystan.chains.ess(sim, var_id))
"""get_kept_samples function in pystan"""
# xtrace = []
# nv = slst[param_names[var_id]]
# for i in range(nv.shape[0] - ns_warmup2[chain]):
#     xtrace.append(nv[ns_warmup2[chain] + i])
# xtrace = np.asarray(xtrace)
varname = param_names[var_id]
xtrace = sim['samples'][chain]['chains'][varname]
pymc_acov = [pm.stats.autocov(xtrace, lag)[0, 1] for lag in range(1, len(xtrace))]
import scipy.fftpack as fft


def autocorr(x):
    """
    Compute autocorrelation using FFT
    https://en.wikipedia.org/wiki/Autocorrelation#Efficient_computation
    """
    x = np.asarray(x)
    x = x - x.mean()
    N = len(x)
    s = fft.fft(x, N * 2 - 1)
    result = np.real(fft.ifft(s * np.conjugate(s), N * 2 - 1))
    result = result[:N]
    result /= (N - np.arange(N))
    result /= result[0]
    return result


def autocov(x):
    acorr = autocorr(x)
    varx = np.var(x) * (len(x) - 1) / len(x)
    return acorr * varx


from scipy.signal import fftconvolve


def autocorr2(x):
    x -= x.mean()
    N = len(x)
    result = fftconvolve(x, x[::-1])
    result = result[len(result) // 2:]
    result /= (N - np.arange(N))
    result /= result[0]
    return result


def autocov2(x):
    acorr = autocorr2(x)
    varx = np.var(x) * (len(x) - 1) / len(x)
    return acorr * varx


acov_ = autocov2(xtrace)
acov_.shape
# plt.plot(autocorr(xtrace))
# plt.plot(autocorr2(xtrace));
plt.plot(autocorr(xtrace)-autocorr2(xtrace));
x = np.random.randn(50000)
%timeit autocorr(x)
%timeit autocorr2(x)
plt.plot(np.array(stan_acov), alpha=.5)
plt.plot(np.array(pymc_acov), alpha=.5)
# plt.plot(autocov(xtrace), alpha=.5)
plt.plot(autocov2(xtrace), alpha=.5)
plt.legend();
varx = np.var(xtrace) * (len(xtrace) - 1) / len(xtrace)
plt.plot(autocorr(xtrace))
plt.acorr(xtrace, maxlags=len(xtrace)-1);
plt.plot(np.array(stan_acov) - acov_);
np.array(stan_acov)[:10]-acov_[:10]
np.array(stan_acov)[:10]-np.array(pymc_acov)[:10]<jupyter_output><empty_output><jupyter_text>PyMC3 translation of PyStan effect_n<jupyter_code>from scipy.signal import fftconvolve


def autocorr(x):
    """
    Compute autocorrelation using FFTfor every lag for the input array
    https://en.wikipedia.org/wiki/Autocorrelation#Efficient_computation

    Parameters
    ----------
    x : Numpy array
        A 1D array containing MCMC samples

    Returns
    -------
    acorr: Numpy array same size as the input array
    """
    y = x - x.mean()
    N = len(y)
    result = fftconvolve(y, y[::-1])
    acorr = result[len(result) // 2:]
    acorr /= (N - np.arange(N))
    acorr /= acorr[0]
    return acorr


def autocov(x):
    """Compute autocovariance estimates for every lag for 
    the input array

    Parameters
    ----------
    x : Numpy array
        A 1D array containing MCMC samples

    Returns
    -------
    acov: Numpy array same size as the input array
    """
    acorr = autocorr(x)
    varx = np.var(x, ddof=1) * (len(x) - 1) / len(x)
    return acorr * varx
var_id = 13 #22
stan_acov_ = deepcopy([pystan._chains._test_autocovariance(sim, im, var_id) for im in range(m)])
trace_values = deepcopy([sim['samples'][im]['chains'][param_names[var_id]]
                for im in range(m)])

x = np.array(trace_values, copy=True)
shape = x.shape

# Make sure to handle scalars correctly, adding extra dimensions if
# needed. We could use np.squeeze here, but we don't want to squeeze
# out dummy dimensions that a user inputs.
if len(shape) == 2:
    x = np.atleast_3d(trace_values)

# Transpose all dimensions, which makes the loop below
# easier by moving the axes of the variable to the front instead
# of the chain and sample axes.
x = x.transpose()
for tup in np.ndindex(*list(x.shape[:-2])):
    print(x[tup].shape)
trace_value = x[0].T
nchain, n_samples = trace_value.shape

acov = np.asarray([autocov(trace_value[chain]) for chain in range(nchain)])
# acov = np.asarray([stan_acov_[chain] for chain in range(nchain)])

chain_mean = np.mean(trace_value, axis=1)
chain_var = acov[:, 0] * n_samples / (n_samples - 1.)
acov_t = acov[:, 1] * n_samples / (n_samples - 1.)
mean_var = np.mean(chain_var)
var_plus = mean_var * (n_samples - 1.) / n_samples
var_plus += np.var(chain_mean, ddof=1)

rho_hat_t = np.zeros(n_samples)
rho_hat_even = 1.
rho_hat_t[0] = rho_hat_even
rho_hat_odd = 1. - (mean_var - np.mean(acov_t)) / var_plus
rho_hat_t[1] = rho_hat_odd
# Geyer's initial positive sequence
max_t = 1
t = 1
while t < (n_samples - 2) and (rho_hat_even + rho_hat_odd) >= 0:
    rho_hat_even = 1. - (mean_var - np.mean(acov[:, t + 1])) / var_plus
    rho_hat_odd = 1. - (mean_var - np.mean(acov[:, t + 2])) / var_plus
    if (rho_hat_even + rho_hat_odd) >= 0:
        rho_hat_t[t + 1] = rho_hat_even
        rho_hat_t[t + 2] = rho_hat_odd
    max_t = t + 2
    t += 2

# Geyer's initial monotone sequence
t = 3
while t <= max_t - 2:
    if (rho_hat_t[t + 1] + rho_hat_t[t + 2]) > (rho_hat_t[t - 1] + rho_hat_t[t]):
        rho_hat_t[t + 1] = (rho_hat_t[t - 1] + rho_hat_t[t]) / 2.
        rho_hat_t[t + 2] = rho_hat_t[t + 1]
    t += 2
ess = nchain * n_samples / (-1. + 2. * np.sum(rho_hat_t))
t = 3

rho_hat_t[t + 1:t+2+1]
ess
pystan._chains.effective_sample_size(sim, var_id)
_, axes = plt.subplots(3, 1, figsize=(15, 10), sharex=True)
axes[0].plot(np.asarray(trace_values).T)
stan_acov = [pystan._chains._test_autocovariance(sim, im, var_id) for im in range(m)]
pymc_acov = [autocov(x) for x in trace_values]
for im in range(m):
    axes[im+1].plot(stan_acov[im])
    axes[im+1].plot(pymc_acov[im])
plt.tight_layout()<jupyter_output><empty_output><jupyter_text>New implementation<jupyter_code>from scipy.signal import fftconvolve


def autocorr(x):
    """
    Compute autocorrelation using FFTfor every lag for the input array
    https://en.wikipedia.org/wiki/Autocorrelation#Efficient_computation

    Parameters
    ----------
    x : Numpy array
        An array containing MCMC samples

    Returns
    -------
    acorr: Numpy array same size as the input array
    """
    y = x - x.mean()
    N = len(y)
    result = fftconvolve(y, y[::-1])
    acorr = result[len(result) // 2:]
    acorr /= (N - np.arange(N))
    acorr /= acorr[0]
    return acorr


def autocov(x):
    """Compute autocovariance estimates for every lag for the input array

    Parameters
    ----------
    x : Numpy array
        An array containing MCMC samples

    Returns
    -------
    acov: Numpy array same size as the input array
    """
    acorr = autocorr(x)
    varx = np.var(x, ddof=1) * (len(x) - 1) / len(x)
    return acorr * varx
def get_neff(x):
    trace_value = x.T
    nchain, n_samples = trace_value.shape

    acov = np.asarray([autocov(trace_value[chain]) for chain in range(nchain)])

    chain_mean = trace_value.mean(axis=1)
    chain_var = acov[:, 0] * n_samples / (n_samples - 1.)
    acov_t = acov[:, 1] * n_samples / (n_samples - 1.)
    mean_var = np.mean(chain_var)
    var_plus = mean_var * (n_samples - 1.) / n_samples
    var_plus += np.var(chain_mean, ddof=1)

    rho_hat_t = np.zeros(n_samples)
    rho_hat_even = 1.
    rho_hat_t[0] = rho_hat_even
    rho_hat_odd = 1. - (mean_var - np.mean(acov_t)) / var_plus
    rho_hat_t[1] = rho_hat_odd
    # Geyer's initial positive sequence
    max_t = 1
    t = 1
    while t < (n_samples - 2) and (rho_hat_even + rho_hat_odd) >= 0.:
        rho_hat_even = 1. - (mean_var - np.mean(acov[:, t + 1])) / var_plus
        rho_hat_odd = 1. - (mean_var - np.mean(acov[:, t + 2])) / var_plus
        if (rho_hat_even + rho_hat_odd) >= 0:
            rho_hat_t[t + 1] = rho_hat_even
            rho_hat_t[t + 2] = rho_hat_odd
        max_t = t + 2
        t += 2

    # Geyer's initial monotone sequence
    t = 3
    while t <= max_t - 2:
        if rho_hat_t[t + 1] + rho_hat_t[t + 2] > rho_hat_t[t - 1] + rho_hat_t[t]:
            rho_hat_t[t + 1] = (rho_hat_t[t - 1] + rho_hat_t[t]) / 2.
            rho_hat_t[t + 2] = rho_hat_t[t + 1]
        t += 2
    ess = nchain * n_samples
    ess = ess / (-1. + 2. * np.sum(rho_hat_t))
    return ess

def generate_neff(trace_values):
    x = np.array(trace_values)
    shape = x.shape

    # Make sure to handle scalars correctly, adding extra dimensions if
    # needed. We could use np.squeeze here, but we don't want to squeeze
    # out dummy dimensions that a user inputs.
    if len(shape) == 2:
        x = np.atleast_3d(x)

    # Transpose all dimensions, which makes the loop below
    # easier by moving the axes of the variable to the front instead
    # of the chain and sample axes.
    x = x.transpose()

    # Get an array the same shape as the var
    _n_eff = np.zeros(x.shape[:-2])

    # Iterate over tuples of indices of the shape of var
    for tup in np.ndindex(*list(x.shape[:-2])):
        _n_eff[tup] = get_neff(x[tup])

    if len(shape) == 2:
        return _n_eff[0]

    return np.transpose(_n_eff)
# trace_values = tr.get_values('x', combine=False)
trace_values = [lst['samples'][0]['chains']['d'], lst['samples'][1]['chains']['d']]
generate_neff(trace_values)
ess3 = []
for varname in param_names:
    trace_values = [sim['samples'][im]['chains'][varname] for im in range(m)]
    ess3.append(generate_neff(trace_values))
df_neff['PyMC3_new'] = pd.Series(ess3)
df_neff<jupyter_output><empty_output><jupyter_text>Current implementation in PyMC3<jupyter_code>from pymc3.stats import autocorr
def get_vhat(x):
    # Chain samples are second to last dim (-2)
    num_samples = x.shape[-2]

    # Calculate between-chain variance
    B = num_samples * np.var(np.mean(x, axis=-2), axis=-1, ddof=1)

    # Calculate within-chain variance
    W = np.mean(np.var(x, axis=-2, ddof=1), axis=-1)

    # Estimate marginal posterior variance
    Vhat = W * (num_samples - 1) / num_samples + B / num_samples

    return Vhat

def get_neff(x, Vhat):
    # Number of chains is last dim (-1)
    num_chains = x.shape[-1]

    # Chain samples are second to last dim (-2)
    num_samples = x.shape[-2]

    negative_autocorr = False

    rho = np.ones(num_samples)
    t = 1

    # Iterate until the sum of consecutive estimates of autocorrelation is
    # negative
    while not negative_autocorr and (t < num_samples):

        variogram = np.mean((x[t:, :] - x[:-t, :])**2)
        rho[t] = 1. - variogram / (2. * Vhat)

        negative_autocorr = sum(rho[t - 1:t + 1]) < 0

        t += 1

    if t % 2:
        t -= 1

    neff = num_chains * num_samples / (1. + 2 * rho[1:t-1].sum())
    return min(num_chains * num_samples, np.floor(neff))

def generate_neff(trace_values):
    x = np.array(trace_values)
    shape = x.shape

    # Make sure to handle scalars correctly, adding extra dimensions if
    # needed. We could use np.squeeze here, but we don't want to squeeze
    # out dummy dimensions that a user inputs.
    if len(shape) == 2:
        x = np.atleast_3d(trace_values)

    # Transpose all dimensions, which makes the loop below
    # easier by moving the axes of the variable to the front instead
    # of the chain and sample axes.
    x = x.transpose()

    Vhat = get_vhat(x)

    # Get an array the same shape as the var
    _n_eff = np.zeros(x.shape[:-2])

    # Iterate over tuples of indices of the shape of var
    for tup in np.ndindex(*list(x.shape[:-2])):
        _n_eff[tup] = get_neff(x[tup], Vhat[tup])

    if len(shape) == 2:
        return _n_eff[0]

    return np.transpose(_n_eff)

#generate_neff(trace_values)
# trace_values = tr.get_values('x', combine=False)
trace_values = [lst['samples'][im]['chains'][varname] for im in range(m)]
generate_neff(trace_values)
ess2 = []
for varname in param_names:
    trace_values = [lst['samples'][im]['chains'][varname] for im in range(m)]
    ess2.append(generate_neff(trace_values))
df_neff['PyMC3_old'] = pd.Series(ess2)
df_neff<jupyter_output><empty_output><jupyter_text>New validation<jupyter_code>cov1 = np.eye(3) * np.diag([.9, 1.2, 1.3])
M = np.random.randn(3, 3)
cov2 = M.dot(M.T)
cov2
with pm.Model() as m:
    x = pm.Normal('x', shape=(3, ))
    pm.MvNormal('x1', mu=x, cov=cov2, shape=(5, 3))
    tr = pm.sample(cores=4)
pm.traceplot(tr);
from pymc3.util import get_default_varnames
from pymc3.backends.base import MultiTrace
from scipy.signal import fftconvolve


def autocorr(x, lag=None):
    """
    Compute autocorrelation using FFT for every lag for the input array
    https://en.wikipedia.org/wiki/Autocorrelation#Efficient_computation

    Parameters
    ----------
    x : Numpy array
        An array containing MCMC samples

    Returns
    -------
    acorr: Numpy array same size as the input array
    """
    y = x - x.mean()
    n = len(y)
    result = fftconvolve(y, y[::-1])
    acorr = result[len(result) // 2:]
    acorr /= np.arange(n, 0, -1)
    acorr /= acorr[0]
    if lag is None:
        return acorr
    else:
        warnings.warn(
            "The `lag` argument has been deprecated. If you want to get "
            "the value of a specific lag please call `autocorr(x)[lag]`.",
            DeprecationWarning)
        return acorr[lag]


def autocov(x, lag=None):
    """Compute autocovariance estimates for every lag for the input array

    Parameters
    ----------
    x : Numpy array
        An array containing MCMC samples

    Returns
    -------
    acov: Numpy array same size as the input array
    """
    acorr = autocorr(x)
    varx = np.var(x, ddof=1) * (len(x) - 1) / len(x)
    acov = acorr * varx
    if lag is None:
        return acov
    else:
        warnings.warn(
            "The `lag` argument has been deprecated. If you want to get "
            "the value of a specific lag please call `autocov(x)[lag]`.",
            DeprecationWarning)
        return acov[lag]
def effective_n(mtrace, varnames=None, include_transformed=False):
    R"""Returns estimate of the effective sample size of a set of traces.

    Parameters
    ----------
    mtrace : MultiTrace or trace object
      A MultiTrace object containing parallel traces (minimum 2)
      of one or more stochastic parameters.
    varnames : list
      Names of variables to include in the effective_n report
    include_transformed : bool
      Flag for reporting automatically transformed variables in addition
      to original variables (defaults to False).

    Returns
    -------
    n_eff : dictionary of floats (MultiTrace) or float (trace object)
        Return the effective sample size, :math:`\hat{n}_{eff}`

    Notes
    -----
    The diagnostic is computed by:

    .. math:: \hat{n}_{eff} = \frac{mn}{1 + 2 \sum_{t=1}^T \hat{\rho}_t}

    where :math:`\hat{\rho}_t` is the estimated autocorrelation at lag t, and T
    is the first odd positive integer for which the sum
    :math:`\hat{\rho}_{T+1} + \hat{\rho}_{T+1}` is negative.

    The current implementation is similar to Stan, which uses Geyer's initial
    monotone sequence criterion (Geyer, 1992; Geyer, 2011).

    References
    ----------
    Gelman et al. BDA (2014)"""

    def get_neff(x):
        """Compute the effective sample size for a 2D array
        """
        trace_value = x.T
        nchain, n_samples = trace_value.shape

        acov = np.asarray([autocov(trace_value[chain])
                           for chain in range(nchain)])

        chain_mean = trace_value.mean(axis=1)
        chain_var = acov[:, 0] * n_samples / (n_samples - 1.)
        acov_t = acov[:, 1] * n_samples / (n_samples - 1.)
        mean_var = np.mean(chain_var)
        var_plus = mean_var * (n_samples - 1.) / n_samples
        var_plus += np.var(chain_mean, ddof=1)

        rho_hat_t = np.zeros(n_samples)
        rho_hat_even = 1.
        rho_hat_t[0] = rho_hat_even
        rho_hat_odd = 1. - (mean_var - np.mean(acov_t)) / var_plus
        rho_hat_t[1] = rho_hat_odd
        # Geyer's initial positive sequence
        max_t = 1
        t = 1
        while t < (n_samples - 2) and (rho_hat_even + rho_hat_odd) >= 0.:
            rho_hat_even = 1. - (mean_var - np.mean(acov[:, t + 1])) / var_plus
            rho_hat_odd = 1. - (mean_var - np.mean(acov[:, t + 2])) / var_plus
            if (rho_hat_even + rho_hat_odd) >= 0:
                rho_hat_t[t + 1] = rho_hat_even
                rho_hat_t[t + 2] = rho_hat_odd
            max_t = t + 2
            t += 2

        # Geyer's initial monotone sequence
        t = 3
        while t <= max_t - 2:
            if (rho_hat_t[t + 1] + rho_hat_t[t + 2]) > (rho_hat_t[t - 1] + rho_hat_t[t]):
                rho_hat_t[t + 1] = (rho_hat_t[t - 1] + rho_hat_t[t]) / 2.
                rho_hat_t[t + 2] = rho_hat_t[t + 1]
            t += 2
        ess = nchain * n_samples
        ess = ess / (-1. + 2. * np.sum(rho_hat_t))
        return ess

    def generate_neff(trace_values):
        x = np.array(trace_values)
        shape = x.shape

        # Make sure to handle scalars correctly, adding extra dimensions if
        # needed. We could use np.squeeze here, but we don't want to squeeze
        # out dummy dimensions that a user inputs.
        if len(shape) == 2:
            x = np.atleast_3d(trace_values)

        # Transpose all dimensions, which makes the loop below
        # easier by moving the axes of the variable to the front instead
        # of the chain and sample axes.
        x = x.transpose()

        # Get an array the same shape as the var
        _n_eff = np.zeros(x.shape[:-2])

        # Iterate over tuples of indices of the shape of var
        for tup in np.ndindex(*list(x.shape[:-2])):
            _n_eff[tup] = get_neff(x[tup])

        if len(shape) == 2:
            return _n_eff[0]

        return np.transpose(_n_eff)

    if not isinstance(mtrace, MultiTrace):
        # Return neff for non-multitrace array
        return generate_neff(mtrace)

    if mtrace.nchains < 2:
        raise ValueError(
            'Calculation of effective sample size requires multiple chains '
            'of the same length.')

    if varnames is None:
        varnames = get_default_varnames(
            mtrace.varnames, include_transformed=include_transformed)

    n_eff = {}

    for var in varnames:
        n_eff[var] = generate_neff(mtrace.get_values(var, combine=False))

    return n_eff
%timeit effective_n(tr)
%timeit pm.effective_n(tr)
effective_n(tr)
pm.effective_n(tr)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13990204003994483, 'fraction_numerical': 0.031337676541918304, 'mean_word_length': 2.8382916590618725, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 34, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '33818', 'license': 'MIT', 'n_tokens_mistral': 8496, 'n_tokens_neox': 7831, 'n_words': 1947, 'path': 'Miscellaneous/Check Effective Sample Size.ipynb', 'repo_name': 'junpenglao/Planet_Sakaar_Data_Science', 'size': '21029'}