diff --git a/.gitignore b/.gitignore index 0ea1e0d..fd98abf 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ dist __pycache__ test.yoda test.yoda.gz +docs/build +.coverage diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 127baac..6d9920c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,8 +20,33 @@ repos: - id: trailing-whitespace - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.6.8" + rev: "v0.7.2" hooks: - id: ruff args: ["--fix", "--show-fixes"] - id: ruff-format + +- repo: https://github.com/adamchainz/blacken-docs + rev: "1.19.1" # replace with latest tag on GitHub + hooks: + - id: blacken-docs +- repo: https://github.com/pre-commit/mirrors-mypy + rev: 'v1.13.0' # Use the sha / tag you want to point at + hooks: + - id: mypy +- repo: https://github.com/codespell-project/codespell + rev: v2.3.0 + hooks: + - id: codespell +- repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.10.0 # Use the ref you want to point at + hooks: + - id: python-use-type-annotations +- repo: https://github.com/executablebooks/mdformat + rev: 0.7.18 # Use the ref you want to point at + hooks: + - id: mdformat + # Optionally add plugins + additional_dependencies: + - mdformat-gfm + - mdformat-black diff --git a/.readthedocs.yaml b/.readthedocs.yaml new file mode 100644 index 0000000..30a46d9 --- /dev/null +++ b/.readthedocs.yaml @@ -0,0 +1,27 @@ +# .readthedocs.yml +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +version: 2 + +sphinx: + configuration: docs/source/conf.py + +formats: all + +submodules: + include: all + +build: + os: "ubuntu-22.04" + tools: + python: "3.10" + + + +python: + install: + - method: pip + path: . + extra_requirements: + - docs \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d298e4f --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2018-2024, Eduardo Rodrigues and Henry Schreiner. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the particle package developers nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/LICENSE.txt b/LICENSE.txt deleted file mode 100644 index d40d159..0000000 --- a/LICENSE.txt +++ /dev/null @@ -1,9 +0,0 @@ -MIT License - -Copyright (c) 2024-present Alexander Puck Neuwirth - -Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..d0c3cbf --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 0000000..747ffb7 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 0000000..779473e --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,32 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = "babyyoda" +copyright = "2024, Alexander Puck Neuwirth" +author = "Alexander Puck Neuwirth" + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", +] + +templates_path = ["_templates"] +exclude_patterns = [ + "**.ipynb_checkpoints", +] + + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output + +html_theme = "alabaster" +html_static_path = ["_static"] diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 0000000..d1f23f9 --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,20 @@ +.. babyyoda documentation master file, created by + sphinx-quickstart on Tue Nov 5 13:28:41 2024. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Welcome to babyyoda's documentation! +==================================== + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/pyproject.toml b/pyproject.toml index 7a946fc..becb787 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,13 +11,15 @@ dynamic = ["version"] description = '' readme = "README.md" requires-python = ">=3.9" -license = "MIT" keywords = [] authors = [ { name = "Alexander Puck Neuwirth", email = "alexander@neuwirth-informatik.de" }, ] classifiers = [ "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: BSD License", "Programming Language :: Python", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", @@ -26,6 +28,7 @@ classifiers = [ "Programming Language :: Python :: 3.13", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", + "Topic :: Scientific/Engineering", ] dependencies = [ "numpy" @@ -49,6 +52,9 @@ develop = [ "pre-commit", "tbump>=6.7.0", ] +docs = [ + "sphinx" +] [tool.hatch.envs.develop] features = [ @@ -93,12 +99,27 @@ exclude_lines = [ [tool.hatch.version] source = "vcs" + + [tool.hatch.version.raw-options] local_scheme = "no-local-version" [tool.hatch.build.hooks.vcs] version-file = "src/babyyoda/_version.py" +[tool.mypy] +warn_unused_configs = true +warn_unused_ignores = true +python_version = "3.9" +files = "src" +strict = true +warn_unreachable = true +enable_error_code = [ + "ignore-without-code", + "redundant-expr", + "truthy-bool", +] + [tool.ruff.lint] extend-select = [ @@ -131,3 +152,19 @@ ignore = [ "RUF012", # TODO: mutable class attributes "SIM115", # TODO: use context manager for opening files ] + +[tool.pytest.ini_options] +minversion = "6.0" +testpaths = ["tests"] +junit_family = "xunit2" +log_cli_level = "info" +xfail_strict = true +addopts = [ + "-ra", + "--strict-markers", + "--strict-config", +] +filterwarnings = [ + "error", + "ignore:Integer weights indicate poissonian data:UserWarning" +] diff --git a/src/babyyoda/analysisobject.py b/src/babyyoda/analysisobject.py index 0b8580a..eb586d6 100644 --- a/src/babyyoda/analysisobject.py +++ b/src/babyyoda/analysisobject.py @@ -1,7 +1,18 @@ class UHIAnalysisObject: - def key(self): + ######## + # Basic needed functions for UHI+YODA + ######## + def path(self) -> str: + err = "UHIAnalysisObject.path() must be implemented by subclass" + raise NotImplementedError(err) + + def setAnnotation(self, key: str, value: str) -> None: + err = "UHIAnalysisObject.setAnnotation() must be implemented by subclass" + raise NotImplementedError(err) + + def key(self) -> str: return self.path() - def setAnnotationsDict(self, d: dict): + def setAnnotationsDict(self, d: dict[str, str]) -> None: for k, v in d.items(): self.setAnnotation(k, v) diff --git a/src/babyyoda/counter.py b/src/babyyoda/counter.py index 10336e6..91858fe 100644 --- a/src/babyyoda/counter.py +++ b/src/babyyoda/counter.py @@ -1,9 +1,11 @@ import contextlib +from typing import Any, Optional +import babyyoda from babyyoda.analysisobject import UHIAnalysisObject -def set_bin0d(target, source): +def set_bin0d(target: Any, source: Any) -> None: if hasattr(target, "set"): target.set( source.numEntries(), @@ -15,24 +17,46 @@ def set_bin0d(target, source): raise NotImplementedError(err) -def Counter(*args, **kwargs): +def Counter(*args: Any, **kwargs: Any) -> "UHICounter": """ Automatically select the correct version of the Histo1D class """ try: from babyyoda import yoda + + return yoda.Counter(*args, **kwargs) except ImportError: - import babyyoda.grogu as yoda - return yoda.Counter(*args, **kwargs) + from babyyoda import grogu + + return grogu.Counter(*args, **kwargs) # TODO make this implementation independent (no V2 or V3...) class UHICounter(UHIAnalysisObject): + ###### + # Minimum required functions + ###### + + def sumW(self) -> float: + raise NotImplementedError + + def sumW2(self) -> float: + raise NotImplementedError + + def numEntries(self) -> float: + raise NotImplementedError + + def annotationsDict(self) -> dict[str, Optional[str]]: + raise NotImplementedError + + def clone(self) -> "UHICounter": + raise NotImplementedError + ###### # BACKENDS ###### - def to_grogu_v2(self): + def to_grogu_v2(self) -> Any: from babyyoda.grogu.counter_v2 import GROGU_COUNTER_V2 return GROGU_COUNTER_V2( @@ -47,7 +71,7 @@ def to_grogu_v2(self): ], ) - def to_grogu_v3(self): + def to_grogu_v3(self) -> "babyyoda.grogu.counter_v3.GROGU_COUNTER_V3": from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3 return GROGU_COUNTER_V3( @@ -62,14 +86,14 @@ def to_grogu_v3(self): ], ) - def to_yoda_v3(self): + def to_yoda_v3(self) -> Any: err = "Not implemented yet" raise NotImplementedError(err) - def to_string(self): + def to_string(self) -> str: # 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() + return str(self.to_grogu_v3().to_string()) ######################################################## # YODA compatibility code (dropped legacy code?) @@ -80,24 +104,24 @@ def to_string(self): ######################################################## @property - def axes(self): + def axes(self) -> list[list[tuple[float, float]]]: return [] @property - def kind(self): + def kind(self) -> str: # TODO reeavaluate this return "COUNT" - def counts(self): + def counts(self) -> float: return self.numEntries() - def values(self): + def values(self) -> float: return self.sumW() - def variances(self): + def variances(self) -> float: return self.sumW2() - def plot(self, *args, **kwargs): + def plot(self, *args: Any, **kwargs: Any) -> None: # TODO check UHI 0D plottable and mplhep hist plot 0D import matplotlib.pyplot as plt @@ -115,7 +139,7 @@ def plot(self, *args, **kwargs): # Show the plot # plt.show() - def _ipython_display_(self): + def _ipython_display_(self) -> "UHICounter": with contextlib.suppress(ImportError): self.plot() return self diff --git a/src/babyyoda/grogu/__init__.py b/src/babyyoda/grogu/__init__.py index 0a46203..062b443 100644 --- a/src/babyyoda/grogu/__init__.py +++ b/src/babyyoda/grogu/__init__.py @@ -1,9 +1,11 @@ +from typing import Any + from babyyoda.grogu.counter_v2 import Counter_v2 -from babyyoda.grogu.counter_v3 import Counter_v3 +from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3, Counter_v3 from babyyoda.grogu.histo1d_v2 import Histo1D_v2 -from babyyoda.grogu.histo1d_v3 import Histo1D_v3 +from babyyoda.grogu.histo1d_v3 import GROGU_HISTO1D_V3, Histo1D_v3 from babyyoda.grogu.histo2d_v2 import Histo2D_v2 -from babyyoda.grogu.histo2d_v3 import Histo2D_v3 +from babyyoda.grogu.histo2d_v3 import GROGU_HISTO2D_V3, Histo2D_v3 from .read import read from .write import write @@ -20,21 +22,19 @@ ] -def Counter(*args, **kwargs): +def Counter(*args: Any, **kwargs: Any) -> GROGU_COUNTER_V3: return Counter_v3(*args, **kwargs) -def Histo1D(*args, **kwargs): +def Histo1D(*args: Any, **kwargs: Any) -> GROGU_HISTO1D_V3: return Histo1D_v3(*args, **kwargs) def Histo2D( - *args, - title=None, - **kwargs, -): + *args: Any, + **kwargs: Any, +) -> GROGU_HISTO2D_V3: return Histo2D_v3( *args, - title=title, **kwargs, ) diff --git a/src/babyyoda/grogu/analysis_object.py b/src/babyyoda/grogu/analysis_object.py index f8ef9fc..e50d94e 100644 --- a/src/babyyoda/grogu/analysis_object.py +++ b/src/babyyoda/grogu/analysis_object.py @@ -1,62 +1,62 @@ import re from dataclasses import dataclass, field +from typing import Optional @dataclass class GROGU_ANALYSIS_OBJECT: - d_annotations: dict = field(default_factory=dict) - # TODO add anotations + d_annotations: dict[str, Optional[str]] = field(default_factory=dict) d_key: str = "" - def __post_init__(self): + def __post_init__(self) -> None: if "Path" not in self.d_annotations: self.d_annotations["Path"] = "/" if "Title" not in self.d_annotations: self.d_annotations["Title"] = "" ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ - def key(self): + def key(self) -> str: return self.d_key - def name(self): + def name(self) -> str: return self.path().split("/")[-1] - def path(self): + def path(self) -> str: p = self.annotation("Path") return p if p else "/" - def title(self): + def title(self) -> Optional[str]: return self.annotation("Title") - def type(self): + def type(self) -> Optional[str]: return self.annotation("Type") - def annotations(self): - return self.d_annotations.keys() + def annotations(self) -> list[str]: + return list(self.d_annotations.keys()) - def annotation(self, k: str, default=None) -> str: + def annotation(self, k: str, default: Optional[str] = None) -> Optional[str]: return self.d_annotations.get(k, default) - def setAnnotation(self, key: str, value: str): + def setAnnotation(self, key: str, value: str) -> None: self.d_annotations[key] = value - def clearAnnotations(self): + def clearAnnotations(self) -> None: self.d_annotations = {} def hasAnnotation(self, key: str) -> bool: return key in self.d_annotations - def annotationsDict(self): + def annotationsDict(self) -> dict[str, Optional[str]]: return self.d_annotations @classmethod def from_string(cls, file_content: str) -> "GROGU_ANALYSIS_OBJECT": lines = file_content.strip().splitlines() # Extract metadata (path, title) - annotations = {"Path": "/"} + annotations: dict[str, Optional[str]] = {"Path": "/"} pattern = re.compile(r"(\S+): (.+)") for line in lines: pattern_match = pattern.match(line) @@ -69,10 +69,10 @@ def from_string(cls, file_content: str) -> "GROGU_ANALYSIS_OBJECT": return cls( d_annotations=annotations, - d_key=annotations.get("Path", ""), + d_key=annotations.get("Path", "") or "", ) - def to_string(self): + def to_string(self) -> str: ret = "" for k, v in self.d_annotations.items(): val = v diff --git a/src/babyyoda/grogu/counter_v2.py b/src/babyyoda/grogu/counter_v2.py index d9d1bab..1fffc0a 100644 --- a/src/babyyoda/grogu/counter_v2.py +++ b/src/babyyoda/grogu/counter_v2.py @@ -1,12 +1,12 @@ import re from dataclasses import dataclass, field -from typing import Union +from typing import Any, Optional, Union from babyyoda.counter import UHICounter from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT -def Counter_v2(title=None, **kwargs): +def Counter_v2(title: Optional[str] = None, **kwargs: Any) -> "GROGU_COUNTER_V2": return GROGU_COUNTER_V2( d_bins=[GROGU_COUNTER_V2.Bin()], d_annotations={"Title": title} if title else {}, @@ -23,23 +23,23 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## - def clone(self): + def clone(self) -> "GROGU_COUNTER_V2.Bin": return GROGU_COUNTER_V2.Bin( d_sumw=self.d_sumw, d_sumw2=self.d_sumw2, d_numentries=self.d_numentries, ) - def fill(self, weight: float = 1.0, fraction: float = 1.0) -> bool: + def fill(self, weight: float = 1.0, fraction: float = 1.0) -> None: sf = fraction * weight self.d_sumw += sf self.d_sumw2 += sf * weight self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: self.d_sumw = bin.sumW() self.d_sumw2 = bin.sumW2() self.d_numentries = bin.numEntries() @@ -49,7 +49,7 @@ def set( numEntries: float, sumW: Union[list[float], float], sumW2: Union[list[float], float], - ): + ) -> None: if isinstance(sumW, float): sumW = [sumW] if isinstance(sumW2, float): @@ -60,13 +60,13 @@ def set( self.d_sumw2 = sumW2[0] self.d_numentries = numEntries - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def variance(self): + def variance(self) -> float: if self.d_sumw**2 - self.d_sumw2 == 0: return 0 return abs( @@ -75,22 +75,22 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> Any: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> Any: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> Any: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> Any: return self.stdDev() / self.effNumEntries() ** 0.5 - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries - def __eq__(self, other): + def __eq__(self, other: object) -> bool: return ( isinstance(other, GROGU_COUNTER_V2.Bin) and self.d_sumw == other.d_sumw @@ -98,7 +98,7 @@ def __eq__(self, other): and self.d_numentries == other.d_numentries ) - def __add__(self, other): + def __add__(self, other: Any) -> "GROGU_COUNTER_V2.Bin": assert isinstance(other, GROGU_COUNTER_V2.Bin) return GROGU_COUNTER_V2.Bin( self.d_sumw + other.d_sumw, @@ -119,39 +119,39 @@ def from_string(cls, string: str) -> "GROGU_COUNTER_V2.Bin": d_bins: list[Bin] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Counter") assert len(self.d_bins) == 1 ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ - def sumW(self): + def sumW(self) -> float: return self.d_bins[0].sumW() - def sumW2(self): + def sumW2(self) -> float: return self.d_bins[0].sumW2() - def numEntries(self): + def numEntries(self) -> float: return self.d_bins[0].numEntries() - def clone(self): + def clone(self) -> "GROGU_COUNTER_V2": return GROGU_COUNTER_V2( d_key=self.d_key, d_annotations=self.annotationsDict(), d_bins=[b.clone() for b in self.d_bins], ) - def fill(self, weight=1.0, fraction=1.0): + def fill(self, weight: float = 1.0, fraction: float = 1.0) -> None: for b in self.bins(): b.fill(weight=weight, fraction=fraction) - def set(self, *args, **kwargs): + def set(self, *args: Any, **kwargs: Any) -> None: self.d_bins[0].set(*args, **kwargs) - def bins(self): + def bins(self) -> list[Bin]: return self.d_bins @classmethod @@ -191,7 +191,7 @@ def from_string(cls, file_content: str) -> "GROGU_COUNTER_V2": d_bins=bins, ) - def to_string(self): + def to_string(self) -> str: """Convert a YODA_COUNTER_V2 object to a formatted string.""" header = ( f"BEGIN YODA_COUNTER_V2 {self.d_key}\n" diff --git a/src/babyyoda/grogu/counter_v3.py b/src/babyyoda/grogu/counter_v3.py index 7fdef75..e095b1f 100644 --- a/src/babyyoda/grogu/counter_v3.py +++ b/src/babyyoda/grogu/counter_v3.py @@ -1,12 +1,12 @@ import re from dataclasses import dataclass, field -from typing import Union +from typing import Any, Optional, Union from babyyoda.counter import UHICounter from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT -def Counter_v3(title=None, **kwargs): +def Counter_v3(title: Optional[str] = None, **kwargs: Any) -> "GROGU_COUNTER_V3": return GROGU_COUNTER_V3( d_bins=[GROGU_COUNTER_V3.Bin()], d_annotations={"Title": title} if title else {}, @@ -23,23 +23,23 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## - def clone(self): + def clone(self) -> "GROGU_COUNTER_V3.Bin": return GROGU_COUNTER_V3.Bin( d_sumw=self.d_sumw, d_sumw2=self.d_sumw2, d_numentries=self.d_numentries, ) - def fill(self, weight: float = 1.0, fraction: float = 1.0) -> bool: + def fill(self, weight: float = 1.0, fraction: float = 1.0) -> None: sf = fraction * weight self.d_sumw += sf self.d_sumw2 += sf * weight self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: self.d_sumw = bin.sumW() self.d_sumw2 = bin.sumW2() self.d_numentries = bin.numEntries() @@ -49,7 +49,7 @@ def set( numEntries: float, sumW: Union[list[float], float], sumW2: Union[list[float], float], - ): + ) -> None: if isinstance(sumW, float): sumW = [sumW] if isinstance(sumW2, float): @@ -60,13 +60,13 @@ def set( self.d_sumw2 = sumW2[0] self.d_numentries = numEntries - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def variance(self): + def variance(self) -> float: if self.d_sumw**2 - self.d_sumw2 == 0: return 0 return abs( @@ -75,22 +75,22 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> Any: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> Any: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> Any: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> Any: return self.stdDev() / self.effNumEntries() ** 0.5 - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries - def __eq__(self, other): + def __eq__(self, other: object) -> bool: return ( isinstance(other, GROGU_COUNTER_V3.Bin) and self.d_sumw == other.d_sumw @@ -98,7 +98,7 @@ def __eq__(self, other): and self.d_numentries == other.d_numentries ) - def __add__(self, other): + def __add__(self, other: Any) -> "GROGU_COUNTER_V3.Bin": assert isinstance(other, GROGU_COUNTER_V3.Bin) return GROGU_COUNTER_V3.Bin( self.d_sumw + other.d_sumw, @@ -119,39 +119,39 @@ def from_string(cls, string: str) -> "GROGU_COUNTER_V3.Bin": d_bins: list[Bin] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Counter") assert len(self.d_bins) == 1 ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ - def sumW(self): + def sumW(self) -> float: return self.d_bins[0].sumW() - def sumW2(self): + def sumW2(self) -> float: return self.d_bins[0].sumW2() - def numEntries(self): + def numEntries(self) -> float: return self.d_bins[0].numEntries() - def clone(self): + def clone(self) -> "GROGU_COUNTER_V3": return GROGU_COUNTER_V3( d_key=self.d_key, d_annotations=self.annotationsDict(), d_bins=[b.clone() for b in self.d_bins], ) - def fill(self, weight=1.0, fraction=1.0): + def fill(self, weight: float = 1.0, fraction: float = 1.0) -> None: for b in self.bins(): b.fill(weight=weight, fraction=fraction) - def set(self, *args, **kwargs): + def set(self, *args: Any, **kwargs: Any) -> None: self.d_bins[0].set(*args, **kwargs) - def bins(self): + def bins(self) -> list[Bin]: return self.d_bins @classmethod @@ -191,7 +191,7 @@ def from_string(cls, file_content: str) -> "GROGU_COUNTER_V3": d_bins=bins, ) - def to_string(self): + def to_string(self) -> str: """Convert a YODA_COUNTER_V3 object to a formatted string.""" header = ( f"BEGIN YODA_COUNTER_V3 {self.d_key}\n" diff --git a/src/babyyoda/grogu/histo1d_v2.py b/src/babyyoda/grogu/histo1d_v2.py index 913ee06..80abca0 100644 --- a/src/babyyoda/grogu/histo1d_v2.py +++ b/src/babyyoda/grogu/histo1d_v2.py @@ -1,13 +1,15 @@ import re from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT from babyyoda.grogu.counter_v2 import Counter_v2 from babyyoda.histo1d import UHIHisto1D -def Histo1D_v2(*args, title=None, **kwargs): +def Histo1D_v2( + *args: Any, title: Optional[str] = None, **kwargs: Any +) -> "GROGU_HISTO1D_V2": edges = [] if isinstance(args[0], list): edges = args[0] @@ -51,16 +53,16 @@ class Bin: d_sumwx2: float = 0.0 d_numentries: float = 0.0 - def __post_init__(self): + def __post_init__(self) -> None: assert ( self.d_xmin is None or self.d_xmax is None or self.d_xmin < self.d_xmax ) ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## - def clone(self): + def clone(self) -> "GROGU_HISTO1D_V2.Bin": return GROGU_HISTO1D_V2.Bin( d_xmin=self.d_xmin, d_xmax=self.d_xmax, @@ -71,7 +73,7 @@ def clone(self): d_numentries=self.d_numentries, ) - def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool: + def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> None: # 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 @@ -80,7 +82,7 @@ def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool: self.d_sumwx2 += sf * x**2 self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: # TODO allow modify those? # self.d_xmin = bin.xMin() # self.d_xmax = bin.xMax() @@ -90,7 +92,12 @@ def set_bin(self, bin): self.d_sumwx2 = bin.sumWX2() self.d_numentries = bin.numEntries() - def set(self, numEntries: float, sumW: list[float], sumW2: list[float]): + def contains(self, x: float) -> bool: + if self.d_xmin is None or self.d_xmax is None: + return False + return x >= self.d_xmin and x < self.d_xmax + + def set(self, numEntries: float, sumW: list[float], sumW2: list[float]) -> None: assert len(sumW) == 2 assert len(sumW2) == 2 self.d_sumw = sumW[0] @@ -99,28 +106,30 @@ def set(self, numEntries: float, sumW: list[float], sumW2: list[float]): self.d_sumwx2 = sumW2[1] self.d_numentries = numEntries - def xMin(self): + def xMin(self) -> Optional[float]: return self.d_xmin - def xMax(self): + def xMax(self) -> Optional[float]: return self.d_xmax - def xMid(self): + def xMid(self) -> Optional[float]: + if self.d_xmin is None or self.d_xmax is None: + return None return (self.d_xmin + self.d_xmax) / 2 - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def sumWX(self): + def sumWX(self) -> float: return self.d_sumwx - def sumWX2(self): + def sumWX2(self) -> float: return self.d_sumwx2 - def variance(self): + def variance(self) -> float: if self.d_sumw**2 - self.d_sumw2 == 0: return 0 return abs( @@ -129,22 +138,24 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> Any: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> Any: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> Any: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> Any: return self.stdDev() / self.effNumEntries() ** 0.5 - def dVol(self): + def dVol(self) -> Optional[float]: + if self.d_xmin is None or self.d_xmax is None: + return None return self.d_xmax - self.d_xmin - def xVariance(self): + def xVariance(self) -> float: # return self.d_sumwx2/self.d_sumw - (self.d_sumwx/self.d_sumw)**2 if self.d_sumw**2 - self.d_sumw2 == 0: return 0 @@ -153,7 +164,7 @@ def xVariance(self): / (self.d_sumw**2 - self.d_sumw2) ) - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries # def __eq__(self, other): @@ -168,7 +179,7 @@ def numEntries(self): # and self.d_numentries == other.d_numentries # ) - def __add__(self, other): + def __add__(self, other: Any) -> "GROGU_HISTO1D_V2.Bin": assert isinstance(other, GROGU_HISTO1D_V2.Bin) return GROGU_HISTO1D_V2.Bin( self.d_xmin, @@ -208,26 +219,26 @@ def from_string(cls, line: str) -> "GROGU_HISTO1D_V2.Bin": float(values[6]), ) - def to_string(bin, label=None) -> str: + def to_string(bin, label: Optional[str] = None) -> str: """Convert a Histo1DBin object to a formatted string.""" if label is None: return f"{bin.d_xmin:<12.6e}\t{bin.d_xmax:<12.6e}\t{bin.d_sumw:<12.6e}\t{bin.d_sumw2:<12.6e}\t{bin.d_sumwx:<12.6e}\t{bin.d_sumwx2:<12.6e}\t{bin.d_numentries:<12.6e}" return f"{label:8}\t{label:8}\t{bin.d_sumw:<12.6e}\t{bin.d_sumw2:<12.6e}\t{bin.d_sumwx:<12.6e}\t{bin.d_sumwx2:<12.6e}\t{bin.d_numentries:<12.6e}" d_bins: list[Bin] = field(default_factory=list) - d_overflow: Optional[Bin] = None - d_underflow: Optional[Bin] = None - d_total: Optional[Bin] = None + d_overflow: Bin = field(default_factory=Bin) + d_underflow: Bin = field(default_factory=Bin) + d_total: Bin = field(default_factory=Bin) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo1D") ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ - def clone(self): + def clone(self) -> "GROGU_HISTO1D_V2": return GROGU_HISTO1D_V2( d_key=self.d_key, d_annotations=self.annotationsDict(), @@ -237,50 +248,54 @@ def clone(self): d_total=self.d_total, ) - def underflow(self): + def underflow(self) -> Bin: return self.d_underflow - def overflow(self): + def overflow(self) -> Bin: return self.d_overflow - def fill(self, x, weight=1.0, fraction=1.0): + def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> None: self.d_total.fill(x, weight, fraction) for b in self.d_bins: - if b.xMin() <= x < b.xMax(): + if b.contains(x): b.fill(x, weight, fraction) - if x >= self.xMax() and self.d_overflow is not None: + xmax = self.xMax() + xmin = self.xMin() + if x >= 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: + if x < 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 xMax(self) -> float: + return max(b.d_xmax for b in self.d_bins if b.d_xmax is not None) - def xMin(self): - return min([b.xMin() for b in self.d_bins]) + def xMin(self) -> float: + return min(b.d_xmin for b in self.d_bins if b.d_xmin is not None) - def bins(self, includeOverflows=False): + def bins(self, includeOverflows: bool = False) -> list[Bin]: if includeOverflows: return [self.d_underflow, *self.d_bins, self.d_overflow] # TODO sorted needed here? - return sorted(self.d_bins, key=lambda b: b.d_xmin) + return sorted(self.d_bins, key=lambda b: b.d_xmin or -float("inf")) - def bin(self, *indices): + def bin(self, *indices: int) -> list[Bin]: return [self.bins()[i] for i in indices] - def binAt(self, x): + def binAt(self, x: float) -> Optional[Bin]: for b in self.bins(): - if b.d_xmin <= x < b.d_xmax: + if b.contains(x): return b return None - def binDim(self): + def binDim(self) -> int: return 1 - def xEdges(self): - return [b.xMin() for b in self.d_bins] + [self.xMax()] + def xEdges(self) -> list[float]: + return list( + {b.d_xmin for b in self.d_bins if b.d_xmin is not None} | {self.xMax()} + ) - def rebinXTo(self, edges: list[float]): + def rebinXTo(self, edges: list[float]) -> None: own_edges = self.xEdges() for e in edges: assert e in own_edges, f"Edge {e} not found in own edges {own_edges}" @@ -289,20 +304,24 @@ def rebinXTo(self, edges: list[float]): for i in range(len(edges) - 1): new_bins.append(GROGU_HISTO1D_V2.Bin(d_xmin=edges[i], d_xmax=edges[i + 1])) for b in self.bins(): - if b.xMid() < min(edges): + bm = b.xMid() + if bm is None: + err = "Bin has no xMid" + raise ValueError(err) + if bm < min(edges): self.d_underflow += b - elif b.xMid() > max(edges): + elif bm > max(edges): self.d_overflow += b else: for i in range(len(edges) - 1): - if edges[i] <= b.xMid() and b.xMid() <= edges[i + 1]: + if edges[i] <= bm <= edges[i + 1]: new_bins[i] += b self.d_bins = new_bins assert len(self.d_bins) == len(self.xEdges()) - 1 # return self - def get_projector(self): + def get_projector(self) -> Any: return Counter_v2 def to_string(histo) -> str: @@ -370,6 +389,9 @@ def from_string(cls, file_content: str) -> "GROGU_HISTO1D_V2": bins.append(cls.Bin.from_string(line)) # Create and return the YODA_HISTO1D_V2 object + if underflow is None or overflow is None or total is None: + err = "Underflow, overflow or total bin not found" + raise ValueError(err) return cls( d_key=key, d_annotations=annotations, diff --git a/src/babyyoda/grogu/histo1d_v3.py b/src/babyyoda/grogu/histo1d_v3.py index d9fd40d..fc40e9d 100644 --- a/src/babyyoda/grogu/histo1d_v3.py +++ b/src/babyyoda/grogu/histo1d_v3.py @@ -1,13 +1,16 @@ import copy import re from dataclasses import dataclass, field +from typing import Any, Optional from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT from babyyoda.grogu.counter_v3 import Counter_v3 from babyyoda.histo1d import UHIHisto1D -def Histo1D_v3(*args, title=None, **kwargs): +def Histo1D_v3( + *args: Any, title: Optional[str] = None, **kwargs: Any +) -> "GROGU_HISTO1D_V3": edges = [] if isinstance(args[0], list): edges = args[0] @@ -45,10 +48,10 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## - def clone(self): + def clone(self) -> "GROGU_HISTO1D_V3.Bin": return GROGU_HISTO1D_V3.Bin( d_sumw=self.d_sumw, d_sumw2=self.d_sumw2, @@ -57,7 +60,7 @@ def clone(self): d_numentries=self.d_numentries, ) - def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool: + def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> None: sf = fraction * weight self.d_sumw += sf self.d_sumw2 += sf * weight @@ -65,14 +68,14 @@ def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> bool: self.d_sumwx2 += sf * x**2 self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: self.d_sumw = bin.sumW() self.d_sumw2 = bin.sumW2() self.d_sumwx = bin.sumWX() self.d_sumwx2 = bin.sumWX2() self.d_numentries = bin.numEntries() - def set(self, numEntries: float, sumW: list[float], sumW2: list[float]): + def set(self, numEntries: float, sumW: list[float], sumW2: list[float]) -> None: assert len(sumW) == 2 assert len(sumW2) == 2 self.d_sumw = sumW[0] @@ -90,19 +93,19 @@ def set(self, numEntries: float, sumW: list[float], sumW2: list[float]): # def xMid(self): # return (self.d_xmin + self.d_xmax) / 2 - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def sumWX(self): + def sumWX(self) -> float: return self.d_sumwx - def sumWX2(self): + def sumWX2(self) -> float: return self.d_sumwx2 - def variance(self): + def variance(self) -> float: if self.d_sumw**2 - self.d_sumw2 == 0: return 0 return abs( @@ -111,31 +114,31 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> Any: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> Any: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> Any: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> Any: return self.stdDev() / self.effNumEntries() ** 0.5 - def xVariance(self): + def xVariance(self) -> float: # return self.d_sumwx2/self.d_sumw - (self.d_sumwx/self.d_sumw)**2 if self.d_sumw**2 - self.d_sumw2 == 0: - return 0 + return 0.0 return abs( (self.d_sumwx2 * self.d_sumw - self.d_sumwx**2) / (self.d_sumw**2 - self.d_sumw2) ) - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries - def __eq__(self, other): + def __eq__(self, other: object) -> bool: return ( isinstance(other, GROGU_HISTO1D_V3.Bin) and self.d_sumw == other.d_sumw @@ -145,7 +148,7 @@ def __eq__(self, other): and self.d_numentries == other.d_numentries ) - def __add__(self, other): + def __add__(self, other: Any) -> "GROGU_HISTO1D_V3.Bin": assert isinstance(other, GROGU_HISTO1D_V3.Bin) return GROGU_HISTO1D_V3.Bin( self.d_sumw + other.d_sumw, @@ -169,7 +172,7 @@ def from_string(cls, string: str) -> "GROGU_HISTO1D_V3.Bin": d_edges: list[float] = field(default_factory=list) d_bins: list[Bin] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo1D") # one more edge than bins, subtract 2 for underflow and overflow @@ -178,10 +181,10 @@ def __post_init__(self): ), f"{len(self.d_edges)} != {len(self.d_bins)} + 1 - 2" ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ - def clone(self): + def clone(self) -> "GROGU_HISTO1D_V3": return GROGU_HISTO1D_V3( d_key=self.d_key, d_annotations=self.annotationsDict(), @@ -189,13 +192,13 @@ def clone(self): d_bins=[b.clone() for b in self.d_bins], ) - def underflow(self): + def underflow(self) -> Bin: return self.bins(includeOverflows=True)[0] - def overflow(self): + def overflow(self) -> Bin: return self.bins(includeOverflows=True)[-1] - def fill(self, x, weight=1.0, fraction=1.0): + def fill(self, x: float, weight: float = 1.0, fraction: float = 1.0) -> None: for i, b in enumerate(self.bins()): if self.xEdges()[i] <= x < self.xEdges()[i + 1]: b.fill(x, weight, fraction) @@ -204,35 +207,35 @@ def fill(self, x, weight=1.0, fraction=1.0): if x < self.xMin(): self.underflow().fill(x, weight, fraction) - def xMax(self): + def xMax(self) -> float: return max(self.xEdges()) - def xMin(self): + def xMin(self) -> float: return min(self.xEdges()) - def bins(self, includeOverflows=False): + def bins(self, includeOverflows: bool = False) -> list[Bin]: return self.d_bins[1:-1] if not includeOverflows else self.d_bins - def bin(self, *indices): + def bin(self, *indices: int) -> list[Bin]: return [self.bins()[i] for i in indices] - def binAt(self, x): + def binAt(self, x: float) -> Optional[Bin]: # TODO add tests for binAt for i, b in enumerate(self.bins()): if self.xEdges()[i] <= x < self.xEdges()[i + 1]: return b return None - def binDim(self): + def binDim(self) -> int: return 1 - def xEdges(self): + def xEdges(self) -> list[float]: return self.d_edges - def xMid(self, i): + def xMid(self, i: int) -> float: return (self.xEdges()[i] + self.xEdges()[i + 1]) / 2 - def rebinXTo(self, edges: list[float]): + def rebinXTo(self, edges: list[float]) -> None: own_edges = self.xEdges() for e in edges: assert e in own_edges, f"Edge {e} not found in own edges {own_edges}" @@ -256,7 +259,7 @@ def rebinXTo(self, edges: list[float]): assert len(self.d_bins) == len(self.xEdges()) - 1 + 2 - def get_projector(self): + def get_projector(self) -> Any: return Counter_v3 @classmethod @@ -304,7 +307,7 @@ def from_string(cls, file_content: str) -> "GROGU_HISTO1D_V3": d_edges=edges, ) - def to_string(self): + def to_string(self) -> str: """Convert a YODA_HISTO1D_V3 object to a formatted string.""" header = ( f"BEGIN YODA_HISTO1D_V3 {self.d_key}\n" diff --git a/src/babyyoda/grogu/histo2d_v2.py b/src/babyyoda/grogu/histo2d_v2.py index 7826129..b38eb43 100644 --- a/src/babyyoda/grogu/histo2d_v2.py +++ b/src/babyyoda/grogu/histo2d_v2.py @@ -1,6 +1,6 @@ import re from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional from babyyoda.grogu.analysis_object import GROGU_ANALYSIS_OBJECT from babyyoda.grogu.histo1d_v2 import Histo1D_v2 @@ -8,10 +8,10 @@ def Histo2D_v2( - *args, - title=None, - **kwargs, -): + *args: Any, + title: Optional[str] = None, + **kwargs: Any, +) -> "GROGU_HISTO2D_V2": xedges = [] yedges = [] if isinstance(args[0], list) and isinstance(args[1], list): @@ -69,10 +69,10 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## - def clone(self): + def clone(self) -> "GROGU_HISTO2D_V2.Bin": return GROGU_HISTO2D_V2.Bin( d_xmin=self.d_xmin, d_xmax=self.d_xmax, @@ -88,7 +88,9 @@ def clone(self): d_numentries=self.d_numentries, ) - def fill(self, x: float, y: float, weight: float = 1.0, fraction=1.0): + def fill( + self, x: float, y: float, weight: float = 1.0, fraction: float = 1.0 + ) -> None: sf = fraction * weight self.d_sumw += sf self.d_sumw2 += sf * weight @@ -99,7 +101,7 @@ def fill(self, x: float, y: float, weight: float = 1.0, fraction=1.0): self.d_sumwxy += sf * x * y self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: self.d_sumw = bin.sumW() self.d_sumw2 = bin.sumW2() self.d_sumwx = bin.sumWX() @@ -115,7 +117,7 @@ def set( sumW: list[float], sumW2: list[float], sumWcross: list[float], - ): + ) -> None: assert len(sumW) == 3 assert len(sumW2) == 3 assert len(sumWcross) == 1 @@ -128,56 +130,77 @@ def set( self.d_sumwxy = sumWcross[0] self.d_numentries = numEntries - def xMin(self): + def contains(self, x: float, y: float) -> bool: + if ( + self.d_xmin is None + or self.d_xmax is None + or self.d_ymin is None + or self.d_ymax is None + ): + return False + return self.d_xmin <= x < self.d_xmax and self.d_ymin <= y < self.d_ymax + + def xMin(self) -> Optional[float]: return self.d_xmin - def xMax(self): + def xMax(self) -> Optional[float]: return self.d_xmax - def xMid(self): + def xMid(self) -> Optional[float]: + if self.d_xmin is None or self.d_xmax is None: + return None return (self.d_xmin + self.d_xmax) / 2 - def yMid(self): + def yMid(self) -> Optional[float]: + if self.d_ymin is None or self.d_ymax is None: + return None return (self.d_ymin + self.d_ymax) / 2 - def yMin(self): + def yMin(self) -> Optional[float]: return self.d_ymin - def yMax(self): + def yMax(self) -> Optional[float]: return self.d_ymax - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def sumWX(self): + def sumWX(self) -> float: return self.d_sumwx - def sumWX2(self): + def sumWX2(self) -> float: return self.d_sumwx2 - def sumWY(self): + def sumWY(self) -> float: return self.d_sumwy - def sumWY2(self): + def sumWY2(self) -> float: return self.d_sumwy2 - def sumWXY(self): + def sumWXY(self) -> float: return self.d_sumwxy - def dVol(self): + def dVol(self) -> Optional[float]: + if ( + self.d_xmin is None + or self.d_xmax is None + or self.d_ymin is None + or self.d_ymax is None + ): + return None return (self.d_xmax - self.d_xmin) * (self.d_ymax - self.d_ymin) - def crossTerm(self, x, y): + def crossTerm(self, x: int, y: int) -> float: assert (x == 0 and y == 1) or (x == 1 and y == 0) return self.sumWXY() - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries - def __add__(self, other): + def __add__(self, other: "GROGU_HISTO2D_V2.Bin") -> "GROGU_HISTO2D_V2.Bin": assert isinstance(other, GROGU_HISTO2D_V2.Bin) return GROGU_HISTO2D_V2.Bin( d_xmin=self.d_xmin, @@ -194,7 +217,7 @@ def __add__(self, other): d_numentries=self.d_numentries + other.d_numentries, ) - def to_string(self, label=None) -> str: + def to_string(self, label: Optional[str] = None) -> str: if label is None: return ( f"{self.d_xmin:<12.6e}\t{self.d_xmax:<12.6e}\t{self.d_ymin:<12.6e}\t{self.d_ymax:<12.6e}\t" @@ -204,17 +227,17 @@ def to_string(self, label=None) -> str: return f"{label:8}\t{label:8}\t{self.d_sumw:<12.6e}\t{self.d_sumw2:<12.6e}\t{self.d_sumwx:<12.6e}\t{self.d_sumwx2:<12.6e}\t{self.d_sumwy:<12.6e}\t{self.d_sumwy2:<12.6e}\t{self.d_sumwxy:<12.6e}\t{self.d_numentries:<12.6e}" d_bins: list[Bin] = field(default_factory=list) - d_total: Optional[Bin] = None + d_total: Bin = field(default_factory=Bin) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo2D") # - # YODA compatibilty code + # YODA compatibility code # - def clone(self): + def clone(self) -> "GROGU_HISTO2D_V2": return GROGU_HISTO2D_V2( d_key=self.d_key, d_annotations=self.annotationsDict(), @@ -222,45 +245,51 @@ def clone(self): d_total=self.d_total.clone(), ) - def fill(self, x, y, weight=1.0, fraction=1.0): + def fill( + self, x: float, y: float, weight: float = 1.0, fraction: float = 1.0 + ) -> None: self.d_total.fill(x, y, weight, fraction) for b in self.d_bins: - if b.d_xmin <= x < b.d_xmax and b.d_ymin <= y < b.d_ymax: + if b.contains(x, y): b.fill(x, y, weight, fraction) - def xEdges(self): + def xEdges(self) -> list[float]: assert all( x == y for x, y in zip( - sorted({b.d_xmin for b in self.d_bins})[1:], - sorted({b.d_xmax for b in self.d_bins})[:-1], + sorted({b.d_xmin for b in self.d_bins if b.d_xmin is not None})[1:], + sorted({b.d_xmax for b in self.d_bins if b.d_xmax is not None})[:-1], ) ) - return sorted({b.d_xmin for b in self.d_bins} | {self.xMax()}) + return sorted( + {b.d_xmin for b in self.d_bins if b.d_xmin is not None} | {self.xMax()} + ) - def yEdges(self): + def yEdges(self) -> list[float]: assert all( x == y for x, y in zip( - sorted({b.d_ymin for b in self.d_bins})[1:], - sorted({b.d_ymax for b in self.d_bins})[:-1], + sorted({b.d_ymin for b in self.d_bins if b.d_ymin is not None})[1:], + sorted({b.d_ymax for b in self.d_bins if b.d_ymax is not None})[:-1], ) ) - return sorted({b.d_ymin for b in self.d_bins} | {self.yMax()}) + return sorted( + {b.d_ymin for b in self.d_bins if b.d_ymin is not None} | {self.yMax()} + ) - def xMin(self): - return min(b.d_xmin for b in self.d_bins) + def xMin(self) -> float: + return min(b.d_xmin for b in self.d_bins if b.d_xmin is not None) - def yMin(self): - return min(b.d_ymin for b in self.d_bins) + def yMin(self) -> float: + return min(b.d_ymin for b in self.d_bins if b.d_ymin is not None) - def xMax(self): - return max(b.d_xmax for b in self.d_bins) + def xMax(self) -> float: + return max(b.d_xmax for b in self.d_bins if b.d_xmax is not None) - def yMax(self): - return max(b.d_ymax for b in self.d_bins) + def yMax(self) -> float: + return max(b.d_ymax for b in self.d_bins if b.d_ymax is not None) - def bins(self, includeOverflows=False): + def bins(self, includeOverflows: bool = False) -> list[Bin]: if includeOverflows: err = "includeFlow=True not supported" raise NotImplementedError(err) @@ -270,16 +299,23 @@ def bins(self, includeOverflows=False): # YODA-2 return sorted(self.d_bins, key=lambda b: (b.d_ymin, b.d_xmin)) - def bin(self, index): + def bin(self, index: int) -> Bin: return self.bins()[index] - def binAt(self, x, y): + def binAt(self, x: float, y: float) -> Optional[Bin]: for b in self.bins(): - if b.d_xmin <= x < b.d_xmax and b.d_ymin <= y < b.d_ymax: + if ( + b.d_xmin is not None + and b.d_xmax is not None + and b.d_ymin is not None + and b.d_ymax is not None + and b.d_xmin <= x < b.d_xmax + and b.d_ymin <= y < b.d_ymax + ): return b return None - def rebinXYTo(self, xedges: list[float], yedges: list[float]): + def rebinXYTo(self, xedges: list[float], yedges: list[float]) -> None: own_xedges = self.xEdges() for e in xedges: assert e in own_xedges, f"Edge {e} not found in own edges {own_xedges}" @@ -301,9 +337,13 @@ def rebinXYTo(self, xedges: list[float], yedges: list[float]): for b in self.bins(): for j in range(len(yedges) - 1): for i in range(len(xedges) - 1): + xm = b.xMid() + ym = b.yMid() if ( - xedges[i] <= b.xMid() < xedges[i + 1] - and yedges[j] <= b.yMid() < yedges[j + 1] + xm + and ym + and xedges[i] <= xm < xedges[i + 1] + and yedges[j] <= ym < yedges[j + 1] ): assert new_bins[i + j * (len(xedges) - 1)].d_xmin == xedges[i] assert ( @@ -313,23 +353,21 @@ def rebinXYTo(self, xedges: list[float], yedges: list[float]): assert ( new_bins[i + j * (len(xedges) - 1)].d_ymax == yedges[j + 1] ) - assert ( - new_bins[i + j * (len(xedges) - 1)].d_xmin - <= b.xMid() - < new_bins[i + j * (len(xedges) - 1)].d_xmax - ) + assert new_bins[i + j * (len(xedges) - 1)].d_xmin is not None + assert new_bins[i + j * (len(xedges) - 1)].d_xmax is not None + assert new_bins[i + j * (len(xedges) - 1)].contains(xm, ym) new_bins[i + j * (len(xedges) - 1)] += b self.d_bins = new_bins assert len(self.d_bins) == (len(self.xEdges()) - 1) * (len(self.yEdges()) - 1) - def rebinXTo(self, xedges: list[float]): + def rebinXTo(self, xedges: list[float]) -> None: self.rebinXYTo(xedges, self.yEdges()) - def rebinYTo(self, yedges: list[float]): + def rebinYTo(self, yedges: list[float]) -> None: self.rebinXYTo(self.xEdges(), yedges) - def get_projector(self): + def get_projector(self) -> Any: return Histo1D_v2 def to_string(self) -> str: @@ -434,10 +472,12 @@ def from_string(cls, file_content: str) -> "GROGU_HISTO2D_V2": numEntries, ) ) - - return cls( - d_key=key, - d_annotations=annotations, - d_bins=bins, - d_total=total, - ) + if total is not None: + return cls( + d_key=key, + d_annotations=annotations, + d_bins=bins, + d_total=total, + ) + err = "Total bin not found in the histogram" + raise ValueError(err) diff --git a/src/babyyoda/grogu/histo2d_v3.py b/src/babyyoda/grogu/histo2d_v3.py index d333fd2..ab1000a 100644 --- a/src/babyyoda/grogu/histo2d_v3.py +++ b/src/babyyoda/grogu/histo2d_v3.py @@ -2,6 +2,7 @@ import re import sys from dataclasses import dataclass, field +from typing import Any, Optional import numpy as np @@ -10,14 +11,14 @@ from babyyoda.histo2d import UHIHisto2D -def to_index(x, y, xedges, yedges): +def to_index(x: float, y: float, xedges: list[float], yedges: list[float]) -> float: # get ix and iy to map to correct bin - fix = None + fix = 0 for ix, xEdge in enumerate([*xedges, sys.float_info.max]): fix = ix if x < xEdge: break - fiy = None + fiy = 0 for iy, yEdge in enumerate([*yedges, sys.float_info.max]): fiy = iy if y < yEdge: @@ -27,10 +28,10 @@ def to_index(x, y, xedges, yedges): def Histo2D_v3( - *args, - title=None, - **kwargs, -): + *args: Any, + title: Optional[str] = None, + **kwargs: Any, +) -> "GROGU_HISTO2D_V3": xedges = [] yedges = [] if isinstance(args[0], list) and isinstance(args[1], list): @@ -81,7 +82,7 @@ class Bin: d_sumwxy: float = 0.0 d_numentries: float = 0.0 - def clone(self): + def clone(self) -> "GROGU_HISTO2D_V3.Bin": return GROGU_HISTO2D_V3.Bin( d_sumw=self.d_sumw, d_sumw2=self.d_sumw2, @@ -93,7 +94,9 @@ def clone(self): d_numentries=self.d_numentries, ) - def fill(self, x: float, y: float, weight: float = 1.0, fraction=1.0): + def fill( + self, x: float, y: float, weight: float = 1.0, fraction: float = 1.0 + ) -> None: sf = fraction * weight self.d_sumw += sf self.d_sumw2 += sf * weight @@ -104,7 +107,7 @@ def fill(self, x: float, y: float, weight: float = 1.0, fraction=1.0): self.d_sumwxy += sf * x * y self.d_numentries += fraction - def set_bin(self, bin): + def set_bin(self, bin: Any) -> None: self.d_sumw = bin.sumW() self.d_sumw2 = bin.sumW2() self.d_sumwx = bin.sumWX() @@ -120,7 +123,7 @@ def set( sumW: list[float], sumW2: list[float], sumWcross: list[float], - ): + ) -> None: assert len(sumW) == 3 assert len(sumW2) == 3 assert len(sumWcross) == 1 @@ -133,32 +136,32 @@ def set( self.d_sumwxy = sumWcross[0] self.d_numentries = numEntries - def sumW(self): + def sumW(self) -> float: return self.d_sumw - def sumW2(self): + def sumW2(self) -> float: return self.d_sumw2 - def sumWX(self): + def sumWX(self) -> float: return self.d_sumwx - def sumWX2(self): + def sumWX2(self) -> float: return self.d_sumwx2 - def sumWY(self): + def sumWY(self) -> float: return self.d_sumwy - def sumWY2(self): + def sumWY2(self) -> float: return self.d_sumwy2 - def sumWXY(self): + def sumWXY(self) -> float: return self.d_sumwxy - def crossTerm(self, x, y): + def crossTerm(self, x: int, y: int) -> float: assert (x == 0 and y == 1) or (x == 1 and y == 0) return self.sumWXY() - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries def __add__(self, other: "GROGU_HISTO2D_V3.Bin") -> "GROGU_HISTO2D_V3.Bin": @@ -191,7 +194,7 @@ def from_string(cls, line: str) -> "GROGU_HISTO2D_V3.Bin": d_bins: list[Bin] = field(default_factory=list) d_edges: list[list[float]] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo2D") @@ -201,10 +204,10 @@ def __post_init__(self): ) # - # YODA compatibilty code + # YODA compatibility code # - def clone(self): + def clone(self) -> "GROGU_HISTO2D_V3": return GROGU_HISTO2D_V3( d_key=self.d_key, d_annotations=self.annotationsDict(), @@ -212,35 +215,37 @@ def clone(self): d_edges=copy.deepcopy(self.d_edges), ) - def xEdges(self): + def xEdges(self) -> list[float]: return self.d_edges[0] - def yEdges(self): + def yEdges(self) -> list[float]: return self.d_edges[1] - def fill(self, x, y, weight=1.0, fraction=1.0): + def fill( + self, x: float, y: float, weight: float = 1.0, fraction: float = 1.0 + ) -> None: # Also fill overflow bins self.bins(True)[to_index(x, y, self.xEdges(), self.yEdges())].fill( x, y, weight, fraction ) - def xMax(self): + def xMax(self) -> float: assert max(self.xEdges()) == self.xEdges()[-1], "xMax is not the last edge" return self.xEdges()[-1] - def xMin(self): + def xMin(self) -> float: assert min(self.xEdges()) == self.xEdges()[0], "xMin is not the first edge" return self.xEdges()[0] - def yMax(self): + def yMax(self) -> float: assert max(self.yEdges()) == self.yEdges()[-1], "yMax is not the last edge" return self.yEdges()[-1] - def yMin(self): + def yMin(self) -> float: assert min(self.yEdges()) == self.yEdges()[0], "yMin is not the first edge" return self.yEdges()[0] - def bins(self, includeOverflows=False): + def bins(self, includeOverflows: bool = False) -> np.ndarray: if includeOverflows: return self.d_bins # TODO consider represent data always as numpy @@ -250,7 +255,7 @@ def bins(self, includeOverflows=False): .flatten() ) - def rebinXYTo(self, xedges: list[float], yedges: list[float]): + def rebinXYTo(self, xedges: list[float], yedges: list[float]) -> None: # print(f"rebinXYTo : {self.xEdges()} -> {xedges}, {self.yEdges()} -> {yedges}") own_xedges = self.xEdges() for e in xedges: @@ -308,13 +313,13 @@ def rebinXYTo(self, xedges: list[float], yedges: list[float]): assert len(self.d_bins) == (len(xedges) + 1) * (len(yedges) + 1) - def rebinXTo(self, xedges: list[float]): + def rebinXTo(self, xedges: list[float]) -> None: self.rebinXYTo(xedges, self.yEdges()) - def rebinYTo(self, yedges: list[float]): + def rebinYTo(self, yedges: list[float]) -> None: self.rebinXYTo(self.xEdges(), yedges) - def get_projector(self): + def get_projector(self) -> Any: return Histo1D_v3 def to_string(self) -> str: diff --git a/src/babyyoda/grogu/read.py b/src/babyyoda/grogu/read.py index ab88373..36916be 100644 --- a/src/babyyoda/grogu/read.py +++ b/src/babyyoda/grogu/read.py @@ -1,5 +1,7 @@ import gzip import re +from io import BufferedReader +from typing import Union from babyyoda.grogu.counter_v2 import GROGU_COUNTER_V2 from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3 @@ -10,7 +12,7 @@ # Copied from pylhe -def _extract_fileobj(filepath): +def _extract_fileobj(filepath: str) -> Union[gzip.GzipFile, BufferedReader]: """ Checks to see if a file is compressed, and if so, extract it with gzip so that the uncompressed file can be returned. @@ -32,36 +34,45 @@ def _extract_fileobj(filepath): ) -def read(file_path: str): +Histograms = Union[ + GROGU_COUNTER_V2, + GROGU_COUNTER_V3, + GROGU_HISTO1D_V2, + GROGU_HISTO1D_V3, + GROGU_HISTO2D_V2, + GROGU_HISTO2D_V3, +] + + +def read( + file_path: str, +) -> dict[ + str, + Histograms, +]: with _extract_fileobj(file_path) as f: - content = f.read() - content = content.decode("utf-8") + bcontent = f.read() + content = bcontent.decode("utf-8") pattern = re.compile( r"(BEGIN (YODA_[A-Z0-9_]+) ([^\n]+)\n(.*?)\nEND \2)", re.DOTALL ) matches = pattern.findall(content) - histograms = {} + histograms: dict[str, Histograms] = {} for full_match, hist_type, name, _body in matches: if hist_type == "YODA_COUNTER_V2": - hist = GROGU_COUNTER_V2.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_COUNTER_V2.from_string(full_match) elif hist_type == "YODA_COUNTER_V3": - hist = GROGU_COUNTER_V3.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_COUNTER_V3.from_string(full_match) elif hist_type == "YODA_HISTO1D_V2": - hist = GROGU_HISTO1D_V2.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_HISTO1D_V2.from_string(full_match) elif hist_type == "YODA_HISTO1D_V3": - hist = GROGU_HISTO1D_V3.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_HISTO1D_V3.from_string(full_match) elif hist_type == "YODA_HISTO2D_V2": - hist = GROGU_HISTO2D_V2.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_HISTO2D_V2.from_string(full_match) elif hist_type == "YODA_HISTO2D_V3": - hist = GROGU_HISTO2D_V3.from_string(full_match) - histograms[name] = hist + histograms[name] = GROGU_HISTO2D_V3.from_string(full_match) else: # Add other parsing logic for different types if necessary print(f"Unknown type: {hist_type}, skipping...") diff --git a/src/babyyoda/grogu/write.py b/src/babyyoda/grogu/write.py index 60b3ace..da2d635 100644 --- a/src/babyyoda/grogu/write.py +++ b/src/babyyoda/grogu/write.py @@ -1,12 +1,38 @@ +from typing import Union + +from babyyoda.counter import UHICounter +from babyyoda.grogu.counter_v2 import GROGU_COUNTER_V2 +from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3 +from babyyoda.grogu.histo1d_v2 import GROGU_HISTO1D_V2 +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 +from babyyoda.histo1d import UHIHisto1D +from babyyoda.histo2d import UHIHisto2D from babyyoda.util import open_write_file +ggHistograms = Union[ + GROGU_COUNTER_V2, + GROGU_COUNTER_V3, + GROGU_HISTO1D_V2, + GROGU_HISTO1D_V3, + GROGU_HISTO2D_V2, + GROGU_HISTO2D_V3, +] +uhiHistograms = Union[UHICounter, UHIHisto1D, UHIHisto2D] +Histograms = Union[uhiHistograms, ggHistograms] + -def write(histograms, file_path: str, gz=False): +def write( + histograms: Union[list[Histograms], dict[str, Histograms]], + file_path: str, + gz: bool = False, +) -> None: """Write multiple histograms to a file in YODA format.""" with open_write_file(file_path, gz=gz) as f: # if dict loop over values if isinstance(histograms, dict): - histograms = histograms.values() + histograms = list(histograms.values()) for histo in histograms: f.write(histo.to_string()) f.write("\n") diff --git a/src/babyyoda/histo1d.py b/src/babyyoda/histo1d.py index fe8e832..1d754f2 100644 --- a/src/babyyoda/histo1d.py +++ b/src/babyyoda/histo1d.py @@ -1,14 +1,15 @@ import contextlib import sys +from typing import Any, Optional, Union import numpy as np +import babyyoda from babyyoda.analysisobject import UHIAnalysisObject -from babyyoda.counter import UHICounter from babyyoda.util import loc, overflow, project, rebin, rebinBy_to_rebinTo, underflow -def set_bin1d(target, source): +def set_bin1d(target: Any, source: Any) -> None: # TODO allow modify those? # self.d_xmin = bin.xMin() # self.d_xmax = bin.xMax() @@ -23,24 +24,49 @@ def set_bin1d(target, source): raise NotImplementedError(err) -def Histo1D(*args, **kwargs): +def Histo1D(*args: Any, **kwargs: Any) -> "UHIHisto1D": """ Automatically select the correct version of the Histo1D class """ try: from babyyoda import yoda + + return yoda.Histo1D(*args, **kwargs) except ImportError: - import babyyoda.grogu as yoda - return yoda.Histo1D(*args, **kwargs) + from babyyoda import grogu + + return grogu.Histo1D(*args, **kwargs) # TODO make this implementation independent (no V2 or V3...) class UHIHisto1D(UHIAnalysisObject): + ###### + # Minimum required functions + ###### + + def bins(self, includeOverflows: bool = False) -> list[Any]: + raise NotImplementedError + + def xEdges(self) -> list[float]: + raise NotImplementedError + + def annotationsDict(self) -> dict[str, Optional[str]]: + raise NotImplementedError + + def clone(self) -> "UHIHisto1D": + raise NotImplementedError + + def get_projector(self) -> Any: + raise NotImplementedError + + def rebinXTo(self, bins: list[float]) -> None: + raise NotImplementedError + ###### # BACKENDS ###### - def to_grogu_v2(self): + def to_grogu_v2(self) -> Any: from babyyoda.grogu.histo1d_v2 import GROGU_HISTO1D_V2 tot = GROGU_HISTO1D_V2.Bin() @@ -87,7 +113,7 @@ def to_grogu_v2(self): ), ) - def to_grogu_v3(self): + def to_grogu_v3(self) -> Any: from babyyoda.grogu.histo1d_v3 import GROGU_HISTO1D_V3 return GROGU_HISTO1D_V3( @@ -124,58 +150,60 @@ def to_grogu_v3(self): ], ) - def to_yoda_v3(self): + def to_yoda_v3(self) -> Any: err = "Not implemented yet" raise NotImplementedError(err) - def to_string(self): + def to_string(self) -> str: # 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() + return str(self.to_grogu_v3().to_string()) ######################################################## # YODA compatibility code (dropped legacy code?) ######################################################## - def overflow(self): + def overflow(self) -> Any: return self.bins(includeOverflows=True)[-1] - def underflow(self): + def underflow(self) -> Any: return self.bins(includeOverflows=True)[0] - def errWs(self): + def errWs(self) -> np.ndarray: return np.sqrt(np.array([b.sumW2() for b in self.bins()])) - def xMins(self): + def xMins(self) -> list[float]: return self.xEdges()[:-1] # return np.array([b.xMin() for b in self.bins()]) - def xMaxs(self): + def xMaxs(self) -> list[float]: return self.xEdges()[1:] # return np.array([b.xMax() for b in self.bins()]) - def sumWs(self): - return np.array([b.sumW() for b in self.bins()]) + def sumWs(self) -> list[float]: + return [b.sumW() for b in self.bins()] - def sumW2s(self): - return np.array([b.sumW2() for b in self.bins()]) + def sumW2s(self) -> list[float]: + return [b.sumW2() for b in self.bins()] - def xMean(self, includeOverflows=True): + def xMean(self, includeOverflows: bool = True) -> float: return sum( - b.sumWX() for b in self.bins(includeOverflows=includeOverflows) - ) / sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows)) + float(b.sumWX()) for b in self.bins(includeOverflows=includeOverflows) + ) / sum(float(b.sumW()) for b in self.bins(includeOverflows=includeOverflows)) - def integral(self, includeOverflows=True): - return sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows)) + def integral(self, includeOverflows: bool = True) -> float: + return sum( + float(b.sumW()) for b in self.bins(includeOverflows=includeOverflows) + ) - def rebinXBy(self, factor: int, begin=1, end=sys.maxsize): + def rebinXBy(self, factor: int, begin: int = 1, end: int = sys.maxsize) -> None: new_edges = rebinBy_to_rebinTo(self.xEdges(), factor, begin, end) self.rebinXTo(new_edges) - def rebinBy(self, *args, **kwargs): + def rebinBy(self, *args: Any, **kwargs: Any) -> None: self.rebinXBy(*args, **kwargs) - def rebinTo(self, *args, **kwargs): + def rebinTo(self, *args: Any, **kwargs: Any) -> None: self.rebinXTo(*args, **kwargs) ######################################################## @@ -183,29 +211,32 @@ def rebinTo(self, *args, **kwargs): ######################################################## @property - def axes(self): + def axes(self) -> list[list[tuple[float, float]]]: return [list(zip(self.xMins(), self.xMaxs()))] @property - def kind(self): + def kind(self) -> str: # TODO reeavaluate this return "COUNT" - def counts(self): + def counts(self) -> np.ndarray: return np.array([b.numEntries() for b in self.bins()]) - def values(self): + def values(self) -> np.ndarray: return np.array([b.sumW() for b in self.bins()]) - def variances(self): + def variances(self) -> np.ndarray: return np.array([(b.sumW2()) for b in self.bins()]) - def __getitem__(self, slices): + def __getitem__( + self, + slices: Union[ + int, loc, slice, type[babyyoda.util.underflow], type[babyyoda.util.overflow] + ], + ) -> Any: index = self.__get_index(slices) # integer index - if isinstance(slices, int): - return self.bins()[index] - if isinstance(slices, loc): + if isinstance(index, int): # loc and int return self.bins()[index] if slices is underflow: return self.underflow() @@ -221,6 +252,11 @@ def __getitem__(self, slices): self.__get_index(item.stop), item.step, ) + if not (start is None or isinstance(start, int)) or not ( + stop is None or isinstance(stop, int) + ): + err = "Invalid argument type" + raise TypeError(err) sc = self.clone() if isinstance(step, rebin): @@ -246,8 +282,17 @@ def __getitem__(self, slices): err = "Invalid argument type" raise TypeError(err) - def __get_index(self, slices): - index = None + def __get_index( + self, + slices: Union[ + int, loc, slice, type[babyyoda.util.underflow], type[babyyoda.util.overflow] + ], + ) -> Optional[ + Union[int, type[babyyoda.util.underflow], type[babyyoda.util.overflow]] + ]: + index: Optional[ + Union[type[Union[babyyoda.util.underflow, babyyoda.util.overflow]], int] + ] = None if isinstance(slices, int): index = slices while index < 0: @@ -261,28 +306,38 @@ def __get_index(self, slices): and slices.value < self.xEdges()[i + 1] ): idx = i - index = idx + slices.offset + if idx is not None: + index = idx + slices.offset if slices is underflow: index = underflow if slices is overflow: index = overflow return index - def __set_by_index(self, index, value): - if index == underflow: + def __set_by_index( + self, + index: Union[type[babyyoda.util.underflow], type[babyyoda.util.overflow], int], + value: Any, + ) -> None: + if index is underflow: set_bin1d(self.underflow(), value) return - if index == overflow: + if index is overflow: set_bin1d(self.overflow(), value) return - set_bin1d(self.bins()[index], value) + if isinstance(index, int): + set_bin1d(self.bins()[index], value) + return + err = "Invalid argument type" + raise TypeError(err) - def __setitem__(self, slices, value): + def __setitem__(self, slices: Any, value: Any) -> None: # integer index index = self.__get_index(slices) - self.__set_by_index(index, value) + if index is not None: + self.__set_by_index(index, value) - def project(self) -> UHICounter: + def project(self) -> Any: # sc = self.clone().rebinTo(self.xEdges()[0], self.xEdges()[-1]) p = self.get_projector()() p.set( @@ -293,7 +348,7 @@ def project(self) -> UHICounter: p.setAnnotationsDict(self.annotationsDict()) return p - def plot(self, *args, binwnorm=1.0, **kwargs): + def plot(self, *args: Any, binwnorm: float = 1.0, **kwargs: Any) -> None: import mplhep as hep hep.histplot( @@ -305,7 +360,7 @@ def plot(self, *args, binwnorm=1.0, **kwargs): **kwargs, ) - def _ipython_display_(self): + def _ipython_display_(self) -> "UHIHisto1D": with contextlib.suppress(ImportError): self.plot() return self diff --git a/src/babyyoda/histo2d.py b/src/babyyoda/histo2d.py index 6e8b92b..a2ba548 100644 --- a/src/babyyoda/histo2d.py +++ b/src/babyyoda/histo2d.py @@ -1,5 +1,6 @@ import contextlib import sys +from typing import Any, Optional, Union import numpy as np @@ -16,7 +17,7 @@ ) -def set_bin2d(target, source): +def set_bin2d(target: Any, source: Any) -> None: # TODO allow modify those? # self.d_xmin = bin.xMin() # self.d_xmax = bin.xMax() @@ -32,23 +33,54 @@ def set_bin2d(target, source): raise NotImplementedError(err) -def Histo2D(*args, **kwargs): +def Histo2D(*args: list[Any], **kwargs: list[Any]) -> "UHIHisto2D": """ Automatically select the correct version of the Histo2D class """ try: from babyyoda import yoda + + return yoda.Histo2D(*args, **kwargs) except ImportError: - import babyyoda.grogu as yoda - return yoda.Histo2D(*args, **kwargs) + from babyyoda import grogu + + return grogu.Histo2D(*args, **kwargs) class UHIHisto2D(UHIAnalysisObject): + ###### + # Minimum required functions + ###### + + def bins(self, includeOverflows: bool = False) -> list[Any]: + raise NotImplementedError + + def xEdges(self) -> list[float]: + raise NotImplementedError + + def yEdges(self) -> list[float]: + raise NotImplementedError + + def rebinXTo(self, edges: list[float]) -> None: + raise NotImplementedError + + def rebinYTo(self, edges: list[float]) -> None: + raise NotImplementedError + + def annotationsDict(self) -> dict[str, Optional[str]]: + raise NotImplementedError + + def clone(self) -> "UHIHisto2D": + raise NotImplementedError + + def get_projector(self) -> Any: + raise NotImplementedError + ##### # BACKENDS ##### - def to_grogu_v2(self): + def to_grogu_v2(self) -> Any: from babyyoda.grogu.histo2d_v2 import GROGU_HISTO2D_V2 tot = GROGU_HISTO2D_V2.Bin() @@ -85,7 +117,7 @@ def to_grogu_v2(self): ], ) - def to_grogu_v3(self): + def to_grogu_v3(self) -> Any: from babyyoda.grogu.histo2d_v3 import GROGU_HISTO2D_V3 bins = [] @@ -122,11 +154,6 @@ def to_grogu_v3(self): ], ) - def to_string(self): - if hasattr(self.target, "to_string"): - return self.target.to_string() - return self.to_grogu_v3().to_string() - # def bins(self, *args, **kwargs): # # fix order # return self.target.bins(*args, **kwargs) @@ -147,50 +174,52 @@ def to_string(self): # # This is a YODA-1 feature that is not present in YODA-2 # return self.bins(includeOverflows=True)[0] - def xMins(self): + def xMins(self) -> list[float]: return self.xEdges()[:-1] # return np.array(sorted(list(set([b.xMin() for b in self.bins()])))) - def xMaxs(self): + def xMaxs(self) -> list[float]: return self.xEdges()[1:] # return np.array(sorted(list(set([b.xMax() for b in self.bins()])))) - def yMins(self): + def yMins(self) -> list[float]: return self.yEdges()[:-1] # return np.array(sorted(list(set([b.yMin() for b in self.bins()])))) - def yMaxs(self): + def yMaxs(self) -> list[float]: return self.yEdges()[1:] # return np.array(sorted(list(set([b.yMax() for b in self.bins()])))) - def sumWs(self): - return np.array([b.sumW() for b in self.bins()]) + def sumWs(self) -> list[float]: + return [b.sumW() for b in self.bins()] - def sumWXYs(self): + def sumWXYs(self) -> list[float]: return [b.crossTerm(0, 1) for b in self.bins()] - def xMean(self, includeOverflows=True): - return sum(b.sumWX() for b in self.d_bins) / sum( - b.sumW() for b in self.bins(includeOverflows=includeOverflows) - ) + def xMean(self, includeOverflows: bool = True) -> Any: + return sum( + float(b.sumWX()) for b in self.bins(includeOverflows=includeOverflows) + ) / sum(float(b.sumW()) for b in self.bins(includeOverflows=includeOverflows)) - def yMean(self, includeOverflows=True): - return sum(b.sumWY() for b in self.d_bins) / sum( - b.sumW() for b in self.bins(includeOverflows=includeOverflows) - ) + def yMean(self, includeOverflows: bool = True) -> Any: + return sum( + float(b.sumWY()) for b in self.bins(includeOverflows=includeOverflows) + ) / sum(float(b.sumW()) for b in self.bins(includeOverflows=includeOverflows)) - def integral(self, includeOverflows=True): - return sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows)) + def integral(self, includeOverflows: bool = True) -> float: + return sum( + float(b.sumW()) for b in self.bins(includeOverflows=includeOverflows) + ) - def rebinXBy(self, factor: int, begin=1, end=sys.maxsize): + def rebinXBy(self, factor: int, begin: int = 1, end: int = sys.maxsize) -> None: new_edges = rebinBy_to_rebinTo(self.xEdges(), factor, begin, end) self.rebinXTo(new_edges) - def rebinYBy(self, factor: int, begin=1, end=sys.maxsize): + def rebinYBy(self, factor: int, begin: int = 1, end: int = sys.maxsize) -> None: new_edges = rebinBy_to_rebinTo(self.yEdges(), factor, begin, end) self.rebinYTo(new_edges) - def dVols(self): + def dVols(self) -> list[float]: ret = [] for iy in range(len(self.yMins())): for ix in range(len(self.xMins())): @@ -198,59 +227,59 @@ def dVols(self): (self.xMaxs()[ix] - self.xMins()[ix]) * (self.yMaxs()[iy] - self.yMins()[iy]) ) - return np.array(ret) + return ret ######################################################## # Generic UHI code ######################################################## @property - def axes(self): + def axes(self) -> list[list[tuple[float, float]]]: return [ list(zip(self.xMins(), self.xMaxs())), list(zip(self.yMins(), self.yMaxs())), ] @property - def kind(self): + def kind(self) -> str: # TODO reeavaluate this return "COUNT" - def values(self): - return self.sumWs().reshape((len(self.axes[1]), len(self.axes[0]))).T + def values(self) -> np.ndarray: + return np.array(self.sumWs()).reshape((len(self.axes[1]), len(self.axes[0]))).T - def variances(self): + def variances(self) -> np.ndarray: return ( np.array([b.sumW2() for b in self.bins()]) .reshape((len(self.axes[1]), len(self.axes[0]))) .T ) - def counts(self): + def counts(self) -> np.ndarray: return ( np.array([b.numEntries() for b in self.bins()]) .reshape((len(self.axes[1]), len(self.axes[0]))) .T ) - def __single_index(self, ix, iy): + def __single_index(self, ix: int, iy: int) -> int: return iy * len(self.axes[0]) + ix # return ix * len(self.axes[1]) + iy - def __get_by_indices(self, ix, iy): + def __get_by_indices(self, ix: int, iy: int) -> Any: return self.bins()[ self.__single_index(ix, iy) ] # THIS is the fault with/without overflows! - def __get_index_by_loc(self, loc, bins): + def __get_index_by_loc(self, oloc: loc, bins: list[tuple[float, float]]) -> int: # find the index in bin where loc is for a, b in bins: - if a <= loc.value and loc.value < b: - return bins.index((a, b)) + loc.offset - err = f"loc {loc.value} is not in the range of {bins}" + if a <= oloc.value and oloc.value < b: + return bins.index((a, b)) + oloc.offset + err = f"loc {oloc.value} is not in the range of {bins}" raise ValueError(err) - def __get_x_index(self, slices): + def __get_x_index(self, slices: Union[int, loc, slice]) -> Optional[int]: ix = None if isinstance(slices, int): ix = slices @@ -258,7 +287,7 @@ def __get_x_index(self, slices): ix = self.__get_index_by_loc(slices, self.axes[0]) return ix - def __get_y_index(self, slices): + def __get_y_index(self, slices: Union[int, loc, slice]) -> Optional[int]: iy = None if isinstance(slices, int): iy = slices @@ -266,21 +295,27 @@ def __get_y_index(self, slices): iy = self.__get_index_by_loc(slices, self.axes[1]) return iy - def __get_indices(self, slices): + def __get_indices( + self, slices: tuple[Union[int, loc, slice], Union[int, loc, slice]] + ) -> tuple[Optional[int], Optional[int]]: return self.__get_x_index(slices[0]), self.__get_y_index(slices[1]) - def __setitem__(self, slices, value): + def __setitem__( + self, slices: tuple[Union[int, slice, loc], Union[int, slice, loc]], value: Any + ) -> None: set_bin2d(self.__getitem__(slices), value) - def __getitem__(self, slices): + def __getitem__( + self, slices: tuple[Union[int, slice, loc], Union[int, slice, loc]] + ) -> Any: # integer index - if slices is underflow: + if slices is underflow: # type: ignore[comparison-overlap] err = "No underflow bin in 2D histogram" raise TypeError(err) - if slices is overflow: + if slices is overflow: # type: ignore[comparison-overlap] err = "No overflow bin in 2D histogram" raise TypeError(err) - if isinstance(slices, tuple) and len(slices) == 2: + if isinstance(slices, tuple) and len(slices) == 2: # type: ignore[redundant-expr] ix, iy = self.__get_indices(slices) if isinstance(ix, int) and isinstance(iy, int): return self.__get_by_indices(ix, iy) @@ -288,18 +323,18 @@ def __getitem__(self, slices): slices = (slices[0], slice(iy, iy + 1, project)) if isinstance(ix, int) and isinstance(slices[1], slice): slices = (slice(ix, ix + 1, project), slices[1]) - ix, iy = slices + s_ix, s_iy = slices sc = self.clone() - if isinstance(ix, slice) and isinstance(iy, slice): + if isinstance(s_ix, slice) and isinstance(s_iy, slice): xstart, xstop, xstep = ( - self.__get_x_index(ix.start), - self.__get_x_index(ix.stop), - ix.step, + self.__get_x_index(s_ix.start), + self.__get_x_index(s_ix.stop), + s_ix.step, ) ystart, ystop, ystep = ( - self.__get_y_index(iy.start), - self.__get_y_index(iy.stop), - iy.step, + self.__get_y_index(s_iy.start), + self.__get_y_index(s_iy.stop), + s_iy.step, ) if isinstance(ystep, rebin): @@ -331,12 +366,11 @@ def __getitem__(self, slices): return sc err = "Slice with Index not implemented" raise NotImplementedError(err) - # TODO implement slice - err = "Invalid argument type" + err = "Invalid argument type" # type: ignore[unreachable] raise TypeError(err) - def projectX(self): + def projectX(self) -> Any: # Sum c = self.clone() c.rebinXTo([self.xEdges()[0], self.xEdges()[-1]]) @@ -347,7 +381,7 @@ def projectX(self): p.setAnnotationsDict(self.annotationsDict()) return p - def projectY(self): + def projectY(self) -> Any: # Sum c = self.clone() c.rebinYTo([self.yEdges()[0], self.yEdges()[-1]]) @@ -359,13 +393,16 @@ def projectY(self): return p # TODO maybe N dim project - def project(self, axis: int = 0): + def project(self, axis: int = 0) -> Any: assert axis in [0, 1] if axis == 0: return self.projectX() return self.projectY() - def plot(self, *args, binwnorm=True, **kwargs): + def to_string(self) -> str: + return str(self.to_grogu_v3().to_string()) + + def plot(self, *args: Any, binwnorm: bool = True, **kwargs: Any) -> None: # # TODO should use histplot # import mplhep as hep @@ -383,7 +420,7 @@ def plot(self, *args, binwnorm=True, **kwargs): # Hack in the temporary division by dVol saved_values = self.values - def temp_values(): + def temp_values() -> np.ndarray: return ( np.array( [b.sumW() / vol for b, vol in zip(self.bins(), self.dVols())] @@ -392,12 +429,12 @@ def temp_values(): .T ) - self.values = temp_values + self.values = temp_values # type: ignore[method-assign] hep.hist2dplot(self, *args, **kwargs) if binwnorm: - self.values = saved_values + self.values = saved_values # type: ignore[method-assign] - def _ipython_display_(self): + def _ipython_display_(self) -> "UHIHisto2D": with contextlib.suppress(ImportError): self.plot() return self diff --git a/src/babyyoda/read.py b/src/babyyoda/read.py index 5ebd7f2..9b5ca3a 100644 --- a/src/babyyoda/read.py +++ b/src/babyyoda/read.py @@ -1,30 +1,30 @@ -import warnings +from typing import Any from babyyoda import grogu -def read(file_path: str): +def read(file_path: str) -> dict[str, Any]: try: return read_yoda(file_path) except ImportError: - warnings.warn( - "yoda is not installed, falling back to python grogu implementation", - stacklevel=2, - ) + # warnings.warn( + # "yoda is not installed, falling back to python grogu implementation", + # stacklevel=2, + # ) return read_grogu(file_path) -def read_yoda(file_path: str): +def read_yoda(file_path: str) -> dict[str, Any]: """ - Wrap yoda histograms in the by HISTO1D_V2 class + Wrap yoda histograms in the yoda classes """ from babyyoda import yoda return yoda.read(file_path) -def read_grogu(file_path: str): +def read_grogu(file_path: str) -> dict[str, Any]: """ - Wrap grogu histograms in the by HISTO1D_V2 class + Wrap grogu histograms in the by classes """ return grogu.read(file_path) diff --git a/src/babyyoda/util.py b/src/babyyoda/util.py index e906d80..8b4d3e5 100644 --- a/src/babyyoda/util.py +++ b/src/babyyoda/util.py @@ -1,27 +1,28 @@ import gzip import inspect import sys +from typing import Any, Optional, TextIO class loc: "When used in the start or stop of a Histogram's slice, x is taken to be the position in data coordinates." - def __init__(self, x, offset=0): + def __init__(self, x: float, offset: int = 0): self.value = x self.offset = offset # add and subtract method - def __add__(self, other): + def __add__(self, other: int) -> "loc": return loc(self.value, self.offset + other) - def __sub__(self, other): + def __sub__(self, other: int) -> "loc": return loc(self.value, self.offset - other) class rebin: "When used in the step of a Histogram's slice, rebin(n) combines bins, scaling their widths by a factor of n. If the number of bins is not divisible by n, the remainder is added to the overflow bin." - def __init__(self, factor): + def __init__(self, factor: int): self.factor = factor @@ -37,23 +38,23 @@ class project: pass -def open_write_file(file_path, gz=False): +def open_write_file(file_path: str, gz: bool = False) -> TextIO: if file_path.endswith((".gz", ".gzip")) or gz: return gzip.open(file_path, "wt") return open(file_path, "w") -def uses_yoda(obj): +def uses_yoda(obj: Any) -> bool: if hasattr(obj, "target"): return uses_yoda(obj.target) return is_yoda(obj) -def is_yoda(obj): +def is_yoda(obj: Any) -> bool: return is_from_package(obj, "yoda.") -def is_from_package(obj, package_name): +def is_from_package(obj: Any, package_name: str) -> bool: # Get the class of the object obj_class = obj.__class__ @@ -68,7 +69,7 @@ def is_from_package(obj, package_name): return False -def has_own_method(cls, method_name): +def has_own_method(cls: Any, method_name: str) -> bool: # Check if the class has the method defined if not hasattr(cls, method_name): return False @@ -77,18 +78,22 @@ def has_own_method(cls, method_name): cls_method = getattr(cls, method_name) parent_method = getattr(cls.__bases__[0], method_name, None) + if cls_method is None: + return False + if parent_method is None: + return True # Compare the underlying function (__func__) if both exist return cls_method.__func__ is not parent_method.__func__ -def rebinBy_to_rebinTo(edges: list[float], factor: int, begin=1, end=sys.maxsize): +def rebinBy_to_rebinTo( + edges: list[float], factor: int, begin: int = 1, end: int = sys.maxsize +) -> list[float]: # Just compute the new edges and call rebinXTo start = begin - 1 stop = end - if start is None: - start = 0 stop = (len(edges) - 1) if stop >= sys.maxsize else stop - 1 - new_edges = [] + new_edges: list[float] = [] # new_bins = [] # new_bins += [self.underflow()] for i in range(start): @@ -110,6 +115,10 @@ def rebinBy_to_rebinTo(edges: list[float], factor: int, begin=1, end=sys.maxsize # add both edges new_edges.append(xmin) new_edges.append(xmax) + if last is None: + # alternatively just return old edges + err = "No bins to rebin" + raise ValueError(err) for j in range(last + 1, (len(edges) - 1)): # new_bins.append(self.bins()[j].clone()) new_edges.append(edges[j]) @@ -118,7 +127,7 @@ def rebinBy_to_rebinTo(edges: list[float], factor: int, begin=1, end=sys.maxsize return list(set(new_edges)) -def shift_rebinby(ystart, ystop): +def shift_rebinby(ystart: Optional[int], ystop: Optional[int]) -> tuple[int, int]: # weird yoda default if ystart is None: ystart = 1 @@ -131,7 +140,9 @@ def shift_rebinby(ystart, ystop): return ystart, ystop -def shift_rebinto(xstart, xstop): +def shift_rebinto( + xstart: Optional[int], xstop: Optional[int] +) -> tuple[Optional[int], Optional[int]]: if xstop is not None: xstop += 1 return xstart, xstop diff --git a/src/babyyoda/write.py b/src/babyyoda/write.py index 9c8a20d..5d11d49 100644 --- a/src/babyyoda/write.py +++ b/src/babyyoda/write.py @@ -1,10 +1,15 @@ +from typing import Any + from babyyoda import grogu -def write(anyhistograms, file_path: str, *args, **kwargs): +def write(anyhistograms: Any, file_path: str, *args: Any, **kwargs: Any) -> None: + listhistograms: list[Any] = [] # if dict loop over values if isinstance(anyhistograms, dict): - listhistograms = anyhistograms.values() + listhistograms = list(anyhistograms.values()) + elif isinstance(anyhistograms, list): + listhistograms = anyhistograms # check if all histograms are yoda => use yoda use_yoda = True try: @@ -24,11 +29,11 @@ def write(anyhistograms, file_path: str, *args, **kwargs): # These functions are just to be similar to the read functions -def write_grogu(histograms, file_path: str, gz=False): +def write_grogu(histograms: Any, file_path: str, gz: bool = False) -> None: grogu.write(histograms, file_path, gz=gz) -def write_yoda(histograms, file_path: str, gz=False): +def write_yoda(histograms: Any, file_path: str, gz: bool = False) -> None: # TODO we could force convert to YODA in Histo{1,2}D here ... from babyyoda import yoda diff --git a/src/babyyoda/yoda/counter.py b/src/babyyoda/yoda/counter.py index 7b27d05..0e28516 100644 --- a/src/babyyoda/yoda/counter.py +++ b/src/babyyoda/yoda/counter.py @@ -1,3 +1,5 @@ +from typing import Any, Optional + import yoda import babyyoda @@ -5,7 +7,7 @@ class Counter(babyyoda.UHICounter): - def __init__(self, *args, **kwargs): + def __init__(self, *args: list[Any], **kwargs: dict[Any, Any]) -> None: """ target is either a yoda or grogu Counter """ @@ -18,11 +20,40 @@ def __init__(self, *args, **kwargs): super().__setattr__("target", target) + ########################## + # Basic needed functions for UHI directly relayed to target + ########################## + + def path(self) -> str: + return str(self.target.path()) + + def sumW(self) -> float: + return float(self.target.sumW()) + + def sumW2(self) -> float: + return float(self.target.sumW2()) + + def numEntries(self) -> float: + return float(self.target.numEntries()) + + def bins(self, *args: Any, **kwargs: Any) -> Any: + return self.target.bins(*args, **kwargs) + + def clone(self) -> "Counter": + return Counter(self.target.clone()) + + # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 + def annotationsDict(self) -> dict[str, Optional[str]]: + d = {} + for k in self.target.annotations(): + d[k] = self.target.annotation(k) + return d + ######################################################## # Relay all attribute access to the target object ######################################################## - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: # if we overwrite it here, use that if has_own_method(Counter, name): return getattr(self, name) @@ -52,16 +83,3 @@ def __getattr__(self, name): # return self.target(*args, **kwargs) # err = f"'{type(self.target).__name__}' object is not callable" # raise TypeError(err) - - def bins(self, *args, **kwargs): - return self.target.bins(*args, **kwargs) - - def clone(self): - return Counter(self.target.clone()) - - # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 - def annotationsDict(self): - d = {} - for k in self.target.annotations(): - d[k] = self.target.annotation(k) - return d diff --git a/src/babyyoda/yoda/histo1d.py b/src/babyyoda/yoda/histo1d.py index 77d1496..53f2e61 100644 --- a/src/babyyoda/yoda/histo1d.py +++ b/src/babyyoda/yoda/histo1d.py @@ -1,3 +1,5 @@ +from typing import Any, Optional, cast + import yoda from packaging import version @@ -7,7 +9,7 @@ class Histo1D(babyyoda.UHIHisto1D): - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: """ target is either a yoda or grogu HISTO1D_V2 """ @@ -21,11 +23,64 @@ def __init__(self, *args, **kwargs): super().__setattr__("target", target) + ########################## + # Basic needed functions for UHI directly relayed to target + ########################## + + def path(self) -> str: + return str(self.target.path()) + + def clone(self) -> "Histo1D": + return Histo1D(self.target.clone()) + + def get_projector(self) -> Any: + return Counter + + def xEdges(self) -> list[float]: + return cast(list[float], self.target.xEdges()) + + def bins(self, includeOverflows: bool = False, *args: Any, **kwargs: Any) -> Any: + import yoda + + if version.parse(yoda.__version__) >= version.parse("2.0.0"): + return self.target.bins(*args, includeOverflows=includeOverflows, **kwargs) + # YODA1 does not offer inlcudeOverflows + if includeOverflows: + return [ + self.target.underflow(), + *self.target.bins(), + self.target.overflow(), + ] + return self.target.bins(*args, **kwargs) + + def rebinXTo(self, *args: Any, **kwargs: Any) -> None: + import yoda + + if version.parse(yoda.__version__) >= version.parse("2.0.0"): + self.target.rebinXTo(*args, **kwargs) + else: + self.target.rebinTo(*args, **kwargs) + + def rebinXBy(self, *args: Any, **kwargs: Any) -> None: + import yoda + + if version.parse(yoda.__version__) >= version.parse("2.0.0"): + self.target.rebinXBy(*args, **kwargs) + else: + self.target.rebinBy(*args, **kwargs) + + # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 + def annotationsDict(self) -> dict[str, Optional[str]]: + d = {} + for k in self.target.annotations(): + d[k] = self.target.annotation(k) + return d + ######################################################## # Relay all attribute access to the target object ######################################################## - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: # if we overwrite it here, use that if has_own_method(Histo1D, name): return getattr(self, name) @@ -56,46 +111,5 @@ def __getattr__(self, name): # err = f"'{type(self.target).__name__}' object is not callable" # raise TypeError(err) - def bins(self, includeOverflows=False, *args, **kwargs): - import yoda - - if version.parse(yoda.__version__) >= version.parse("2.0.0"): - return self.target.bins(*args, includeOverflows=includeOverflows, **kwargs) - # YODA1 does not offer inlcudeOverflows - if includeOverflows: - return [ - self.target.underflow(), - *self.target.bins(), - self.target.overflow(), - ] - return self.target.bins(*args, **kwargs) - - def rebinXTo(self, *args, **kwargs): - import yoda - - if version.parse(yoda.__version__) >= version.parse("2.0.0"): - return self.target.rebinXTo(*args, **kwargs) - return self.target.rebinTo(*args, **kwargs) - - def rebinXBy(self, *args, **kwargs): - import yoda - - if version.parse(yoda.__version__) >= version.parse("2.0.0"): - return self.target.rebinXBy(*args, **kwargs) - return self.target.rebinBy(*args, **kwargs) - - def __getitem__(self, slices): + def __getitem__(self, slices: Any) -> Any: return super().__getitem__(slices) - - def clone(self): - return Histo1D(self.target.clone()) - - def get_projector(self): - return Counter - - # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 - def annotationsDict(self): - d = {} - for k in self.target.annotations(): - d[k] = self.target.annotation(k) - return d diff --git a/src/babyyoda/yoda/histo2d.py b/src/babyyoda/yoda/histo2d.py index d318f55..e7832f0 100644 --- a/src/babyyoda/yoda/histo2d.py +++ b/src/babyyoda/yoda/histo2d.py @@ -1,3 +1,5 @@ +from typing import Any, Optional, cast + import yoda from packaging import version @@ -7,7 +9,7 @@ class Histo2D(babyyoda.UHIHisto2D): - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: """ target is either a yoda or grogu HISTO2D_V2 """ @@ -21,11 +23,56 @@ def __init__(self, *args, **kwargs): # Store the target object where calls and attributes will be forwarded super().__setattr__("target", target) + ########################## + # Basic needed functions for UHI directly relayed to target + ########################## + + def clone(self) -> "Histo2D": + return Histo2D(self.target.clone()) + + def path(self) -> str: + return str(self.target.path()) + + def bins(self, includeOverflows: bool = False, *args: Any, **kwargs: Any) -> Any: + import yoda + + if version.parse(yoda.__version__) >= version.parse("2.0.0"): + return self.target.bins(*args, includeOverflows=includeOverflows, **kwargs) + if not includeOverflows: + # YODA1 bins are not sorted than YODA2 + return sorted( + self.target.bins(*args, **kwargs), key=lambda b: (b.yMin(), b.xMin()) + ) + err = "YODA1 backend can not include overflows" + raise NotImplementedError(err) + + def xEdges(self) -> list[float]: + return cast(list[float], self.target.xEdges()) + + def yEdges(self) -> list[float]: + return cast(list[float], self.target.yEdges()) + + def rebinXTo(self, bins: list[float]) -> None: + self.target.rebinXTo(bins) + + def rebinYTo(self, bins: list[float]) -> None: + self.target.rebinYTo(bins) + + def get_projector(self) -> Any: + return Histo1D + + # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 + def annotationsDict(self) -> dict[str, Optional[str]]: + d = {} + for k in self.target.annotations(): + d[k] = self.target.annotation(k) + return d + ######################################################## # Relay all attribute access to the target object ######################################################## - def __getattr__(self, name): + def __getattr__(self, name: Any) -> Any: # if we overwrite it here, use that if has_own_method(Histo2D, name): return getattr(self, name) @@ -56,31 +103,5 @@ def __getattr__(self, name): # err = f"'{type(self.target).__name__}' object is not callable" # raise TypeError(err) - def bins(self, includeOverflows=False, *args, **kwargs): - import yoda - - if version.parse(yoda.__version__) >= version.parse("2.0.0"): - return self.target.bins(*args, includeOverflows=includeOverflows, **kwargs) - if not includeOverflows: - # YODA1 bins are not sorted than YODA2 - return sorted( - self.target.bins(*args, **kwargs), key=lambda b: (b.yMin(), b.xMin()) - ) - err = "YODA1 backend can not include overflows" - raise NotImplementedError(err) - - def __getitem__(self, slices): + def __getitem__(self, slices: Any) -> Any: return super().__getitem__(slices) - - def clone(self): - return Histo2D(self.target.clone()) - - def get_projector(self): - return Histo1D - - # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 - def annotationsDict(self): - d = {} - for k in self.target.annotations(): - d[k] = self.target.annotation(k) - return d diff --git a/src/babyyoda/yoda/read.py b/src/babyyoda/yoda/read.py index 45979db..1291afa 100644 --- a/src/babyyoda/yoda/read.py +++ b/src/babyyoda/yoda/read.py @@ -1,3 +1,5 @@ +from typing import Any + import yoda as yd from babyyoda.yoda.counter import Counter @@ -5,12 +7,12 @@ from babyyoda.yoda.histo2d import Histo2D -def read(file_path: str): +def read(file_path: str) -> dict[str, Any]: """ Wrap yoda histograms in the by HISTO1D_V2 class """ - ret = {} + ret: dict[str, Any] = {} for k, v in yd.read(file_path).items(): if isinstance(v, yd.Histo1D): ret[k] = Histo1D(v) diff --git a/src/babyyoda/yoda/write.py b/src/babyyoda/yoda/write.py index c592792..d70a1ac 100644 --- a/src/babyyoda/yoda/write.py +++ b/src/babyyoda/yoda/write.py @@ -1,9 +1,22 @@ import warnings +from typing import Any, Union import yoda as yd +import babyyoda.yoda as by -def write(anyhistograms, file_path: str, *args, gz=False, **kwargs): +byHistograms = Union[by.Counter, by.Histo1D, by.Histo2D] +ydHistograms = Union[yd.Counter, yd.Histo1D, yd.Histo2D] +Histograms = Union[ydHistograms, byHistograms] + + +def write( + anyhistograms: Union[list[Histograms], dict[str, Histograms]], + file_path: str, + *args: list[Any], + gz: bool = False, + **kwargs: dict[Any, Any], +) -> None: if gz and not file_path.endswith((".gz", ".gzip")): warnings.warn( "gz is True but file_path does not end with .gz or .gzip", stacklevel=2 @@ -17,6 +30,3 @@ def write(anyhistograms, file_path: str, *args, gz=False, **kwargs): # replace every value of list by value.target anyhistograms = [v.target for v in anyhistograms] yd.write(anyhistograms, file_path, *args, **kwargs) - else: - err = "anyhistograms should be a dict or a list of histograms" - raise ValueError(err)