Rebin tests

src/babyyoda/grogu/histo2d_v2.py

@@ -117,6 +117,12 @@ def xMin(self):
         def xMax(self):
             return self.d_xmax
 
+        def xMid(self):
+            return (self.d_xmin + self.d_xmax) / 2
+
+        def yMid(self):
+            return (self.d_ymin + self.d_ymax) / 2
+
         def yMin(self):
             return self.d_ymin
 
@@ -154,6 +160,23 @@ def crossTerm(self, x, y):
         def numEntries(self):
             return self.d_numentries
 
+        def __add__(self, other):
+            assert isinstance(other, GROGU_HISTO2D_V2.Bin)
+            return GROGU_HISTO2D_V2.Bin(
+                d_xmin=self.d_xmin,
+                d_xmax=self.d_xmax,
+                d_ymin=self.d_ymin,
+                d_ymax=self.d_ymax,
+                d_sumw=self.d_sumw + other.d_sumw,
+                d_sumw2=self.d_sumw2 + other.d_sumw2,
+                d_sumwx=self.d_sumwx + other.d_sumwx,
+                d_sumwx2=self.d_sumwx2 + other.d_sumwx2,
+                d_sumwy=self.d_sumwy + other.d_sumwy,
+                d_sumwy2=self.d_sumwy2 + other.d_sumwy2,
+                d_sumwxy=self.d_sumwxy + other.d_sumwxy,
+                d_numentries=self.d_numentries + other.d_numentries,
+            )
+
         def to_string(self, label=None) -> str:
             if label is None:
                 return (
@@ -239,6 +262,45 @@ def binAt(self, x, y):
                 return b
         return None
 
+    def rebinXYTo(self, xedges: list[float], yedges: list[float]):
+        own_xedges = self.xEdges()
+        for e in xedges:
+            assert e in own_xedges, f"Edge {e} not found in own edges {own_xedges}"
+        own_yedges = self.yEdges()
+        for e in yedges:
+            assert e in own_yedges, f"Edge {e} not found in own edges {own_yedges}"
+
+        new_bins = []
+        for j in range(len(yedges) - 1):
+            for i in range(len(xedges) - 1):
+                new_bins.append(
+                    GROGU_HISTO2D_V2.Bin(
+                        d_xmin=xedges[i],
+                        d_xmax=xedges[i + 1],
+                        d_ymin=yedges[j],
+                        d_ymax=yedges[j + 1],
+                    )
+                )
+        for b in self.bins():
+            for j in range(len(yedges) - 1):
+                for i in range(len(xedges) - 1):
+                    if (
+                        xedges[i] <= b.xMid()
+                        and b.xMid() <= xedges[i + 1]
+                        and yedges[j] <= b.yMid()
+                        and b.yMid() <= yedges[j + 1]
+                    ):
+                        new_bins[i + j * (len(xedges) - 1)] += b
+        self.d_bins = new_bins
+
+        assert len(self.d_bins) == (len(self.xEdges()) - 1) * (len(self.yEdges()) - 1)
+
+    def rebinXTo(self, xedges: list[float]):
+        self.rebinXYTo(xedges, self.yEdges())
+
+    def rebinYTo(self, yedges: list[float]):
+        self.rebinXYTo(self.xEdges(), yedges)
+
     def get_projector(self):
         return Histo1D_v2
 

src/babyyoda/grogu/histo2d_v3.py

@@ -6,9 +6,26 @@
 import numpy as np
 
 from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT
+from babyyoda.grogu.histo1d_v3 import Histo1D_v3
 from babyyoda.histo2d import UHIHisto2D
 
 
+def to_index(x, y, xedges, yedges):
+    # get ix and iy to map to correct bin
+    fix = None
+    for ix, xEdge in enumerate([*xedges, sys.float_info.max]):
+        fix = ix
+        if x < xEdge:
+            break
+    fiy = None
+    for iy, yEdge in enumerate([*yedges, sys.float_info.max]):
+        fiy = iy
+        if y < yEdge:
+            break
+    # Also fill overflow bins
+    return fiy * (len(xedges) + 1) + fix
+
+
 def Histo2D_v3(
     nxbins: int,
     xstart: float,
@@ -126,6 +143,19 @@ def crossTerm(self, x, y):
         def numEntries(self):
             return self.d_numentries
 
+        def __add__(self, other: "GROGU_HISTO2D_V3.Bin") -> "GROGU_HISTO2D_V3.Bin":
+            assert isinstance(other, GROGU_HISTO2D_V3.Bin)
+            return GROGU_HISTO2D_V3.Bin(
+                d_sumw=self.d_sumw + other.d_sumw,
+                d_sumw2=self.d_sumw2 + other.d_sumw2,
+                d_sumwx=self.d_sumwx + other.d_sumwx,
+                d_sumwx2=self.d_sumwx2 + other.d_sumwx2,
+                d_sumwy=self.d_sumwy + other.d_sumwy,
+                d_sumwy2=self.d_sumwy2 + other.d_sumwy2,
+                d_sumwxy=self.d_sumwxy + other.d_sumwxy,
+                d_numentries=self.d_numentries + other.d_numentries,
+            )
+
         def to_string(self) -> str:
             return (
                 f"{self.d_sumw:<13.6e}\t{self.d_sumw2:<13.6e}\t{self.d_sumwx:<13.6e}\t{self.d_sumwx2:<13.6e}\t"
@@ -171,19 +201,8 @@ def yEdges(self):
         return self.d_edges[1]
 
     def fill(self, x, y, weight=1.0, fraction=1.0):
-        # get ix and iy to map to correct bin
-        fix = None
-        for ix, xEdge in enumerate([*self.xEdges(), sys.float_info.max]):
-            fix = ix
-            if x < xEdge:
-                break
-        fiy = None
-        for iy, yEdge in enumerate([*self.yEdges(), sys.float_info.max]):
-            fiy = iy
-            if y < yEdge:
-                break
         # Also fill overflow bins
-        self.bins(True)[fiy * (len(self.xEdges()) + 1) + fix].fill(
+        self.bins(True)[to_index(x, y, self.xEdges(), self.yEdges())].fill(
             x, y, weight, fraction
         )
 
@@ -213,6 +232,72 @@ def bins(self, includeOverflows=False):
             .flatten()
         )
 
+    def rebinXYTo(self, xedges: list[float], yedges: list[float]):
+        own_xedges = self.xEdges()
+        for e in xedges:
+            assert e in own_xedges, f"Edge {e} not found in own edges {own_xedges}"
+        own_yedges = self.yEdges()
+        for e in yedges:
+            assert e in own_yedges, f"Edge {e} not found in own edges {own_yedges}"
+
+        # new bins inclusive of overflow and underflow
+        new_bins = []
+        for _ in range((len(xedges) + 1) * (len(yedges) + 1)):
+            new_bins.append(GROGU_HISTO2D_V3.Bin())
+        new_hist = np.array(new_bins).reshape((len(yedges) + 1, len(xedges) + 1))
+        old_hist = np.array(self.d_bins).reshape(
+            (len(self.yEdges()) + 1, len(self.xEdges()) + 1)
+        )
+        old_xedges = np.array(self.xEdges())
+        old_yedges = np.array(self.yEdges())
+        new_xedges = np.array(xedges)
+        new_yedges = np.array(yedges)
+
+        for i in range(len(xedges) + 1):
+            for j in range(len(yedges) + 1):
+                if i == 0:
+                    x_mask = old_xedges <= new_xedges[0]
+                    x_mask = np.concatenate((x_mask, [False]))  # skip overflow
+                elif i == len(xedges):
+                    x_mask = old_xedges > new_xedges[-1]
+                    x_mask = np.concatenate((x_mask, [True]))  # keep underflow
+                    # x_mask = x_mask + [True] # keep overflow
+                else:
+                    x_mask = (old_xedges > new_xedges[i - 1]) & (
+                        old_xedges <= new_xedges[i]
+                    )
+                    x_mask = np.concatenate((x_mask, [False]))  # skip overflow
+                    # x_mask = x_mask + [False] # skip overflow
+                if j == 0:
+                    y_mask = old_yedges <= new_yedges[0]
+                    y_mask = np.concatenate((y_mask, [False]))  # skip overflow
+                    # y_mask = y_mask + [False] # skip overflow
+                elif j == len(yedges):
+                    y_mask = old_yedges > new_yedges[-1]
+                    y_mask = np.concatenate((y_mask, [True]))  # keep underflow
+                    # y_mask = y_mask + [True] # keep overflow
+                else:
+                    y_mask = (old_yedges > new_yedges[j - 1]) & (
+                        old_yedges <= new_yedges[j]
+                    )
+                    y_mask = np.concatenate((y_mask, [False]))  # skip overflow
+                    # y_mask = y_mask + [False] # skip overflow
+                new_hist[j, i] = old_hist[y_mask, :][:, x_mask].sum()
+
+        self.d_bins = new_hist.flatten()
+        self.d_edges = [xedges, yedges]
+
+        assert len(self.d_bins) == (len(xedges) + 1) * (len(yedges) + 1)
+
+    def rebinXTo(self, xedges: list[float]):
+        self.rebinXYTo(xedges, self.yEdges())
+
+    def rebinYTo(self, yedges: list[float]):
+        self.rebinXYTo(self.xEdges(), yedges)
+
+    def get_projector(self):
+        return Histo1D_v3
+
     def to_string(self) -> str:
         """Convert a YODA_HISTO2D_V3 object to a formatted string."""
         header = (

src/babyyoda/histo1d.py

@@ -5,7 +5,7 @@
 
 from babyyoda.analysisobject import UHIAnalysisObject
 from babyyoda.counter import UHICounter
-from babyyoda.util import loc, overflow, project, rebin, underflow
+from babyyoda.util import loc, overflow, project, rebin, rebinBy_to_rebinTo, underflow
 
 
 def set_bin1d(target, source):
@@ -160,42 +160,8 @@ def integral(self, includeOverflows=True):
         return sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows))
 
     def rebinXBy(self, factor: int, begin=1, end=sys.maxsize):
-        # Just compute the new edges and call rebinXTo
-        start = begin - 1
-        stop = end
-        if start is None:
-            start = 0
-        stop = len(self.bins()) if stop >= sys.maxsize else stop - 1
-        new_edges = []
-        # new_bins = []
-        # new_bins += [self.underflow()]
-        for i in range(start):
-            # new_bins.append(self.bins()[i].clone())
-            new_edges.append(self.xEdges()[i])
-            new_edges.append(self.xEdges()[i + 1])
-        last = None
-        for i in range(start, stop, factor):
-            if i + factor <= len(self.bins()):
-                xmin = self.xEdges()[i]
-                xmax = self.xEdges()[i + 1]
-                # nb = GROGU_HISTO1D_V3.Bin()
-                for j in range(factor):
-                    last = i + j
-                    # nb += self.bins()[i + j]
-                    xmin = min(xmin, self.xEdges()[i + j])
-                    xmax = max(xmax, self.xEdges()[i + j + 1])
-                # new_bins.append(nb)
-                # add both edges
-                new_edges.append(xmin)
-                new_edges.append(xmax)
-        for j in range(last + 1, len(self.bins())):
-            # new_bins.append(self.bins()[j].clone())
-            new_edges.append(self.xEdges()[j])
-            new_edges.append(self.xEdges()[j + 1])
-        # new_bins += [self.overflow()]
-        # self.d_bins = new_bins
-        # drop duplicate edges
-        self.rebinXTo(list(set(new_edges)))
+        new_edges = rebinBy_to_rebinTo(self.xEdges(), factor, begin, end)
+        self.rebinXTo(new_edges)
 
     def rebinBy(self, *args, **kwargs):
         self.rebinXBy(*args, **kwargs)

src/babyyoda/histo2d.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from babyyoda.analysisobject import UHIAnalysisObject
-from babyyoda.util import loc, overflow, project, rebin, underflow
+from babyyoda.util import loc, overflow, project, rebin, rebinBy_to_rebinTo, underflow
 
 
 def set_bin2d(target, source):
@@ -161,6 +161,24 @@ def yMean(self, includeOverflows=True):
     def integral(self, includeOverflows=True):
         return sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows))
 
