<commit_before>from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

from nose.tools import raises

import numpy as np
from numpy.random import RandomState
from numpy.testing import assert_equal
from cryptorandom.cryptorandom import SHA256
from cryptorandom.sample import random_sample, random_permutation

from ..utils import (binom_conf_interval,
                     hypergeom_conf_interval,
                     get_prng,
                     permute,
                     permute_rows,
                     permute_within_groups,
                     permute_incidence_fixed_sums,
                     potential_outcomes)


def test_binom_conf_interval():
    res = binom_conf_interval(10, 3)
    expected = (0.05154625578928545, 0.6915018049393984)
    np.testing.assert_equal(res, expected)

    res2 = binom_conf_interval(10, 5, cl=0.95, alternative="upper")
    expected2 = (0.0, 0.7775588989918742)
    np.testing.assert_equal(res2, expected2)

    res3 = binom_conf_interval(10, 5, cl=0.95, alternative="lower")
    expected3 = (0.22244110100812578, 1.0)
    np.testing.assert_equal(res3, expected3)

    res4 = binom_conf_interval(10, 5, cl=0.95, alternative="upper", p=1)
    expected4 = (0.0, 0.7775588989918742)
    np.testing.assert_equal(res4, expected4)

    res5 = binom_conf_interval(10, 5, cl=0.95, alternative="lower", p=0)
    expected5 = (0.22244110100812578, 1.0)
    np.testing.assert_equal(res5, expected5)


def test_hypergeom_conf_interval():
    res = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="two-sided")
    expected = (1.0, 4.0)
    np.testing.assert_equal(res, expected)

    res2 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="upper")
    expected2 = (0.0, 4.0)
    np.testing.assert_equal(res2, expected2)

    res3 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="lower")
    expected3 = (1.0, 5.0)
    np.testing.assert_equal(res3, expected3)

    res4 = hypergeom_conf_interval(2, 2, 5, cl=0.95, alternative="two-sided")
    expected4 = (2.0, 5.0)
    np.testing.assert_equal(res4, expected4)

    cl = 0.95
    n = 10
    x = 5
    N = 20
    [lot, hit] = [6, 14]
    alternative = "two-sided"
    [lo, hi] = hypergeom_conf_interval(n, x, N, cl=cl, alternative=alternative)
    np.testing.assert_equal(lo, lot)
    np.testing.assert_equal(hi, hit)

    res5 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="upper", G=5)
    np.testing.assert_equal(res5, expected2)

    res6 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="lower", G=0)
    np.testing.assert_equal(res6, expected3)


def test_get_random_state():
    prng1 = RandomState(42)
    prng2 = get_prng(42)
    prng3 = get_prng(prng1)
    prng4 = get_prng(prng2)
    prng5 = get_prng()
    prng6 = get_prng(None)
    prng7 = get_prng(np.random)
    prng8 = get_prng(SHA256(42))
    assert(isinstance(prng1, RandomState))
    assert(isinstance(prng2, SHA256))
    assert(isinstance(prng3, RandomState))
    assert(isinstance(prng4, SHA256))
    assert(isinstance(prng5, SHA256))
    assert(isinstance(prng6, SHA256))
    assert(isinstance(prng7, RandomState))
    x1 = prng1.randint(0, 5, size=10)
    x2 = prng2.randint(0, 5, size=10)
    x3 = prng3.randint(0, 5, size=10)
    x4 = prng4.randint(0, 5, size=10)
    x5 = prng5.randint(0, 5, size=10)
    x6 = prng6.randint(0, 5, size=10)
    x7 = prng7.randint(0, 5, size=10)
    x8 = prng8.randint(0, 5, size=10)
    assert_equal(x2, x8)
    assert_equal(prng2.counter, 1)
    assert_equal(prng2.baseseed, 42)
    assert_equal(prng2.baseseed, prng4.baseseed)
    assert_equal(len(x5), 10)
    assert_equal(len(x6), 10)
    assert_equal(len(x7), 10)


@raises(ValueError)
def test_get_random_state_error():
    get_prng([1, 1.11])


