atlas-schema 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

atlas_schema/__init__.py

@@ -10,7 +10,8 @@ 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__ = ["ParticleOrigin", "PhotonID", "__version__"]
+__all__ = ["ParticleOrigin", "PhotonID", "__version__", "isin"]

atlas_schema/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)

atlas_schema/enums.py

@@ -1,9 +1,31 @@
 from __future__ import annotations
 
-from enum import IntEnum
+import sys
+from enum import Enum, IntEnum
 
+if sys.version_info >= (3, 11):
+    from enum import EnumType
+else:
+    from enum import EnumMeta as EnumType
 
-class ParticleType(IntEnum):
+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>`_.
     """
@@ -50,7 +72,7 @@ class ParticleType(IntEnum):
     UnknownJet = 38
 
 
-class ParticleOrigin(IntEnum):
+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>`_.
     """
@@ -105,7 +127,7 @@ class ParticleOrigin(IntEnum):
     QCD = 45
 
 
-class PhotonID(IntEnum):
+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>`_.
     """

atlas_schema/schema.py

@@ -49,7 +49,7 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
     }
 
     # These are stored as length-1 vectors unnecessarily
-    singletons: ClassVar[list[str]] = []
+    singletons: ClassVar[set[str]] = set()
 
     docstrings: ClassVar[dict[str, str]] = {
         "charge": "charge",
@@ -127,8 +127,8 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
 
         output = {}
 
-        # first, register the event-level stuff directly
-        for name in self.event_ids:
+        # first, register singletons (event-level, others)
+        for name in {*self.event_ids, *self.singletons}:
             if name in missing_event_ids:
                 continue
             output[name] = branch_forms[name]
@@ -163,7 +163,17 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
                 }
             )
 
-            output[name] = zip_forms(content, name, record_name=mixin)
+            if not used and not content:
+                warnings.warn(
+                    f"I identified a branch that likely does not have any leaves: '{name}'. I will treat this as a 'singleton'. To suppress this warning next time, please define your singletons explicitly.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                self.singletons.add(name)
+                output[name] = branch_forms[name]
+
+            else:
+                output[name] = zip_forms(content, name, record_name=mixin)
 
             output[name].setdefault("parameters", {})
             output[name]["parameters"].update({"collection_name": name})
@@ -174,6 +184,9 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
             elif output[name]["class"] == "RecordArray":
                 parameters = output[name]["fields"]
                 contents = output[name]["contents"]
+            elif output[name]["class"] == "NumpyArray":
+                # these are singletons that we just pass through
+                continue
             else:
                 msg = f"Unhandled class {output[name]['class']}"
                 raise RuntimeError(msg)
@@ -191,11 +204,6 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
                     ),
                 )
 
-            if name in self.singletons:
-                # flatten! this 'promotes' the content of an inner dimension
-                # upwards, effectively hiding one nested dimension
-                output[name] = output[name]["content"]
-
         return output.keys(), output.values()
 
     @classmethod

atlas_schema/utils.py

@@ -0,0 +1,49 @@
+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(element: Array, test_elements: dak.Array | ak.Array, axis: int = -1) -> Array:
+    """
+    Find test_elements in element. Similar in API as :func:`numpy.isin`.
+
+    Calculates `element in test_elements`, broadcasting over *element elements only*. Returns a boolean array of the same shape as *element* that is `True` where an element of *element* is in *test_elements* and `False` otherwise.
+
+    This works by first transforming *test_elements* to an array with one more
+    dimension than the *element*, placing the *test_elements* at *axis*, and then doing a
+    comparison.
+
+    Args:
+        element (dak.Array or ak.Array): input array of values.
+        test_elements (dak.Array or ak.Array): one-dimensional set of values against which to test each value of *element*.
+        axis (int): the axis along which the comparison is performed
+
+    Returns:
+        dak.Array or ak.Array: result of comparison for test_elements in *element*
+
+    Example:
+        >>> import awkward as ak
+        >>> import atlas_schema as ats
+        >>> truth_origins = ak.Array([[1, 2, 3], [4], [5, 6, 7], [1]])
+        >>> prompt_origins = ak.Array([1, 2, 7])
+        >>> ats.isin(truth_origins, prompt_origins).to_list()
+        [[True, True, False], [False], [False, False, True], [True]]
+    """
+    assert test_elements.ndim == 1, "test_elements must be one-dimensional"
+    assert axis >= -1, "axis must be -1 or positive-valued"
+    assert axis < element.ndim + 1, "axis too large for the element"
+
+    # First, build up the transformation, with slice(None) indicating where to stick the test_elements
+    reshaper: list[None | slice] = [None] * element.ndim
+    axis = element.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(element == test_elements[tuple(reshaper)], axis=-1))

