setting works with yoda2 too

src/babyyoda/Histo1D_v2.py

@@ -2,6 +2,18 @@
 from babyyoda.util import loc, overflow, rebin, underflow
 
 
+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()],
+    )
+
+
+# TODO make this implementation independent (no V2 or V3...)
 class HISTO1D_V2:
     def __init__(self, target):
         """
@@ -92,20 +104,6 @@ def values(self):
     def variances(self):
         return np.array([(b.sumW2()) for b in self.bins()])
 
-    def __setitem__(self, slices, value):
-        # integer index
-        index = self.__get_index(slices)
-        self.__set_by_index(index, value)
-
-    def __set_by_index(self, index, value):
-        if index == underflow:
-            self.underflow = value
-            return
-        if index == overflow:
-            self.overflow = value
-            return
-        self.bins()[index] = value
-
     def __getitem__(self, slices):
         index = self.__get_index(slices)
         # integer index
@@ -157,6 +155,20 @@ def __get_index(self, slices):
             index = overflow
         return index
 
+    def __set_by_index(self, index, value):
+        if index == underflow:
+            set_bin(self.underflow(), value)
+            return
+        if index == overflow:
+            set_bin(self.overflow(), value)
+            return
+        set_bin(self.bins()[index], value)
+
+    def __setitem__(self, slices, value):
+        # integer index
+        index = self.__get_index(slices)
+        self.__set_by_index(index, value)
+
     def plot(self, *args, binwnorm=1.0, **kwargs):
         import mplhep as hep
 

src/babyyoda/grogu/__init__.py

@@ -1,3 +1,22 @@
 from .read import read
+from .histo1d_v2 import GROGU_HISTO1D_V2 as Histo1D_v2
+
 
 __all__ = ["read"]
+
+
+# TODO same function for Hist1D in babyyoda.Histo1D_v2, but how pick backend? probably just yoda if yoda available
+def Histo1D(nbins: int, start: float, end: float, title=None, **kwargs):
+    return Histo1D_v2(
+        d_bins=[
+            Histo1D_v2.Bin(
+                d_xmin=start + i * (end - start) / nbins,
+                d_xmax=start + (i + 1) * (end - start) / nbins,
+            )
+            for i in range(nbins)
+        ],
+        d_overflow=Histo1D_v2.Bin(),
+        d_underflow=Histo1D_v2.Bin(),
+        d_title=title,
+        **kwargs,
+    )

src/babyyoda/grogu/histo1d_v2.py

@@ -1,5 +1,5 @@
 import re
-from typing import Optional
+from typing import List, Optional
 from dataclasses import dataclass, field
 
 from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT
@@ -30,6 +30,25 @@ def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool:
             self.d_sumwx2 += sf * x**2
             self.d_numentries += fraction
 
+        def set_bin(self, bin):
+            # TODO allow modify those?
+            # self.d_xmin = bin.xMin()
+            # self.d_xmax = bin.xMax()
+            self.d_sumw = bin.sumW()
+            self.d_sumw2 = bin.sumW2()
+            self.d_sumwx = bin.sumWX()
+            self.d_sumwx2 = bin.sumWX2()
+            self.d_numentries = bin.numEntries()
+
+        def set(self, numEntries: float, sumW: List[float], sumW2: List[float]):
+            assert len(sumW) == 2
+            assert len(sumW2) == 2
+            self.d_sumw = sumW[0]
+            self.d_sumw2 = sumW2[0]
+            self.d_sumwx = sumW[1]
+            self.d_sumwx2 = sumW2[1]
+            self.d_numentries = numEntries
+
         def xMin(self):
             return self.d_xmin
 

src/babyyoda/grogu/histo2d_v2.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional
 
 from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT
 
@@ -37,6 +37,40 @@ def fill(self, x: float, y: float, weight: float = 1.0, fraction=1.0):
             self.d_sumwxy += sf * x * y
             self.d_numentries += fraction
 
+        def set_bin(self, bin):
+            # TODO allow modify those?
+            # self.d_xmin = bin.xMin()
+            # self.d_xmax = bin.xMax()
+            # self.d_ymin = bin.yMin()
+            # self.d_ymax = bin.yMax()
+            self.d_sumw = bin.sumW()
+            self.d_sumw2 = bin.sumW2()
+            self.d_sumwx = bin.sumWX()
+            self.d_sumwx2 = bin.sumWX2()
+            self.d_sumwy = bin.sumWY()
+            self.d_sumwy2 = bin.sumWY2()
+            self.d_sumwxy = bin.sumWXY()
+            self.d_numentries = bin.numEntries()
+
+        def set(
+            self,
+            numEntries: float,
+            sumW: List[float],
+            sumW2: List[float],
+            sumWcross: List[float],
+        ):
+            assert len(sumW) == 3
+            assert len(sumW2) == 3
+            assert len(sumWcross) == 1
+            self.d_sumw = sumW[0]
+            self.d_sumw2 = sumW2[0]
+            self.d_sumwx = sumW[1]
+            self.d_sumwx2 = sumW2[1]
+            self.d_sumwy = sumW[2]
+            self.d_sumwy2 = sumW2[2]
+            self.d_sumwxy = sumWcross[0]
+            self.d_numentries = numEntries
+
         def xMin(self):
             return self.d_xmin
 

tests/grogu/uhi/test_gg_histo1d_setting.py

@@ -0,0 +1,49 @@
+from babyyoda.Histo1D_v2 import HISTO1D_V2
+import babyyoda.grogu as yoda
+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_setting_index():
+    yuhi1d = get_histo1d()
+    yuhi1d[1] = yuhi1d[2]
+    yuhi1d[7] = yuhi1d[8]
+    yuhi1d[3] = yuhi1d[5]
+    assert yuhi1d[1].sumW() == yuhi1d[2].sumW()
+    assert yuhi1d[7].sumW2() == yuhi1d[8].sumW2()
+    assert yuhi1d[3].numEntries() == yuhi1d[5].numEntries()
+
+
+def test_setting_loc():
+    yuhi1d = get_histo1d()
+    yuhi1d[loc(1)] = yuhi1d[2]
+    yuhi1d[loc(7)] = yuhi1d[8]
+    yuhi1d[loc(3)] = yuhi1d[5]
+    assert yuhi1d[1].sumW() == yuhi1d[2].sumW()
+    assert yuhi1d[7].sumW2() == yuhi1d[8].sumW2()
+    assert yuhi1d[3].numEntries() == yuhi1d[5].numEntries()
+
+
+def test_setting_underflow():
+    yuhi1d = get_histo1d()
+    yuhi1d[underflow] = yuhi1d[overflow]
+    assert yuhi1d[underflow].sumW() == yuhi1d[overflow].sumW()
+    assert yuhi1d[underflow].sumW2() == yuhi1d[overflow].sumW2()
+    assert yuhi1d[underflow].numEntries() == yuhi1d[overflow].numEntries()
+
+
+def test_setting_overflow():
+    yuhi1d = get_histo1d()
+    yuhi1d[overflow] = yuhi1d[1]
+    assert yuhi1d[overflow].sumW() == yuhi1d[1].sumW()
+    assert yuhi1d[overflow].sumW2() == yuhi1d[1].sumW2()
+    assert yuhi1d[overflow].numEntries() == yuhi1d[1].numEntries()

tests/yoda/uhi/test_gg_histo1d_setting.py

@@ -0,0 +1,50 @@
+from babyyoda.Histo1D_v2 import HISTO1D_V2
+import yoda
+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 = HISTO1D_V2(h)
+    h.underflow().fill(-1)
+    h.overflow().fill(10)
+    return h
+
+
+def test_setting_index():
+    yuhi1d = get_histo1d()
+    yuhi1d[1] = yuhi1d[2]
+    yuhi1d[7] = yuhi1d[8]
+    yuhi1d[3] = yuhi1d[5]
+    assert yuhi1d[1].sumW() == yuhi1d[2].sumW()
+    assert yuhi1d[7].sumW2() == yuhi1d[8].sumW2()
+    assert yuhi1d[3].numEntries() == yuhi1d[5].numEntries()
+
+
+def test_setting_loc():
+    yuhi1d = get_histo1d()
+    yuhi1d[loc(1)] = yuhi1d[2]
+    yuhi1d[loc(7)] = yuhi1d[8]
+    yuhi1d[loc(3)] = yuhi1d[5]
+    assert yuhi1d[1].sumW() == yuhi1d[2].sumW()
+    assert yuhi1d[7].sumW2() == yuhi1d[8].sumW2()
+    assert yuhi1d[3].numEntries() == yuhi1d[5].numEntries()
+
+
+def test_setting_underflow():
+    yuhi1d = get_histo1d()
+    yuhi1d[underflow] = yuhi1d[overflow]
+    assert yuhi1d[underflow].sumW() == yuhi1d[overflow].sumW()
+    assert yuhi1d[underflow].sumW2() == yuhi1d[overflow].sumW2()
+    assert yuhi1d[underflow].numEntries() == yuhi1d[overflow].numEntries()
+
+
+def test_setting_overflow():
+    yuhi1d = get_histo1d()
+    yuhi1d[overflow] = yuhi1d[1]
+    assert yuhi1d[overflow].sumW() == yuhi1d[1].sumW()
+    assert yuhi1d[overflow].sumW2() == yuhi1d[1].sumW2()
+    assert yuhi1d[overflow].numEntries() == yuhi1d[1].numEntries()