def test_permute_within_group():
    x = np.repeat([1, 2, 3] * 3, 3)
    group = np.repeat([1, 2, 3], 9)
    #response = np.zeros_like(group)
    #response[[0,  1,  3,  9, 10, 11, 18, 19, 20]] = 1
    
    res1 = permute_within_groups(x, group, seed=42)
    res2 = permute_within_groups(x, group, seed=SHA256(42))
    np.testing.assert_equal(res1, res2)

    res3 = permute_within_groups(x, group)
    np.testing.assert_equal(res3.max(), 3)
    res3.sort()
    np.testing.assert_equal(group, res3)


def test_permute():
    prng = RandomState(42)

    x = prng.randint(10, size=20)
    permute(x, prng)
    expected = np.array([3, 2, 7, 9, 6, 5, 1, 6, 4, 2,
                         7, 7, 7, 4, 4, 3, 1, 5, 7, 6])
    np.testing.assert_array_equal(x, expected)

    permute(x)
    np.testing.assert_equal(x.max(), 9)
    np.testing.assert_equal(x.min(), 1)


def test_permute_rows():
    prng = RandomState(42)

    x = prng.randint(10, size=20).reshape(2, 10)
    permute_rows(x, prng)
    expected = np.array([[2, 7, 7, 6, 4, 9, 3, 4, 6, 6],
                         [7, 4, 5, 5, 3, 7, 1, 2, 7, 1]])
    np.testing.assert_array_equal(x, expected)

    permute_rows(x)
    np.testing.assert_equal(x.max(), 9)
    np.testing.assert_equal(x.min(), 1)


def test_permute_incidence_fixed_sums():
    prng = RandomState(42)
    x0 = prng.randint(2, size=80).reshape((8, 10))
    x1 = permute_incidence_fixed_sums(x0)

    K = 5

    m = []
    for i in range(1000):
        x2 = permute_incidence_fixed_sums(x0, k=K)
        m.append(np.sum(x0 != x2))
<commit_msg>TST: fix permute.utils tests to work with updated permutation helpers<commit_after>from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

from nose.tools import raises

import numpy as np
from numpy.random import RandomState
from numpy.testing import assert_equal
from cryptorandom.cryptorandom import SHA256
from cryptorandom.sample import random_sample, random_permutation

from ..utils import (binom_conf_interval,
                     hypergeom_conf_interval,
                     get_prng,
                     permute,
                     permute_rows,
                     permute_within_groups,
                     permute_incidence_fixed_sums,
                     potential_outcomes)


def test_binom_conf_interval():
    res = binom_conf_interval(10, 3)
    expected = (0.05154625578928545, 0.6915018049393984)
    np.testing.assert_equal(res, expected)

    res2 = binom_conf_interval(10, 5, cl=0.95, alternative="upper")
    expected2 = (0.0, 0.7775588989918742)
    np.testing.assert_equal(res2, expected2)

    res3 = binom_conf_interval(10, 5, cl=0.95, alternative="lower")
    expected3 = (0.22244110100812578, 1.0)
    np.testing.assert_equal(res3, expected3)

    res4 = binom_conf_interval(10, 5, cl=0.95, alternative="upper", p=1)
    expected4 = (0.0, 0.7775588989918742)
    np.testing.assert_equal(res4, expected4)

    res5 = binom_conf_interval(10, 5, cl=0.95, alternative="lower", p=0)
    expected5 = (0.22244110100812578, 1.0)
    np.testing.assert_equal(res5, expected5)


def test_hypergeom_conf_interval():
    res = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="two-sided")
    expected = (1.0, 4.0)
    np.testing.assert_equal(res, expected)

    res2 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="upper")
    expected2 = (0.0, 4.0)
    np.testing.assert_equal(res2, expected2)

    res3 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="lower")
    expected3 = (1.0, 5.0)
    np.testing.assert_equal(res3, expected3)

    res4 = hypergeom_conf_interval(2, 2, 5, cl=0.95, alternative="two-sided")
    expected4 = (2.0, 5.0)
    np.testing.assert_equal(res4, expected4)

    cl = 0.95
    n = 10
    x = 5
    N = 20
    [lot, hit] = [6, 14]
    alternative = "two-sided"
    [lo, hi] = hypergeom_conf_interval(n, x, N, cl=cl, alternative=alternative)
    np.testing.assert_equal(lo, lot)
    np.testing.assert_equal(hi, hit)

    res5 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="upper", G=5)
    np.testing.assert_equal(res5, expected2)

    res6 = hypergeom_conf_interval(2, 1, 5, cl=0.95, alternative="lower", G=0)
    np.testing.assert_equal(res6, expected3)


