<commit_before>import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels.imputation.bayes_mi import BayesMI, MI
from numpy.testing import assert_allclose


def test_bmi():

    x = np.asarray([[1, np.nan, 3], [np.nan, 2, np.nan], [3, np.nan, 0],
                    [np.nan, 1, np.nan], [3, 2, 1]])
    bm = BayesMI(x)
    assert_allclose(bm.patterns[0], np.r_[0, 2])
    assert_allclose(bm.patterns[1], np.r_[1, 3])


def test_2x2():

    # Generate correlated data with mean and variance
    np.random.seed(3434)
    x = np.random.normal(size=(1000, 2))
    r = 0.5
    x[:, 1] = r*x[:, 0] + np.sqrt(1-r**2)*x[:, 1]
    x[:, 0] *= 2
    x[:, 1] *= 3
    x[:, 0] += 1
    x[:, 1] -= 2

    # Introduce some missing values
    u = np.random.normal(size=x.shape[0])
    x[u > 1, 0] = np.nan
    u = np.random.normal(size=x.shape[0])
    x[u > 1, 1] = np.nan

    bm = BayesMI(x)

    # Burn-in
    for k in range(500):
        bm.update()

    # Estimate the posterior mean
    mean = 0
    cov = 0
    dmean = 0
    dcov = 0
    for k in range(500):
        bm.update()
        mean += bm.mean
        cov += bm.cov
        dmean += bm.data.mean(0)
        dcov += np.cov(bm.data.T)
    mean /= 500
    cov /= 500
    dmean /= 500
    dcov /= 500

    assert_allclose(mean, np.r_[1, -2], 0.1)
    assert_allclose(dmean, np.r_[1, -2], 0.1)
    assert_allclose(cov, np.asarray([[4, 6*r], [6*r, 9]]), 0.1)
    assert_allclose(dcov, np.asarray([[4, 6*r], [6*r, 9]]), 0.1)


def test_MI():

    np.random.seed(414)
    x = np.random.normal(size=(200, 4))
    x[[1, 3, 9], 0] = np.nan
    x[[1, 4, 3], 1] = np.nan
    x[[2, 11, 21], 2] = np.nan
    x[[11, 22, 99], 3] = np.nan

    def model_args(x):
        # Return endog, exog
        # Regress x0 on x1 and x2
        return (x[:, 0], x[:, 1:])

    for j in (0, 1):
        np.random.seed(2342)
        imp = BayesMI(x.copy())
        mi = MI(imp, sm.OLS, model_args, burn=0)
        r = mi.fit()
        r.summary()  # smoke test
        assert_allclose(r.params, np.r_[
            -0.05347919, -0.02479701, 0.10075517], 1e-4, 1e-4)

        c = np.asarray([[0.00418232, 0.00029746, -0.00035057],
                        [0.00029746, 0.00407264, 0.00019496],
                        [-0.00035057, 0.00019496, 0.00509413]])
        assert_allclose(r.cov_params(), c, 1e-4, 1e-4)

        # Test with ndarray and pandas input
        x = pd.DataFrame(x)


def test_MI_stat():
    # Test for MI where we know statistically what should happen. The
    # analysis model is x0 ~ x1 with standard error 1/sqrt(n) for the
    # slope parameter.  The nominal n is 1000, but half of the cases<commit_msg>test_bayes_mi.py<commit_after>import numpy as np
import pandas as pd
import statsmodels.api as sm
from statsmodels.imputation.bayes_mi import BayesMI, MI
from numpy.testing import assert_allclose


def test_pat():

    x = np.asarray([[1, np.nan, 3], [np.nan, 2, np.nan], [3, np.nan, 0],
                    [np.nan, 1, np.nan], [3, 2, 1]])
    bm = BayesMI(x)
    assert_allclose(bm.patterns[0], np.r_[0, 2])
    assert_allclose(bm.patterns[1], np.r_[1, 3])


def test_2x2():

    # Generate correlated data with mean and variance
    np.random.seed(3434)
    x = np.random.normal(size=(1000, 2))
    r = 0.5
    x[:, 1] = r*x[:, 0] + np.sqrt(1-r**2)*x[:, 1]
    x[:, 0] *= 2
    x[:, 1] *= 3
    x[:, 0] += 1
    x[:, 1] -= 2

    # Introduce some missing values
    u = np.random.normal(size=x.shape[0])
    x[u > 1, 0] = np.nan
    u = np.random.normal(size=x.shape[0])
    x[u > 1, 1] = np.nan

    bm = BayesMI(x)

    # Burn-in
    for k in range(500):
        bm.update()

    # Estimate the posterior mean
    mean = 0
    cov = 0
    dmean = 0
    dcov = 0
    for k in range(500):
        bm.update()
        mean += bm.mean
        cov += bm.cov
        dmean += bm.data.mean(0)
        dcov += np.cov(bm.data.T)
    mean /= 500
    cov /= 500
    dmean /= 500
    dcov /= 500

    assert_allclose(mean, np.r_[1, -2], 0.1)
    assert_allclose(dmean, np.r_[1, -2], 0.1)
    assert_allclose(cov, np.asarray([[4, 6*r], [6*r, 9]]), 0.1)
    assert_allclose(dcov, np.asarray([[4, 6*r], [6*r, 9]]), 0.1)


def test_MI():

    np.random.seed(414)
    x = np.random.normal(size=(200, 4))
    x[[1, 3, 9], 0] = np.nan
    x[[1, 4, 3], 1] = np.nan
    x[[2, 11, 21], 2] = np.nan
    x[[11, 22, 99], 3] = np.nan

    def model_args(x):
        # Return endog, exog
        # Regress x0 on x1 and x2
        return (x[:, 0], x[:, 1:])

    for j in (0, 1):
        np.random.seed(2342)
        imp = BayesMI(x.copy())
        mi = MI(imp, sm.OLS, model_args, burn=0)
        r = mi.fit()
        r.summary()  # smoke test
        assert_allclose(r.params, np.r_[
            -0.05347919, -0.02479701, 0.10075517], 1e-3, 1e-2)

        c = np.asarray([[0.00418232, 0.00029746, -0.00035057],
                        [0.00029746, 0.00407264, 0.00019496],
                        [-0.00035057, 0.00019496, 0.00509413]])
        assert_allclose(r.cov_params(), c, 1e-2, 1e-3)

        # Test with ndarray and pandas input
        x = pd.DataFrame(x)


def test_MI_stat():
    # Test for MI where we know statistically what should happen. The
    # analysis model is x0 ~ x1 with standard error 1/sqrt(n) for the
    # slope parameter.  The nominal n is 1000, but half of the cases
{'fraction_non_alphanumeric': 0.16850570413315877, 'fraction_numerical': 0.09463250420796708, 'mean_word_length': 2.0806451612903225, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '5208160', 'n_tokens_mistral': 2506, 'n_tokens_neox': 2146, 'n_words': 563}