atlas-schema 0.2.0__tar.gz → 0.2.2__tar.gz

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/.gitignore

@@ -70,6 +70,7 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+docs/_generated/
 
 # PyBuilder
 .pybuilder/

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/PKG-INFO

@@ -1,11 +1,13 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: atlas-schema
-Version: 0.2.0
+Version: 0.2.2
 Summary: Helper python package for ATLAS Common NTuple Analysis work.
 Project-URL: Homepage, https://github.com/scipp-atlas/atlas-schema
 Project-URL: Bug Tracker, https://github.com/scipp-atlas/atlas-schema/issues
 Project-URL: Discussions, https://github.com/scipp-atlas/atlas-schema/discussions
-Project-URL: Changelog, https://github.com/scipp-atlas/atlas-schema/releases
+Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.2.2/
+Project-URL: Releases, https://github.com/scipp-atlas/atlas-schema/releases
+Project-URL: Release Notes, https://atlas-schema.readthedocs.io/en/latest/history.html
 Author-email: Giordon Stark <kratsg@gmail.com>
 License: 
                                          Apache License
@@ -209,6 +211,7 @@ License:
            WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
            See the License for the specific language governing permissions and
            limitations under the License.
+License-File: LICENSE
 Classifier: Development Status :: 1 - Planning
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
@@ -225,14 +228,19 @@ Classifier: Topic :: Scientific/Engineering
 Classifier: Typing :: Typed
 Requires-Python: >=3.9
 Requires-Dist: coffea[dask]>=2024.4.1
+Requires-Dist: particle>=0.25.0
 Provides-Extra: dev
 Requires-Dist: pytest-cov>=3; extra == 'dev'
 Requires-Dist: pytest>=6; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: furo>=2023.08.17; extra == 'docs'
+Requires-Dist: intersphinx-registry>=0.2411.17; extra == 'docs'
+Requires-Dist: ipywidgets; extra == 'docs'
 Requires-Dist: myst-parser>=0.13; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
-Requires-Dist: sphinx-copybutton; extra == 'docs'
+Requires-Dist: sphinx-click; extra == 'docs'
+Requires-Dist: sphinx-copybutton!=0.5.1,>=0.3.2; extra == 'docs'
+Requires-Dist: sphinx-issues; extra == 'docs'
 Requires-Dist: sphinx>=7.0; extra == 'docs'
 Provides-Extra: test
 Requires-Dist: build; extra == 'test'
@@ -243,7 +251,7 @@ Requires-Dist: tbump>=6.7.0; extra == 'test'
 Requires-Dist: twine; extra == 'test'
 Description-Content-Type: text/markdown
 
-# atlas-schema v0.2.0
+# atlas-schema v0.2.2
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/README.md

@@ -1,4 +1,4 @@
-# atlas-schema v0.2.0
+# atlas-schema v0.2.2
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
   "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["coffea[dask] >= 2024.4.1"]
