<commit_before>import json
from .variables import load_variables, SimpleColumn, BIDSVariableCollection
from . import transform
from collections import namedtuple
from six import string_types
import pandas as pd
import numpy as np


DesignMatrix = namedtuple('DesignMatrix', ('data', 'image', 'entities'))


class Analysis(object):

    ''' Represents an entire BIDS-Model analysis.
    Args:
        layouts (BIDSLayout or list): One or more BIDSLayout objects to pull
            variables from.
        model (str or dict): a BIDS model specification. Can either be a
            string giving the path of the JSON model spec, or an already-loaded
            dict containing the model info.
        collection (BIDSVariableCollection): Optional BIDSVariableCollection
            object used to load/manage variables. If None, a new collection is
            initialized from the provided layouts.
        selectors (dict): Optional keyword arguments to pass onto the
            collection; these will be passed on to the Layout's .get() method,
            and can be used to restrict variables.
    '''

    def __init__(self, model, collection=None, layouts=None, **selectors):

        if layouts is None and collection is None:
            raise ValueError("At least one of the 'layout' and 'collection' "
                             "arguments must be passed.")

        if isinstance(model, str):
            model = json.load(open(model))
        self.model = model

        if 'input' in model:
            selectors.update(model['input'])

        if collection is None:
            collection = load_variables(layouts, **selectors)

        self.collection = collection

        self._load_blocks(model['blocks'])

    @property
    def layout(self):
        return self.collection.layout  # for convenience

    def __iter__(self):
        for b in self.blocks:
            yield b

    def __getitem__(self, index):
        if isinstance(index, int):
            return self.blocks[index]
        return list(filter(lambda x: x.name == index, self.blocks))[0]

    def _load_blocks(self, blocks):
        self.blocks = []
        for i, b in enumerate(blocks):
            self.blocks.append(Block(self, index=i, **b))

    def setup(self, apply_transformations=True):
        ''' Set up the sequence of blocks--applying transformations,
        generating outputs, etc. '''

        # pass a copy of the collection through the pipeline (columns mutate)
        _collection = self.collection.clone()

        for b in self.blocks:
            b.setup(_collection, apply_transformations=apply_transformations)
            _collection = b.output_collection.clone()


