added draft and first tests for vectorized Histogram collection

cmsl1t/collections/__init__.py

@@ -5,10 +5,12 @@
 from .by_pileup import HistogramsByPileUpCollection
 from .resolution import ResolutionCollection
 from .efficiency import EfficiencyCollection
+from .vectorized import VectorizedHistCollection
 
 __all__ = [
     'BaseHistCollection',
     'HistogramsByPileUpCollection',
     'ResolutionCollection',
     'EfficiencyCollection',
+    'VectorizedHistCollection',
 ]

cmsl1t/collections/vectorized.py

@@ -0,0 +1,28 @@
+import numbda
+
+from . import BaseHistCollection
+
+
+@numba.jit(nopython=True)
+def extend(arr1, starts, stops):
+    repeat = stops - starts
+    return np.repeat(arr1, repeat, axis=0)
+
+
+class VectorizedHistCollection(object):
+
+    def __init__(self, innerBins):
+        self._innerBins = innerBins
+        self._innerHist = Hist(100, 0, 100, name='inner')
+
+    def _get_inner_indices(self, values):
+        '''
+            Returns the pileup bin corresponding to the provided pileup value.
+             - bin 0 is underflow
+             - bin len(innerBins) is overflow
+
+            :Example:
+                >>> hists = VectorizedHistCollection(innerBins=[0,10,15,20,30,999])
+                >>> hists._get_inner_indices([1, 11, 1111]) # returns [0, 1, 5]
+        '''
+        return np.digitize(values, self._innerBins)

test/collections/test_vectorized.py

@@ -0,0 +1,19 @@
+import pytest
+import numpy as np
+from rootpy.plotting import Hist
+
+from cmsl1t.collections import VectorizedHistCollection
+
+
+@pytest.mark.parametrize(
+    "values,expected",
+    [
+        ([1, 12, 1, 50], [1, 2, 1, 5]),
+        ([1, 11, 1111], [1, 2, 6]),
+        ([-10, 1111, 20], [0, 6, 4]),
+    ])
+def test_inner_index(values, expected):
+    innerBins = np.array([0, 10, 15, 20, 30, 999])
+    coll = VectorizedHistCollection(innerBins)
+
+    np.testing.assert_array_equal(coll._get_inner_indices(values), expected)