+dependencies = ["coffea[dask] >= 2024.4.1", "particle >= 0.25.0"]
 
 [project.optional-dependencies]
 test = [
@@ -47,16 +47,22 @@ dev = [
 docs = [
   "sphinx>=7.0",
   "myst_parser>=0.13",
-  "sphinx_copybutton",
+  "sphinx-copybutton>=0.3.2,!=0.5.1",
   "sphinx_autodoc_typehints",
   "furo>=2023.08.17",
+  "sphinx-click",
+  "ipywidgets",
+  "intersphinx_registry>=0.2411.17",
+  "sphinx-issues",
 ]
 
 [project.urls]
 Homepage = "https://github.com/scipp-atlas/atlas-schema"
 "Bug Tracker" = "https://github.com/scipp-atlas/atlas-schema/issues"
 Discussions = "https://github.com/scipp-atlas/atlas-schema/discussions"
-Changelog = "https://github.com/scipp-atlas/atlas-schema/releases"
+Documentation = "https://atlas-schema.readthedocs.io/en/v0.2.2/"
+Releases = "https://github.com/scipp-atlas/atlas-schema/releases"
+"Release Notes" = "https://atlas-schema.readthedocs.io/en/latest/history.html"
 
 
 [tool.hatch]
@@ -90,7 +96,19 @@ packages = ["src/atlas_schema"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
+addopts = [
+  "-ra",
+  "--showlocals",
+  "--strict-markers",
+  "--strict-config",
+  "--doctest-modules",
+  "--doctest-glob=*.rst",
+  "--cov",
+  "--cov-report=xml",
+  "--cov-report=term",
+  "--durations=20",
+  "--ignore=docs/conf.py",
+]
 xfail_strict = true
 filterwarnings = [
   "error",
@@ -98,9 +116,9 @@ filterwarnings = [
 log_cli_level = "INFO"
 testpaths = [
   "tests",
+  "docs",
 ]
 
-
 [tool.coverage]
 run.source = ["atlas_schema"]
 report.exclude_also = [
@@ -109,6 +127,7 @@ report.exclude_also = [
 ]
 
 [tool.mypy]
+mypy_path = ["src"]
 files = ["src", "tests"]
 python_version = "3.9"
 warn_unused_configs = true
@@ -129,6 +148,7 @@ disallow_incomplete_defs = true
 module = [
   'awkward.*',
   'coffea.*',
+  'dask_awkward.*',
 ]
 ignore_missing_imports = true
 

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/src/atlas_schema/__init__.py

@@ -6,7 +6,12 @@ atlas_schema: Collection of utilities and helper functions for HEP ATLAS analyse
 
 from __future__ import annotations
 
+import warnings
+
 from atlas_schema._version import version as __version__
 from atlas_schema.enums import ParticleOrigin, PhotonID
+from atlas_schema.utils import isin
+
+warnings.filterwarnings("ignore", module="coffea.*")
 
-__all__ = ["__version__", "ParticleOrigin", "PhotonID"]
+__all__ = ["ParticleOrigin", "PhotonID", "__version__", "isin"]

{atlas_schema-0.2.0 → atlas_schema-0.2.2}/src/atlas_schema/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '0.2.2'
+__version_tuple__ = version_tuple = (0, 2, 2)

atlas_schema-0.2.2/src/atlas_schema/enums.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import sys
+from enum import Enum, IntEnum
+
+if sys.version_info >= (3, 11):
+    from enum import EnumType
+else:
+    from enum import EnumMeta as EnumType
+
+from typing import Callable, TypeVar, cast
+
+_E = TypeVar("_E", bound=Enum)
+
+
+class MultipleEnumAccessMeta(EnumType):
+    """
+    Enum Metaclass to provide a way to access multiple values all at once.
+    """
+
+    def __getitem__(self: type[_E], key: str | tuple[str]) -> _E | list[_E]:  # type:ignore[misc,override]
+        getitem = cast(Callable[[str], _E], super().__getitem__)  # type:ignore[misc]
+        if isinstance(key, tuple):
+            return [getitem(name) for name in key]
+        return getitem(key)
+
+
+class ParticleType(IntEnum, metaclass=MultipleEnumAccessMeta):
+    """
+    Taken from `ATLAS Truth Utilities for ParticleType <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L8-49>`_.
+    """
+
+    Unknown = 0
+    UnknownElectron = 1
+    IsoElectron = 2
+    NonIsoElectron = 3
+    BkgElectron = 4
+    UnknownMuon = 5
+    IsoMuon = 6
+    NonIsoMuon = 7
+    BkgMuon = 8
+    UnknownTau = 9
+    IsoTau = 10
+    NonIsoTau = 11
+    BkgTau = 12
+    UnknownPhoton = 13
+    IsoPhoton = 14
+    NonIsoPhoton = 15
+    BkgPhoton = 16
+    Hadron = 17
+    Neutrino = 18
+    NuclFrag = 19
+    NonPrimary = 20
+    GenParticle = 21
+    SUSYParticle = 22
+    OtherBSMParticle = 39
+    BBbarMesonPart = 23
+    BottomMesonPart = 24
+    CCbarMesonPart = 25
+    CharmedMesonPart = 26
+    BottomBaryonPart = 27
+    CharmedBaryonPart = 28
+    StrangeBaryonPart = 29
+    LightBaryonPart = 30
+    StrangeMesonPart = 31
+    LightMesonPart = 32
+    BJet = 33
+    CJet = 34
+    LJet = 35
+    GJet = 36
+    TauJet = 37
+    UnknownJet = 38
+
+
+class ParticleOrigin(IntEnum, metaclass=MultipleEnumAccessMeta):
+    """
+    Taken from `ATLAS Truth Utilities for ParticleOrigin <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103>`_.
+    """
+
+    NonDefined = 0
+    SingleElec = 1
+    SingleMuon = 2
+    SinglePhot = 3
+    SingleTau = 4
+    PhotonConv = 5
+    DalitzDec = 6
+    ElMagProc = 7
+    Mu = 8
+    TauLep = 9
+    top = 10
+    QuarkWeakDec = 11
+    WBoson = 12
+    ZBoson = 13
+    Higgs = 14
+    HiggsMSSM = 15
+    HeavyBoson = 16
+    WBosonLRSM = 17
+    NuREle = 18
+    NuRMu = 19
+    NuRTau = 20
+    LQ = 21
+    SUSY = 22
+    OtherBSM = 46
+    LightMeson = 23
+    StrangeMeson = 24
+    CharmedMeson = 25
+    BottomMeson = 26
+    CCbarMeson = 27
+    JPsi = 28
+    BBbarMeson = 29
+    LightBaryon = 30
+    StrangeBaryon = 31
+    CharmedBaryon = 32
+    BottomBaryon = 33
+    PionDecay = 34
+    KaonDecay = 35
+    BremPhot = 36
+    PromptPhot = 37
+    UndrPhot = 38
+    ISRPhot = 39
+    FSRPhot = 40
+    NucReact = 41
+    PiZero = 42
+    DiBoson = 43
+    ZorHeavyBoson = 44
+    MultiBoson = 47
+    QCD = 45
+
+
+class PhotonID(IntEnum, metaclass=MultipleEnumAccessMeta):
+    """
+    Taken from the `EGamma Identification CP group's twiki <https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/EGammaIdentificationRun2#Photon_isEM_word>`_.
+    """
+
+    Rhad = 10  # ClusterHadronicLeakage_Photon
+    E277 = 11  # ClusterMiddleEnergy_Photon
+    Reta = 12  # ClusterMiddleEratio37_Photon
+    Rphi = 13  # ClusterMiddleEratio33_Photon
+    Weta2 = 14  # ClusterMiddleWidth_Photon
+    f1 = 15  # ClusterStripsEratio_Photon
+    DeltaE = 17  # ClusterStripsDeltaE_Photon
+    Wstot = 18  # ClusterStripsWtot_Photon
+    fside = 19  # ClusterStripsFracm_Photon
+    Ws3 = 20  # ClusterStripsWeta1c_Photon
+    ERatio = 21  # ClusterStripsDEmaxs1_Photon

atlas_schema-0.2.2/src/atlas_schema/methods.py

@@ -0,0 +1,241 @@
+"""Mixins for the Ntuple schema"""
+
+from __future__ import annotations
+
+from functools import reduce
+from operator import ior
+
+import awkward
+import particle
+from coffea.nanoevents.methods import base, candidate, vector
+from dask_awkward import dask_method
+
+from atlas_schema.enums import PhotonID
+from atlas_schema.typing_compat import Behavior
+
+behavior: Behavior = {}
+behavior.update(base.behavior)
+# vector behavior is included in candidate behavior
+behavior.update(candidate.behavior)
+
+
+class NtupleEvents(behavior["NanoEvents"]):  # type: ignore[misc, valid-type, name-defined]
+    def __repr__(self):
+        return f"<event {getattr(self, 'runNumber', '??')}:\
+                {getattr(self, 'eventNumber', '??')}:\
+                {getattr(self, 'mcChannelNumber', '??')}>"
+
+
+behavior["NanoEvents"] = NtupleEvents
+
+
+def _set_repr_name(classname):
+    def namefcn(_self):
+        return classname
+
+    behavior[("__typestr__", classname)] = classname[0].lower() + classname[1:]
+    behavior[classname].__repr__ = namefcn
+
+
+@awkward.mixin_class(behavior)
+class Weight(base.NanoCollection, base.Systematic): ...
+
+
+_set_repr_name("Weight")
+
+
+@awkward.mixin_class(behavior)
+class Pass(base.NanoCollection, base.Systematic): ...
+
+
+_set_repr_name("Pass")
+
+behavior.update(
+    awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior)
+)
+
+
+@awkward.mixin_class(behavior)
+class Particle(vector.PtEtaPhiMLorentzVector):
+    """Generic particle collection that has Lorentz vector properties
+
+    Also handles the following additional branches:
+    - '{obj}_select'
+    """
+
+    @property
+    def mass(self):
+        r"""Invariant mass (+, -, -, -)
+
+        :math:`\sqrt{t^2-x^2-y^2-z^2}`
+        """
+        return self["mass"] / 1.0e3
+
+    @dask_method
+    def passes(self, name):
+        return self[f"select_{name}"] == 1
+
+    @passes.dask
+    def passes(self, dask_array, name):
+        return dask_array[f"select_{name}"] == 1
+
+    # NB: fields with the name 'pt' take precedence over this
+    # @dask_property
+    # def pt(self):
+    #     print('inside non-dask prop')
+    #     return self["pt_NOSYS"]
+
+    # @pt.dask
+    # def pt(self, dask_array):
+    #     branch = 'pt'
+    #     print('inside dask prop')
+    #     variation = dask_array._events().metadata.get("systematic", "NOSYS")
+    #     with contextlib.suppress(Exception):
+    #         return dask_array[f"{branch}_{variation}"]
+
+    #     if variation != "NOSYS":
+    #         with contextlib.suppress(Exception):
+    #             return dask_array[f"{branch}_NOSYS"]
+
+    #     return dask_array[branch]
+
+
+_set_repr_name("Particle")
+
+ParticleArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+ParticleArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+ParticleArray.ProjectionClass4D = ParticleArray  # noqa: F821
+ParticleArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+
+behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior))
+
+
+@awkward.mixin_class(behavior)
+class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic):
+    @property
+    def r(self):
+        """Distance from origin in XY plane"""
+        return self["met"]
+
+
+_set_repr_name("MissingET")
+
+MissingETArray.ProjectionClass2D = MissingETArray  # noqa: F821
+MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray  # noqa: F821
+MissingETArray.ProjectionClass4D = vector.LorentzVectorArray  # noqa: F821
+MissingETArray.MomentumClass = MissingETArray  # noqa: F821
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Photon", behavior))
+
+
+@awkward.mixin_class(behavior)
+class Photon(Particle, base.NanoCollection, base.Systematic):
+    @property
+    def mass(self):
+        """Return zero mass for photon."""
+        return awkward.zeros_like(self.pt)
+
+    @property
+    def charge(self):
+        """Return zero charge for photon."""
+        return awkward.zeros_like(self.pt)
+
+    @property
+    def isEM(self):
+        return self.isEM_syst.NOSYS == 0
+
+    def pass_isEM(self, words: list[PhotonID]):
+        # 0 is pass, 1 is fail
+        return (
+            self.isEM_syst.NOSYS & reduce(ior, (1 << word.value for word in words))
+        ) == 0
+
+
+_set_repr_name("Photon")
+
+PhotonArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+PhotonArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+PhotonArray.ProjectionClass4D = PhotonArray  # noqa: F821
+PhotonArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior))
+
+
+@awkward.mixin_class(behavior)
+class Electron(Particle, base.NanoCollection, base.Systematic):
+    @property
+    def mass(self):
+        """Electron mass in GeV"""
+        return particle.literals.e_minus.mass / 1.0e3
+
+
+_set_repr_name("Electron")
+
+ElectronArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+ElectronArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+ElectronArray.ProjectionClass4D = ElectronArray  # noqa: F821
+ElectronArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior))
+
+
+@awkward.mixin_class(behavior)
+class Muon(Particle, base.NanoCollection, base.Systematic):
+    @property
+    def mass(self):
+        """Muon mass in GeV"""
+        return particle.literals.mu_minus.mass / 1.0e3
+
+
+_set_repr_name("Muon")
+
+MuonArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+MuonArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+MuonArray.ProjectionClass4D = MuonArray  # noqa: F821
+MuonArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Tau", behavior))
+
+
+@awkward.mixin_class(behavior)
+class Tau(Particle, base.NanoCollection, base.Systematic):
+    @property
+    def mass(self):
+        """Tau mass in GeV"""
+        return particle.literals.tau_minus.mass / 1.0e3
+
+
+_set_repr_name("Tau")
+
+TauArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+TauArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+TauArray.ProjectionClass4D = TauArray  # noqa: F821
+TauArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior))
+
+
+@awkward.mixin_class(behavior)
+class Jet(Particle, base.NanoCollection, base.Systematic): ...
+
+
+_set_repr_name("Jet")
+
+JetArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821
+JetArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821
+JetArray.ProjectionClass4D = JetArray  # noqa: F821
+JetArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821
+
+__all__ = [
+    "Electron",
+    "Jet",
+    "MissingET",
+    "Muon",
+    "NtupleEvents",
+    "Particle",
+    "Pass",
+    "Photon",
+    "Weight",
+]

