Add slicing

src/babyyoda/Histo1D_v2.py

@@ -6,11 +6,14 @@ def set_bin(target, source):
     # TODO allow modify those?
     # self.d_xmin = bin.xMin()
     # self.d_xmax = bin.xMax()
-    target.set(
-        source.numEntries(),
-        [source.sumW(), source.sumWX()],
-        [source.sumW2(), source.sumWX2()],
-    )
+    if hasattr(target, "set"):
+        target.set(
+            source.numEntries(),
+            [source.sumW(), source.sumWX()],
+            [source.sumW2(), source.sumWX2()],
+        )
+    else:
+        raise NotImplementedError("YODA1 backend can not set bin values")
 
 
 # TODO make this implementation independent (no V2 or V3...)
@@ -59,6 +62,9 @@ def __call__(self, *args, **kwargs):
     # YODA compatibility code (dropped legacy code?)
     ########################################################
 
+    def copy(self):
+        return HISTO1D_V2(self.target.copy())
+
     def overflow(self):
         # if target has overflow method, call it
         if hasattr(self.target, "overflow"):
@@ -71,6 +77,9 @@ def underflow(self):
             return self.target.underflow()
         return self.bins(includeOverflows=True)[0]
 
+    def errWs(self):
+        return np.sqrt(np.array([b.sumW2() for b in self.bins()]))
+
     def xMins(self):
         return np.array([b.xMin() for b in self.bins()])
 
@@ -83,6 +92,12 @@ def sumWs(self):
     def sumW2s(self):
         return np.array([b.sumW2() for b in self.bins()])
 
+    def rebinBy(self, *args, **kwargs):
+        self.rebinXBy(*args, **kwargs)
+
+    def rebinTo(self, *args, **kwargs):
+        self.rebinXTo(*args, **kwargs)
+
     ########################################################
     # Generic UHI code
     ########################################################
@@ -128,11 +143,16 @@ def __getitem__(self, slices):
             if isinstance(step, rebin):
                 if start is None:
                     start = 0
-                return self.rebinBy(step.factor, start, stop)
-
-            # print(f" {start} {stop} {step}")
+                cs = self.copy()
+                cs.rebinBy(step.factor, start, stop)
+                return cs
 
-            return self.rebinTo(self.bins()[start:stop])
+            print(f" {start} {stop} {step}")
+            if stop is not None:
+                stop += 1
+            sc = self.copy()
+            sc.rebinTo(self.xEdges()[start:stop])
+            return sc
 
         raise TypeError("Invalid argument type")
 
@@ -141,7 +161,7 @@ def __get_index(self, slices):
         if isinstance(slices, int):
             index = slices
             while index < 0:
-                index = len(self.bins) + index
+                index = len(self.bins()) + index
         if isinstance(slices, loc):
             # TODO cyclic maybe
             idx = None

src/babyyoda/grogu/histo1d_v2.py

@@ -21,6 +21,17 @@ class Bin:
         # YODA compatibilty code
         ########################################################
 
+        def copy(self):
+            return GROGU_HISTO1D_V2.Bin(
+                d_xmin=self.d_xmin,
+                d_xmax=self.d_xmax,
+                d_sumw=self.d_sumw,
+                d_sumw2=self.d_sumw2,
+                d_sumwx=self.d_sumwx,
+                d_sumwx2=self.d_sumwx2,
+                d_numentries=self.d_numentries,
+            )
+
         def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool:
             # if (self.d_xmin is None or x > self.d_xmin) and (self.d_xmax is None or x < self.d_xmax):
             sf = fraction * weight
@@ -55,6 +66,9 @@ def xMin(self):
         def xMax(self):
             return self.d_xmax
 
+        def xMid(self):
+            return (self.d_xmin + self.d_xmax) / 2
+
         def sumW(self):
             return self.d_sumw
 
@@ -117,18 +131,16 @@ def __eq__(self, other):
 
         def __add__(self, other):
             assert isinstance(other, GROGU_HISTO1D_V2.Bin)
