<commit_before>"""Functions to implement convergence testing."""

import numpy as np
from nose.tools import eq_

# Get the accepted positions
#mixed_start = opts.nsteps / 2
#mixed_positions = mcmc.positions[mixed_start:,:]
#mixed_accepts = mcmc.accepts[mixed_start:]
#mixed_accept_positions = mixed_positions[mixed_accepts]

def calculate_within_chain_variances(chain_set):
    """
    Takes a list of chains (each expressed as an array of positions,
    each transformed to linear, not log space. 

    Returns W

    Note that all chains should be the same length, n."""
    # TODO: Check/get chain lengths
    # Make sure that we have more than one chain!
    if len(chain_set) < 2:
        raise Exception("To calculate the convergence criterion, there must " +
                        "be more than one chain.")

    # Calculate each within-chain variance (s_j^2):
    chain_variances = []
    for chain in chain_set:
        chain_variances.append(np.var(chain, axis=0, ddof=1))

    # Calculate the average within-chain variance (W):
    W = np.mean(chain_variances)

    return W


def calculate_between_chain_variances(chain_set):
    """
    Takes a list of chains (each expressed as an array of positions,
    each transformed to linear, not log space. 

    Returns B

    Note that all chains should be the same length, n."""
    # TODO: Check/get chain lengths
    # Make sure that we have more than one chain!
    if len(chain_set) < 2:
        raise Exception("To calculate the convergence criterion, there must " +
                        "be more than one chain.")

    # Calculate each within-chain average, that is, psi_{.j} :
    chain_averages = []
    for chain in chain_set:
        chain_averages.append(np.mean(chain, axis=0))

    # Calculate the between-chain variance (B):
    n = len(chain_set[0])
    B = n * np.var(chain_averages, axis=0, ddof=1)

    return B

def calculate_parameter_variance_estimates(chain_set):
    # TODO: Check/get chain lengths
    n = float(len(chain_set[0]))

    W = calculate_within_chain_variances(chain_set)
    B = calculate_between_chain_variances(chain_set)

    return (((n-1)/n) * W) + ((1/n) * B)

def convergence_criterion(chain_set):
    W = calculate_within_chain_variances(chain_set)
    var = calculate_parameter_variance_estimates(chain_set)

    return np.sqrt(var / W)<commit_msg>Initial commit of Gelman-Rubin convergence criterion<commit_after>"""Functions to implement convergence testing."""

import numpy as np
from nose.tools import eq_

def within_chain_variances(chain_set):
    """
    Takes a list of chains (each expressed as an array of positions,
    each transformed to linear, not log space. 

    Returns W

    Note that all chains should be the same length, n."""
    # TODO: Check/get chain lengths
    # Make sure that we have more than one chain!
    if len(chain_set) < 2:
        raise Exception("To calculate the convergence criterion, there must " +
                        "be more than one chain.")

    # Calculate each within-chain variance (s_j^2):
    chain_variances = []
    for chain in chain_set:
        chain_variances.append(np.var(chain, axis=0, ddof=1))

    # Calculate the average within-chain variance (W):
    W = np.mean(chain_variances, axis=0)

    return W

def between_chain_variances(chain_set):
    """
    Takes a list of chains (each expressed as an array of positions,
    each transformed to linear, not log space. 

    Returns B

    Note that all chains should be the same length, n."""
    # TODO: Check/get chain lengths
    # Make sure that we have more than one chain!
    if len(chain_set) < 2:
        raise Exception("To calculate the convergence criterion, there must " +
                        "be more than one chain.")

    # Calculate each within-chain average, that is, psi_{.j} :
    chain_averages = []
    for chain in chain_set:
        chain_averages.append(np.mean(chain, axis=0))

    # Calculate the between-chain variance (B):
    n = len(chain_set[0])
    B = n * np.var(chain_averages, axis=0, ddof=1)

    return B

def parameter_variance_estimates(chain_set):
    # TODO: Check/get chain lengths
    n = float(len(chain_set[0]))

    W = within_chain_variances(chain_set)
    B = between_chain_variances(chain_set)

    return (((n-1)/n) * W) + ((1/n) * B)

def convergence_criterion(chain_set):
    W = within_chain_variances(chain_set)
    var = parameter_variance_estimates(chain_set)

    return np.sqrt(var / W)
{'fraction_non_alphanumeric': 0.09200535953550693, 'fraction_numerical': 0.005806163465832961, 'mean_word_length': 3.1510658016682114, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, '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}}
{'dir': 'git-commits-cleaned', 'id': '4768177', 'n_tokens_mistral': 1409, 'n_tokens_neox': 1324, 'n_words': 446}