atlas_schema-0.2.2/src/atlas_schema/utils.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import TypeVar, Union, cast
+
+import awkward as ak
+import dask_awkward as dak
+
+Array = TypeVar("Array", bound=Union[dak.Array, ak.Array])
+_E = TypeVar("_E", bound=Enum)
+
+
+def isin(haystack: Array, needles: dak.Array | ak.Array, axis: int = -1) -> Array:
+    """
+    Find needles in haystack.
+
+    This works by first transforming needles to an array with one more
+    dimension than the haystack, placing the needles at axis, and then doing a
+    comparison.
+
+    Args:
+        haystack (dak.Array or ak.Array): haystack of values.
+        needles (dak.Array or ak.Array): one-dimensional set of needles to find in haystack.
+        axis (int): the axis along which the comparison is performed
+
+    Returns:
+        dak.Array or ak.Array: result of comparison for needles in haystack
+    """
+    assert needles.ndim == 1, "Needles must be one-dimensional"
+    assert axis >= -1, "axis must be -1 or positive-valued"
+    assert axis < haystack.ndim + 1, "axis too large for the haystack"
+
+    # First, build up the transformation, with slice(None) indicating where to stick the needles
+    reshaper: list[None | slice] = [None] * haystack.ndim
+    axis = haystack.ndim if axis == -1 else axis
+    reshaper.insert(axis, slice(None))
+
+    # Note: reshaper needs to be a tuple for indexing purposes
+    return cast(Array, ak.any(haystack == needles[tuple(reshaper)], axis=-1))

