Returns estimate of the effective sample size of a set of traces.

    Parameters
    ----------
    x : array-like
      An array containing the 2 or more traces of a stochastic parameter. That is, an array of dimension m x n x k, where m is the number of traces, n the number of samples, and k the dimension of the stochastic.
    
    Returns
    -------
    n_eff : float
      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.

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

def effective_n(x):
    """ Returns estimate of the effective sample size of a set of traces.

    Parameters
    ----------
    x : array-like
      An array containing the 2 or more traces of a stochastic parameter. That is, an array of dimension m x n x k, where m is the number of traces, n the number of samples, and k the dimension of the stochastic.
    
    Returns
    -------
    n_eff : float
      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.

    References
    ----------
    Gelman et al. (2014)"""
    
    if np.shape(x) < (2,):
        raise ValueError(
            'Calculation of effective sample size requires multiple chains of the same length.')

    try:
        m, n = np.shape(x)
    except ValueError:
        return [effective_n(np.transpose(y)) for y in np.transpose(x)]
        
    s2 = gelman_rubin(x, return_var=True)
    
    negative_autocorr = False
    t = 1
    
    variogram = lambda t: (sum(sum((x[j][i] - x[j][i-t])**2 for i in range(t,n)) for j in range(m)) 
                                / (m*(n - t)))
    rho = np.ones(n)
    # Iterate until the sum of consecutive estimates of autocorrelation is negative
    while not negative_autocorr and (t < n):
        
        rho[t] = 1. - variogram(t)/(2.*s2)
        
        if not t % 2:
            negative_autocorr = sum(rho[t-1:t+1]) < 0
        
        t += 1
        
    return int(m*n / (1 + 2*rho[1:t].sum()))