{atlas_schema-0.2.1.dist-info → atlas_schema-0.2.3.dist-info}/METADATA

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: atlas-schema
-Version: 0.2.1
+Version: 0.2.3
 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.3/
+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
@@ -249,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.1
+# atlas-schema v0.2.3
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]
@@ -277,6 +279,129 @@ Description-Content-Type: text/markdown
 
 <!-- prettier-ignore-end -->
 
+This is the python package containing schemas and helper functions enabling
+analyzers to work with ATLAS datasets (Monte Carlo and Data), using
+[coffea](https://coffea-hep.readthedocs.io/en/latest/).
+
+## Hello World
+
+The simplest example is to just get started processing the file as expected:
+
+```python
+from atlas_schema.schema import NtupleSchema
+from coffea import dataset_tools
+import awkward as ak
+
+fileset = {"ttbar": {"files": {"path/to/ttbar.root": "tree_name"}}}
+samples, report = dataset_tools.preprocess(fileset)
+
+
+def noop(events):
+    return ak.fields(events)
+
+
+fields = dataset_tools.apply_to_fileset(noop, samples, schemaclass=NtupleSchema)
+print(fields)
+```
+
+which produces something similar to
+
+```python
+{
+    "ttbar": [
+        "dataTakingYear",
+        "mcChannelNumber",
+        "runNumber",
+        "eventNumber",
+        "lumiBlock",
+        "actualInteractionsPerCrossing",
+        "averageInteractionsPerCrossing",
+        "truthjet",
+        "PileupWeight",
+        "RandomRunNumber",
+        "met",
+        "recojet",
+        "truth",
+        "generatorWeight",
+        "beamSpotWeight",
+        "trigPassed",
+        "jvt",
+    ]
+}
+```
+
+However, a more involved example to apply a selection and fill a histogram looks
+like below:
+
+```python
+import awkward as ak
+import dask
+import hist.dask as had
+import matplotlib.pyplot as plt
+from coffea import processor
+from coffea.nanoevents import NanoEventsFactory
+from distributed import Client
+
+from atlas_schema.schema import NtupleSchema
+
+
+class MyFirstProcessor(processor.ProcessorABC):
+    def __init__(self):
+        pass
+
+    def process(self, events):
+        dataset = events.metadata["dataset"]
+        h_ph_pt = (
+            had.Hist.new.StrCat(["all", "pass", "fail"], name="isEM")
+            .Regular(200, 0.0, 2000.0, name="pt", label="$pt_{\gamma}$ [GeV]")
+            .Int64()
+        )
+
+        cut = ak.all(events.ph.isEM, axis=1)
+        h_ph_pt.fill(isEM="all", pt=ak.firsts(events.ph.pt / 1.0e3))
+        h_ph_pt.fill(isEM="pass", pt=ak.firsts(events[cut].ph.pt / 1.0e3))
+        h_ph_pt.fill(isEM="fail", pt=ak.firsts(events[~cut].ph.pt / 1.0e3))
+
+        return {
+            dataset: {
+                "entries": ak.num(events, axis=0),
+                "ph_pt": h_ph_pt,
+            }
+        }
+
+    def postprocess(self, accumulator):
+        pass
+
+
+if __name__ == "__main__":
+    client = Client()
+
+    fname = "ntuple.root"
+    events = NanoEventsFactory.from_root(
+        {fname: "analysis"},
+        schemaclass=NtupleSchema,
+        metadata={"dataset": "700352.Zqqgamma.mc20d.v1"},
+    ).events()
+
+    p = MyFirstProcessor()
+    out = p.process(events)
+    (computed,) = dask.compute(out)
+    print(computed)
+
+    fig, ax = plt.subplots()
+    computed["700352.Zqqgamma.mc20d.v1"]["ph_pt"].plot1d(ax=ax)
+    ax.set_xscale("log")
+    ax.legend(title="Photon pT for Zqqgamma")
+
+    fig.savefig("ph_pt.pdf")
+```
+
+which produces
+
+<img src="https://raw.githubusercontent.com/scipp-atlas/atlas-schema/main/docs/_static/img/ph_pt.png" alt="three stacked histograms of photon pT, with each stack corresponding to: no selection, requiring the isEM flag, and inverting the isEM requirement" width="500" style="display: block; margin-left: auto; margin-right: auto;">
+
+<!-- SPHINX-END -->
+
 ## Developer Notes
 
 ### Converting Enums from C++ to Python

atlas_schema-0.2.3.dist-info/RECORD

@@ -0,0 +1,13 @@
+atlas_schema/__init__.py,sha256=ebY-rTiwSGnfvt1yWATze2GE7K3fVgJj6fT64Sl4sH8,469
+atlas_schema/_version.py,sha256=AaQEeqeDwmZAHoPuwg2C0ulADePbIYLSFanZzt0cytQ,411
+atlas_schema/_version.pyi,sha256=j5kbzfm6lOn8BzASXWjGIA1yT0OlHTWqlbyZ8Si_o0E,118
+atlas_schema/enums.py,sha256=hwgOvFBmITNxL0MQkrNpbiPv9VMezFoE-eyGgjzem8E,3688
+atlas_schema/methods.py,sha256=K7u6HGKXrtpMg7jjCjKPwIEnknOShUH4HQ1ibKBzkZ0,6832
+atlas_schema/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+atlas_schema/schema.py,sha256=s3bcSa5DH5iILPD_2BD3co8MSTsXNs1rYBmn44388Kc,8082
+atlas_schema/typing_compat.py,sha256=RwkxiiYbXO9yxkeaL8CdRaOHH7wq6vO_epg1YD7RbRs,439
+atlas_schema/utils.py,sha256=spk7KIMBbXSPpZBTltyxaHWvyitkEGSVldfuKFoyavk,2137
+atlas_schema-0.2.3.dist-info/METADATA,sha256=xtXdXa-9ra8TTtZSXLycUKPvBLPzXCIPa5cueuv0w90,20107
+atlas_schema-0.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlas_schema-0.2.3.dist-info/licenses/LICENSE,sha256=snem82NV8fgAi4DKaaUIfReaM5RqIWbH5OOXOvy40_w,11344
+atlas_schema-0.2.3.dist-info/RECORD,,

atlas_schema-0.2.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-atlas_schema/__init__.py,sha256=mwY_EsW32pdZxihzpAg_enYPl7S-_d27idXKIlYvVqE,425
-atlas_schema/_version.py,sha256=MxUhzLJIZQfEpDTTcKSxciTGrMLd5v2VmMlHa2HGeo0,411
-atlas_schema/_version.pyi,sha256=j5kbzfm6lOn8BzASXWjGIA1yT0OlHTWqlbyZ8Si_o0E,118
-atlas_schema/enums.py,sha256=a44N1UbUA4K1y6wzK7zBUxbw4xT02J7vLeSqT03dFaU,2941
-atlas_schema/methods.py,sha256=K7u6HGKXrtpMg7jjCjKPwIEnknOShUH4HQ1ibKBzkZ0,6832
-atlas_schema/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-atlas_schema/schema.py,sha256=YRVaiDa5Evl2HZ9CzH23d0-TLkvxqyvFQhn0ixyWCcw,7668
-atlas_schema/typing_compat.py,sha256=RwkxiiYbXO9yxkeaL8CdRaOHH7wq6vO_epg1YD7RbRs,439
-atlas_schema-0.2.1.dist-info/METADATA,sha256=KFUTH5W2HUx8wJOzznaalEqLc9frqeMcbPJ1TmNyTYc,16662
-atlas_schema-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-atlas_schema-0.2.1.dist-info/licenses/LICENSE,sha256=snem82NV8fgAi4DKaaUIfReaM5RqIWbH5OOXOvy40_w,11344
-atlas_schema-0.2.1.dist-info/RECORD,,