<commit_before>import functools

import numpy as np
from scipy.stats import norm as ndist

import regreg.api as rr

from selection.tests.instance import gaussian_instance
from selection.algorithms.lasso import ROSI

from learn_selection.core import (infer_full_target,
                                  split_sampler, 
                                  normal_sampler,
                                  logit_fit,
                                  repeat_selection,
                                  probit_fit)
from learn_selection.keras_fit import keras_fit

def simulate(n=1000, p=100, s=10, signal=(0.5, 1), sigma=2, alpha=0.1, seed=0):

    # description of statistical problem

    np.random.seed(seed)
    X, y, truth = gaussian_instance(n=n,
                                    p=p, 
                                    s=s,
                                    equicorrelated=False,
                                    rho=0.5, 
                                    sigma=sigma,
                                    signal=signal,
                                    random_signs=True,
                                    scale=False,
                                    center=False)[:3]

    dispersion = sigma**2

    S = X.T.dot(y)
    covS = dispersion * X.T.dot(X)
    smooth_sampler = normal_sampler(S, covS)

    def meta_algorithm(X, XTXi, resid, sampler):

        n, p = X.shape

        rho = 0.8
        S = sampler(scale=0.) # deterministic with scale=0
        ynew = X.dot(XTXi).dot(S) + resid # will be ok for n>p and non-degen X
        Xnew = rho * X + np.sqrt(1 - rho**2) * np.random.standard_normal(X.shape)

        X_full = np.hstack([X, Xnew])
        beta_full = np.linalg.pinv(X_full).dot(ynew)
        winners = np.fabs(beta_full)[:p] > np.fabs(beta_full)[p:]
        return set(np.nonzero(winners)[0])

    XTX = X.T.dot(X)
    XTXi = np.linalg.inv(XTX)
    resid = y - X.dot(XTXi.dot(X.T.dot(y)))
    dispersion = np.linalg.norm(resid)**2 / (n-p)
                         
    selection_algorithm = functools.partial(meta_algorithm, X, XTXi, resid)

    # run selection algorithm

    success_params = (8, 10)

    observed_set = repeat_selection(selection_algorithm, smooth_sampler, *success_params)

    # find the target, based on the observed outcome

    # find the target, based on the observed outcome

    pivots, covered, lengths, pvalues = [], [], [], []
    lower, upper = [], []
    naive_pvalues, naive_pivots, naive_covered, naive_lengths =  [], [], [], []

    targets = []
    true_target = truth[sorted(observed_set)]

    if len(observed_set) > 0:
        results = infer_full_target(selection_algorithm,
                                    observed_set,
                                    sorted(observed_set),
                                    smooth_sampler,
                                    dispersion,
                                    hypothesis=true_target,
                                    fit_probability=keras_fit,
                                    fit_args={'epochs':20, 'sizes':[100]*5, 'dropout':0., 'activation':'relu'},
                                    success_params=success_params,
                                    alpha=alpha,
                                    B=3000)

        for i, result in enumerate(results):

            (pivot, 
             interval,
             pvalue,
             _) = result

            pvalues.append(pvalue)
            pivots.append(pivot)
            covered.append((interval[0] < true_target[i]) * (interval[1] > true_target[i]))
            lengths.append(interval[1] - interval[0])

        for i, idx in enumerate(sorted(observed_set)):
            target_sd = np.sqrt(dispersion * XTXi[idx, idx])
            observed_target = np.squeeze(XTXi[idx].dot(X.T.dot(y)))
            quantile = ndist.ppf(1 - 0.5 * alpha)
            naive_interval = (observed_target - quantile * target_sd, observed_target + quantile * target_sd)

            naive_pivot = (1 - ndist.cdf((observed_target - true_target[i]) / target_sd))
            naive_pivot = 2 * min(naive_pivot, 1 - naive_pivot)
            naive_pivots.append(naive_pivot)

            naive_pvalue = (1 - ndist.cdf(observed_target / target_sd))
            naive_pvalue = 2 * min(naive_pivot, 1 - naive_pivot)
            naive_pvalues.append(naive_pvalue)

            naive_covered.append((naive_interval[0] < true_target[i]) * (naive_interval[1] > true_target[i]))
            naive_lengths.append(naive_interval[1] - naive_interval[0])
            lower.append(interval[0])
            upper.append(interval[1])

    if len(pvalues) > 0:
        return pd.DataFrame({'pivot':pivots,
                             'pvalue':pvalues,
                             'coverage':(np.array(lower) < true_target) * (np.array(upper) > true_target),
                             'length':lengths,
                             'naive_pivot':naive_pivots,
                             'naive_coverage':naive_covered,
                             'naive_length':naive_lengths,
                             'upper':upper,
                             'lower':lower,
                             'target':true_target
                             })


if __name__ == "__main__":
    import statsmodels.api as sm
    import matplotlib.pyplot as plt
    import pandas as pd

    U = np.linspace(0, 1, 101)
    plt.clf()

    for i in range(500):
        df = simulate()
        csvfile = 'knockoff_kernel_multi.csv'<commit_msg>make sure to use the seed<commit_after>import functools

import numpy as np
from scipy.stats import norm as ndist

import regreg.api as rr

from selection.tests.instance import gaussian_instance
from selection.algorithms.lasso import ROSI

