<commit_before>

"""
Test Cython and weave mixing code.
"""


import scipy.stats as stats
import Yank.mixing as mixing
import numpy as np

def mix_replicas(n_swaps=100, n_states=16, u_kl=None):
    """
    Utility function to generate "replicas" and call the mixing function a certain number of times

    Arguments
    ---------
    n_swaps : int (optional)
        The number of times to call the mixing code (default 100)
    n_states : int (optional)
        The number of replica states to include (default 16)
    u_kl : n_states x n_states ndarray of float64 (optional)
        Energies for each state. If None, will be initialized to zeros

    Returns
    -------
    Nij_proposed : n_states x n_states ndarray of np.int64
        Contains the number of times a given swap was proposed
    Nij_accepted : n_states x n_states ndarray of np.int64
        Contains the number of times a given swap was accepted
    """

    if u_kl is None:
        u_kl = np.zeros([n_states, n_states], dtype=np.float64)
    replica_states = range(n_states)
    Nij_proposed =  np.zeros([n_states,n_states], dtype=np.int64)
    Nij_accepted = np.zeros([n_states,n_states], dtype=np.int64)
    for i in range(n_swaps):
        mixing._mix_replicas_cython(n_states, replica_states, u_kl, Nij_proposed, Nij_accepted)
    return Nij_proposed, Nij_accepted




def test_even_mixing():
    pass

def test_general_mixing():<commit_msg>Use chi-square test to check that replicas with equal 0 energies are sampled equally<commit_after>

"""
Test Cython and weave mixing code.
"""


import scipy.stats as stats
import Yank.mixing as mixing
import numpy as np

def mix_replicas(n_swaps=100, n_states=16, u_kl=None):
    """
    Utility function to generate "replicas" and call the mixing function a certain number of times

    Arguments
    ---------
    n_swaps : int (optional)
        The number of times to call the mixing code (default 100)
    n_states : int (optional)
        The number of replica states to include (default 16)
    u_kl : n_states x n_states ndarray of float64 (optional)
        Energies for each state. If None, will be initialized to zeros

    Returns
    -------
    Nij_proposed : n_states x n_states ndarray of np.int64
        Contains the number of times a given swap was proposed
    Nij_accepted : n_states x n_states ndarray of np.int64
        Contains the number of times a given swap was accepted
    """

    if u_kl is None:
        u_kl = np.zeros([n_states, n_states], dtype=np.float64)
    replica_states = range(n_states)
    Nij_proposed =  np.zeros([n_states,n_states], dtype=np.int64)
    Nij_accepted = np.zeros([n_states,n_states], dtype=np.int64)
    for i in range(n_swaps):
        mixing._mix_replicas_cython(n_states, replica_states, u_kl, Nij_proposed, Nij_accepted)
    return Nij_proposed, Nij_accepted




def test_even_mixing():
    n_swaps = 100
    n_states = 16
    Nij_proposed, Nij_accepted = mix_replicas(n_swaps=n_swaps, n_states=n_states)
    for state in range(n_states):
        chi_sq, p_val = stats.chisquare(Nij_accepted[state,:])
        if p_val > 0.05:
            return 1
    return 0







def test_general_mixing():
{'fraction_non_alphanumeric': 0.07960981749528005, 'fraction_numerical': 0.017306482064191317, 'mean_word_length': 3.1125485122897802, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, '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': '6168002', 'n_tokens_mistral': 1119, 'n_tokens_neox': 1048, 'n_words': 337}