class Block(object):

    ''' Represents a single analysis block from a BIDS-Model specification.
    Args:
        analysis (Analysis): The parent Analysis this Block belongs to.
        level (str): The BIDS keyword to use as the grouping variable; must be
            one of ['run', 'session', 'subject', or 'dataset'].
        index (int): The numerical index of the current Block within the
            sequence of blocks.
        name (str): Optional name to assign to the block. Must be specified
            in order to enable name-based indexing in the parent Analysis.
        transformations (list): List of BIDS-Model transformations to apply.
        model (dict): The 'model' part of the BIDS-Model block specification.
        contrasts (list): List of contrasts to apply to the parameter estimates
            generated when the model is fit.
    '''

    def __init__(self, analysis, level, index, name=None, transformations=None,
                 model=None, contrasts=None):

        self.analysis = analysis
        self.level = level
        self.index = index
        self.name = name
        self.transformations = transformations or []
        self.model = model or None
        self.contrasts = contrasts or []
        self._design_matrix = None
        self.input_collection = None
        self.output_collection = None

    def _get_design_matrix(self, **selectors):
        if self._design_matrix is None:
            raise ValueError("Block hasn't been set up yet; please call "
                             "setup() before you try to retrieve the DM.")
        # subset the data if needed
        data = self._design_matrix
        if selectors:
            # TODO: make sure this handles constraints on int columns properly
            bad_keys = list(set(selectors.keys()) - set(data.columns))
            if bad_keys:
                raise ValueError("The following query constraints do not map "
                                 "onto existing columns: %s." % bad_keys)
            query = ' and '.join(["{} in {}".format(k, v)
                                  for k, v in selectors.items()])
            data = data.query(query)
        return data

    def _drop_columns(self, data, drop_entities=True, drop_timing=False):

        entities = []
        if drop_entities:
            entities += ['run', 'session', 'subject', 'task']
        if drop_timing:
            entities += ['onset', 'duration']
        common_ents = list(set(entities) & set(data.columns))
        return data.drop(common_ents, axis=1)

    def _get_groupby_cols(self):
        # Get a list of keywords that define the grouping at the current level.
        # Note that we need to include all entities *above* the current one--
        # e.g., if the block level is 'run', this means we actually want to
        # groupby(['subject', 'session', 'run']), otherwise we would only
        # end up with n(runs) groups.
        hierarchy = ['subject', 'session', 'run']

        if self.level is None:
            return None
        elif self.level not in hierarchy:  # e.g., for 'dataset'
            return []
        pos = hierarchy.index(self.level)
        return hierarchy[:(pos + 1)]

    def apply_transformations(self):
        ''' Apply all transformations to the variables in the collection.
        '''
        for t in self.transformations:
            kwargs = dict(t)
            func = kwargs.pop('name')
            cols = kwargs.pop('input', None)

            if isinstance(func, string_types):
                if not hasattr(transform, func):
                    raise ValueError("No transformation '%s' found!" % func)
                func = getattr(transform, func)
                func(self.input_collection, cols, **kwargs)

    def generate_output_collection(self, keep_input_columns=True):

        data = self._get_design_matrix()

        # Figure out which column names to keep
        contrast_names = [c['name'] for c in self.contrasts]
        if keep_input_columns:
            sel_cols = list(data['condition'].unique())
            if self.model is not None and 'variables' in self.model:
                sel_cols = list(set(sel_cols) and set(self.model['variables']))
            contrast_names = sel_cols + contrast_names

        ent_cols = self._get_groupby_cols()
        ent_cols = list(set(ent_cols) & set(data.columns))

        contrast_names = list(set(contrast_names) -
                              set(self.input_collection.entities))

        # Set the same (unit) weights for all rows
        unit_weights = pd.Series(np.ones(len(contrast_names)),
                                 index=contrast_names)

        if ent_cols:
            data = data.groupby(ent_cols).apply(lambda x: unit_weights)

        # Generate a new BIDSVariableCollection to pass to next block
        ic = self.input_collection

        collection = BIDSVariableCollection(ic.layout, entities=ent_cols)

        for col_name in data.columns:
            col_data = data[col_name].reset_index()
            col_data = col_data.rename(columns={col_name: 'amplitude'})
            col = SimpleColumn(collection, col_name, col_data)
            collection.columns[col_name] = col

        self.output_collection = collection
        for k, v in collection.columns.items():
            pass

    def get_contrasts(self, format='matrix', **selectors):
        ''' Return contrast information for the current block.
        Args:
            format (str): What format to return the contrast specifications in.
                Valid options are:
                    'matrix' or 'df': Returns a pandas DataFrame with each
                        contrast as a row and each of the existing design
                        matrix columns as a column. This format makes it easy
                        to matrix-multiply existing in-memory images by the
                        contrast definition matrix in one shot.
                    'patsy': Returns a list of strings, where each string gives
                        a patsy-compatible definition of the contrast. E.g.,
                        if there are conditions 'A' and 'B', and the weights
                        are [1, -1], the returned string would be "A-B".
                    'dict': Returns the BIDS-Model contrast specification
                        as a dict loaded from the original json.<commit_msg>filter columns if variables field set in model<commit_after>import json
from .variables import load_variables, SimpleColumn, BIDSVariableCollection
from . import transform
from collections import namedtuple
from six import string_types
import pandas as pd
import numpy as np


DesignMatrix = namedtuple('DesignMatrix', ('data', 'image', 'entities'))


class Analysis(object):

    ''' Represents an entire BIDS-Model analysis.
    Args:
        layouts (BIDSLayout or list): One or more BIDSLayout objects to pull
            variables from.
        model (str or dict): a BIDS model specification. Can either be a
            string giving the path of the JSON model spec, or an already-loaded
            dict containing the model info.
        collection (BIDSVariableCollection): Optional BIDSVariableCollection
            object used to load/manage variables. If None, a new collection is
            initialized from the provided layouts.
        selectors (dict): Optional keyword arguments to pass onto the
            collection; these will be passed on to the Layout's .get() method,
            and can be used to restrict variables.
    '''

    def __init__(self, model, collection=None, layouts=None, **selectors):

        if layouts is None and collection is None:
            raise ValueError("At least one of the 'layout' and 'collection' "
                             "arguments must be passed.")

        if isinstance(model, str):
            model = json.load(open(model))
        self.model = model

        if 'input' in model:
            selectors.update(model['input'])

        if collection is None:
            collection = load_variables(layouts, **selectors)

        self.collection = collection

        self._load_blocks(model['blocks'])

    @property
    def layout(self):
        return self.collection.layout  # for convenience

    def __iter__(self):
        for b in self.blocks:
            yield b

    def __getitem__(self, index):
        if isinstance(index, int):
            return self.blocks[index]
        return list(filter(lambda x: x.name == index, self.blocks))[0]

    def _load_blocks(self, blocks):
        self.blocks = []
        for i, b in enumerate(blocks):
            self.blocks.append(Block(self, index=i, **b))

    def setup(self, apply_transformations=True):
        ''' Set up the sequence of blocks--applying transformations,
        generating outputs, etc. '''

        # pass a copy of the collection through the pipeline (columns mutate)
        _collection = self.collection.clone()

        for b in self.blocks:
            b.setup(_collection, apply_transformations=apply_transformations)
            _collection = b.output_collection.clone()


class Block(object):

    ''' Represents a single analysis block from a BIDS-Model specification.
    Args:
        analysis (Analysis): The parent Analysis this Block belongs to.
        level (str): The BIDS keyword to use as the grouping variable; must be
            one of ['run', 'session', 'subject', or 'dataset'].
        index (int): The numerical index of the current Block within the
            sequence of blocks.
        name (str): Optional name to assign to the block. Must be specified
            in order to enable name-based indexing in the parent Analysis.
        transformations (list): List of BIDS-Model transformations to apply.
        model (dict): The 'model' part of the BIDS-Model block specification.
        contrasts (list): List of contrasts to apply to the parameter estimates
            generated when the model is fit.
    '''

    def __init__(self, analysis, level, index, name=None, transformations=None,
                 model=None, contrasts=None):

        self.analysis = analysis
        self.level = level
        self.index = index
        self.name = name
        self.transformations = transformations or []
        self.model = model or None
        self.contrasts = contrasts or []
        self._design_matrix = None
        self.input_collection = None
        self.output_collection = None

    def _get_design_matrix(self, **selectors):
        if self._design_matrix is None:
            raise ValueError("Block hasn't been set up yet; please call "
                             "setup() before you try to retrieve the DM.")
        # subset the data if needed
        data = self._design_matrix
        if selectors:
            # TODO: make sure this handles constraints on int columns properly
            bad_keys = list(set(selectors.keys()) - set(data.columns))
            if bad_keys:
                raise ValueError("The following query constraints do not map "
                                 "onto existing columns: %s." % bad_keys)
            query = ' and '.join(["{} in {}".format(k, v)
                                  for k, v in selectors.items()])
            data = data.query(query)
        return data

    def _drop_columns(self, data, drop_entities=True, drop_timing=False):

        entities = []
        if drop_entities:
            entities += ['run', 'session', 'subject', 'task']
        if drop_timing:
            entities += ['onset', 'duration']
        common_ents = list(set(entities) & set(data.columns))
        return data.drop(common_ents, axis=1)

    def _get_groupby_cols(self):
        # Get a list of keywords that define the grouping at the current level.
        # Note that we need to include all entities *above* the current one--
        # e.g., if the block level is 'run', this means we actually want to
        # groupby(['subject', 'session', 'run']), otherwise we would only
        # end up with n(runs) groups.
        hierarchy = ['subject', 'session', 'run']

        if self.level is None:
            return None
        elif self.level not in hierarchy:  # e.g., for 'dataset'
            return []
        pos = hierarchy.index(self.level)
        return hierarchy[:(pos + 1)]

    def apply_transformations(self):
        ''' Apply all transformations to the variables in the collection.
        '''
        for t in self.transformations:
            kwargs = dict(t)
            func = kwargs.pop('name')
            cols = kwargs.pop('input', None)

            if isinstance(func, string_types):
                if not hasattr(transform, func):
                    raise ValueError("No transformation '%s' found!" % func)
                func = getattr(transform, func)
                func(self.input_collection, cols, **kwargs)

        # Also apply variable selection even if it's not represented in
        # its own transformation.
        if self.model is not None and 'variables' in self.model:
            # TODO: add a "select" transformation and construct a call to it
            # here instead of filtering directly on the collection.
            coll = self.input_collection
            coll.columns = {k: v for k, v in coll.columns.items()
                            if k in self.model['variables']}

    def generate_output_collection(self, keep_input_columns=True):

        data = self._get_design_matrix()

        # Figure out which column names to keep
        contrast_names = [c['name'] for c in self.contrasts]
        if keep_input_columns:
            sel_cols = list(data['condition'].unique())
            contrast_names = sel_cols + contrast_names

        ent_cols = self._get_groupby_cols()
        ent_cols = list(set(ent_cols) & set(data.columns))

        contrast_names = list(set(contrast_names) -
                              set(self.input_collection.entities))

        # Set the same (unit) weights for all rows
        unit_weights = pd.Series(np.ones(len(contrast_names)),
                                 index=contrast_names)

        if ent_cols:
            data = data.groupby(ent_cols).apply(lambda x: unit_weights)

        # Generate a new BIDSVariableCollection to pass to next block
        ic = self.input_collection

        collection = BIDSVariableCollection(ic.layout, entities=ent_cols)

        for col_name in data.columns:
            col_data = data[col_name].reset_index()
            col_data = col_data.rename(columns={col_name: 'amplitude'})
            col = SimpleColumn(collection, col_name, col_data)
            collection.columns[col_name] = col

        self.output_collection = collection

    def get_contrasts(self, format='matrix', **selectors):
        ''' Return contrast information for the current block.
        Args:
            format (str): What format to return the contrast specifications in.
                Valid options are:
                    'matrix' or 'df': Returns a pandas DataFrame with each
                        contrast as a row and each of the existing design
                        matrix columns as a column. This format makes it easy
                        to matrix-multiply existing in-memory images by the
                        contrast definition matrix in one shot.
                    'patsy': Returns a list of strings, where each string gives
                        a patsy-compatible definition of the contrast. E.g.,
                        if there are conditions 'A' and 'B', and the weights
                        are [1, -1], the returned string would be "A-B".
                    'dict': Returns the BIDS-Model contrast specification
                        as a dict loaded from the original json.
{'fraction_non_alphanumeric': 0.08291811091854419, 'fraction_numerical': 0.0005415944540727903, 'mean_word_length': 2.1096328730212193, '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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1418877', 'n_tokens_mistral': 4643, 'n_tokens_neox': 4506, 'n_words': 1706}