<commit_before>import pandas as pd
import numpy as np
from six import string_types
from collections import OrderedDict, defaultdict
from bambi.utils import listify
from patsy import dmatrices, dmatrix
import warnings
from bambi.priors import default_priors
import re


def listify(obj):
    return obj if isinstance(obj, (list, tuple)) else [obj]


class Model(object):

    def __init__(self, data=None, intercept=True, backend='pymc3'):
        '''
        Args:
            dataset (DataFrame): the pandas DF containing the data to use.
        '''

        obj_cols = data.select_dtypes(['object']).columns
        data[obj_cols] = data[obj_cols].apply(lambda x: x.astype('category'))
        self.data = data
        # Some random effects stuff later requires us to make guesses about
        # column groupings into terms based on patsy's naming scheme.
        if re.search("[\[\]]+", ''.join(data.columns)):
            warnings.warn("At least one of the column names in the specified "
                          "dataset contain square brackets ('[' or ']')."
                          "This may cause unexpected behavior if you specify "
                          "models with random effects. You are encouraged to "
                          "rename your columns to avoid square brackets.")

        self.reset()

        if backend.lower() == 'pymc3':
            from bambi.backends import PyMC3BackEnd
            self.backend = PyMC3BackEnd()
        else:
            raise ValueError("At the moment, only the PyMC3 backend is supported.")

    def reset(self):
        # self.cache = OrderedDict()
        self.contrasts = OrderedDict()
        self.terms = OrderedDict()
        self.y = None

    def build(self):
        ''' Build the PyMC3 model. '''
        if self.y is None:
            raise ValueError("No outcome (y) variable is set! Please call "
                             "set_y() before build() or fit().")
        self.backend.build(self)
        self.built = True

    def fit(self, fixed=None, random=None, **kwargs):
        if fixed is not None:
            self.add_formula(fixed, random=random, append=False)
        ''' Run the BackEnd to fit the model. '''
        if not self.built:
            warnings.warn("Current Bayesian model has not been built yet; "
              "building it first before sampling begins.")
            self.build()
        self.backend.run(**kwargs)

    def add_intercept(self):
        n = len(self.data)
        df = pd.DataFrame(np.ones((n, 1)), columns=['Intercept'])
        self.add_term('Intercept', df)

    def add_formula(self, f, random=None, append=False, priors=None,
                    categorical=None):

        data = self.data

        if not append:
            self.reset()

        # Explicitly convert columns to category if desired--though this can
        # also be done within the formula using C().
        if categorical is not None:
            data = data.copy()
            categorical = listify(categorical)
            data[categorical] = data[categorical].apply(lambda x: x.astype('category'))

        if '~' in f:
            y, X = dmatrices(f, data=data)
            y_label = y.design_info.term_names[0]
            self.add_y(y_label)
        else:
            X = dmatrix(f, data=data)

        # Loop over predictor terms
        for _name, _slice in X.design_info.term_name_slices.items():
            cols = X.design_info.column_names[_slice]
            term_data = pd.DataFrame(X[:, _slice], columns=cols)
            self.add_term(_name, data=term_data)

        # Random effects
        if random is not None:
            random = listify(random)
            for f in random:
                kwargs = {'random': True, 'categorical': False}
                # '1|factor' is considered identical to 'factor'
                f = re.sub(r'^1\s+\|(.*)', r'\1', f).strip()
                if '|' not in f:
                    variable = f
                else:
                    variable, split_by = re.split('\s*\|\s*', f)
                    kwargs['split_by'] = split_by
                self.add_term(variable=variable, **kwargs)

    def set_y(self, label):
        ''' Set the outcome variable. '''
        if self.y is not None:
            self.terms[self.y.label] = self.y
        self.y = self.terms.pop(label)
        self.built = False

    def add_y(self, *args, **kwargs):
        self.add_term(*args, **kwargs)
        name = list(self.terms.values())[-1].name # last-added term
        self.set_y(name)

    def add_term(self, variable, data=None, label=None, categorical=False,
                 random=False, split_by=None, prior=None):
        ''' Create a new Term and add it to the current Model. All positional
        and keyword arguments are passed directly to the Term initializer. '''

        if data is None:
            data = self.data.copy()

        if categorical:
            data[variable] = data[variable].astype('category')

        # Extract splitting variable
        if split_by is not None:
            data[split_by] = data[split_by].astype('category')
            split_by = listify(split_by)
            group_term = ':'.join(split_by)
            f = '0 + %s : (%s)' % (variable, group_term)
            data = dmatrix(f, data=data)
            cols = data.design_info.column_names
            data = pd.DataFrame(data, columns=cols)

            # For categorical random effects, one variance term per split_by level
            if random and categorical:
                # if group_term not in self.terms:
                #     raise ValueError("The variable '%s' cannot be nested in or"
                #                  " crossed with '%s', because the latter does "
                #                  "not exist yet. Please make sure that you "
                #                  "explicitly add all terms to the model before"
                #                  " crossing or nesting with other terms." %
                #                  (variable, split_by))
                split_data = {}
                groups = list(set([re.sub(r'^.*?\:', '', c) for c in cols]))
                for g in groups:
                    patt = re.escape(r':%s' % g) + '$'
                    level_data = data.filter(regex=patt)
                    level_data.columns = [c.split(':')[0] for c in level_data.columns]
                    level_data = level_data.loc[:, (level_data!=0).any(axis=0)]
                    split_data[g] = level_data.values
                data = split_data

        elif categorical or data[variable].dtype.name in ['object', 'category']:
            data = pd.get_dummies(data[variable])
        else:
            # If all columns have identical names except for levels in [],
            # assume they've already been contrast-coded, and pass data as-is<commit_msg>fix random effects naming bug<commit_after>import pandas as pd