+    def rebinXBy(self, factor: int, begin=1, end=sys.maxsize):
+        new_edges = rebinBy_to_rebinTo(self.xEdges(), factor, begin, end)
+        self.rebinXTo(new_edges)
+
+    def rebinYBy(self, factor: int, begin=1, end=sys.maxsize):
+        new_edges = rebinBy_to_rebinTo(self.yEdges(), factor, begin, end)
+        self.rebinYTo(new_edges)
+
+    def dVols(self):
+        ret = []
+        for iy in range(len(self.yMins())):
+            for ix in range(len(self.xMins())):
+                ret.append(
+                    (self.xMaxs()[ix] - self.xMins()[ix])
+                    * (self.yMaxs()[iy] - self.yMins()[iy])
+                )
+        return np.array(ret)
+
     ########################################################
     # Generic UHI code
     ########################################################
@@ -359,7 +377,9 @@ def plot(self, *args, binwnorm=True, **kwargs):
 
             def temp_values():
                 return (
-                    np.array([b.sumW() / b.dVol() for b in self.bins()])
+                    np.array(
+                        [b.sumW() / vol for b, vol in zip(self.bins(), self.dVols())]
+                    )
                     .reshape((len(self.axes[1]), len(self.axes[0])))
                     .T
                 )

src/babyyoda/test.py

