<commit_before>"""
This module should handle:
* conversion of TTrees into NumPy arrays
* conversion of TTrees into carrays (http://pypi.python.org/pypi/carray)
"""
from .types import Variable, convert
from .utils import asrootpy
import numpy as np

def to_numpy_array(trees, branches=None,
                   use_cache=False, cache_size=1000000,
                   include_weight=False,
                   weight_dtype='f4'):
    
    if type(trees) not in (list, tuple):
        trees = [trees]
    trees = [asrootpy(tree) for tree in trees]
    # if branches is None then select only branches with basic types
    # i.e. no vectors or other special objects
    tree = trees[0]
    _branches = {}
    if branches is None:
        for name, value in tree.buffer.items():
            if isinstance(value, Variable):
                _branches[name] = value
    else:
        for branch in branches:
            if branch not in tree.buffer:
                raise ValueError("Branch %s does not exist in tree" % branch)
            value = tree.buffer[branch]
            if not isinstance(value, Variable):
                raise TypeError("Branch %s is not a basic type: %s" % 
                                (branch, type(value)))
            _branches[branch] = tree.buffer[branch]
    if not _branches:
        return None
    dtype = [(name, convert('ROOTCODE', 'NUMPY', value.type)) for name, value in _branches.items()]
    if include_weight:
        if 'weight' not in _branches.keys():
            dtype.append(('weight', weight_dtype))
        else:
            include_weight = False
    total_entries = sum([tree.GetEntries() for tree in trees])
    array = np.recarray(shape=(total_entries,), dtype=dtype)
    i = 0
    for tree in trees:
        tree.use_cache(use_cache, cache_size=cache_size, learn_entries=1)
        if use_cache:
            tree.always_read(_branches.keys())
        weight = tree.GetWeight()
        for entry in tree:
            for j, (branch, value) in enumerate(_branches.items()):
                array[i][j] = entry[branch].value
            if include_weight:
                array[i][-1] = weight
            i += 1
    return array
<commit_msg>updates<commit_after>"""
This module should handle:
* conversion of TTrees into NumPy arrays
* conversion of TTrees into carrays (http://pypi.python.org/pypi/carray)
"""
from .types import Variable, convert
from .utils import asrootpy
import numpy as np

def to_numpy_array(trees, branches=None,
                   use_cache=False, cache_size=1000000,
                   include_weight=False,
                   weight_dtype='f4'):
    
    if type(trees) not in (list, tuple):
        trees = [trees]
    trees = [asrootpy(tree) for tree in trees]
    # if branches is None then select only branches with basic types
    # i.e. no vectors or other special objects
    tree = trees[0]
    _branches = {}
    if branches is None:
        for name, value in tree.buffer.items():
            if isinstance(value, Variable):
                _branches[name] = value
    else:
        if len(set(branches)) != len(branches):
            raise ValueError('branches contains duplicates')
        for branch in branches:
            if branch not in tree.buffer:
                raise ValueError("Branch %s does not exist in tree" % branch)
            value = tree.buffer[branch]
            if not isinstance(value, Variable):
                raise TypeError("Branch %s is not a basic type: %s" % 
                                (branch, type(value)))
            _branches[branch] = tree.buffer[branch]
    if not _branches:
        return None
    dtype = [(name, convert('ROOTCODE', 'NUMPY', value.type)) for name, value in _branches.items()]
    print dtype
    if include_weight:
        if 'weight' not in _branches.keys():
            dtype.append(('weight', weight_dtype))
        else:
            include_weight = False
    total_entries = sum([tree.GetEntries() for tree in trees])
    array = np.recarray(shape=(total_entries,), dtype=dtype)
    i = 0
    for tree in trees:
        tree.use_cache(use_cache, cache_size=cache_size, learn_entries=1)
        if use_cache:
            tree.always_read(_branches.keys())
        weight = tree.GetWeight()
        for entry in tree:
            for j, branch in enumerate(_branches.keys()):
                array[i][j] = entry[branch].value
            if include_weight:
                array[i][-1] = weight
            i += 1
    return array

{'fraction_non_alphanumeric': 0.09895252952975261, 'fraction_numerical': 0.005794517494985513, 'mean_word_length': 2.1384615384615384, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4186232', 'n_tokens_mistral': 1300, 'n_tokens_neox': 1256, 'n_words': 396}