-            nxhigh = None
-            nxlow = None
-            # combine if the bins are adjacent
-            if self.d_xmax == other.d_xmin:
-                nxlow = self.d_xmin
-                nxhigh = other.d_xmax
-            elif self.d_xmin == other.d_xmax:
-                nxlow = other.d_xmin
-                nxhigh = self.d_xmax
+            ## combine if the bins are adjacent
+            # if self.d_xmax == other.d_xmin:
+            #    nxlow = self.d_xmin
+            #    nxhigh = other.d_xmax
+            # elif self.d_xmin == other.d_xmax:
+            #    nxlow = other.d_xmin
+            #    nxhigh = self.d_xmax
             return GROGU_HISTO1D_V2.Bin(
-                nxlow,
-                nxhigh,
+                self.d_xmin,
+                self.d_xmax,
                 self.d_sumw + other.d_sumw,
                 self.d_sumw2 + other.d_sumw2,
                 self.d_sumwx + other.d_sumwx,
@@ -147,6 +159,16 @@ def __post_init__(self):
     # YODA compatibilty code
     ############################################
 
+    def copy(self):
+        return GROGU_HISTO1D_V2(
+            d_key=self.d_key,
+            d_path=self.d_path,
+            d_title=self.d_title,
+            d_bins=[b.copy() for b in self.d_bins],
+            d_underflow=self.d_underflow,
+            d_overflow=self.d_overflow,
+        )
+
     def underflow(self):
         return self.d_underflow
 
@@ -183,28 +205,57 @@ def binAt(self, x):
     def binDim(self):
         return 1
 
-    def rebinBy(self, factor: int, start: None, stop: None):
+    def rebinXBy(self, factor: int, start: None, stop: None):
+        # TODO what about not fitting start stop with factor?!
         if start is None:
             start = 0
         if stop is None:
             stop = len(self.bins())
         new_bins = []
         for i in range(start, stop, factor):
-            nb = self.bins[i]
-            for j in range(1, factor):
+            nb = GROGU_HISTO1D_V2.Bin(
+                d_xmin=self.bins[i].xMin(), d_xmax=self.bins[i].xMax()
+            )
+            for j in range(0, factor):
                 nb += self.bins[i + j]
+                nb.d_xmin = min(nb.d_xmin, self.bins[i + j].xMin())
+                nb.d_xmax = max(nb.d_xmax, self.bins[i + j].xMax())
+
             new_bins.append(nb)
-        return GROGU_HISTO1D_V2(
-            d_key=self.d_key,
-            d_path=self.d_path,
-            d_title=self.d_title,
-            d_bins=new_bins,
-            d_underflow=self.d_underflow,
-            d_overflow=self.d_overflow,
-        )
+        self.d_bins = new_bins
+        # return self
+        # not inplace
+        # return GROGU_HISTO1D_V2(
+        #    d_key=self.d_key,
+        #    d_path=self.d_path,
+        #    d_title=self.d_title,
+        #    d_bins=new_bins,
+        #    d_underflow=self.d_underflow,
+        #    d_overflow=self.d_overflow,
+        # )
+
+    def xEdges(self):
+        return [b.xMin() for b in self.d_bins] + [self.xMax()]
+
+    def rebinXTo(self, edges: List[float]):
+        own_edges = self.xEdges()
+        for e in edges:
+            assert e in own_edges, f"Edge {e} not found in own edges {own_edges}"
 
-    def rebinTo(bins):
-        raise NotImplementedError
+        new_bins = []
+        for i in range(len(edges) - 1):
+            new_bins.append(GROGU_HISTO1D_V2.Bin(d_xmin=edges[i], d_xmax=edges[i + 1]))
+        for b in self.bins():
+            if b.xMid() < min(edges):
+                self.d_underflow += b
+            elif b.xMid() > max(edges):
+                self.d_overflow += b
+            else:
+                for i in range(len(edges) - 1):
+                    if edges[i] <= b.xMid() and b.xMid() <= edges[i + 1]:
+                        new_bins[i] += b
+        self.d_bins = new_bins
+        # return self
 
 
 def parse_histo1d_v2(file_content: str, key: str = "") -> GROGU_HISTO1D_V2:

tests/grogu/uhi/test_gg_histo1d_slicing.py

@@ -0,0 +1,78 @@
+from babyyoda.Histo1D_v2 import HISTO1D_V2
+import babyyoda.grogu as yoda
+from babyyoda.test import assert_histo1d
+from babyyoda.util import loc, overflow, underflow
+
+
+def get_histo1d():
+    h = yoda.Histo1D(10, 0, 10, title="test")
+    for i in range(12):
+        for _ in range(i):
+            h.fill(i)
+    h.underflow().fill(-1)
+    h.overflow().fill(10)
+    return HISTO1D_V2(h)
+
+
+def test_slicing_everything():
+    yuhi1d = get_histo1d()
+    assert yuhi1d.copy() != yuhi1d
+    assert_histo1d(yuhi1d.copy(), yuhi1d)
+    assert yuhi1d[:] != yuhi1d
+    assert_histo1d(yuhi1d[:], yuhi1d)
+    assert yuhi1d.copy()[:] != yuhi1d
+
+
+def test_slicing_subset():
+    yuhi1d = get_histo1d()
+    assert yuhi1d.copy()[1:3] != yuhi1d
+    assert yuhi1d[1:3] != yuhi1d[1:3]
+    assert_histo1d(yuhi1d[1:3], yuhi1d[1:3])
+    assert yuhi1d[1:3][0].sumW() == yuhi1d[1].sumW()
+
+
+def test_slicing_upper_bound():
+    yuhi1d = get_histo1d()
+    print(yuhi1d[: loc(5)])
+    assert yuhi1d[: loc(5)][-1].sumW() == yuhi1d[loc(4)].sumW()
+    assert yuhi1d[: loc(5)][-1].sumW() == yuhi1d[4].sumW()
+
+
+def test_slicing_lower_bound():
+    yuhi1d = get_histo1d()
+    assert yuhi1d[loc(5) :][0].sumW() == yuhi1d[5].sumW()
+
+
+def test_slicing_mixed_bound():
+    yuhi1d = get_histo1d()
+    assert yuhi1d[1:9] != yuhi1d[1:-1]
+    assert_histo1d(yuhi1d[1:9], yuhi1d[1:-1])
+    assert yuhi1d[1:] != yuhi1d[1:10]
+    assert_histo1d(yuhi1d[1:], yuhi1d[1:10])
+    assert yuhi1d[1:][2:-1] != yuhi1d[1:10][2:8]
+    assert_histo1d(yuhi1d[1:][2:-1], yuhi1d[1:10][2:8])
+    assert yuhi1d[:3][2:] != yuhi1d[:3][2:]
+    assert_histo1d(yuhi1d[:3][2:], yuhi1d[:3][2:])
+
+    assert yuhi1d[:3][overflow].sumW() == yuhi1d[2:3][overflow].sumW()
+    assert yuhi1d[:3][2:][overflow].sumW() == yuhi1d[2:3][overflow].sumW()
+    assert yuhi1d[2:][:3] != yuhi1d[2:5]
+    assert_histo1d(yuhi1d[2:][:3], yuhi1d[2:5])
+
+
+def test_slicing_overflow():
+    yuhi1d = get_histo1d()
+    assert (yuhi1d)[overflow].sumW() == yuhi1d[overflow].sumW()
+    assert (yuhi1d[:3])[overflow].sumW() == yuhi1d[:3][overflow].sumW()
+    assert (yuhi1d)[overflow].sumW() != yuhi1d[:3][overflow].sumW()
+    assert (yuhi1d)[overflow].sumW() == yuhi1d[3:][overflow].sumW()
+    assert (yuhi1d[:3])[2:][overflow].sumW() == yuhi1d[:3][overflow].sumW()
+
+
+def test_slicing_underflow():
+    yuhi1d = get_histo1d()
+    assert (yuhi1d)[underflow].sumW() == yuhi1d[underflow].sumW()
+    assert (yuhi1d[3:])[underflow].sumW() == yuhi1d[3:][underflow].sumW()
+    assert (yuhi1d)[underflow].sumW() == yuhi1d[:3][underflow].sumW()
+    assert (yuhi1d)[underflow].sumW() != yuhi1d[3:][underflow].sumW()
+    assert (yuhi1d[3:])[:2][underflow].sumW() == yuhi1d[3:][underflow].sumW()

tests/yoda/uhi/test_gg_histo1d_setting.py → tests/yoda/uhi/test_yd_histo1d_setting.py

@@ -1,9 +1,14 @@
+import pytest
 from babyyoda.Histo1D_v2 import HISTO1D_V2
-import yoda
 from babyyoda.util import loc, overflow, underflow
 
+# YODA1 does not support setting
+pytest.importorskip("yoda", minversion="2.0.0")
+
 
 def get_histo1d():
+    import yoda
+
     h = yoda.Histo1D(10, 0, 10, title="test")
     for i in range(12):
         for _ in range(i):