<commit_before>import numpy as np
from collections import OrderedDict


class LabeledArray(np.ndarray):
    """
    Each rows corresponds to labels, each columns corresponds to cells.
    Underlying data structure can be N-dimensional array. First dimension will be used for labeled array.
    
    Examples:
        >> arr = np.arange(12).reshape((3, 2, 2))
        >> labelarr = np.array([['a1' ,'b1', ''], 
                                ['a1' ,'b2' , 'c1'], 
                                ['a1' ,'b2' , 'c2']], dtype=object)
        >> darr = DArray(arr, labelarr)
        >> assert darr['a1'].shape
        (3, 2, 2)
        >> darr['a1', 'b1'].shape
        (2, 2)
        >> darr['a1', 'b2', 'c1']
        DArray([[4, 5],
               [6, 7]])
    """

    idx = None
    label = None
    
    def __new__(cls, arr=None, label=None, idx=None):
        obj = np.asarray(arr).view(cls)
        obj.label = label
        obj.idx = idx
        return obj
    
    def __array_finalize__(self, obj):
        if obj is None: return
        self.label = getattr(obj, 'label', None)
        if hasattr(obj, 'idx') and np.any(self.label):
            self.label = self.label[obj.idx]
            if isinstance(self.label, str):
                return
            if self.label.ndim > 1:
                f_leftshift = lambda a1:all(x>=y for x, y in zip(a1, a1[1:]))
                all_column = np.all(self.label == self.label[0,:], axis=0)
                sl = 0 if not f_leftshift(all_column) else all_column.sum()
                self.label = self.label[:, slice(sl, None)]


    def __getitem__(self, item):
        if isinstance(item, str):
            item = self._label2idx(item)
        if isinstance(item, tuple):
            if isinstance(item[0], str):
                item = self._label2idx(item)
        self.idx = item
        ret = super(LabeledArray, self).__getitem__(item)
        return ret.squeeze()

    def _label2idx(self, item):
        item = (item, ) if not isinstance(item, tuple) else item
        boolarr = np.ones(self.label.shape[0], dtype=bool)<commit_msg>fix bugs for 3d indexing for larger arrays<commit_after>import numpy as np
from collections import OrderedDict


class LabeledArray(np.ndarray):
    """
    Each rows corresponds to labels, each columns corresponds to cells.
    Underlying data structure can be N-dimensional array. First dimension will be used for labeled array.
    
    Examples:
        >> arr = np.arange(12).reshape((3, 2, 2))
        >> labelarr = np.array([['a1' ,'b1', ''], 
                                ['a1' ,'b2' , 'c1'], 
                                ['a1' ,'b2' , 'c2']], dtype=object)
        >> darr = DArray(arr, labelarr)
        >> assert darr['a1'].shape
        (3, 2, 2)
        >> darr['a1', 'b1'].shape
        (2, 2)
        >> darr['a1', 'b2', 'c1']
        DArray([[4, 5],
               [6, 7]])
    """

    idx = None
    label = None
    
    def __new__(cls, arr=None, label=None, idx=None):
        obj = np.asarray(arr).view(cls)
        obj.label = label
        obj.idx = idx
        return obj
    
    def __array_finalize__(self, obj):
        if obj is None: return
        self.label = getattr(obj, 'label', None)
        if hasattr(obj, 'idx') and np.any(self.label) and self.ndim > 1:
            self.label = self.label[obj.idx]
            if isinstance(self.label, str):
                return
            if self.label.ndim > 1:
                f_leftshift = lambda a1:all(x>=y for x, y in zip(a1, a1[1:]))
                all_column = np.all(self.label == self.label[0,:], axis=0)
                sl = 0 if not f_leftshift(all_column) else all_column.sum()
                self.label = self.label[:, slice(sl, None)]


    def __getitem__(self, item):
        if isinstance(item, str):
            item = self._label2idx(item)
        if isinstance(item, tuple):
            if isinstance(item[0], str):
                item = self._label2idx(item)
        self.idx = item
        ret = super(LabeledArray, self).__getitem__(item)
        return ret.squeeze()

    def _label2idx(self, item):
        item = (item, ) if not isinstance(item, tuple) else item
        boolarr = np.ones(self.label.shape[0], dtype=bool)
{'fraction_non_alphanumeric': 0.1372131535367873, 'fraction_numerical': 0.0198722498225692, 'mean_word_length': 1.8742352141400407, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 32, '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': '491305', 'n_tokens_mistral': 1388, 'n_tokens_neox': 1337, 'n_words': 355}