from learn_selection.core import (infer_full_target,
                                  split_sampler, 
                                  normal_sampler,
                                  logit_fit,
                                  repeat_selection,
                                  probit_fit)
from learn_selection.keras_fit import keras_fit

def simulate(n=1000, p=100, s=10, signal=(0.5, 1), sigma=2, alpha=0.1, seed=0):

    # description of statistical problem

    np.random.seed(seed)
    X, y, truth = gaussian_instance(n=n,
                                    p=p, 
                                    s=s,
                                    equicorrelated=False,
                                    rho=0.5, 
                                    sigma=sigma,
                                    signal=signal,
                                    random_signs=True,
                                    scale=False,
                                    center=False)[:3]

    dispersion = sigma**2

    S = X.T.dot(y)
    covS = dispersion * X.T.dot(X)
    smooth_sampler = normal_sampler(S, covS)

    def meta_algorithm(X, XTXi, resid, sampler):

        n, p = X.shape

        rho = 0.8
        S = sampler(scale=0.) # deterministic with scale=0
        ynew = X.dot(XTXi).dot(S) + resid # will be ok for n>p and non-degen X
        Xnew = rho * X + np.sqrt(1 - rho**2) * np.random.standard_normal(X.shape)

        X_full = np.hstack([X, Xnew])
        beta_full = np.linalg.pinv(X_full).dot(ynew)
        winners = np.fabs(beta_full)[:p] > np.fabs(beta_full)[p:]
        return set(np.nonzero(winners)[0])

    XTX = X.T.dot(X)
    XTXi = np.linalg.inv(XTX)
    resid = y - X.dot(XTXi.dot(X.T.dot(y)))
    dispersion = np.linalg.norm(resid)**2 / (n-p)
                         
    selection_algorithm = functools.partial(meta_algorithm, X, XTXi, resid)

    # run selection algorithm

    success_params = (8, 10)

    observed_set = repeat_selection(selection_algorithm, smooth_sampler, *success_params)

    # find the target, based on the observed outcome

    # find the target, based on the observed outcome

    pivots, covered, lengths, pvalues = [], [], [], []
    lower, upper = [], []
    naive_pvalues, naive_pivots, naive_covered, naive_lengths =  [], [], [], []

    targets = []
    true_target = truth[sorted(observed_set)]

    if len(observed_set) > 0:
        results = infer_full_target(selection_algorithm,
                                    observed_set,
                                    sorted(observed_set),
                                    smooth_sampler,
                                    dispersion,
                                    hypothesis=true_target,
                                    fit_probability=keras_fit,
                                    fit_args={'epochs':20, 'sizes':[100]*5, 'dropout':0., 'activation':'relu'},
                                    success_params=success_params,
                                    alpha=alpha,
                                    B=3000)

        for i, result in enumerate(results):

            (pivot, 
             interval,
             pvalue,
             _) = result

            pvalues.append(pvalue)
            pivots.append(pivot)
            covered.append((interval[0] < true_target[i]) * (interval[1] > true_target[i]))
            lengths.append(interval[1] - interval[0])

        for i, idx in enumerate(sorted(observed_set)):
            target_sd = np.sqrt(dispersion * XTXi[idx, idx])
            observed_target = np.squeeze(XTXi[idx].dot(X.T.dot(y)))
            quantile = ndist.ppf(1 - 0.5 * alpha)
            naive_interval = (observed_target - quantile * target_sd, observed_target + quantile * target_sd)

            naive_pivot = (1 - ndist.cdf((observed_target - true_target[i]) / target_sd))
            naive_pivot = 2 * min(naive_pivot, 1 - naive_pivot)
            naive_pivots.append(naive_pivot)

            naive_pvalue = (1 - ndist.cdf(observed_target / target_sd))
            naive_pvalue = 2 * min(naive_pvalue, 1 - naive_pvalue)
            naive_pvalues.append(naive_pvalue)

            naive_covered.append((naive_interval[0] < true_target[i]) * (naive_interval[1] > true_target[i]))
            naive_lengths.append(naive_interval[1] - naive_interval[0])
            lower.append(interval[0])
            upper.append(interval[1])

    if len(pvalues) > 0:
        return pd.DataFrame({'pivot':pivots,
                             'pvalue':pvalues,
                             'coverage':(np.array(lower) < true_target) * (np.array(upper) > true_target),
                             'length':lengths,
                             'naive_pivot':naive_pivots,
                             'naive_coverage':naive_covered,
                             'naive_length':naive_lengths,
                             'upper':upper,
                             'lower':lower,
                             'target':true_target
                             })


if __name__ == "__main__":
    import statsmodels.api as sm
    import matplotlib.pyplot as plt
    import pandas as pd

    U = np.linspace(0, 1, 101)
    plt.clf()

    iseed = int(np.fabs(np.random.standard_normal() * 50000))
    for i in range(500):
        df = simulate(seed=i + iseed)
        csvfile = 'knockoff_kernel_multi.csv'
{'fraction_non_alphanumeric': 0.10295975455693918, 'fraction_numerical': 0.013264753654574987, 'mean_word_length': 1.5029358626919602, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 48, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2001049', 'n_tokens_mistral': 3455, 'n_tokens_neox': 3234, 'n_words': 610}