@@ -111,7 +111,7 @@ def assert_bin2d(gb, yb):
 
 
 def assert_value2d(gb, yb):
-    assert gb.sumW() == yb.sumW()
+    assert gb.sumW() == yb.sumW(), f"{gb.sumW()} != {yb.sumW()}"
     assert gb.sumW2() == yb.sumW2()
     assert gb.sumWX() == yb.sumWX(), f"{gb.sumWX()} != {yb.sumWX()}"
     assert gb.sumWX2() == yb.sumWX2()

src/babyyoda/util.py

@@ -1,5 +1,6 @@
 import gzip
 import inspect
+import sys
 
 
 class loc:
@@ -78,3 +79,40 @@ def has_own_method(cls, method_name):
 
     # Compare the underlying function (__func__) if both exist
     return cls_method.__func__ is not parent_method.__func__
+
+
+def rebinBy_to_rebinTo(edges: list[float], factor: int, begin=1, end=sys.maxsize):
+    # Just compute the new edges and call rebinXTo
+    start = begin - 1
+    stop = end
+    if start is None:
+        start = 0
+    stop = (len(edges) - 1) if stop >= sys.maxsize else stop - 1
+    new_edges = []
+    # new_bins = []
+    # new_bins += [self.underflow()]
+    for i in range(start):
+        # new_bins.append(self.bins()[i].clone())
+        new_edges.append(edges[i])
+        new_edges.append(edges[i + 1])
+    last = None
+    for i in range(start, stop, factor):
+        if i + factor <= (len(edges) - 1):
+            xmin = edges[i]
+            xmax = edges[i + 1]
+            # nb = GROGU_HISTO1D_V3.Bin()
+            for j in range(factor):
+                last = i + j
+                # nb += self.bins()[i + j]
+                xmin = min(xmin, edges[i + j])
+                xmax = max(xmax, edges[i + j + 1])
+            # new_bins.append(nb)
+            # add both edges
+            new_edges.append(xmin)
+            new_edges.append(xmax)
+    for j in range(last + 1, (len(edges) - 1)):
+        # new_bins.append(self.bins()[j].clone())
+        new_edges.append(edges[j])
+        new_edges.append(edges[j + 1])
+    # no duplicates
+    return list(set(new_edges))