Fix grogu histo2d_V3 fill

src/babyyoda/grogu/__init__.py

@@ -51,7 +51,7 @@ def Histo2D(
     title=None,
     **kwargs,
 ):
-    return Histo2D_v2(
+    return Histo2D_v3(
         nxbins=nxbins,
         xstart=xstart,
         xend=xend,
@@ -106,7 +106,7 @@ def Histo2D_v3(
         ],
         d_bins=[
             GROGU_HISTO2D_V3.Bin()
-            for _ in range((nxbins + 1) * (nybins + 1))  # add overflow and underflow
+            for _ in range((nxbins + 2) * (nybins + 2))  # add overflow and underflow
         ],
         d_title=title,
         **kwargs,

src/babyyoda/grogu/histo2d_v2.py

@@ -168,7 +168,9 @@ def xMax(self):
     def yMax(self):
         return max(b.d_ymax for b in self.d_bins)
 
-    def bins(self):
+    def bins(self, includeFlow=False):
+        if includeFlow:
+            raise NotImplementedError("includeFlow=True not supported")
         # sort the bins by xlow, then ylow
         # YODA-1
         # return sorted(self.d_bins, key=lambda b: (b.d_xmin, b.d_ymin))

src/babyyoda/grogu/histo2d_v3.py

@@ -1,5 +1,6 @@
 import re
 from dataclasses import dataclass, field
+import sys
 from typing import List
 
 import numpy as np
@@ -124,20 +125,15 @@ def yEdges(self):
         return self.d_edges[1]
 
     def fill(self, x, y, weight=1.0, fraction=1.0):
-        for i, b in enumerate(self.bins()):
-            if (
-                self.xEdges()[i] <= x < self.xEdges()[i + 1]
-                and self.yEdges()[i] <= y < self.yEdges()[i + 1]
-            ):
-                b.fill(x, y, weight, fraction)
-                return  # should be done here
-        raise ValueError(
-            f"fill({x}, {y}, ...) out of range, Overflow not implemented yet"
-        )
-        if x >= self.xMax() and self.d_overflow is not None:
-            self.d_overflow.fill(x, y, weight, fraction)
-        if x < self.xMin() and self.d_underflow is not None:
-            self.d_underflow.fill(x, y, weight, fraction)
+        # get ix and iy to map to correct bin
+        for ix, xEdge in enumerate(self.xEdges() + [sys.float_info.max]):
+            if x < xEdge:
+                break
+        for iy, yEdge in enumerate(self.yEdges() + [sys.float_info.max]):
+            if y < yEdge:
+                break
+        # Also fill overflow bins
+        self.bins(True)[iy * (len(self.xEdges()) + 1) + ix].fill(x, y, weight, fraction)
 
     def xMax(self):
         assert max(self.xEdges()) == self.xEdges()[-1], "xMax is not the last edge"

src/babyyoda/test.py

@@ -96,7 +96,7 @@ def assert_value2d(gb, yb):
     assert gb.numEntries() == yb.numEntries()
 
 
-def assert_histo2d(gh1, yh1):
+def assert_histo2d(gh1, yh1, includeFlow=True):
     assert_ao(gh1, yh1)
 
     assert len(gh1.bins()) == len(yh1.bins())
@@ -107,8 +107,9 @@ def assert_histo2d(gh1, yh1):
     for ge, ye in zip(gh1.yEdges(), yh1.yEdges()):
         assert ge == ye, f"{gh1.yEdges()} != {yh1.yEdges()}"
 
-    for gb, yb in zip(gh1.bins(True), yh1.bins(True)):
-        assert_value2d(gb, yb)
+    if includeFlow:
+        for gb, yb in zip(gh1.bins(True), yh1.bins(True)):
+            assert_value2d(gb, yb)
 
     for gb, yb in zip(gh1.bins(), yh1.bins()):
         assert_value2d(gb, yb)

tests/babyyoda/test_histo2d.py

@@ -29,23 +29,32 @@ def create_histo(backend):
 
 
 @pytest.mark.parametrize(
-    "factory", [grogu.Histo1D, grogu.Histo1D_v2, grogu.Histo1D_v3, yoda.Histo1D]
+    "factory", [grogu.Histo2D, grogu.Histo2D_v2, grogu.Histo2D_v3, yoda.Histo2D]
 )
 def test_create_histo(factory):
     create_histo(factory)
 
 
 @pytest.mark.parametrize(
-    "factory1", [grogu.Histo1D, grogu.Histo1D_v2, grogu.Histo1D_v3, yoda.Histo1D]
+    "factory1", [grogu.Histo2D, grogu.Histo2D_v2, grogu.Histo2D_v3, yoda.Histo2D]
 )
 @pytest.mark.parametrize(
-    "factory2", [grogu.Histo1D, grogu.Histo1D_v2, grogu.Histo1D_v3, yoda.Histo1D]
+    "factory2", [grogu.Histo2D, grogu.Histo2D_v2, grogu.Histo2D_v3, yoda.Histo2D]
 )
 def test_histos_equal(factory1, factory2):
     h1 = create_histo(factory1)
     h2 = create_histo(factory2)
 
-    assert_histo2d(h1, h2)
+    assert_histo2d(h1, h2, includeFlow=False)
+
+
+@pytest.mark.parametrize("factory1", [grogu.Histo2D, grogu.Histo2D_v3, yoda.Histo2D])
+@pytest.mark.parametrize("factory2", [grogu.Histo2D, grogu.Histo2D_v3, yoda.Histo2D])
+def test_histos_overflow_equal(factory1, factory2):
+    h1 = create_histo(factory1)
+    h2 = create_histo(factory2)
+
+    assert_histo2d(h1, h2, includeFlow=False)
 
 
 # TODO more like in 1d

tests/babyyoda/test_histo2d_overflow.py

@@ -0,0 +1,42 @@
+import pytest
+from babyyoda.histo2D import Histo2D
+from babyyoda.test import assert_histo2d
+
+import babyyoda.grogu as grogu
+
+
+# YODA1 does not support histo2d overflows
+pytest.importorskip("yoda", minversion="2.0.0")
+
+try:
+    import yoda
+
+    yoda_available = True
+    # version dependence possible here
+except ImportError:
+    import babyyoda.grogu as yoda
+
+    yoda_available = False
+
+
+# TODO use metafunction fixtures instead fo many pytest.mark
+
+
+def create_histo(backend):
+    h = Histo2D(10, 0, 10, 10, 0, 10, title="test", backend=backend)
+    w = 0
+    for i in range(-10, 12):
+        for j in range(-10, 12):
+            w += 1
+            h.fill(i, j, w)
+    # do we already want to use HISTO1D here?
+    return h
+
+
+@pytest.mark.parametrize("factory1", [grogu.Histo2D, grogu.Histo2D_v3, yoda.Histo2D])
+@pytest.mark.parametrize("factory2", [grogu.Histo2D, grogu.Histo2D_v3, yoda.Histo2D])
+def test_histos_overflow_equal(factory1, factory2):
+    h1 = create_histo(factory1)
+    h2 = create_histo(factory2)
+
+    assert_histo2d(h1, h2, includeFlow=True)

tests/yoda/test_yoda_vs_grogu.py

@@ -154,4 +154,4 @@ def test_create_histo2d():
                 h.fill(i, j)
                 g.fill(i, j)
 
-    assert_histo2d(Histo2D(g), Histo2D(h))
+    assert_histo2d(Histo2D(g), Histo2D(h), includeFlow=False)