atlas-schema 0.2.1__py3-none-any.whl → 0.2.2__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.2'
+__version_tuple__ = version_tuple = (0, 2, 2)

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/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.1.dist-info → atlas_schema-0.2.2.dist-info}/METADATA

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: atlas-schema
-Version: 0.2.1
+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
@@ -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.2
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]

atlas_schema-0.2.2.dist-info/RECORD

@@ -0,0 +1,13 @@
+atlas_schema/__init__.py,sha256=ebY-rTiwSGnfvt1yWATze2GE7K3fVgJj6fT64Sl4sH8,469
+atlas_schema/_version.py,sha256=RrHB9KG1O3GPm--rbTedqmZbdDrbgeRLXBmT4OBUqqI,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=YRVaiDa5Evl2HZ9CzH23d0-TLkvxqyvFQhn0ixyWCcw,7668
+atlas_schema/typing_compat.py,sha256=RwkxiiYbXO9yxkeaL8CdRaOHH7wq6vO_epg1YD7RbRs,439
+atlas_schema/utils.py,sha256=Oe2G3pe009Uhawsdk9e0MuqOHbAa5vZ8F2F9pOmz_Ok,1442
+atlas_schema-0.2.2.dist-info/METADATA,sha256=QeHezHbhZY-hA2xdVlrQNeZN2OSCA8hn24jzoMUZDX8,16823
+atlas_schema-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlas_schema-0.2.2.dist-info/licenses/LICENSE,sha256=snem82NV8fgAi4DKaaUIfReaM5RqIWbH5OOXOvy40_w,11344
+atlas_schema-0.2.2.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,,