atlas-ftag-tools 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

{atlas_ftag_tools-0.2.6.dist-info → atlas_ftag_tools-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atlas-ftag-tools
-Version: 0.2.6
+Version: 0.2.7
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT

{atlas_ftag_tools-0.2.6.dist-info → atlas_ftag_tools-0.2.7.dist-info}/RECORD

@@ -1,14 +1,14 @@
-ftag/__init__.py,sha256=edvBW3ySbO10OzSeqTQd9HPuXqiosdWibAhi3ZJFT4g,680
+ftag/__init__.py,sha256=k5qBmtC7Ieh0trgm2Ba9Qj_6A2wQSpmAmXo2iIOAaI0,737
 ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
 ftag/cuts.py,sha256=9_ooLZHaO3SnIQBNxwbaPZn-qptGdKnB27FdKQGTiTY,2933
-ftag/flavour.py,sha256=MUij7lknepnkidErDsyNsJDbAPotAEwS-uMaTSYCet4,3800
+ftag/flavours.py,sha256=ShH4M2UjQZpZ_NlCctTm2q1tJbzYxjmGteioQ2GcqEU,114
 ftag/flavours.yaml,sha256=E_vpn38qJ3-Tygg2aHlH4wkn_rR1On_lMeaG8OemHCQ,8285
 ftag/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
-ftag/labeller.py,sha256=iyDjRkOaly48jTaaFimA5eZ5YySp1lrLVEt3GyDZW-8,2553
+ftag/labeller.py,sha256=IXUgU9UBir39PxVWRKs5r5fqI66Tv0x7nJD3-RYpbrg,2780
+ftag/labels.py,sha256=C7IylPTnc32dFXq8C2Ks2wuljYK3WaY2EsPLGrhtXy8,3932
 ftag/mock.py,sha256=Eyj3tkkaSSnqvS3G6NS7fq8sB__Nx8YE9-OM2_lpdoQ,4992
 ftag/region.py,sha256=ANv0dGI2W6NJqD9fp7EfqAUReH4FOjc1gwl_Qn8llcM,360
-ftag/sample.py,sha256=TFXMhDkbPmjkms9-b-bINJ32T3bO86JcU70C0nY7wa8,2500
-ftag/test_cli_utils.py,sha256=xa08vf6SEOow58SSFagYdAselb-dkNOVvWsWheMnW-g,1001
+ftag/sample.py,sha256=3N0FrRcu9l1sX8ohuGOHuMYGD0See6gMO4--7NzR2tE,2538
 ftag/track_selector.py,sha256=fJNk_kIBQriBqV4CPT_3ReJbOUnavDDzO-u3EQlRuyk,2654
 ftag/transform.py,sha256=uEGGJSnqoKOzLYQv650XdK_kDNw4Aw-5dc60z9Dp_y0,3963
 ftag/vds.py,sha256=nRViQZQIORB95nC7NZsW3KsSoGkLzEdOsuCViH5h8-U,3296
@@ -19,10 +19,10 @@ ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
 ftag/hdf5/h5utils.py,sha256=-4zKTMtNCrDZr_9Ww7uzfsB7M7muBKpmm_1IkKJnHOI,3222
 ftag/hdf5/h5writer.py,sha256=9FkClV__UbBqmFsq_h2jwiZnbWVm8QFRL_4mDZZBbTs,5316
 ftag/wps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ftag/wps/discriminant.py,sha256=rTASnt9s6xEHK5hLKHgvc9c7TiyyX6Upbp4LeTh6TB0,4270
-ftag/wps/working_points.py,sha256=fPfSZoTp-2qxL3YB7HdPxIB-7aSR_HozEtWWIA0BUtg,9749
-atlas_ftag_tools-0.2.6.dist-info/METADATA,sha256=gmEDL518a_xrRBsPsRYCFt-C06BD3K9bH6HCStV6mAU,5169
-atlas_ftag_tools-0.2.6.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-atlas_ftag_tools-0.2.6.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
-atlas_ftag_tools-0.2.6.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.2.6.dist-info/RECORD,,
+ftag/wps/discriminant.py,sha256=VJdZlJJUwaTeyxmIDEk23rQSAuvWs6wDA3XRjDI6-_c,4277
+ftag/wps/working_points.py,sha256=cvStSpP8Cbb_FWM8v59tFsscUvdeqi831tLn5BiHUEg,9741
+atlas_ftag_tools-0.2.7.dist-info/METADATA,sha256=oo5m85dK467AWuK-L8xaIbLDVmRUO3r7vA_J1vgR5b8,5169
+atlas_ftag_tools-0.2.7.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+atlas_ftag_tools-0.2.7.dist-info/entry_points.txt,sha256=LfVLsZHQolqbPnwPgtmc5IQTh527BKkN2v-IpXWTNHw,137
+atlas_ftag_tools-0.2.7.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.7.dist-info/RECORD,,

{atlas_ftag_tools-0.2.6.dist-info → atlas_ftag_tools-0.2.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.2.0)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ftag/__init__.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.6"
-
+__version__ = "v0.2.7"
 
 from ftag import hdf5
 from ftag.cuts import Cuts
-from ftag.flavour import Flavour, Flavours
+from ftag.flavours import Flavours
 from ftag.labeller import Labeller
+from ftag.labels import Label, LabelContainer
 from ftag.mock import get_mock_file
 from ftag.sample import Sample
 from ftag.transform import Transform
@@ -17,8 +17,9 @@ from ftag.wps.working_points import get_working_points
 
 __all__ = [
     "Cuts",
-    "Flavour",
     "Flavours",
+    "Label",
+    "LabelContainer",
     "Labeller",
     "Sample",
     "Transform",

ftag/flavours.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from ftag.labels import LabelContainer
+
+Flavours = LabelContainer.from_yaml()

ftag/labeller.py

@@ -4,8 +4,9 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from ftag.flavour import Flavour, FlavourContainer, Flavours
+from ftag import Flavours
 from ftag.hdf5 import join_structured_arrays, structured_from_dict
+from ftag.labels import Label, LabelContainer
 
 
 @dataclass
@@ -13,28 +14,36 @@ class Labeller:
     """
     Defines a labelling scheme.
 
-    Labels are [0, ..., n] and are assigned using pre-defined selections.
+    Classes are assigned integer labels in [0, ..., n] based on pre-defined selections.
 
     Parameters
     ----------
-    labels : FlavourContainer | list[str | Flavour]
+    labels : LabelContainer | list[str | Label]
         The labels to be use.
     require_labels : bool
         Whether to require that all objects are labelled.
     """
 
-    labels: FlavourContainer | list[str | Flavour]
+    labels: LabelContainer | list[str | Label]
     require_labels: bool = True
 
-    @property
-    def variables(self):
-        return sum((label.cuts.variables for label in self.labels), [])
-
     def __post_init__(self) -> None:
-        if isinstance(self.labels, FlavourContainer):
+        if isinstance(self.labels, LabelContainer):
             self.labels = list(self.labels)
         self.labels = sorted([Flavours[label] for label in self.labels])
 
+    @property
+    def variables(self) -> list[str]:
+        """
+        Returns the variables used for labelling.
+
+        Returns
+        -------
+        list[str]
+            The variables used for labelling.
+        """
+        return sum((label.cuts.variables for label in self.labels), [])  # type: ignore[union-attr]
+
     def get_labels(self, array: np.ndarray) -> np.ndarray:
         """
         Returns the labels for the given array.

ftag/labels.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+from ftag.cuts import Cuts
+
+
+def remove_suffix(string: str, suffix: str) -> str:
+    if string.endswith(suffix):
+        return string[: -len(suffix)]
+    return string
+
+
+@dataclass(frozen=True)
+class Label:
+    name: str
+    label: str
+    cuts: Cuts
+    colour: str
+    category: str
+    _px: str | None = None
+
+    @property
+    def px(self) -> str:
+        return self._px or f"p{remove_suffix(self.name, 'jets')}"
+
+    @property
+    def eff_str(self) -> str:
+        return self.label.replace("jets", "jet") + " efficiency"
+
+    @property
+    def rej_str(self) -> str:
+        return self.label.replace("jets", "jet") + " rejection"
+
+    @property
+    def frac_str(self) -> str:
+        return "f" + remove_suffix(self.name, "jets")
+
+    def __str__(self) -> str:
+        return self.name
+
+    def __lt__(self, other) -> bool:
+        return self.name < other.name
+
+
+@dataclass
+class LabelContainer:
+    labels: dict[str, Label]
+
+    def __iter__(self) -> Iterator:
+        yield from self.labels.values()
+
+    def __getitem__(self, key) -> Label:
+        if isinstance(key, Label):
+            key = key.name
+        try:
+            return self.labels[key]
+        except KeyError as e:
+            raise KeyError(f"Label '{key}' not found") from e
+
+    def __getattr__(self, name) -> Label:
+        return self[name]
+
+    def __contains__(self, label: str | Label) -> bool:
+        if isinstance(label, Label):
+            label = label.name
+        return label in self.labels
+
+    def __eq__(self, other) -> bool:
+        if isinstance(other, LabelContainer):
+            return self.labels == other.labels
+        if isinstance(other, list) and all(isinstance(f, str) for f in other):
+            return {f.name for f in self} == set(other)
+        return False
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({', '.join([f.name for f in self])})"
+
+    @property
+    def categories(self) -> list[str]:
+        return list(dict.fromkeys(f.category for f in self))
+
+    def by_category(self, category: str) -> LabelContainer:
+        f = LabelContainer({k: v for k, v in self.labels.items() if v.category == category})
+        if not f.labels:
+            raise KeyError(f"No labels with category '{category}' found")
+        return f
+
+    def from_cuts(self, cuts: list | Cuts) -> Label:
+        if isinstance(cuts, list):
+            cuts = Cuts.from_list(cuts)
+        for label in self:
+            if label.cuts == cuts:
+                return label
+        raise KeyError(f"Label with {cuts} not found")
+
+    @classmethod
+    def from_yaml(cls, yaml_path: Path | None = None) -> LabelContainer:
+        if yaml_path is None:
+            yaml_path = Path(__file__).parent / "flavours.yaml"
+        with open(yaml_path) as f:
+            config = yaml.safe_load(f)
+
+        # sanity checks
+        cuts = [Cuts.from_list(f["cuts"]) for f in config]
+        if duplicates := [c for c in cuts if cuts.count(c) > 1]:
+            raise ValueError(f"Duplicate label definitions detected: {duplicates}")
+        names = [f["name"] for f in config]
+        if duplicates := [n for n in names if names.count(n) > 1]:
+            raise ValueError(f"Duplicate label names detected: {duplicates}")
+
+        labels = {f["name"]: Label(cuts=Cuts.from_list(f.pop("cuts")), **f) for f in config}
+        return cls(labels)
+
+    @classmethod
+    def from_list(cls, labels: list[Label]) -> LabelContainer:
+        return cls({f.name: f for f in labels})
+
+    def backgrounds(self, label: Label, only_signals: bool = True) -> LabelContainer:
+        bkg = [f for f in self if f.category == label.category and f != label]
+        if not only_signals:
+            bkg = [f for f in bkg if f.name not in {"ujets", "qcd"}]
+        return LabelContainer.from_list(bkg)

ftag/sample.py

@@ -4,7 +4,7 @@ import glob
 from dataclasses import dataclass
 from pathlib import Path
 
-from ftag.flavour import remove_suffix
+from ftag.labels import remove_suffix
 from ftag.vds import create_virtual_file
 
 
@@ -13,6 +13,7 @@ class Sample:
     pattern: Path | str | tuple[Path | str, ...]
     ntuple_dir: Path | str | None = None
     name: str | None = None
+    weights: list[float] | None = None
 
     def __post_init__(self) -> None:
         if not self.pattern:

ftag/wps/discriminant.py

@@ -4,13 +4,14 @@ from typing import Callable
 
 import numpy as np
 
-from ftag.flavour import Flavour, Flavours, remove_suffix
+from ftag import Flavours
+from ftag.labels import Label, remove_suffix
 
 
 def discriminant(
     jets: np.ndarray,
     tagger: str,
-    signal: Flavour,
+    signal: Label,
     fxs: dict[str, float],
     epsilon: float = 1e-10,
 ) -> np.ndarray:
@@ -87,7 +88,7 @@ def hcc_discriminant(jets, tagger, ftop=0.25, fhbb=0.3, epsilon=1e-10):
 
 
 def get_discriminant(
-    jets: np.ndarray, tagger: str, signal: Flavour | str, epsilon: float = 1e-10, **fxs
+    jets: np.ndarray, tagger: str, signal: Label | str, epsilon: float = 1e-10, **fxs
 ):
     """Calculate the b-tag or c-tag discriminant for a given tagger.
 
@@ -97,7 +98,7 @@ def get_discriminant(
         Structured array of jets containing tagger outputs
     tagger : str
         Name of the tagger
-    signal : Flavour
+    signal : Label
         Signal flavour (bjets/cjets or hbb/hcc)
     epsilon : float, optional
         Small number to avoid division by zero, by default 1e-10

ftag/wps/working_points.py

@@ -8,9 +8,9 @@ from pathlib import Path
 import numpy as np
 import yaml
 
+from ftag import Flavours
 from ftag.cli_utils import HelpFormatter
 from ftag.cuts import Cuts
-from ftag.flavour import Flavours
 from ftag.hdf5 import H5Reader
 from ftag.wps.discriminant import get_discriminant
 

ftag/flavour.py

@@ -1,127 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import Iterator
-from dataclasses import dataclass
-from pathlib import Path
-
-import yaml
-
-from ftag.cuts import Cuts
-
-
-def remove_suffix(string: str, suffix: str) -> str:
-    if string.endswith(suffix):
-        return string[: -len(suffix)]
-    return string
-
-
-@dataclass(frozen=True)
-class Flavour:
-    name: str
-    label: str
-    cuts: Cuts
-    colour: str
-    category: str
-    _px: str | None = None
-
-    @property
-    def px(self) -> str:
-        return self._px or f"p{remove_suffix(self.name, 'jets')}"
-
-    @property
-    def eff_str(self) -> str:
-        return self.label.replace("jets", "jet") + " efficiency"
-
-    @property
-    def rej_str(self) -> str:
-        return self.label.replace("jets", "jet") + " rejection"
-
-    @property
-    def frac_str(self) -> str:
-        return "f" + remove_suffix(self.name, "jets")
-
-    def __str__(self) -> str:
-        return self.name
-
-    def __lt__(self, other) -> bool:
-        return self.name < other.name
-
-
-@dataclass
-class FlavourContainer:
-    flavours: dict[str, Flavour]
-
-    def __iter__(self) -> Iterator:
-        yield from self.flavours.values()
-
-    def __getitem__(self, key) -> Flavour:
-        if isinstance(key, Flavour):
-            key = key.name
-        try:
-            return self.flavours[key]
-        except KeyError as e:
-            raise KeyError(f"Flavour '{key}' not found") from e
-
-    def __getattr__(self, name) -> Flavour:
-        return self[name]
-
-    def __contains__(self, flavour: str | Flavour) -> bool:
-        if isinstance(flavour, Flavour):
-            flavour = flavour.name
-        return flavour in self.flavours
-
-    def __eq__(self, other) -> bool:
-        if isinstance(other, FlavourContainer):
-            return self.flavours == other.flavours
-        if isinstance(other, list) and all(isinstance(f, str) for f in other):
-            return {f.name for f in self} == set(other)
-        return False
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({', '.join([f.name for f in self])})"
-
-    @property
-    def categories(self) -> list[str]:
-        return list(dict.fromkeys(f.category for f in self))
-
-    def by_category(self, category: str) -> FlavourContainer:
-        f = FlavourContainer({k: v for k, v in self.flavours.items() if v.category == category})
-        if not f.flavours:
-            raise KeyError(f"No flavours with category '{category}' found")
-        return f
-
-    def from_cuts(self, cuts: list | Cuts) -> Flavour:
-        if isinstance(cuts, list):
-            cuts = Cuts.from_list(cuts)
-        for flavour in self:
-            if flavour.cuts == cuts:
-                return flavour
-        raise KeyError(f"Flavour with {cuts} not found")
-
-    @classmethod
-    def from_yaml(cls, yaml_path: Path | None = None) -> FlavourContainer:
-        if yaml_path is None:
-            yaml_path = Path(__file__).parent / "flavours.yaml"
-
-        with open(yaml_path) as f:
-            flavours_yaml = yaml.safe_load(f)
-
-        flavours_dict = {
-            f["name"]: Flavour(cuts=Cuts.from_list(f.pop("cuts")), **f) for f in flavours_yaml
-        }
-        assert len(flavours_dict) == len(flavours_yaml), "Duplicate flavour names detected"
-
-        return cls(flavours_dict)
-
-    @classmethod
-    def from_list(cls, flavours: list[Flavour]) -> FlavourContainer:
-        return cls({f.name: f for f in flavours})
-
-    def backgrounds(self, flavour: Flavour, keep_possible_signals: bool = True) -> FlavourContainer:
-        bkg = [f for f in self if f.category == flavour.category and f != flavour]
-        if not keep_possible_signals:
-            bkg = [f for f in bkg if f.name not in {"ujets", "qcd"}]
-        return FlavourContainer.from_list(bkg)
-
-
-Flavours = FlavourContainer.from_yaml()

ftag/test_cli_utils.py

@@ -1,34 +0,0 @@
-from __future__ import annotations
-
-from pathlib import Path
-from tempfile import NamedTemporaryFile, TemporaryDirectory
-
-import pytest
-
-from ftag.cli_utils import valid_path
-
-
-def test_valid_path_existing_file():
-    # Test when the input path is an existing file
-    # get a temp directory
-    with TemporaryDirectory() as tmpdir, NamedTemporaryFile(dir=tmpdir) as f:
-        input_path = f.name
-        expected_output = Path(f.name)
-        result = valid_path(input_path)
-        assert result == expected_output
-
-
-def test_valid_path_non_existing_file():
-    # Test when the input path is a non-existing file
-    input_path = "non_existing_file.txt"
-    with pytest.raises(FileNotFoundError) as e:
-        valid_path(input_path)
-    assert str(e.value) == input_path
-
-
-def test_valid_path_directory():
-    # Test when the input path is a directory
-    input_path = "directory/"
-    with pytest.raises(FileNotFoundError) as e:
-        valid_path(input_path)
-    assert str(e.value) == input_path