atlas-schema 0.3.0__tar.gz → 0.4.0__tar.gz

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: atlas-schema
-Version: 0.3.0
+Version: 0.4.0
 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: Documentation, https://atlas-schema.readthedocs.io/en/v0.3.0/
+Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.4.0/
 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>
@@ -251,7 +251,7 @@ Requires-Dist: tbump>=6.7.0; extra == 'test'
 Requires-Dist: twine; extra == 'test'
 Description-Content-Type: text/markdown
 
-# atlas-schema v0.3.0
+# atlas-schema v0.4.0
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]
@@ -399,6 +399,57 @@ 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;">
 
+## Processing with Systematic Variations
+
+For analyses requiring systematic uncertainty evaluation, you can easily iterate
+over all systematic variations using the new `events["NOSYS"]` alias and
+`systematic_names` property:
+
+```python
+import awkward as ak
+from hist import Hist
+from coffea import processor
+from atlas_schema.schema import NtupleSchema
+
+
+class SystematicsProcessor(processor.ProcessorABC):
+    def __init__(self):
+        self.h = (
+            Hist.new.StrCat([], name="variation", growth=True)
+            .Regular(50, 0.0, 500.0, name="jet_pt", label="Leading Jet $p_T$ [GeV]")
+            .Int64()
+        )
+
+    def process(self, events):
+        dsid = events.metadata["dataset"]
+
+        # Process all systematic variations including nominal ("NOSYS")
+        for variation in events.systematic_names:
+            event_view = events[variation]
+
+            # Fill histogram with leading jet pT for this systematic variation
+            leading_jet_pt = event_view.jet.pt[:, 0] / 1_000  # Convert MeV to GeV
+            weights = (
+                event_view.weight.mc
+                if hasattr(event_view, "weight")
+                else ak.ones_like(leading_jet_pt)
+            )
+
+            self.h.fill(variation=variation, jet_pt=leading_jet_pt, weight=weights)
+
+        return {
+            "hist": self.h,
+            "meta": {"sumw": {dsid: {(events.metadata["fileuuid"], ak.sum(weights))}}},
+        }
+
+    def postprocess(self, accumulator):
+        return accumulator
+```
+
+This approach allows you to seamlessly process both nominal and systematic
+variations in a single loop, eliminating the need for special-case handling of
+the nominal variation.
+
 <!-- SPHINX-END -->
 
 ## Developer Notes

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/README.md

@@ -1,4 +1,4 @@
-# atlas-schema v0.3.0
+# atlas-schema v0.4.0
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]
@@ -146,6 +146,57 @@ 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;">
 
+## Processing with Systematic Variations
+
+For analyses requiring systematic uncertainty evaluation, you can easily iterate
+over all systematic variations using the new `events["NOSYS"]` alias and
+`systematic_names` property:
+
+```python
+import awkward as ak
+from hist import Hist
+from coffea import processor
+from atlas_schema.schema import NtupleSchema
+
+
+class SystematicsProcessor(processor.ProcessorABC):
+    def __init__(self):
+        self.h = (
+            Hist.new.StrCat([], name="variation", growth=True)
+            .Regular(50, 0.0, 500.0, name="jet_pt", label="Leading Jet $p_T$ [GeV]")
+            .Int64()
+        )
+
+    def process(self, events):
+        dsid = events.metadata["dataset"]
+
+        # Process all systematic variations including nominal ("NOSYS")
+        for variation in events.systematic_names:
+            event_view = events[variation]
+
+            # Fill histogram with leading jet pT for this systematic variation
+            leading_jet_pt = event_view.jet.pt[:, 0] / 1_000  # Convert MeV to GeV
+            weights = (
+                event_view.weight.mc
+                if hasattr(event_view, "weight")
+                else ak.ones_like(leading_jet_pt)
+            )
+
+            self.h.fill(variation=variation, jet_pt=leading_jet_pt, weight=weights)
+
+        return {
+            "hist": self.h,
+            "meta": {"sumw": {dsid: {(events.metadata["fileuuid"], ak.sum(weights))}}},
+        }
+
+    def postprocess(self, accumulator):
+        return accumulator
+```
+
+This approach allows you to seamlessly process both nominal and systematic
+variations in a single loop, eliminating the need for special-case handling of
+the nominal variation.
+
 <!-- SPHINX-END -->
 
 ## Developer Notes

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/pyproject.toml

