From 5946a249d3c2cd437515ddeff33517a059ae1294 Mon Sep 17 00:00:00 2001
From: kreczko <lkreczko@googlemail.com>
Date: Tue, 9 Jul 2019 10:49:19 +0100
Subject: [PATCH] added draft and first tests for vectorized Histogram
 collection

---
 cmsl1t/collections/__init__.py      |  2 ++
 cmsl1t/collections/vectorized.py    | 28 ++++++++++++++++++++++++++++
 test/collections/test_vectorized.py | 19 +++++++++++++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 cmsl1t/collections/vectorized.py
 create mode 100644 test/collections/test_vectorized.py

diff --git a/cmsl1t/collections/__init__.py b/cmsl1t/collections/__init__.py
index daab1251d3d..30722f35d3e 100644
--- a/cmsl1t/collections/__init__.py
+++ b/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',
 ]
diff --git a/cmsl1t/collections/vectorized.py b/cmsl1t/collections/vectorized.py
new file mode 100644
index 00000000000..58f9ce82b26
--- /dev/null
+++ b/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)
diff --git a/test/collections/test_vectorized.py b/test/collections/test_vectorized.py
new file mode 100644
index 00000000000..a0eef2bfcb9
--- /dev/null
+++ b/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)