atlas_schema-0.2.0/src/atlas_schema/enums.py

@@ -1,116 +0,0 @@
-from __future__ import annotations
-
-from enum import IntEnum
-
-from atlas_schema.typing_compat import Annotated
-
-
-# https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103
-class ParticleType(IntEnum):
-    Unknown: Annotated[int, "Unknown"] = 0
-    UnknownElectron: Annotated[int, "UnknownElectron"] = 1
-    IsoElectron: Annotated[int, "IsoElectron"] = 2
-    NonIsoElectron: Annotated[int, "NonIsoElectron"] = 3
-    BkgElectron: Annotated[int, "BkgElectron"] = 4
-    UnknownMuon: Annotated[int, "UnknownMuon"] = 5
-    IsoMuon: Annotated[int, "IsoMuon"] = 6
-    NonIsoMuon: Annotated[int, "NonIsoMuon"] = 7
-    BkgMuon: Annotated[int, "BkgMuon"] = 8
-    UnknownTau: Annotated[int, "UnknownTau"] = 9
-    IsoTau: Annotated[int, "IsoTau"] = 10
-    NonIsoTau: Annotated[int, "NonIsoTau"] = 11
-    BkgTau: Annotated[int, "BkgTau"] = 12
-    UnknownPhoton: Annotated[int, "UnknownPhoton"] = 13
-    IsoPhoton: Annotated[int, "IsoPhoton"] = 14
-    NonIsoPhoton: Annotated[int, "NonIsoPhoton"] = 15
-    BkgPhoton: Annotated[int, "BkgPhoton"] = 16
-    Hadron: Annotated[int, "Hadron"] = 17
-    Neutrino: Annotated[int, "Neutrino"] = 18
-    NuclFrag: Annotated[int, "NuclFrag"] = 19
-    NonPrimary: Annotated[int, "NonPrimary"] = 20
-    GenParticle: Annotated[int, "GenParticle"] = 21
-    SUSYParticle: Annotated[int, "SUSYParticle"] = 22
-    OtherBSMParticle: Annotated[int, "OtherBSMParticle"] = 39
-    BBbarMesonPart: Annotated[int, "BBbarMesonPart"] = 23
-    BottomMesonPart: Annotated[int, "BottomMesonPart"] = 24
-    CCbarMesonPart: Annotated[int, "CCbarMesonPart"] = 25
-    CharmedMesonPart: Annotated[int, "CharmedMesonPart"] = 26
-    BottomBaryonPart: Annotated[int, "BottomBaryonPart"] = 27
-    CharmedBaryonPart: Annotated[int, "CharmedBaryonPart"] = 28
-    StrangeBaryonPart: Annotated[int, "StrangeBaryonPart"] = 29
-    LightBaryonPart: Annotated[int, "LightBaryonPart"] = 30
-    StrangeMesonPart: Annotated[int, "StrangeMesonPart"] = 31
-    LightMesonPart: Annotated[int, "LightMesonPart"] = 32
-    BJet: Annotated[int, "BJet"] = 33
-    CJet: Annotated[int, "CJet"] = 34
-    LJet: Annotated[int, "LJet"] = 35
-    GJet: Annotated[int, "GJet"] = 36
-    TauJet: Annotated[int, "TauJet"] = 37
-    UnknownJet: Annotated[int, "UnknownJet"] = 38
-
-
-# https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103
-class ParticleOrigin(IntEnum):
-    NonDefined: Annotated[int, "NonDefined"] = 0
-    SingleElec: Annotated[int, "SingleElec"] = 1
-    SingleMuon: Annotated[int, "SingleMuon"] = 2
-    SinglePhot: Annotated[int, "SinglePhot"] = 3
-    SingleTau: Annotated[int, "SingleTau"] = 4
-    PhotonConv: Annotated[int, "PhotonConv"] = 5
-    DalitzDec: Annotated[int, "DalitzDec"] = 6
-    ElMagProc: Annotated[int, "ElMagProc"] = 7
-    Mu: Annotated[int, "Mu"] = 8
-    TauLep: Annotated[int, "TauLep"] = 9
-    top: Annotated[int, "top"] = 10
-    QuarkWeakDec: Annotated[int, "QuarkWeakDec"] = 11
-    WBoson: Annotated[int, "WBoson"] = 12
-    ZBoson: Annotated[int, "ZBoson"] = 13
-    Higgs: Annotated[int, "Higgs"] = 14
-    HiggsMSSM: Annotated[int, "HiggsMSSM"] = 15
-    HeavyBoson: Annotated[int, "HeavyBoson"] = 16
-    WBosonLRSM: Annotated[int, "WBosonLRSM"] = 17
-    NuREle: Annotated[int, "NuREle"] = 18
-    NuRMu: Annotated[int, "NuRMu"] = 19
-    NuRTau: Annotated[int, "NuRTau"] = 20
-    LQ: Annotated[int, "LQ"] = 21
-    SUSY: Annotated[int, "SUSY"] = 22
-    OtherBSM: Annotated[int, "OtherBSM"] = 46
-    LightMeson: Annotated[int, "LightMeson"] = 23
-    StrangeMeson: Annotated[int, "StrangeMeson"] = 24
-    CharmedMeson: Annotated[int, "CharmedMeson"] = 25
-    BottomMeson: Annotated[int, "BottomMeson"] = 26
-    CCbarMeson: Annotated[int, "CCbarMeson"] = 27
-    JPsi: Annotated[int, "JPsi"] = 28
-    BBbarMeson: Annotated[int, "BBbarMeson"] = 29
-    LightBaryon: Annotated[int, "LightBaryon"] = 30
-    StrangeBaryon: Annotated[int, "StrangeBaryon"] = 31
-    CharmedBaryon: Annotated[int, "CharmedBaryon"] = 32
-    BottomBaryon: Annotated[int, "BottomBaryon"] = 33
-    PionDecay: Annotated[int, "PionDecay"] = 34
-    KaonDecay: Annotated[int, "KaonDecay"] = 35
-    BremPhot: Annotated[int, "BremPhot"] = 36
-    PromptPhot: Annotated[int, "PromptPhot"] = 37
-    UndrPhot: Annotated[int, "UndrPhot"] = 38
-    ISRPhot: Annotated[int, "ISRPhot"] = 39
-    FSRPhot: Annotated[int, "FSRPhot"] = 40
-    NucReact: Annotated[int, "NucReact"] = 41
-    PiZero: Annotated[int, "PiZero"] = 42
-    DiBoson: Annotated[int, "DiBoson"] = 43
-    ZorHeavyBoson: Annotated[int, "ZorHeavyBoson"] = 44
-    MultiBoson: Annotated[int, "MultiBoson"] = 47
-    QCD: Annotated[int, "QCD"] = 45
-
-
-# https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/EGammaIdentificationRun2#Photon_isEM_word
-class PhotonID(IntEnum):
-    Rhad: Annotated[int, "ClusterHadronicLeakage_Photon"] = 10
-    E277: Annotated[int, "ClusterMiddleEnergy_Photon"] = 11
-    Reta: Annotated[int, "ClusterMiddleEratio37_Photon"] = 12
-    Rphi: Annotated[int, "ClusterMiddleEratio33_Photon"] = 13
-    Weta2: Annotated[int, "ClusterMiddleWidth_Photon"] = 14
-    f1: Annotated[int, "ClusterStripsEratio_Photon"] = 15
-    DeltaE: Annotated[int, "ClusterStripsDeltaE_Photon"] = 17
-    Wstot: Annotated[int, "ClusterStripsWtot_Photon"] = 18
-    fside: Annotated[int, "ClusterStripsFracm_Photon"] = 19
-    Ws3: Annotated[int, "ClusterStripsWeta1c_Photon"] = 20
-    ERatio: Annotated[int, "ClusterStripsDEmaxs1_Photon"] = 21

