Add counter (#31)

* Add counter

* Add AO and annotations

debug/.ipynb_checkpoints/Counter-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

debug/Counter.ipynb

@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "538a4798-95f5-4116-ab6f-f411f9db9fa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import yoda"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "698a0a2e-b833-405f-a0fc-b53cf1138653",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "c = yoda.Counter(title=\"hi\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "1a8a28c1-7f1b-4f67-8bd8-cbd93ff3a981",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "yoda.write([c], \"test_counter_v2.yoda\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "ee895f9f-c60e-4a79-bd2e-e4e474cb7ace",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'path': '/', 'title': 'hi', 'type': 'Counter'}"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "c.annotationsDict()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "019e131e-b9ec-409d-aeb3-2800e3c01f77",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['Path', 'Title', 'Type']"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "c.annotations()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "82a60bfd-dbae-4ea5-9983-418c9c38acf4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "c.annotation(\"path\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3a7db0c-4307-40ef-98dd-3c9683c8d99e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

debug/test_counter.yoda

@@ -0,0 +1,8 @@
+BEGIN YODA_COUNTER_V3 /
+Path: /
+Title: hi
+Type: Counter
+---
+# sumW       	sumW2        	numEntries
+0.000000e+00 	0.000000e+00 	0.000000e+00
+END YODA_COUNTER_V3

debug/test_counter_v2.yoda

@@ -0,0 +1,8 @@
+BEGIN YODA_COUNTER_V3 /
+Path: /
+Title: hi
+Type: Counter
+---
+# sumW       	sumW2        	numEntries
+0.000000e+00 	0.000000e+00 	0.000000e+00
+END YODA_COUNTER_V3

debug/test_counter_v3.yoda

@@ -0,0 +1,8 @@
+BEGIN YODA_COUNTER_V3 /
+Path: /
+Title: hi
+Type: Counter
+---
+# sumW       	sumW2        	numEntries
+0.000000e+00 	0.000000e+00 	0.000000e+00
+END YODA_COUNTER_V3

src/babyyoda/__init__.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MIT
 
 from ._version import version as __version__
+from .counter import UHICounter
 from .histo1d import UHIHisto1D
 from .histo2d import UHIHisto2D
 from .read import read, read_grogu, read_yoda
@@ -11,6 +12,7 @@
 
 __all__ = [
     "__version__",
+    "UHICounter",
     "UHIHisto1D",
     "UHIHisto2D",
     "read",

src/babyyoda/analysisobject.py

@@ -0,0 +1,3 @@
+class UHIAnalysisObject:
+    def key(self):
+        return self.path()

src/babyyoda/counter.py

@@ -0,0 +1,115 @@
+import contextlib
+
+from babyyoda.analysisobject import UHIAnalysisObject
+
+
+def set_bin0d(target, source):
+    if hasattr(target, "set"):
+        target.set(
+            source.numEntries(),
+            [source.sumW()],
+            [source.sumW2()],
+        )
+    else:
+        err = "YODA1 backend can not set bin values"
+        raise NotImplementedError(err)
+
+
+def Counter(*args, **kwargs):
+    """
+    Automatically select the correct version of the Histo1D class
+    """
+    try:
+        from babyyoda import yoda
+    except ImportError:
+        import babyyoda.grogu as yoda
+    return yoda.Counter(*args, **kwargs)
+
+
+# TODO make this implementation independent (no V2 or V3...)
+class UHICounter(UHIAnalysisObject):
+    ######
+    # BACKENDS
+    ######
+
+    def to_grogu_v2(self):
+        from babyyoda.grogu.counter_v2 import GROGU_COUNTER_V2
+
+        return GROGU_COUNTER_V2(
+            d_key=self.key(),
+            d_annotations=self.annotationsDict(),
+            d_bins=[
+                GROGU_COUNTER_V2.Bin(
+                    d_sumw=self.sumW(),
+                    d_sumw2=self.sumW2(),
+                    d_numentries=self.numEntries(),
+                )
+            ],
+        )
+
+    def to_grogu_v3(self):
+        from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3
+
+        return GROGU_COUNTER_V3(
+            d_key=self.key(),
+            d_annotations=self.annotationsDict(),
+            d_bins=[
+                GROGU_COUNTER_V3.Bin(
+                    d_sumw=self.sumW(),
+                    d_sumw2=self.sumW2(),
+                    d_numentries=self.numEntries(),
+                )
+            ],
+        )
+
+    def to_yoda_v3(self):
+        err = "Not implemented yet"
+        raise NotImplementedError(err)
+
+    def to_string(self):
+        # Now we need to map YODA to grogu and then call to_string
+        # TODO do we want to hardcode v3 here?
+        return self.to_grogu_v3().to_string()
+
+    ########################################################
+    # YODA compatibility code (dropped legacy code?)
+    ########################################################
+
+    ########################################################
+    # Generic UHI code
+    ########################################################
+
+    @property
+    def axes(self):
+        return []
+
+    @property
+    def kind(self):
+        # TODO reeavaluate this
+        return "COUNT"
+
+    def counts(self):
+        return self.numEntries()
+
+    def values(self):
+        return self.sumW()
+
+    def variances(self):
+        return self.sumW2()
+
+    def plot(self, *args, binwnorm=1.0, **kwargs):
+        import mplhep as hep
+
+        hep.histplot(
+            self,
+            *args,
+            yerr=self.variances() ** 0.5,
+            w2method="sqrt",
+            binwnorm=binwnorm,
+            **kwargs,
+        )
+
+    def _ipython_display_(self):
+        with contextlib.suppress(ImportError):
+            self.plot()
+        return self

src/babyyoda/grogu/__init__.py

@@ -1,3 +1,5 @@
+from babyyoda.grogu.counter_v2 import GROGU_COUNTER_V2
+from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3
 from babyyoda.grogu.histo1d_v3 import GROGU_HISTO1D_V3
 from babyyoda.grogu.histo2d_v2 import GROGU_HISTO2D_V2
 from babyyoda.grogu.histo2d_v3 import GROGU_HISTO2D_V3
@@ -9,6 +11,26 @@
 __all__ = ["read", "write"]
 
 
+def Counter(title=None, **kwargs):
+    return Counter_v3(title=title, **kwargs)
+
+
+def Counter_v3(title=None, **kwargs):
+    return GROGU_COUNTER_V3(
+        d_bins=[GROGU_COUNTER_V3.Bin()],
+        d_annotations={"Title": title} if title else {},
+        **kwargs,
+    )
+
+
+def Counter_v2(title=None, **kwargs):
+    return GROGU_COUNTER_V2(
+        d_bins=[GROGU_COUNTER_V2.Bin()],
+        d_annotations={"Title": title} if title else {},
+        **kwargs,
+    )
+
+
 def Histo1D(nbins: int, start: float, end: float, title=None, **kwargs):
     return Histo1D_v2(nbins=nbins, start=start, end=end, title=title, **kwargs)
 
@@ -25,7 +47,7 @@ def Histo1D_v2(nbins: int, start: float, end: float, title=None, **kwargs):
         d_overflow=GROGU_HISTO1D_V2.Bin(),
         d_underflow=GROGU_HISTO1D_V2.Bin(),
         d_total=GROGU_HISTO1D_V2.Bin(),
-        d_title=title,
+        d_annotations={"Title": title} if title else {},
         **kwargs,
     )
 
@@ -37,7 +59,7 @@ def Histo1D_v3(nbins: int, start: float, end: float, title=None, **kwargs):
             GROGU_HISTO1D_V3.Bin()
             for i in range(nbins + 2)  # add overflow and underflow
         ],
-        d_title=title,
+        d_annotations={"Title": title} if title else {},
         **kwargs,
     )
 
@@ -86,7 +108,7 @@ def Histo2D_v2(
             for j in range(nybins)
         ],
         d_total=GROGU_HISTO2D_V2.Bin(),
-        d_title=title,
+        d_annotations={"Title": title} if title else {},
         **kwargs,
     )
 
@@ -110,6 +132,6 @@ def Histo2D_v3(
             GROGU_HISTO2D_V3.Bin()
             for _ in range((nxbins + 2) * (nybins + 2))  # add overflow and underflow
         ],
-        d_title=title,
+        d_annotations={"Title": title} if title else {},
         **kwargs,
     )