From 910e6539a8430807009b78bbcf37f5042820c4cb Mon Sep 17 00:00:00 2001 From: Alexander Puck Neuwirth Date: Sat, 12 Oct 2024 22:04:58 +0200 Subject: [PATCH] Fix grogu histo2d_V3 fill --- src/babyyoda/grogu/__init__.py | 4 +-- src/babyyoda/grogu/histo2d_v2.py | 4 ++- src/babyyoda/grogu/histo2d_v3.py | 24 ++++++-------- src/babyyoda/test.py | 7 +++-- tests/babyyoda/test_histo2d.py | 17 +++++++--- tests/babyyoda/test_histo2d_overflow.py | 42 +++++++++++++++++++++++++ tests/yoda/test_yoda_vs_grogu.py | 2 +- 7 files changed, 75 insertions(+), 25 deletions(-) create mode 100644 tests/babyyoda/test_histo2d_overflow.py diff --git a/src/babyyoda/grogu/__init__.py b/src/babyyoda/grogu/__init__.py index b9e86ef..6b70ac5 100644 --- a/src/babyyoda/grogu/__init__.py +++ b/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, diff --git a/src/babyyoda/grogu/histo2d_v2.py b/src/babyyoda/grogu/histo2d_v2.py index 6ea98b4..92626f6 100644 --- a/src/babyyoda/grogu/histo2d_v2.py +++ b/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)) diff --git a/src/babyyoda/grogu/histo2d_v3.py b/src/babyyoda/grogu/histo2d_v3.py index b7c9349..c59b4c1 100644 --- a/src/babyyoda/grogu/histo2d_v3.py +++ b/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" diff --git a/src/babyyoda/test.py b/src/babyyoda/test.py index ff5eade..14d3ec2 100644 --- a/src/babyyoda/test.py +++ b/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) diff --git a/tests/babyyoda/test_histo2d.py b/tests/babyyoda/test_histo2d.py index 9674003..a8fb90c 100644 --- a/tests/babyyoda/test_histo2d.py +++ b/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 diff --git a/tests/babyyoda/test_histo2d_overflow.py b/tests/babyyoda/test_histo2d_overflow.py new file mode 100644 index 0000000..78905c5 --- /dev/null +++ b/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) diff --git a/tests/yoda/test_yoda_vs_grogu.py b/tests/yoda/test_yoda_vs_grogu.py index f491b17..3dae613 100644 --- a/tests/yoda/test_yoda_vs_grogu.py +++ b/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)