import numpy as np
from six import string_types
from collections import OrderedDict, defaultdict
from bambi.utils import listify
from patsy import dmatrices, dmatrix
import warnings
from bambi.priors import default_priors
import re


def listify(obj):
    return obj if isinstance(obj, (list, tuple)) else [obj]


class Model(object):

    def __init__(self, data=None, intercept=True, backend='pymc3'):
        '''
        Args:
            dataset (DataFrame): the pandas DF containing the data to use.
        '''

        obj_cols = data.select_dtypes(['object']).columns
        data[obj_cols] = data[obj_cols].apply(lambda x: x.astype('category'))
        self.data = data
        # Some random effects stuff later requires us to make guesses about
        # column groupings into terms based on patsy's naming scheme.
        if re.search("[\[\]]+", ''.join(data.columns)):
            warnings.warn("At least one of the column names in the specified "
                          "dataset contain square brackets ('[' or ']')."
                          "This may cause unexpected behavior if you specify "
                          "models with random effects. You are encouraged to "
                          "rename your columns to avoid square brackets.")

        self.reset()

        if backend.lower() == 'pymc3':
            from bambi.backends import PyMC3BackEnd
            self.backend = PyMC3BackEnd()
        else:
            raise ValueError("At the moment, only the PyMC3 backend is supported.")

    def reset(self):
        # self.cache = OrderedDict()
        self.contrasts = OrderedDict()
        self.terms = OrderedDict()
        self.y = None

    def build(self):
        ''' Build the PyMC3 model. '''
        if self.y is None:
            raise ValueError("No outcome (y) variable is set! Please call "
                             "set_y() before build() or fit().")
        self.backend.build(self)
        self.built = True

    def fit(self, fixed=None, random=None, **kwargs):
        if fixed is not None:
            self.add_formula(fixed, random=random, append=False)
        ''' Run the BackEnd to fit the model. '''
        if not self.built:
            warnings.warn("Current Bayesian model has not been built yet; "
              "building it first before sampling begins.")
            self.build()
        self.backend.run(**kwargs)

    def add_intercept(self):
        n = len(self.data)
        df = pd.DataFrame(np.ones((n, 1)), columns=['Intercept'])
        self.add_term('Intercept', df)

    def add_formula(self, f, random=None, append=False, priors=None,
                    categorical=None):

        data = self.data

        if not append:
            self.reset()

        # Explicitly convert columns to category if desired--though this can
        # also be done within the formula using C().
        if categorical is not None:
            data = data.copy()
            categorical = listify(categorical)
            data[categorical] = data[categorical].apply(lambda x: x.astype('category'))

        if '~' in f:
            y, X = dmatrices(f, data=data)
            y_label = y.design_info.term_names[0]
            self.add_y(y_label)
        else:
            X = dmatrix(f, data=data)

        # Loop over predictor terms
        for _name, _slice in X.design_info.term_name_slices.items():
            cols = X.design_info.column_names[_slice]
            term_data = pd.DataFrame(X[:, _slice], columns=cols)
            self.add_term(_name, data=term_data)

        # Random effects
        if random is not None:
            random = listify(random)
            for f in random:
                kwargs = {'random': True, 'categorical': False}
                # '1|factor' is considered identical to 'factor'
                f = re.sub(r'^1\s+\|(.*)', r'\1', f).strip()
                if '|' not in f:
                    variable = f
                else:
                    variable, split_by = re.split('\s*\|\s*', f)
                    kwargs['split_by'] = split_by
                self.add_term(variable=variable, label=f, **kwargs)

    def set_y(self, label):
        ''' Set the outcome variable. '''
        if self.y is not None:
            self.terms[self.y.label] = self.y
        self.y = self.terms.pop(label)
        self.built = False

    def add_y(self, *args, **kwargs):
        self.add_term(*args, **kwargs)
        name = list(self.terms.values())[-1].name # last-added term
        self.set_y(name)

    def add_term(self, variable, data=None, label=None, categorical=False,
                 random=False, split_by=None, prior=None):
        ''' Create a new Term and add it to the current Model. All positional
        and keyword arguments are passed directly to the Term initializer. '''

        if data is None:
            data = self.data.copy()

        if categorical:
            data[variable] = data[variable].astype('category')

        # Extract splitting variable
        if split_by is not None:
            data[split_by] = data[split_by].astype('category')
            split_by = listify(split_by)
            group_term = ':'.join(split_by)
            f = '0 + %s : (%s)' % (variable, group_term)
            data = dmatrix(f, data=data)
            cols = data.design_info.column_names
            data = pd.DataFrame(data, columns=cols)

            # For categorical random effects, one variance term per split_by level
            if random and categorical:
                # if group_term not in self.terms:
                #     raise ValueError("The variable '%s' cannot be nested in or"
                #                  " crossed with '%s', because the latter does "
                #                  "not exist yet. Please make sure that you "
                #                  "explicitly add all terms to the model before"
                #                  " crossing or nesting with other terms." %
                #                  (variable, split_by))
                split_data = {}
                groups = list(set([re.sub(r'^.*?\:', '', c) for c in cols]))
                for g in groups:
                    patt = re.escape(r':%s' % g) + '$'
                    level_data = data.filter(regex=patt)
                    level_data.columns = [c.split(':')[0] for c in level_data.columns]
                    level_data = level_data.loc[:, (level_data!=0).any(axis=0)]
                    split_data[g] = level_data.values
                data = split_data

        elif categorical or (variable in data.columns and \
                             data[variable].dtype.name in ['object', 'category']):
            data = pd.get_dummies(data[variable])
        else:
            # If all columns have identical names except for levels in [],
            # assume they've already been contrast-coded, and pass data as-is
{'fraction_non_alphanumeric': 0.10653893606212352, 'fraction_numerical': 0.0023223746280571886, 'mean_word_length': 1.8076609616951915, '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': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '178320', 'n_tokens_mistral': 3869, 'n_tokens_neox': 3703, 'n_words': 1154}