<commit_before>"""
Class and standalone app for querying
"""

import os
import argparse
import glob
from numpy import zeros
from scipy.io import loadmat
from pyspark import SparkContext
from thunder.utils import load, getdims, subtoind
from thunder.utils import save


class Query(object):
    """Class for querying time series data by averaging together
    data points with the specified indices

    Parameters
    ----------
    indsfile : str, or array-like (2D)
        Array of indices, each an array-like of integer indices, or
        filename of a MAT file containing a set of indices as a cell array
        stored in the variable inds
    """

    def __init__(self, indsfile):
        if type(indsfile) is str:
            inds = loadmat(indsfile)['inds'][0]
        else:
            inds = indsfile
        self.inds = inds
        self.n = len(inds)

    def select(self, data, i):
        """Subselect rows of data with a given set of integer indices

        Parameters
        ----------
        data : RDD of (int, array) pairs, each array of shape (ncols,)
            Data to subselect rows from

        i : int
            Which set of indices to use

        Returns
        -------
        subset : RDD of (int, array) pairs
            Data with subset of rows
        """

        inds_set = set(self.inds[i].flat)
        inds_bc = data.context.broadcast(inds_set)
        subset = data.filter(lambda (k, _): k in inds_bc.value)
        return subset

    def calc(self, data):
        """Calculate averages. Keys (tuples) are converted
        into linear indices based on their dimensions

        Parameters
        ----------
        data : RDD of (tuple, array) pairs, each array of shape (ncols,)
            Data to compute averages from

        Returns
        -------
        ts : array, shape (n, ncols)
        """

        dims = getdims(data)
        data = subtoind(data, dims.max)

        # loop over indices, averaging time series
        ts = zeros((self.n, len(data.first()[1])))
        for i in range(0, self.n):
            if len(self.inds[i]) > 1:
                ts[i, :] = self.select(data, i).map(lambda (k, x): x).sum() / len(self.inds[i])<commit_msg>Fixed index bound<commit_after>"""
Class and standalone app for querying
"""

import os
import argparse
import glob
from numpy import zeros
from scipy.io import loadmat
from pyspark import SparkContext
from thunder.utils import load, getdims, subtoind
from thunder.utils import save


class Query(object):
    """Class for querying time series data by averaging together
    data points with the specified indices

    Parameters
    ----------
    indsfile : str, or array-like (2D)
        Array of indices, each an array-like of integer indices, or
        filename of a MAT file containing a set of indices as a cell array
        stored in the variable inds
    """

    def __init__(self, indsfile):
        if type(indsfile) is str:
            inds = loadmat(indsfile)['inds'][0]
        else:
            inds = indsfile
        self.inds = inds
        self.n = len(inds)

    def select(self, data, i):
        """Subselect rows of data with a given set of integer indices

        Parameters
        ----------
        data : RDD of (int, array) pairs, each array of shape (ncols,)
            Data to subselect rows from

        i : int
            Which set of indices to use

        Returns
        -------
        subset : RDD of (int, array) pairs
            Data with subset of rows
        """

        inds_set = set(self.inds[i].flat)
        inds_bc = data.context.broadcast(inds_set)
        subset = data.filter(lambda (k, _): k in inds_bc.value)
        return subset

    def calc(self, data):
        """Calculate averages. Keys (tuples) are converted
        into linear indices based on their dimensions

        Parameters
        ----------
        data : RDD of (tuple, array) pairs, each array of shape (ncols,)
            Data to compute averages from

        Returns
        -------
        ts : array, shape (n, ncols)
        """

        dims = getdims(data)
        data = subtoind(data, dims.max)

        # loop over indices, averaging time series
        ts = zeros((self.n, len(data.first()[1])))
        for i in range(0, self.n):
            if len(self.inds[i]) > 0:
                ts[i, :] = self.select(data, i).map(lambda (k, x): x).sum() / len(self.inds[i])
{'fraction_non_alphanumeric': 0.1044776119402985, 'fraction_numerical': 0.002261420171867933, 'mean_word_length': 1.9194719471947195, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3379107', 'n_tokens_mistral': 1295, 'n_tokens_neox': 1246, 'n_words': 471}