From d2b96b7b6f83f9568b10ffc96709b807942dc069 Mon Sep 17 00:00:00 2001 From: Alexander Neuwirth Date: Tue, 5 Nov 2024 16:32:47 +0000 Subject: [PATCH] Remove warning on grogu use --- .pre-commit-config.yaml | 2 +- src/babyyoda/analysisobject.py | 12 ++- src/babyyoda/counter.py | 39 ++++++-- src/babyyoda/grogu/__init__.py | 20 ++-- src/babyyoda/grogu/analysis_object.py | 36 ++++---- src/babyyoda/grogu/counter_v2.py | 6 +- src/babyyoda/grogu/counter_v3.py | 32 +++---- src/babyyoda/grogu/histo1d_v2.py | 48 +++++----- src/babyyoda/grogu/histo1d_v3.py | 22 ++--- src/babyyoda/grogu/histo2d_v2.py | 18 ++-- src/babyyoda/grogu/histo2d_v3.py | 30 +++--- src/babyyoda/grogu/write.py | 26 +++++- src/babyyoda/histo1d.py | 18 ++-- src/babyyoda/histo2d.py | 126 ++++++++++++++++---------- src/babyyoda/read.py | 10 +- src/babyyoda/util.py | 41 ++++++--- src/babyyoda/yoda/counter.py | 12 ++- src/babyyoda/yoda/histo1d.py | 2 +- src/babyyoda/yoda/histo2d.py | 2 +- src/babyyoda/yoda/read.py | 6 +- src/babyyoda/yoda/write.py | 18 +++- 21 files changed, 322 insertions(+), 204 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 54faf08..6d9920c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,7 +20,7 @@ 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"] diff --git a/src/babyyoda/analysisobject.py b/src/babyyoda/analysisobject.py index 0b8580a..ad5aba6 100644 --- a/src/babyyoda/analysisobject.py +++ b/src/babyyoda/analysisobject.py @@ -1,7 +1,15 @@ class UHIAnalysisObject: - def key(self): + def key(self) -> str: return self.path() - def setAnnotationsDict(self, d: dict): + def path(self) -> str: + err = "UHIAnalysisObject.path() must be implemented by subclass" + raise NotImplementedError(err) + + def setAnnotationsDict(self, d: dict[str, str]) -> None: for k, v in d.items(): self.setAnnotation(k, v) + + def setAnnotation(self, key: str, value: str) -> None: + err = "UHIAnalysisObject.setAnnotation() must be implemented by subclass" + raise NotImplementedError(err) diff --git a/src/babyyoda/counter.py b/src/babyyoda/counter.py index 10336e6..ac6e68e 100644 --- a/src/babyyoda/counter.py +++ b/src/babyyoda/counter.py @@ -1,9 +1,10 @@ import contextlib +from typing import Any, Optional 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 +16,46 @@ def set_bin0d(target, source): raise NotImplementedError(err) -def Counter(*args, **kwargs): +def Counter(*args, **kwargs) -> "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 +70,7 @@ def to_grogu_v2(self): ], ) - def to_grogu_v3(self): + def to_grogu_v3(self) -> Any: from babyyoda.grogu.counter_v3 import GROGU_COUNTER_V3 return GROGU_COUNTER_V3( @@ -62,11 +85,11 @@ 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() 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..98249d9 100644 --- a/src/babyyoda/grogu/counter_v2.py +++ b/src/babyyoda/grogu/counter_v2.py @@ -23,7 +23,7 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## def clone(self): @@ -119,13 +119,13 @@ 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): diff --git a/src/babyyoda/grogu/counter_v3.py b/src/babyyoda/grogu/counter_v3.py index 7fdef75..b88fdaa 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 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) -> "GROGU_COUNTER_V3": return GROGU_COUNTER_V3( d_bins=[GROGU_COUNTER_V3.Bin()], d_annotations={"Title": title} if title else {}, @@ -23,7 +23,7 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## def clone(self): @@ -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,19 +75,19 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> float: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> float: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> float: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> float: return self.stdDev() / self.effNumEntries() ** 0.5 - def numEntries(self): + def numEntries(self) -> float: return self.d_numentries def __eq__(self, other): @@ -119,22 +119,22 @@ 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): diff --git a/src/babyyoda/grogu/histo1d_v2.py b/src/babyyoda/grogu/histo1d_v2.py index 913ee06..cffa081 100644 --- a/src/babyyoda/grogu/histo1d_v2.py +++ b/src/babyyoda/grogu/histo1d_v2.py @@ -57,7 +57,7 @@ def __post_init__(self): ) ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## def clone(self): @@ -99,28 +99,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,19 +131,21 @@ def variance(self): ) # return self.d_sumw2/self.d_numentries - (self.d_sumw/self.d_numentries)**2 - def errW(self): + def errW(self) -> float: return self.d_sumw2**0.5 - def stdDev(self): + def stdDev(self) -> float: return self.variance() ** 0.5 - def effNumEntries(self): + def effNumEntries(self) -> float: return self.sumW() ** 2 / self.sumW2() - def stdErr(self): + def stdErr(self) -> float: 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): @@ -215,16 +219,16 @@ def to_string(bin, label=None) -> str: 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): GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo1D") ############################################ - # YODA compatibilty code + # YODA compatibility code ############################################ def clone(self): @@ -253,10 +257,10 @@ def fill(self, x, weight=1.0, fraction=1.0): if x < self.xMin() and self.d_underflow is not None: self.d_underflow.fill(x, weight, fraction) - def xMax(self): + def xMax(self) -> Optional[float]: return max([b.xMax() for b in self.d_bins]) - def xMin(self): + def xMin(self) -> Optional[float]: return min([b.xMin() for b in self.d_bins]) def bins(self, includeOverflows=False): @@ -274,10 +278,10 @@ def binAt(self, x): return b return None - def binDim(self): + def binDim(self) -> int: return 1 - def xEdges(self): + def xEdges(self) -> list[float]: return [b.xMin() for b in self.d_bins] + [self.xMax()] def rebinXTo(self, edges: list[float]): diff --git a/src/babyyoda/grogu/histo1d_v3.py b/src/babyyoda/grogu/histo1d_v3.py index d9fd40d..043ab49 100644 --- a/src/babyyoda/grogu/histo1d_v3.py +++ b/src/babyyoda/grogu/histo1d_v3.py @@ -45,7 +45,7 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## def clone(self): @@ -169,7 +169,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 +178,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 +189,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,10 +204,10 @@ 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): @@ -223,10 +223,10 @@ def binAt(self, x): 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): diff --git a/src/babyyoda/grogu/histo2d_v2.py b/src/babyyoda/grogu/histo2d_v2.py index 7826129..5865a68 100644 --- a/src/babyyoda/grogu/histo2d_v2.py +++ b/src/babyyoda/grogu/histo2d_v2.py @@ -69,7 +69,7 @@ class Bin: d_numentries: float = 0.0 ######################################################## - # YODA compatibilty code + # YODA compatibility code ######################################################## def clone(self): @@ -204,14 +204,14 @@ 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): GROGU_ANALYSIS_OBJECT.__post_init__(self) self.setAnnotation("Type", "Histo2D") # - # YODA compatibilty code + # YODA compatibility code # def clone(self): @@ -228,7 +228,7 @@ def fill(self, x, y, weight=1.0, fraction=1.0): if b.d_xmin <= x < b.d_xmax and b.d_ymin <= y < b.d_ymax: b.fill(x, y, weight, fraction) - def xEdges(self): + def xEdges(self) -> list[float]: assert all( x == y for x, y in zip( @@ -238,7 +238,7 @@ def xEdges(self): ) return sorted({b.d_xmin for b in self.d_bins} | {self.xMax()}) - def yEdges(self): + def yEdges(self) -> list[float]: assert all( x == y for x, y in zip( @@ -248,16 +248,16 @@ def yEdges(self): ) return sorted({b.d_ymin for b in self.d_bins} | {self.yMax()}) - def xMin(self): + def xMin(self) -> float: return min(b.d_xmin for b in self.d_bins) - def yMin(self): + def yMin(self) -> float: return min(b.d_ymin for b in self.d_bins) - def xMax(self): + def xMax(self) -> float: return max(b.d_xmax for b in self.d_bins) - def yMax(self): + def yMax(self) -> float: return max(b.d_ymax for b in self.d_bins) def bins(self, includeOverflows=False): diff --git a/src/babyyoda/grogu/histo2d_v3.py b/src/babyyoda/grogu/histo2d_v3.py index d333fd2..38b7927 100644 --- a/src/babyyoda/grogu/histo2d_v3.py +++ b/src/babyyoda/grogu/histo2d_v3.py @@ -30,7 +30,7 @@ def Histo2D_v3( *args, title=None, **kwargs, -): +) -> "GROGU_HISTO2D_V3": xedges = [] yedges = [] if isinstance(args[0], list) and isinstance(args[1], list): @@ -133,32 +133,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, y) -> 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 +191,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,7 +201,7 @@ def __post_init__(self): ) # - # YODA compatibilty code + # YODA compatibility code # def clone(self): @@ -212,13 +212,13 @@ 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, y, 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 diff --git a/src/babyyoda/grogu/write.py b/src/babyyoda/grogu/write.py index 60b3ace..a03a7f7 100644 --- a/src/babyyoda/grogu/write.py +++ b/src/babyyoda/grogu/write.py @@ -1,7 +1,31 @@ +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=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 diff --git a/src/babyyoda/histo1d.py b/src/babyyoda/histo1d.py index fe8e832..ca2297b 100644 --- a/src/babyyoda/histo1d.py +++ b/src/babyyoda/histo1d.py @@ -1,5 +1,6 @@ import contextlib import sys +from typing import Any import numpy as np @@ -8,7 +9,7 @@ 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() @@ -29,9 +30,12 @@ def Histo1D(*args, **kwargs): """ 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...) @@ -40,7 +44,7 @@ class UHIHisto1D(UHIAnalysisObject): # 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 +91,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,11 +128,11 @@ 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() diff --git a/src/babyyoda/histo2d.py b/src/babyyoda/histo2d.py index 6e8b92b..6513269 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,50 @@ 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) -> float: + return sum( + b.sumWX() for b in self.bins(includeOverflows=includeOverflows) + ) / sum(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) -> float: + return sum( + b.sumWY() for b in self.bins(includeOverflows=includeOverflows) + ) / sum(b.sumW() for b in self.bins(includeOverflows=includeOverflows)) - def integral(self, includeOverflows=True): + def integral(self, includeOverflows: bool = True) -> float: return sum(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 +225,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) -> int: 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]) -> int: ix = None if isinstance(slices, int): ix = slices @@ -258,7 +285,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]) -> int: iy = None if isinstance(slices, int): iy = slices @@ -266,10 +293,12 @@ 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], Union[int, loc]] + ) -> tuple[int, int]: return self.__get_x_index(slices[0]), self.__get_y_index(slices[1]) - def __setitem__(self, slices, value): + def __setitem__(self, slices, value: Any) -> None: set_bin2d(self.__getitem__(slices), value) def __getitem__(self, slices): @@ -365,6 +394,9 @@ def project(self, axis: int = 0): return self.projectX() return self.projectY() + def to_string(self) -> str: + return self.to_grogu_v3().to_string() + def plot(self, *args, binwnorm=True, **kwargs): # # TODO should use histplot # import mplhep as hep diff --git a/src/babyyoda/read.py b/src/babyyoda/read.py index 5ebd7f2..3def87f 100644 --- a/src/babyyoda/read.py +++ b/src/babyyoda/read.py @@ -1,5 +1,3 @@ -import warnings - from babyyoda import grogu @@ -7,10 +5,10 @@ def read(file_path: str): 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) 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/yoda/counter.py b/src/babyyoda/yoda/counter.py index 7b27d05..02b6176 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 """ @@ -22,7 +24,7 @@ def __init__(self, *args, **kwargs): # Relay all attribute access to the target object ######################################################## - def __getattr__(self, name): + def __getattr__(self, name: str): # if we overwrite it here, use that if has_own_method(Counter, name): return getattr(self, name) @@ -53,14 +55,14 @@ def __getattr__(self, name): # err = f"'{type(self.target).__name__}' object is not callable" # raise TypeError(err) - def bins(self, *args, **kwargs): + def bins(self, *args, **kwargs) -> Any: return self.target.bins(*args, **kwargs) - def clone(self): + def clone(self) -> "Counter": return Counter(self.target.clone()) # Fix https://gitlab.com/hepcedar/yoda/-/issues/101 - def annotationsDict(self): + def annotationsDict(self) -> dict[str, Optional[str]]: d = {} for k in self.target.annotations(): d[k] = self.target.annotation(k) diff --git a/src/babyyoda/yoda/histo1d.py b/src/babyyoda/yoda/histo1d.py index 77d1496..24d77bd 100644 --- a/src/babyyoda/yoda/histo1d.py +++ b/src/babyyoda/yoda/histo1d.py @@ -7,7 +7,7 @@ class Histo1D(babyyoda.UHIHisto1D): - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: """ target is either a yoda or grogu HISTO1D_V2 """ diff --git a/src/babyyoda/yoda/histo2d.py b/src/babyyoda/yoda/histo2d.py index d318f55..e07ff93 100644 --- a/src/babyyoda/yoda/histo2d.py +++ b/src/babyyoda/yoda/histo2d.py @@ -7,7 +7,7 @@ class Histo2D(babyyoda.UHIHisto2D): - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: """ target is either a yoda or grogu HISTO2D_V2 """ 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)