<commit_before>"""
In this module, we implement forward stepwise model selection for $K$ steps.

The main goal of this is to produce a set of linear inequality constraints satisfied by
$y$ after $K$ steps.

"""

import warnings
from copy import copy

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

# local imports 

from ..constraints.affine import constraints, gibbs_test, stack
from ..distributions.chisq import quadratic_test
from ..distributions.discrete_family import discrete_family
from .projection import projection

DEBUG = False

class forward_step(object):

    """
    Centers columns of X!
    """

    def __init__(self, X, Y, 
                 subset=[],
                 covariance=None):
        self.subset = subset
        self.X, self.Y = X, Y
        if subset != []:
            self.adjusted_X = X.copy()[subset]
            self.subset_subset = Y.copy()[subset]
            self.subset_selector = np.identity(self.X.shape[0])[subset]
        else:
            self.adjusted_X = X.copy()
            self.subset_Y = Y.copy()

        self.P = [None] # residual forming projections
        self.variables = []
        self.Z = []
        self.Zfunc = []
        self.signs = []
        self.covariance = covariance
        self._resid_vector = self.subset_Y.copy() 

    def __iter__(self):
        n, p = self.X.shape
        self.identity_cone = []
        self.inactive = range(p)
        self.offset = [[np.ones(p) * np.inf, np.ones(p) * np.inf]]
        return self

    def next(self, compute_pval=False,
             use_identity=False,
             burnin=2000,
             ndraw=8000):
        """
        """
        
        adjusted_X, Y, inactive = self.adjusted_X, self.subset_Y, self.inactive
        resid_vector = self._resid_vector
        n, p = adjusted_X.shape

        # up to now inactive
        inactive = sorted(set(range(p)).difference(self.variables))
        scale = np.sqrt(np.sum(adjusted_X**2, 0))
        
        linear_part = []

        Zstat = np.dot(adjusted_X.T[inactive], Y) / scale[inactive]
        idx = np.argmax(np.fabs(Zstat))
        next_var = inactive[idx]
        next_sign = np.sign(Zstat[idx])

        realized_Z = Zstat[idx]

        # keep track of identity for testing
        # variables other than the last one added

        keep = np.zeros(p, np.bool)
        keep[inactive] = True
        keep[next_var] = False
        identity_linpart = np.vstack([adjusted_X[:,keep].T 
                                      / scale[keep,None] -
                                      next_sign * adjusted_X[:,next_var] / scale[next_var],
                                      -adjusted_X[:,keep].T 
                                      / scale[keep,None] -
                                      next_sign * adjusted_X[:,next_var] / scale[next_var],
                                      -next_sign * adjusted_X[:,next_var].reshape((1,-1))])

        identity_con = constraints(identity_linpart,
                                   np.zeros(identity_linpart.shape[0]))

        self.identity_cone.append(identity_linpart)

        eta = adjusted_X[:,next_var]

        if compute_pval:

            XI = self.X[:,inactive]
            linear_part = np.vstack([XI.T, -XI.T])
            offset = np.array(self.offset)
            offset = offset[:,:,inactive]
            offset_pos = np.min(offset[:,0], 0)
            offset_neg = np.min(offset[:,1], 0)
            offset = np.hstack([offset_pos, offset_neg])
            con = constraints(linear_part, offset,
                              covariance=self.covariance)

            if use_identity:
                con = stack(con, identity_con)
                con.covariance = self.covariance
            if self.variables:
                XA = self.X[:,self.variables]
                self.sequential_con = con.conditional(XA.T,
                                                      np.dot(XA.T, Y))
            else:
                self.sequential_con = con

            def maxT(Z, L=adjusted_X[:,inactive], S=scale[inactive]):
                Tstat = np.fabs(np.dot(Z, L) / S[None,:]).max(1)
                return Tstat

            pval = gibbs_test(self.sequential_con,
                              Y,
                              eta,
                              sigma_known=True,
                              white=False,
                              ndraw=ndraw,
                              burnin=burnin,
                              how_often=-1,
                              UMPU=False,
                              use_random_directions=False,
                              tilt=None,
                              alternative='greater',
                              test_statistic=maxT,
                              )[0]

        # now update state for next step

        inactive.pop(idx)
        self.variables.append(next_var); self.signs.append(next_sign)

        realized_Z_adjusted = np.fabs(realized_Z) * scale
        offset_shift = np.dot(self.X.T, Y - resid_vector)
        self.offset.append([realized_Z_adjusted + offset_shift,
                            realized_Z_adjusted - offset_shift])

        resid_vector -= realized_Z * adjusted_X[:,next_var] / scale[next_var]
        adjusted_X -= (np.multiply.outer(eta, 
                                         np.dot(eta,
                                                adjusted_X)) / 
                       (eta**2).sum())
        if compute_pval:
            return pval

    def constraints(self, step=np.inf, identify_last_variable=True):
        default_step = len(self.variables)
        if default_step > 0 and not identify_last_variable:
            default_step -= 1
        step = min(step, default_step)
        A = np.vstack(self.identity_cone[:step])

        if self.subset == []:
            return constraints(A, 
                               np.zeros(A.shape[0]), <commit_msg>using subset selector<commit_after>"""
In this module, we implement forward stepwise model selection for $K$ steps.

The main goal of this is to produce a set of linear inequality constraints satisfied by
$y$ after $K$ steps.

"""

import warnings
from copy import copy

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

# local imports 

from ..constraints.affine import constraints, gibbs_test, stack
from ..distributions.chisq import quadratic_test
from ..distributions.discrete_family import discrete_family
from .projection import projection

DEBUG = False

class forward_step(object):

    """
    Centers columns of X!
    """

    def __init__(self, X, Y, 
                 subset=[],
                 covariance=None):
        self.subset = subset
        self.X, self.Y = X, Y

        if subset != []:
            self.adjusted_X = X.copy()[subset]
            self.subset_X = X.copy()[subset]
            self.subset_Y = Y.copy()[subset]
            self.subset_selector = np.identity(self.X.shape[0])[subset]
        else:
            self.adjusted_X = X.copy()
            self.subset_Y = Y.copy()

        self.P = [None] # residual forming projections
        self.variables = []
        self.Z = []
        self.Zfunc = []
        self.signs = []
        self.covariance = covariance
        self._resid_vector = self.subset_Y.copy() 

    def __iter__(self):
        n, p = self.X.shape
        self.identity_cone = []
        self.inactive = range(p)
        self.offset = [[np.ones(p) * np.inf, np.ones(p) * np.inf]]
        return self

    def next(self, compute_pval=False,
             use_identity=False,
             burnin=2000,
             ndraw=8000):
        """
        """
        
        adjusted_X, Y, inactive = self.adjusted_X, self.subset_Y, self.inactive
        resid_vector = self._resid_vector
        n, p = adjusted_X.shape

        # up to now inactive
        inactive = sorted(set(range(p)).difference(self.variables))
        scale = np.sqrt(np.sum(adjusted_X**2, 0))
        
        linear_part = []

        Zstat = np.dot(adjusted_X.T[inactive], Y) / scale[inactive]
        idx = np.argmax(np.fabs(Zstat))
        next_var = inactive[idx]
        next_sign = np.sign(Zstat[idx])

        realized_Z = Zstat[idx]

        # keep track of identity for testing
        # variables other than the last one added

        keep = np.zeros(p, np.bool)
        keep[inactive] = True
        keep[next_var] = False
        identity_linpart = np.vstack([adjusted_X[:,keep].T 
                                      / scale[keep,None] -
                                      next_sign * adjusted_X[:,next_var] / scale[next_var],
                                      -adjusted_X[:,keep].T 
                                      / scale[keep,None] -
                                      next_sign * adjusted_X[:,next_var] / scale[next_var],
                                      -next_sign * adjusted_X[:,next_var].reshape((1,-1))])

        if self.subset != []:
            identity_linpart = np.dot(identity_linpart, 
                                      self.subset_selector)

        identity_con = constraints(identity_linpart,
                                   np.zeros(identity_linpart.shape[0]))

        self.identity_cone.append(identity_linpart)

        eta = adjusted_X[:,next_var]

        if compute_pval:

            XI = self.subset_X[:,inactive]
            linear_part = np.vstack([XI.T, -XI.T])
            offset = np.array(self.offset)
            offset = offset[:,:,inactive]
            offset_pos = np.min(offset[:,0], 0)
            offset_neg = np.min(offset[:,1], 0)
            offset = np.hstack([offset_pos, offset_neg])
            con = constraints(linear_part, offset,
                              covariance=self.covariance)

            if use_identity:
                con = stack(con, identity_con)
                con.covariance = self.covariance
            if self.variables:
                XA = self.subset_X[:,self.variables]
                self.sequential_con = con.conditional(XA.T,
                                                      np.dot(XA.T, Y))
            else:
                self.sequential_con = con

            def maxT(Z, L=adjusted_X[:,inactive], S=scale[inactive]):
                Tstat = np.fabs(np.dot(Z, L) / S[None,:]).max(1)
                return Tstat

            pval = gibbs_test(self.sequential_con,
                              Y,
                              eta,
                              sigma_known=True,
                              white=False,
                              ndraw=ndraw,
                              burnin=burnin,
                              how_often=-1,
                              UMPU=False,
                              use_random_directions=False,
                              tilt=None,
                              alternative='greater',
                              test_statistic=maxT,
                              )[0]

        # now update state for next step

        inactive.pop(idx)
        self.variables.append(next_var); self.signs.append(next_sign)

        realized_Z_adjusted = np.fabs(realized_Z) * scale
        offset_shift = np.dot(self.subset_X.T, Y - resid_vector)
        self.offset.append([realized_Z_adjusted + offset_shift,
                            realized_Z_adjusted - offset_shift])

        resid_vector -= realized_Z * adjusted_X[:,next_var] / scale[next_var]
        adjusted_X -= (np.multiply.outer(eta, 
                                         np.dot(eta,
                                                adjusted_X)) / 
                       (eta**2).sum())
        if compute_pval:
            return pval

    def constraints(self, step=np.inf, identify_last_variable=True):
        default_step = len(self.variables)
        if default_step > 0 and not identify_last_variable:
            default_step -= 1
        step = min(step, default_step)
        A = np.vstack(self.identity_cone[:step])

        if self.subset == []:
            return constraints(A, 
                               np.zeros(A.shape[0]), 
{'fraction_non_alphanumeric': 0.10157545605306799, 'fraction_numerical': 0.0041459369817578775, 'mean_word_length': 1.4326341266639775, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 27, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6184678', 'n_tokens_mistral': 3482, 'n_tokens_neox': 3303, 'n_words': 673}