atlas_schema-0.2.0/src/atlas_schema/methods.py

@@ -1,181 +0,0 @@
-"""Mixins for the Ntuple schema"""
-
-from __future__ import annotations
-
-from functools import reduce
-from operator import ior
-
-import awkward
-from coffea.nanoevents.methods import base, candidate, vector
-from dask_awkward import dask_method
-
-from atlas_schema.enums import PhotonID
-from atlas_schema.typing_compat import Behavior
-
-behavior: Behavior = {}
-behavior.update(base.behavior)
-# vector behavior is included in candidate behavior
-behavior.update(candidate.behavior)
-
-
-class NtupleEvents(behavior["NanoEvents"]):  # type: ignore[misc, valid-type, name-defined]
-    def __repr__(self):
-        return f"<event {getattr(self,'runNumber','??')}:\
-                {getattr(self,'eventNumber','??')}:\
-                {getattr(self,'mcChannelNumber','??')}>"
-
-
-behavior["NanoEvents"] = NtupleEvents
-
-
-def _set_repr_name(classname):
-    def namefcn(_self):
-        return classname
-
-    behavior[("__typestr__", classname)] = classname[0].lower() + classname[1:]
-    behavior[classname].__repr__ = namefcn
-
-
-@awkward.mixin_class(behavior)
-class Weight(base.NanoCollection, base.Systematic): ...
-
-
-_set_repr_name("Weight")
-
-
-@awkward.mixin_class(behavior)
-class Pass(base.NanoCollection, base.Systematic): ...
-
-
-_set_repr_name("Pass")
-
-
-@awkward.mixin_class(behavior)
-class Particle(vector.PtEtaPhiMLorentzVector):
-    """Generic particle collection that has Lorentz vector properties
-
-    Also handles the following additional branches:
-    - '{obj}_select'
-    """
-
-    @property
-    def mass(self):
-        r"""Invariant mass (+, -, -, -)
-
-        :math:`\sqrt{t^2-x^2-y^2-z^2}`
-        """
-        return self["mass"] / 1.0e3
-
-    @dask_method
-    def passes(self, name):
-        return self[f"select_{name}"] == 1
-
-    @passes.dask
-    def passes(self, dask_array, name):
-        return dask_array[f"select_{name}"] == 1
-
-    # NB: fields with the name 'pt' take precedence over this
-    # @dask_property
-    # def pt(self):
-    #     print('inside non-dask prop')
-    #     return self["pt_NOSYS"]
-
-    # @pt.dask
-    # def pt(self, dask_array):
-    #     branch = 'pt'
-    #     print('inside dask prop')
-    #     variation = dask_array._events().metadata.get("systematic", "NOSYS")
-    #     with contextlib.suppress(Exception):
-    #         return dask_array[f"{branch}_{variation}"]
-
-    #     if variation != "NOSYS":
-    #         with contextlib.suppress(Exception):
-    #             return dask_array[f"{branch}_NOSYS"]
-
-    #     return dask_array[branch]
-
-
-_set_repr_name("Particle")
-
-
-@awkward.mixin_class(behavior)
-class MasslessParticle(Particle, base.NanoCollection):
-    @property
-    def mass(self):
-        r"""Invariant mass (+, -, -, -)
-
-        :math:`\sqrt{t^2-x^2-y^2-z^2}`
-        """
-        return 0.0 * self.pt
-
-
-_set_repr_name("MasslessParticle")
-
-
-@awkward.mixin_class(behavior)
-class MissingET(MasslessParticle, base.NanoCollection, base.Systematic):
-    @property
-    def pt(self):
-        """Alias for `r`"""
-        return self["met"] / 1.0e3
-
-    @property
-    def eta(self):
-        r"""Pseudorapidity
-
-        :math:`-\ln\tan(\theta/2) = \text{arcsinh}(z/r)`
-        """
-        return 0.0 * self.pt
-
-
-_set_repr_name("MissingET")
-
-
-@awkward.mixin_class(behavior)
-class Photon(MasslessParticle, base.NanoCollection, base.Systematic):
-    @property
-    def isEM(self):
-        return self.isEM_syst.NOSYS == 0
-
-    def pass_isEM(self, words: list[PhotonID]):
-        # 0 is pass, 1 is fail
-        return (
-            self.isEM_syst.NOSYS & reduce(ior, (1 << word.value for word in words))
-        ) == 0
-
-
-_set_repr_name("Photon")
-
-
-@awkward.mixin_class(behavior)
-class Electron(MasslessParticle, base.NanoCollection, base.Systematic): ...
-
-
-_set_repr_name("Electron")
-
-
-@awkward.mixin_class(behavior)
-class Muon(MasslessParticle, base.NanoCollection, base.Systematic): ...
-
-
-_set_repr_name("Muon")
-
-
-@awkward.mixin_class(behavior)
-class Jet(Particle, base.NanoCollection, base.Systematic): ...
-
-
-_set_repr_name("Jet")
-
-
-__all__ = [
-    "NtupleEvents",
-    "Weight",
-    "Pass",
-    "MissingET",
-    "Particle",
-    "Photon",
-    "Electron",
-    "Muon",
-    "Jet",
-]