From 4a613173da88e5208b5e3814d6bd5aeb2699c570 Mon Sep 17 00:00:00 2001
From: Alexander Puck Neuwirth <alexander@neuwirth-informatik.de>
Date: Wed, 9 Oct 2024 09:09:21 +0200
Subject: [PATCH] Add yoda vs grogu tests

Closes: #2
---
 .github/workflows/ci.yml               |  2 +-
 src/babyyoda/grogu/grogu_histo1d_v2.py | 34 +++++++++++++++---
 tests/test_yoda_vs_grogu.py            | 49 ++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0f9869f..3b65588 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -3,7 +3,7 @@ name: CI/CD
 on:
   push:
     branches:
-      - main
+      - master
   pull_request:
   # Run daily at 0:01 UTC
   schedule:
diff --git a/src/babyyoda/grogu/grogu_histo1d_v2.py b/src/babyyoda/grogu/grogu_histo1d_v2.py
index 2a4bdd9..0e1b644 100644
--- a/src/babyyoda/grogu/grogu_histo1d_v2.py
+++ b/src/babyyoda/grogu/grogu_histo1d_v2.py
@@ -11,16 +11,25 @@ class GROGU_HISTO1D_V2(GROGU_ANALYSIS_OBJECT):
     class Bin:
         d_xmin: Optional[float]
         d_xmax: Optional[float]
-        d_sumw: float
-        d_sumw2: float
-        d_sumwx: float
-        d_sumwx2: float
-        d_numentries: float
+        d_sumw: float = 0.0
+        d_sumw2: float = 0.0
+        d_sumwx: float = 0.0
+        d_sumwx2: float = 0.0
+        d_numentries: float = 0.0
 
         ########################################################
         # YODA compatibilty code
         ########################################################
 
+        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
+            self.d_sumw += sf
+            self.d_sumw2 += sf * weight
+            self.d_sumwx += sf * x
+            self.d_sumwx2 += sf * x**2
+            self.d_numentries += fraction
+
         def xMin(self):
             return self.d_xmin
 
@@ -74,6 +83,21 @@ def __post_init__(self):
     # YODA compatibilty code
     ############################################
 
+    def fill(self, x, weight=1.0, fraction=1.0):
+        for b in self.d_bins:
+            if b.xMin() <= x < b.xMax():
+                b.fill(x, weight, fraction)
+        if x > self.xMax() and self.d_overflow is not None:
+            self.d_overflow.fill(x, weight, fraction)
+        if x < self.xMin() and self.d_underflow is not None:
+            self.d_underflow.fill(x, weight, fraction)
+
+    def xMax(self):
+        return max([b.xMax() for b in self.d_bins])
+
+    def xMin(self):
+        return min([b.xMin() for b in self.d_bins])
+
     def bins(self):
         return sorted(self.d_bins, key=lambda b: b.d_xmin)
 
diff --git a/tests/test_yoda_vs_grogu.py b/tests/test_yoda_vs_grogu.py
index 8719865..9508371 100644
--- a/tests/test_yoda_vs_grogu.py
+++ b/tests/test_yoda_vs_grogu.py
@@ -1,4 +1,25 @@
 import babyyoda as by
+from babyyoda.grogu.grogu_histo1d_v2 import GROGU_HISTO1D_V2
+
+
+def compare_histo1d(gh1, yh1):
+    assert gh1.name() == yh1.name()
+    assert gh1.path() == yh1.path()
+    assert gh1.title() == yh1.title()
+    assert gh1.type() == yh1.type()
+
+    assert len(gh1.bins()) == len(yh1.bins())
+
+    for gb, yb in zip(gh1.bins(), yh1.bins()):
+        assert gb.xMin() == yb.xMin()
+        assert gb.xMax() == yb.xMax()
+        assert gb.sumW() == yb.sumW()
+        assert gb.sumW2() == yb.sumW2()
+        assert gb.sumWX() == yb.sumWX()
+        assert gb.sumWX2() == yb.sumWX2()
+        assert gb.numEntries() == yb.numEntries()
+
+    # TODO test overflow and underflow
 
 
 def test_ao():
@@ -34,6 +55,8 @@ def test_histo1():
         assert gb.sumWX2() == yb.sumWX2()
         assert gb.numEntries() == yb.numEntries()
 
+    # TODO test overflow and underflow
+
 
 def test_histo2():
     gh2 = next(iter(by.read_grogu("tests/test2d.yoda").values()))
@@ -50,3 +73,29 @@ def test_histo2():
         assert gb.sumW() == yb.sumW()
         assert gb.sumW2() == yb.sumW2()
         assert gb.numEntries() == yb.numEntries()
+
+    # TODO test overflow and underflow
+
+
+def test_create_histo1d():
+    import yoda
+
+    h = yoda.Histo1D(10, 0, 10, title="test")
+
+    g = GROGU_HISTO1D_V2(
+        d_title="test",
+        d_bins=[
+            GROGU_HISTO1D_V2.Bin(d_xmin=hb.xMin(), d_xmax=hb.xMax()) for hb in h.bins()
+        ],
+    )
+
+    for i in range(12):
+        for _ in range(i):
+            h.fill(i)
+            g.fill(i)
+    h.fill(-1)
+    g.fill(-1)
+    h.fill(10)
+    g.fill(10)
+
+    compare_histo1d(g, h)