Test backends same

src/babyyoda/histo2D.py

@@ -1,6 +1,8 @@
 import sys
 import numpy as np
 import babyyoda
+from babyyoda.grogu.histo2d_v2 import GROGU_HISTO2D_V2
+from babyyoda.grogu.histo2d_v3 import GROGU_HISTO2D_V3
 from babyyoda.util import loc, overflow, rebin, underflow
 
 
@@ -242,6 +244,81 @@ def __getitem__(self, slices):
         # TODO implement slice
         raise TypeError("Invalid argument type")
 
+    def key(self):
+        if hasattr(self.target, "key"):
+            return self.target.key()
+        return self.path()
+
+    def to_grogu_v2(self):
+        return Histo2D(
+            GROGU_HISTO2D_V2(
+                d_key=self.key(),
+                d_path=self.path(),
+                d_title=self.title(),
+                d_bins=[
+                    GROGU_HISTO2D_V2.Bin(
+                        d_xmin=self.xEdges()[i % len(self.xEdges())],
+                        d_xmax=self.xEdges()[i % len(self.xEdges()) + 1],
+                        d_ymin=self.yEdges()[i // len(self.xEdges())],
+                        d_ymax=self.yEdges()[i // len(self.xEdges()) + 1],
+                        d_sumw=b.sumW(),
+                        d_sumw2=b.sumW2(),
+                        d_sumwx=b.sumWX(),
+                        d_sumwx2=b.sumWX2(),
+                        d_sumwy=b.sumWY(),
+                        d_sumwy2=b.sumWY2(),
+                        d_sumwxy=b.sumWXY(),
+                        d_numentries=b.numEntries(),
+                    )
+                    for i, b in enumerate(self.bins())
+                ],
+            )
+        )
+
+    def to_grogu_v3(self):
+        bins = []
+        try:
+            bins = self.bins(True)
+        except NotImplementedError:
+            nobins = self.bins()
+            bins += [GROGU_HISTO2D_V3.Bin()] * (len(self.xEdges()))
+            for j in range(len(nobins)):
+                if j % (len(self.xEdges()) - 1) == 0:
+                    bins += [GROGU_HISTO2D_V3.Bin()]  # overflow
+                    bins += [GROGU_HISTO2D_V3.Bin()]  # underflow
+                bins += [nobins[j]]
+            bins += [GROGU_HISTO2D_V3.Bin()]  # overflow
+            bins += [GROGU_HISTO2D_V3.Bin()]  # underflow
+            bins += [GROGU_HISTO2D_V3.Bin()] * (len(self.xEdges()))
+
+            # Fill up with empty overflow bins
+        return Histo2D(
+            GROGU_HISTO2D_V3(
+                d_key=self.key(),
+                d_path=self.path(),
+                d_title=self.title(),
+                d_edges=[self.xEdges(), self.yEdges()],
+                d_bins=[
+                    GROGU_HISTO2D_V3.Bin(
+                        d_sumw=b.sumW(),
+                        d_sumw2=b.sumW2(),
+                        d_sumwx=b.sumWX(),
+                        d_sumwx2=b.sumWX2(),
+                        d_sumwy=b.sumWY(),
+                        d_sumwy2=b.sumWY2(),
+                        d_sumwxy=b.crossTerm(0, 1),
+                        d_numentries=b.numEntries(),
+                    )
+                    for b in bins
+                ],
+            )
+        )
+
+    def to_string(self):
+        if hasattr(self.target, "to_string"):
+            return self.target.to_string()
+        return self.to_grogu_v3().to_string()
+
     def plot(self, *args, binwnorm=True, **kwargs):
         import mplhep as hep
 

tests/babyyoda/test_backend.py

@@ -0,0 +1,65 @@
+import pytest
+from babyyoda.histo1D import Histo1D
+from babyyoda.test import assert_histo1d
+
+import babyyoda.grogu as grogu
+
+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 = Histo1D(10, 0, 10, title="test", backend=backend)
+    for i in range(12):
+        for _ in range(i):
+            h.fill(i)
+    # do we already want to use HISTO1D here?
+    h.underflow().fill(-1)
+    h.overflow().fill(10)
+
+    return h
+
+
+@pytest.mark.parametrize(
+    "factory1",
+    [
+        None,
+        # babyyoda.Histo1D,
+        grogu.Histo1D,
+        grogu.Histo1D_v2,
+        grogu.Histo1D_v3,
+        yoda.Histo1D,
+    ],
+)
+@pytest.mark.parametrize(
+    "factory2",
+    [
+        None,
+        # babyyoda.Histo1D,
+        grogu.Histo1D,
+        grogu.Histo1D_v2,
+        grogu.Histo1D_v3,
+        yoda.Histo1D,
+    ],
+)
+def test_histo1d_backends(factory1, factory2):
+    h = create_histo(factory1)
+    g = create_histo(factory2)
+    assert_histo1d(h, g)
+    assert_histo1d(h.to_grogu_v2(), g)
+    assert_histo1d(h, g.to_grogu_v2())
+    assert_histo1d(h.to_grogu_v2(), g.to_grogu_v2())
+    assert_histo1d(h.to_grogu_v3(), g)
+    assert_histo1d(h, g.to_grogu_v3())
+    assert_histo1d(h.to_grogu_v3(), g.to_grogu_v3())
+    assert_histo1d(h.to_grogu_v2(), g.to_grogu_v3())
+    # TODO add to_yoda when available

tests/babyyoda/test_write.py

@@ -2,7 +2,7 @@
 import babyyoda
 import babyyoda.read
 import babyyoda.grogu as grogu
-from babyyoda.test import assert_histo1d
+from babyyoda.test import assert_histo1d, assert_histo2d
 from babyyoda.util import is_yoda, uses_yoda
 
 try:
@@ -56,3 +56,96 @@ def test_write_histo1d_v2(read, write, reread, filename):
     h1 = next(iter(hs1.values()))
     h2 = next(iter(hs2.values()))
     assert_histo1d(h1, h2)
+
+
+@pytest.mark.parametrize(
+    "read, write, reread",
+    [
+        (yoda.read, yoda.write, yoda.read),
+        (babyyoda.read, babyyoda.write, babyyoda.read),
+        (babyyoda.read_yoda, babyyoda.write, babyyoda.read_yoda),
+        (babyyoda.read_grogu, babyyoda.write_grogu, babyyoda.read_grogu),
+        (grogu.read, grogu.write, grogu.read),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_grogu),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_yoda),
+        (yoda.read, yoda.write, babyyoda.read),
+        (babyyoda.read, babyyoda.write, yoda.read),
+    ],
+)
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "test.yoda",
+        "test.yoda.gz",
+    ],
+)
+def test_write_histo1d_v3(read, write, reread, filename):
+    hs1 = read("tests/test_histo1d_v3.yoda")
+    write(hs1, filename)
+    hs2 = reread(filename)
+
+    h1 = next(iter(hs1.values()))
+    h2 = next(iter(hs2.values()))
+    assert_histo1d(h1, h2)
+
+
+@pytest.mark.parametrize(
+    "read, write, reread",
+    [
+        (yoda.read, yoda.write, yoda.read),
+        (babyyoda.read, babyyoda.write, babyyoda.read),
+        (babyyoda.read_yoda, babyyoda.write, babyyoda.read_yoda),
+        (babyyoda.read_grogu, babyyoda.write_grogu, babyyoda.read_grogu),
+        (grogu.read, grogu.write, grogu.read),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_grogu),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_yoda),
+        (yoda.read, yoda.write, babyyoda.read),
+        (babyyoda.read, babyyoda.write, yoda.read),
+    ],
+)
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "test.yoda",
+        "test.yoda.gz",
+    ],
+)
+def test_write_histo2d_v2(read, write, reread, filename):
+    hs1 = read("tests/test_histo2d_v2.yoda")
+    write(hs1, filename)
+    hs2 = reread(filename)
+
+    h1 = next(iter(hs1.values()))
+    h2 = next(iter(hs2.values()))
+    assert_histo2d(h1, h2, includeFlow=False)
+
+
+@pytest.mark.parametrize(
+    "read, write, reread",
+    [
+        (yoda.read, yoda.write, yoda.read),
+        (babyyoda.read, babyyoda.write, babyyoda.read),
+        (babyyoda.read_yoda, babyyoda.write, babyyoda.read_yoda),
+        (babyyoda.read_grogu, babyyoda.write_grogu, babyyoda.read_grogu),
+        (grogu.read, grogu.write, grogu.read),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_grogu),
+        (babyyoda.read_yoda, babyyoda.write_grogu, babyyoda.read_yoda),
+        (yoda.read, yoda.write, babyyoda.read),
+        (babyyoda.read, babyyoda.write, yoda.read),
+    ],
+)
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "test.yoda",
+        "test.yoda.gz",
+    ],
+)
+def test_write_histo2d_v3(read, write, reread, filename):
+    hs1 = read("tests/test_histo2d_v3.yoda")
+    write(hs1, filename)
+    hs2 = reread(filename)
+
+    h1 = next(iter(hs1.values()))
+    h2 = next(iter(hs2.values()))
+    assert_histo2d(h1, h2)