@@ -60,7 +60,7 @@ docs = [
 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"
-Documentation = "https://atlas-schema.readthedocs.io/en/v0.3.0/"
+Documentation = "https://atlas-schema.readthedocs.io/en/v0.4.0/"
 Releases = "https://github.com/scipp-atlas/atlas-schema/releases"
 "Release Notes" = "https://atlas-schema.readthedocs.io/en/latest/history.html"
 

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/src/atlas_schema/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.0'
-__version_tuple__ = version_tuple = (0, 3, 0)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)
+
+__commit_id__ = commit_id = None

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/src/atlas_schema/enums.py

@@ -18,7 +18,7 @@ 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]
+    def __getitem__(cls: 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]

atlas_schema-0.4.0/src/atlas_schema/methods.py

@@ -0,0 +1,363 @@
+"""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 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)
+
+
+def _set_repr_name(classname):
+    def namefcn(_self):
+        return classname
+
+    behavior[("__typestr__", classname)] = classname[0].lower() + classname[1:]
+    behavior[classname].__repr__ = namefcn
+
+
+class NtupleEvents(behavior["NanoEvents"]):  # type: ignore[misc, valid-type, name-defined]
+    """Individual systematic variation of events."""
+
+    def __repr__(self):
+        return f"<event {getattr(self, 'runNumber', '??')}:{getattr(self, 'eventNumber', '??')}:{getattr(self, 'mcChannelNumber', '??')}>"
+
+    def __getitem__(self, key):
+        """Support accessing systematic variations via bracket notation.
+
+        Args:
+            key: The systematic variation name. "NOSYS" returns the nominal events.
+
+        Returns:
+            The requested systematic variation or nominal events for "NOSYS".
+        """
+        if key == "NOSYS":
+            return self
+        return super().__getitem__(key)
+
+    @property
+    def systematic(self):
+        """Get the systematic variation name for this event collection."""
+        return "nominal"
+
+    @property
+    def systematic_names(self):
+        """Get all systematic variations available in this event collection.
+
+        Returns a list of systematic variation names, including 'NOSYS' for nominal.
+        """
+        # Get systematics from metadata stored during schema building
+        systematics = self.metadata.get("systematics", [])
+        return ["NOSYS", *systematics]
+
+    @property
+    def systematics(self):
+        """Get all systematic variations available in this event collection.
+
+        Returns a list of systematic variation names, excluding 'nominal'.
+        """
+        # Get systematics from metadata stored during schema building
+        return [
+            getattr(self, systematic)
+            for systematic in self.systematic_names
+            if systematic != "NOSYS"
+        ]
+
+
+behavior["NtupleEvents"] = NtupleEvents
+
+
+class NtupleEventsArray(behavior[("*", "NanoEvents")]):  # type: ignore[misc, valid-type, name-defined]
+    """Collection of NtupleEvents objects, one for each systematic variation."""
+
+    def __getitem__(self, key):
+        """Support accessing systematic variations via bracket notation.
+
+        Args:
+            key: The systematic variation name. "NOSYS" returns the nominal events.
+
+        Returns:
+            The requested systematic variation or nominal events for "NOSYS".
+        """
+        if key == "NOSYS":
+            return self
+        return super().__getitem__(key)
+
+    @property
+    def systematic_names(self):
+        """Get all systematic variations available in this event collection.
+
+        Returns a list of systematic variation names, including 'NOSYS' for nominal.
+        """
+        # Get systematics from metadata stored during schema building
+        systematics = self.metadata.get("systematics", [])
+        return ["NOSYS", *systematics]
+
+    @property
+    def systematics(self):
+        """Get all systematic variations available in this event collection.
+
+        Returns a list of systematic variation names, excluding 'nominal'.
+        """
+        # Get systematics from metadata stored during schema building
+        return [
+            getattr(self, systematic)
+            for systematic in self.systematic_names
+            if systematic != "NOSYS"
+        ]
+
+
+behavior[("*", "NtupleEvents")] = NtupleEventsArray
+
+
+@awkward.mixin_class(behavior)
+class Systematic(base.NanoCollection, base.Systematic):
+    """Base class for systematic variations."""
+
+    @property
+    def metadata(self):
+        """Arbitrary metadata"""
+        return self.layout.purelist_parameter("metadata")  # pylint: disable=no-member
+
+    @property
+    def systematic(self):
+        """Get the systematic variation name for this event collection."""
+        return self.metadata["systematic"]
+
+    def __repr__(self):
+        return f"<event {self.systematic}>"
+
+
+_set_repr_name("Systematic")
+
+
+@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)  # pylint: disable=protected-access
+)
+
+
+@awkward.mixin_class(behavior)
+class Particle(vector.PtEtaPhiMLorentzVector):
+    """Generic particle collection that has Lorentz vector properties
+
+    Also handles the following additional branches:
+    - '{obj}_select'
+    """
+
+    def passes(self, name):
+        return self[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  # pylint: disable=undefined-variable,no-member
+ParticleArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+ParticleArray.ProjectionClass4D = ParticleArray  # noqa: F821  # pylint: disable=undefined-variable
+ParticleArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+
+behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic):
+    """Missing transverse energy collection."""
+
+    @property
+    def r(self):
+        """Distance from origin in XY plane"""
+        return self["met"]
+
+
+_set_repr_name("MissingET")
+
+MissingETArray.ProjectionClass2D = MissingETArray  # noqa: F821  # pylint: disable=undefined-variable
+MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+MissingETArray.ProjectionClass4D = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+MissingETArray.MomentumClass = MissingETArray  # noqa: F821  # pylint: disable=undefined-variable
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Photon", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class Photon(Particle, base.NanoCollection, base.Systematic):
+    """Photon particle collection."""
+
+    @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  # pylint: disable=no-member
+
+    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))  # pylint: disable=no-member
+        ) == 0
+
+
+_set_repr_name("Photon")
+
+PhotonArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+PhotonArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+PhotonArray.ProjectionClass4D = PhotonArray  # noqa: F821  # pylint: disable=undefined-variable
+PhotonArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class Electron(Particle, base.NanoCollection, base.Systematic):
+    """Electron particle collection."""
+
+    @property
+    def mass(self):
+        """Electron mass in MeV"""
+        return awkward.ones_like(self.pt) * particle.literals.e_minus.mass  # pylint: disable=no-member
+
+
+_set_repr_name("Electron")
+
+ElectronArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+ElectronArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+ElectronArray.ProjectionClass4D = ElectronArray  # noqa: F821  # pylint: disable=undefined-variable
+ElectronArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class Muon(Particle, base.NanoCollection, base.Systematic):
+    """Muon particle collection."""
+
+    @property
+    def mass(self):
+        """Muon mass in MeV"""
+        return awkward.ones_like(self.pt) * particle.literals.mu_minus.mass  # pylint: disable=no-member
+
+
+_set_repr_name("Muon")
+
+MuonArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+MuonArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+MuonArray.ProjectionClass4D = MuonArray  # noqa: F821  # pylint: disable=undefined-variable
+MuonArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Tau", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class Tau(Particle, base.NanoCollection, base.Systematic):
+    """Tau particle collection."""
+
+    @property
+    def mass(self):
+        """Tau mass in MeV"""
+        return awkward.ones_like(self.pt) * particle.literals.tau_minus.mass  # pylint: disable=no-member
+
+
+_set_repr_name("Tau")
+
+TauArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+TauArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+TauArray.ProjectionClass4D = TauArray  # noqa: F821  # pylint: disable=undefined-variable
+TauArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+
+behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior))  # pylint: disable=protected-access
+
+
+@awkward.mixin_class(behavior)
+class Jet(Particle, base.NanoCollection, base.Systematic):
+    """Jet particle collection."""
+
+    @property
+    def mass(self):
+        r"""Invariant mass (+, -, -, -)
+
+        :math:`\sqrt{t^2-x^2-y^2-z^2}`
+        """
+        return self["m"]
+
+
+_set_repr_name("Jet")
+
+JetArray.ProjectionClass2D = vector.TwoVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+JetArray.ProjectionClass3D = vector.ThreeVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+JetArray.ProjectionClass4D = JetArray  # noqa: F821  # pylint: disable=undefined-variable
+JetArray.MomentumClass = vector.LorentzVectorArray  # noqa: F821  # pylint: disable=undefined-variable,no-member
+
+__all__ = [
+    "Electron",
+    "ElectronArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "ElectronRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "Jet",
+    "JetArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "JetRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "MissingET",
+    "MissingETArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "MissingETRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "Muon",
+    "MuonArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "MuonRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "NtupleEvents",
+    "Particle",
+    "ParticleArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "ParticleRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "Pass",
+    "Photon",
+    "PhotonArray",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "PhotonRecord",  # noqa: F822  # pylint: disable=undefined-all-variable
+    "Weight",
+]

{atlas_schema-0.3.0 → atlas_schema-0.4.0}/src/atlas_schema/schema.py

@@ -171,10 +171,12 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
             pass
         else:
             pass
-        self._form["fields"], self._form["contents"] = self._build_collections(
-            self._form["fields"], self._form["contents"]
+        self._form["fields"], self._form["contents"], discovered_systematics = (
+            self._build_collections(self._form["fields"], self._form["contents"])
         )
         self._form["parameters"]["metadata"]["version"] = self._version
+        self._form["parameters"]["metadata"]["systematics"] = discovered_systematics
+        self._form["parameters"]["__record__"] = "NtupleEvents"
 
     @classmethod
     def v1(cls, base_form: dict[str, Any]) -> Self:
@@ -187,7 +189,7 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
 
     def _build_collections(
         self, field_names: list[str], input_contents: list[Any]
-    ) -> tuple[KeysView[str], ValuesView[dict[str, Any]]]:
+    ) -> tuple[KeysView[str], ValuesView[dict[str, Any]], list[str]]:
         branch_forms = dict(zip(field_names, input_contents))
 
         # parse into high-level records (collections, list collections, and singletons)
@@ -240,6 +242,19 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
             msg = "One of the branches does not follow the assumed pattern for this schema. [invalid-branch-name]"
             raise RuntimeError(msg) from exc
 
+        all_systematics = self._discover_systematics(
+            branch_forms, collections, subcollections
+        )
+
+        # Pre-compute systematic branch patterns for O(1) lookups
+        # This replaces the expensive O(m*s) nested condition checks
+        systematic_branch_patterns = set()
+        for collection in collections:
+            for subcoll in subcollections:
+                for sys in all_systematics:
+                    if sys != "NOSYS":
+                        systematic_branch_patterns.add(f"{collection}_{subcoll}_{sys}")
+
         # Check the presence of the event_ids
         missing_event_ids = [
             event_id for event_id in self.event_ids if event_id not in branch_forms
@@ -285,91 +300,206 @@ class NtupleSchema(BaseSchema):  # type: ignore[misc]
 
             output[name] = branch_forms[name]
 
-        # next, go through and start grouping up collections
-        for name in collections:
-            content = {}
+        # First, build nominal collections the traditional way
+        nominal_collections = {}
+        for collection_name in collections:
+            collection_content = {}
             used = set()
 
+            # Process subcollections with NOSYS variations
             for subname in subcollections:
-                prefix = f"{name}_{subname}_"
-                used.update({k for k in branch_forms if k.startswith(prefix)})
-                subcontent = {
-                    k[len(prefix) :]: branch_forms[k]
-                    for k in branch_forms
-                    if k.startswith(prefix)
-                }
-                if subcontent:
-                    # create the nominal version
-                    content[subname] = branch_forms[f"{prefix}NOSYS"]
-                    # create a collection of the systematic variations for the given variable
-                    content[f"{subname}_syst"] = zip_forms(
-                        subcontent, f"{name}_syst", record_name="NanoCollection"
+                prefix = f"{collection_name}_{subname}_"
+                nosys_branch = f"{prefix}NOSYS"
+
+                if nosys_branch in branch_forms:
+                    collection_content[subname] = branch_forms[nosys_branch]
+                    used.add(nosys_branch)
+
+            # Add non-systematic branches (like eta, phi that don't vary)
+            for k, form in branch_forms.items():
+                if (
+                    k.startswith(collection_name + "_")
+                    and k not in used
+                    and "_NOSYS" not in k
+                    and k
+                    not in systematic_branch_patterns  # O(1) lookup instead of O(m*s)
+                ):
+                    field_name = k[len(collection_name) + 1 :]
+                    if field_name not in collection_content:
+                        collection_content[field_name] = form
+
+            if collection_content:
+                behavior = self.mixins.get(collection_name, "")
+                if not behavior:
+                    behavior = self.suggested_behavior(collection_name)
+                    warnings.warn(
+                        f"I found a collection with no defined mixin: '{collection_name}'. I will assume behavior: '{behavior}'. To suppress this warning next time, please define mixins for your custom collections. [mixin-undefined]",
+                        RuntimeWarning,
+                        stacklevel=2,
                     )
+                nominal_collections[collection_name] = zip_forms(
+                    collection_content, collection_name, record_name=behavior
+                )
+                nominal_collections[collection_name].setdefault("parameters", {})
+                nominal_collections[collection_name]["parameters"].update(
+                    {"collection_name": collection_name}
+                )
 
-            content.update(
-                {
-                    k[len(name) + 1 :]: branch_forms[k]
-                    for k in branch_forms
-                    if k.startswith(name + "_") and k not in used
+        # Add nominal collections to output
+        output.update(nominal_collections)
+
+        # Now build systematic event structures
+        for systematic in all_systematics:
+            if systematic == "NOSYS":
+                continue
+
+            # Check which collections actually have this systematic variation
+            systematic_collections = {}
+
+            for collection_name in collections:
+                # Check if this collection has any systematic branches for this systematic
+                has_systematic_data = False
+                collection_content = {}
+                used = set()
+
+                # Process subcollections with systematic variations
+                for subname in subcollections:
+                    prefix = f"{collection_name}_{subname}_"
+                    target_branch = f"{prefix}{systematic}"
+                    fallback_branch = f"{prefix}NOSYS"
+
+                    if target_branch in branch_forms:
+                        # Use the systematic variation
+                        collection_content[subname] = branch_forms[target_branch]
+                        used.add(target_branch)
+                        has_systematic_data = True
+                    elif fallback_branch in branch_forms:
+                        # Fall back to nominal
+                        collection_content[subname] = branch_forms[fallback_branch]
+                        used.add(fallback_branch)
+
+                # Add non-systematic branches
+                for k, form in branch_forms.items():
+                    if (
+                        k.startswith(collection_name + "_")
+                        and k not in used
+                        and "_NOSYS" not in k
+                        and k
+                        not in systematic_branch_patterns  # O(1) lookup instead of O(m*s)
+                    ):
+                        field_name = k[len(collection_name) + 1 :]
+                        if field_name not in collection_content:
+                            collection_content[field_name] = form
+
+                # If this collection has systematic data or fallback data, include it
+                if collection_content:
+                    behavior = self.mixins.get(collection_name, "")
+                    if not behavior:
+                        behavior = self.suggested_behavior(collection_name)
+                        # Only warn once (for nominal collections)
+
+                    # If no systematic data, use the nominal collection directly
+                    if (
+                        not has_systematic_data
+                        and collection_name in nominal_collections
+                    ):
+                        systematic_collections[collection_name] = nominal_collections[
+                            collection_name
+                        ]
+                    else:
+                        # Build the systematic collection
+                        systematic_collections[collection_name] = zip_forms(
+                            collection_content, collection_name, record_name=behavior
+                        )
+                        systematic_collections[collection_name].setdefault(
+                            "parameters", {}
+                        )
+                        systematic_collections[collection_name]["parameters"].update(
+                            {"collection_name": collection_name}
+                        )
+
+            # Only create systematic event if there are collections for it
+            if systematic_collections:
+                output[systematic] = {
+                    "class": "RecordArray",
+                    "contents": list(systematic_collections.values()),
+                    "fields": list(systematic_collections.keys()),
+                    "form_key": f"%21invalid%2C{systematic}",
+                    "parameters": {
+                        "__record__": "Systematic",
+                        "metadata": {"systematic": systematic},
+                    },
                 }
-            )
 
-            if not used and not content:
+        # Handle any remaining unrecognized branches as singletons
+        processed_branches = set()
+        # Add event IDs and explicit singletons
+        processed_branches.update(self.event_ids)
+        processed_branches.update(self.singletons)
+        # Add collection-related branches
+        for collection_name in collections:
+            for branch_name in branch_forms:
+                if branch_name.startswith(collection_name + "_"):
+                    processed_branches.add(branch_name)
+
+        # Find unrecognized branches
+        for branch_name, form in branch_forms.items():
+            if branch_name not in processed_branches:
+                # This is an unrecognized branch - treat as singleton with warning
                 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. [singleton-undefined]",
+                    f"I identified a branch that likely does not have any leaves: '{branch_name}'. I will treat this as a 'singleton'. To suppress this warning, add this branch to the singletons set. [singleton-undefined]",
                     RuntimeWarning,
                     stacklevel=2,
                 )
-                self.singletons.add(name)
-                output[name] = branch_forms[name]
+                output[branch_name] = form
 
-            else:
-                behavior = self.mixins.get(name, "")
-                if not behavior:
-                    behavior = self.suggested_behavior(name)
-                    warnings.warn(
-                        f"I found a collection with no defined mixin: '{name}'. I will assume behavior: '{behavior}'. To suppress this warning next time, please define mixins for your custom collections. [mixin-undefined]",
-                        RuntimeWarning,
-                        stacklevel=2,
-                    )
+        # Return discovered systematics (excluding NOSYS/nominal)
+        discovered_systematics = sorted([s for s in all_systematics if s != "NOSYS"])
 
-                output[name] = zip_forms(content, name, record_name=behavior)
-
-            output[name].setdefault("parameters", {})
-            output[name]["parameters"].update({"collection_name": name})
-
-            if output[name]["class"] == "ListOffsetArray":
-                if output[name]["class"] == "RecordArray":
-                    parameters = output[name]["content"]["fields"]
-                    contents = output[name]["content"]["contents"]
-                else:
-                    # these are also singletons of another kind that we just pass through
-                    continue
-            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)
+        return output.keys(), output.values(), discovered_systematics
 
-            # update docstrings as needed
-            # NB: must be before flattening for easier logic
-            for index, parameter in enumerate(parameters):
-                if "parameters" not in contents[index]:
-                    continue
-
-                parsed_name = parameter.replace("_NOSYS", "")
-                contents[index]["parameters"]["__doc__"] = self.docstrings.get(
-                    parsed_name,
-                    contents[index]["parameters"].get(
-                        "__doc__", "no docstring available"
-                    ),
-                )
+    def _discover_systematics(
+        self,
+        branch_forms: dict[str, Any],
+        collections: set[str],
+        subcollections: set[str],
+    ) -> set[str]:
+        """Extract systematic variations from branch names.
+
+        Returns:
+            set: Set of all systematic variation names found in branches
+        """
+        # Optimize systematic discovery: pre-index branches by pattern
+        # This avoids O(n*m) nested loops in systematic discovery
+        subcoll_patterns = {f"{subcoll}_" for subcoll in subcollections}
+
+        all_systematics = set()
+        for k in branch_forms:
+            if not ("_" in k and k not in self.singletons):
+                continue
+            # Handle the pattern: collection_subcollection_systematic
+            # where systematic can contain double underscores like "JET_EnergyResolution__1up"
+            parts = k.split("_")
+            if len(parts) < 3:
+                continue
+            # Find the collection and subcollection parts
+            collection = parts[0]
+            if collection not in collections:
+                continue
+            # Find where the subcollection ends by looking for a known pattern
+            # The systematic starts after the subcollection
+            remaining = "_".join(parts[1:])
+            # Use optimized lookup instead of iterating all subcollections
+            for pattern in subcoll_patterns:
+                if remaining.startswith(pattern):
+                    systematic = remaining[len(pattern) :]
+                    if systematic and systematic != "NOSYS":
+                        all_systematics.add(systematic)
+                    break
 
-        return output.keys(), output.values()
+        # Always include NOSYS as the nominal case
+        all_systematics.add("NOSYS")
+        return all_systematics
 
     @classmethod
     def behavior(cls) -> Behavior:

atlas_schema-0.3.0/src/atlas_schema/methods.py

@@ -1,246 +0,0 @@
-"""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 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'
-    """
-
-    def passes(self, name):
-        return self[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 MeV"""
-        return awkward.ones_like(self.pt) * particle.literals.e_minus.mass
-
-
-_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 MeV"""
-        return awkward.ones_like(self.pt) * particle.literals.mu_minus.mass
-
-
-_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 MeV"""
-        return awkward.ones_like(self.pt) * particle.literals.tau_minus.mass
-
-
-_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):
-    @property
-    def mass(self):
-        r"""Invariant mass (+, -, -, -)
-
-        :math:`\sqrt{t^2-x^2-y^2-z^2}`
-        """
-        return self["m"]
-
-
-_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",
-    "ElectronArray",  # noqa: F822
-    "ElectronRecord",  # noqa: F822
-    "Jet",
-    "JetArray",  # noqa: F822
-    "JetRecord",  # noqa: F822
-    "MissingET",
-    "MissingETArray",  # noqa: F822
-    "MissingETRecord",  # noqa: F822
-    "Muon",
-    "MuonArray",  # noqa: F822
-    "MuonRecord",  # noqa: F822
-    "NtupleEvents",
-    "Particle",
-    "ParticleArray",  # noqa: F822
-    "ParticleRecord",  # noqa: F822
-    "Pass",
-    "Photon",
-    "PhotonArray",  # noqa: F822
-    "PhotonRecord",  # noqa: F822
-    "Weight",
-]