def test_get_random_state():
    prng1 = RandomState(42)
    prng2 = get_prng(42)
    prng3 = get_prng(prng1)
    prng4 = get_prng(prng2)
    prng5 = get_prng()
    prng6 = get_prng(None)
    prng7 = get_prng(np.random)
    prng8 = get_prng(SHA256(42))
    assert(isinstance(prng1, RandomState))
    assert(isinstance(prng2, SHA256))
    assert(isinstance(prng3, RandomState))
    assert(isinstance(prng4, SHA256))
    assert(isinstance(prng5, SHA256))
    assert(isinstance(prng6, SHA256))
    assert(isinstance(prng7, RandomState))
    x1 = prng1.randint(0, 5, size=10)
    x2 = prng2.randint(0, 5, size=10)
    x3 = prng3.randint(0, 5, size=10)
    x4 = prng4.randint(0, 5, size=10)
    x5 = prng5.randint(0, 5, size=10)
    x6 = prng6.randint(0, 5, size=10)
    x7 = prng7.randint(0, 5, size=10)
    x8 = prng8.randint(0, 5, size=10)
    assert_equal(x2, x8)
    assert_equal(prng2.counter, 1)
    assert_equal(prng2.baseseed, 42)
    assert_equal(prng2.baseseed, prng4.baseseed)
    assert_equal(len(x5), 10)
    assert_equal(len(x6), 10)
    assert_equal(len(x7), 10)


@raises(ValueError)
def test_get_random_state_error():
    get_prng([1, 1.11])


def test_permute_within_group():
    x = np.repeat([1, 2, 3] * 3, 3)
    group = np.repeat([1, 2, 3], 9)
    #response = np.zeros_like(group)
    #response[[0,  1,  3,  9, 10, 11, 18, 19, 20]] = 1
    
    res1 = permute_within_groups(x, group, seed=42)
    res2 = permute_within_groups(x, group, seed=SHA256(42))
    np.testing.assert_equal(res1, res2)

    res3 = permute_within_groups(x, group)
    np.testing.assert_equal(res3.max(), 3)
    res3.sort()
    np.testing.assert_equal(group, res3)


def test_permute():
    prng = SHA256(42)

    x = prng.randint(0, 10, size=20)
    actual = permute(x, prng)
    expected = np.array([6, 9, 5, 1, 3, 1, 4, 7, 6, 9, 
                         8, 7, 2, 1, 9, 7, 8, 1, 8, 1])
    np.testing.assert_array_equal(actual, expected)

    actual = permute(x)
    np.testing.assert_equal(actual.max(), 9)
    np.testing.assert_equal(actual.min(), 1)


def test_permute_rows():
    prng = SHA256(42)

    x = prng.randint(0, 10, size=20).reshape(2, 10)
    actual = permute_rows(x, prng)
    expected = np.array([[9, 1, 8, 6, 9, 1, 5, 4, 3, 6], 
                         [1, 7, 2, 1, 9, 7, 8, 7, 8, 1]])
    np.testing.assert_array_equal(actual, expected)

    a = permute_rows(x)
    np.testing.assert_equal(a.max(), 9)
    np.testing.assert_equal(a.min(), 1)


def test_permute_incidence_fixed_sums():
    prng = RandomState(42)
    x0 = prng.randint(2, size=80).reshape((8, 10))
    x1 = permute_incidence_fixed_sums(x0)

    K = 5

    m = []
    for i in range(1000):
        x2 = permute_incidence_fixed_sums(x0, k=K)
        m.append(np.sum(x0 != x2))

{'fraction_non_alphanumeric': 0.13545816733067728, 'fraction_numerical': 0.08802001297137033, 'mean_word_length': 3.1659590891547666, 'pattern_counts': {'":': 0, '<': 3, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '940096', 'n_tokens_mistral': 4776, 'n_tokens_neox': 4219, 'n_words': 829}