import awkward
import numpy as np
from awkward_zipper.behaviors import base, vector
behavior = {}
behavior.update(base.behavior)
def _copy_behaviors(from_name, to_name, behavior):
"""Copy cross-class behaviors from ``from_name`` onto ``to_name`` in place.
Uses ``setdefault`` rather than ``dict.update``: these copies run *after* the
target class is defined, and ``awkward.mixin_class`` registers a class's own
record via plain assignment (overwrite). A blanket ``update`` would therefore
clobber the subclass's own record/array registration with the parent's,
silently downcasting e.g. an ``MCParticle`` to a ``MomentumCandidate`` (which
then lacks methods like ``get_daughters``). ``setdefault`` fills in only the
genuinely-new cross-class keys and leaves the subclass's own entries intact.
See scikit-hep/coffea#1594.
"""
for key, value in awkward._util.copy_behaviors(
from_name, to_name, behavior
).items():
behavior.setdefault(key, value)
class _EDM4HEPEvents(behavior["NanoEvents"]):
"""EDM4HEP Events"""
def __repr__(self):
return "EDM4HEP-Event"
behavior["NanoEvents"] = _EDM4HEPEvents
def _set_repr_name(classname):
def namefcn(self):
return classname
behavior[classname].__repr__ = namefcn
[docs]
@awkward.mixin_class(behavior)
class edm4hep_nanocollection(base.NanoCollection):
"""Modified NanoCollection for EDM4HEP"""
def _apply_nested_global_index(self, index):
"""Similar to _apply_global_index but the indexes are twice nested"""
# extract the shape of the index
counts1 = awkward.num(index, axis=1)
counts2 = awkward.flatten(awkward.num(index, axis=2), axis=1)
if index.ndim == 4:
counts3 = awkward.ravel(awkward.num(index, axis=3))
# Apply the flat index
out = self._apply_global_index(awkward.ravel(index))
# Rebuild the shape of the index
if index.ndim == 4:
out = awkward.unflatten(out, counts3, axis=0)
out = awkward.unflatten(out, counts2, axis=0)
return awkward.unflatten(out, counts1, axis=0)
@property
def List_Relations(self):
"""List all the branches that are for OneToOneRelations or OneToManyRelations"""
return {name for name in self.fields if "_idx_" in name}
[docs]
def Map_Relation(self, generic_name, target_name):
"""Map a OneToOne or a OneToMany Relation to it's target collection
Options: generic_name | str | name of the relation in EDM4HEP.yaml
target_name | str | name of the target collection
Note: All relation branches are named as,
genericname_idx_targetname_index
or
genericname_idx_targetname_index_Global
or
genericname_idx_targetname_collectionID
One can list out all such branches with the method List_Relations
"""
name_struct = generic_name + "_idx_" + target_name + "_index_Global"
idx_field_names = [name for name in self.fields if name_struct in name]
if len(idx_field_names) == 0:
msg = f"*{name_struct} not found in the current collection"
raise FileNotFoundError(msg)
if len(idx_field_names) > 1:
msg = f"More than one field available for *{name_struct}!"
raise RuntimeError(msg)
index = self[idx_field_names[0]]
if index.ndim == 2:
return self._events()[target_name]._apply_global_index(index)
if index.ndim == 3 or index.ndim == 4:
return self._events()[target_name]._apply_nested_global_index(index)
msg = f"Index is highly nested!\n{index}"
raise RuntimeError(msg)
@property
def List_Links(self):
"""List all the branches that are Links"""
return {
name
for name in self.fields
if (("Link_from" in name) or ("Link_to" in name))
}
[docs]
def Map_Link(self, generic_name, target_name):
"""Map a Link to it's target collection
Note: This method only works when copy_links_to_target_datatype is set to true in the schema
Options: generic_name | str | name of the Link in EDM4HEP.yaml
target_name | str | name of the target collection
Note: All link branches are named as,
Link_from_genericname_idx_targetname_index or Link_to_genericname_idx_targetname_index
or
Link_from_genericname_idx_targetname_index_Global or Link_to_genericname_idx_targetname_index_Global
or
Link_from_genericname_idx_targetname_collectionID or Link_to_genericname_idx_targetname_collectionID
One can list out all such branches with the method List_Links
"""
idx_field_name = "Link_" + generic_name + "_" + target_name
if idx_field_name not in self.fields:
msg = f"{idx_field_name} not found in the current collection"
raise FileNotFoundError(msg)
return self._events()[target_name]._apply_global_index(
self[idx_field_name]["index_Global"]
)
_copy_behaviors("NanoCollection", "edm4hep_nanocollection", behavior)
[docs]
@awkward.mixin_class(behavior)
class MomentumCandidate(vector.LorentzVector):
"""A Lorentz vector with charge
This mixin class requires the parent class to provide the items ``px``, ``py``, ``pz``, `E`, and ``charge``.
"""
[docs]
@awkward.mixin_class_method(np.add, {"MomentumCandidate"})
def add(self, other):
"""Add two candidates together elementwise using ``px``, ``py``, ``pz``, `E`, and ``charge`` components"""
return awkward.zip(
{
"px": self.px + other.px,
"py": self.py + other.py,
"pz": self.pz + other.pz,
"E": self.E + other.E,
"charge": self.charge + other.charge,
},
with_name="MomentumCandidate",
behavior=self.behavior,
)
[docs]
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using ``px``, ``py``, ``pz``, `E`, and ``charge`` components"""
return awkward.zip(
{
"px": awkward.sum(self.px, axis=axis),
"py": awkward.sum(self.py, axis=axis),
"pz": awkward.sum(self.pz, axis=axis),
"E": awkward.sum(self.E, axis=axis),
"charge": awkward.sum(self.charge, axis=axis),
},
with_name="MomentumCandidate",
behavior=self.behavior,
)
@property
def absolute_mass(self):
"""The absolute value of the mass"""
return np.sqrt(np.abs(self.mass2))
def __awkward_validation__(self):
if "charge" not in self.fields:
msg = f"{type(self).__name__} requires the 'charge' field"
raise ValueError(msg)
parent = super()
if hasattr(parent, "__awkward_validation__"):
parent.__awkward_validation__()
_copy_behaviors("LorentzVector", "MomentumCandidate", behavior)
MomentumCandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
MomentumCandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
MomentumCandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821
MomentumCandidateArray.MomentumClass = MomentumCandidateArray # noqa: F821
MomentumCandidateRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
MomentumCandidateRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
MomentumCandidateRecord.ProjectionClass4D = vector.LorentzVectorRecord # noqa: F821
MomentumCandidateRecord.MomentumClass = MomentumCandidateRecord # noqa: F821
#########################################################################################
# Selected Components #
#########################################################################################
[docs]
@awkward.mixin_class(behavior)
class TrackState(base.NanoCollection):
"""EDM4HEP Component: TrackState"""
_set_repr_name("TrackState")
TrackState.ProjectionClass2D = vector.TwoVectorArray
TrackState.ProjectionClass3D = vector.ThreeVectorArray
TrackState.ProjectionClass4D = TrackState
TrackState.MomentumClass = vector.LorentzVectorArray
TrackStateRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
TrackStateRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
TrackStateRecord.ProjectionClass4D = TrackStateRecord # noqa: F821
TrackStateRecord.MomentumClass = vector.LorentzVectorRecord # noqa: F821
[docs]
@awkward.mixin_class(behavior)
class Quantity(base.NanoCollection):
"""EDM4HEP Component: Quantity"""
_set_repr_name("Quantity")
[docs]
@awkward.mixin_class(behavior)
class covMatrix(base.NanoCollection):
"""EDM4HEP Component: covMatrix"""
_set_repr_name("covMatrix")
#########################################################################################
# Datatypes #
#########################################################################################
_set_repr_name("EventHeader")
[docs]
@awkward.mixin_class(behavior)
class MCParticle(MomentumCandidate, edm4hep_nanocollection):
"""EDM4HEP Datatype: MCParticle"""
_set_repr_name("MCParticle")
_copy_behaviors("MomentumCandidate", "MCParticle", behavior)
MCParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
MCParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
MCParticleArray.ProjectionClass4D = MCParticleArray # noqa: F821
MCParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
MCParticleRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
MCParticleRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
MCParticleRecord.ProjectionClass4D = MCParticleRecord # noqa: F821
MCParticleRecord.MomentumClass = vector.LorentzVectorRecord # noqa: F821
[docs]
@awkward.mixin_class(behavior)
class SimTrackerHit(edm4hep_nanocollection):
"""EDM4HEP Datatype: SimTrackerHit"""
_set_repr_name("SimTrackerHit")
[docs]
@awkward.mixin_class(behavior)
class CaloHitContribution(edm4hep_nanocollection):
"""EDM4HEP Datatype: CaloHitContribution"""
_set_repr_name("CaloHitContribution")
[docs]
@awkward.mixin_class(behavior)
class SimCalorimeterHit(edm4hep_nanocollection):
"""EDM4HEP Datatype: SimCalorimeterHit"""
_set_repr_name("SimCalorimeterHit")
[docs]
@awkward.mixin_class(behavior)
class RawCalorimeterHit(edm4hep_nanocollection):
"""EDM4HEP Datatype: RawCalorimeterHit"""
_set_repr_name("RawCalorimeterHit")
[docs]
@awkward.mixin_class(behavior)
class CalorimeterHit(edm4hep_nanocollection):
"""EDM4HEP Datatype: CalorimeterHit"""
_set_repr_name("CalorimeterHit")
[docs]
@awkward.mixin_class(behavior)
class ParticleID(edm4hep_nanocollection):
"""EDM4HEP Datatype: ParticleID"""
_set_repr_name("ParticleID")
[docs]
@awkward.mixin_class(behavior)
class Cluster(edm4hep_nanocollection):
"""EDM4HEP Datatype: Cluster"""
_set_repr_name("Cluster")
[docs]
@awkward.mixin_class(behavior)
class TrackerHit3D(edm4hep_nanocollection):
"""EDM4HEP Datatype: TrackerHit3D"""
_set_repr_name("TrackerHit3D")
[docs]
@awkward.mixin_class(behavior)
class TrackerHitPlane(edm4hep_nanocollection):
"""EDM4HEP Datatype: TrackerHitPlane"""
_set_repr_name("TrackerHitPlane")
[docs]
@awkward.mixin_class(behavior)
class RawTimeSeries(edm4hep_nanocollection):
"""EDM4HEP Datatype: RawTimeSeries"""
_set_repr_name("RawTimeSeries")
[docs]
@awkward.mixin_class(behavior)
class Track(edm4hep_nanocollection):
"""EDM4HEP Datatype: Track"""
_set_repr_name("Track")
[docs]
@awkward.mixin_class(behavior)
class Vertex(edm4hep_nanocollection):
"""EDM4HEP Datatype: Vertex"""
_set_repr_name("Vertex")
[docs]
@awkward.mixin_class(behavior)
class ReconstructedParticle(MomentumCandidate, edm4hep_nanocollection):
"""EDM4HEP Datatype: Reconstructed particle"""
_set_repr_name("ReconstructedParticle")
_copy_behaviors("MomentumCandidate", "ReconstructedParticle", behavior)
ReconstructedParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
ReconstructedParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
ReconstructedParticleArray.ProjectionClass4D = ReconstructedParticleArray # noqa: F821
ReconstructedParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
ReconstructedParticleRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
ReconstructedParticleRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
ReconstructedParticleRecord.ProjectionClass4D = ( # noqa: F821
ReconstructedParticleRecord # noqa: F821
)
ReconstructedParticleRecord.MomentumClass = vector.LorentzVectorRecord # noqa: F821
[docs]
@awkward.mixin_class(behavior)
class TimeSeries(edm4hep_nanocollection):
"""EDM4HEP Datatype: TimeSeries"""
_set_repr_name("TimeSeries")
[docs]
@awkward.mixin_class(behavior)
class RecDqdx(edm4hep_nanocollection):
"""EDM4HEP Datatype: RecDqdx"""
_set_repr_name("RecDqdx")
#########################################################################################
# Selected Links #
#########################################################################################
[docs]
@awkward.mixin_class(behavior)
class RecoMCParticleLink(edm4hep_nanocollection):
"""EDM4HEP Link: MCRecoParticleLink"""
_set_repr_name("RecoMCParticleLink")
RecoMCParticleLinkArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821
RecoMCParticleLinkArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821
RecoMCParticleLinkArray.ProjectionClass4D = RecoMCParticleLinkArray # noqa: F821
RecoMCParticleLinkArray.MomentumClass = vector.LorentzVectorArray # noqa: F821
RecoMCParticleLinkRecord.ProjectionClass2D = vector.TwoVectorRecord # noqa: F821
RecoMCParticleLinkRecord.ProjectionClass3D = vector.ThreeVectorRecord # noqa: F821
RecoMCParticleLinkRecord.ProjectionClass4D = RecoMCParticleLinkRecord # noqa: F821
RecoMCParticleLinkRecord.MomentumClass = vector.LorentzVectorRecord # noqa: F821
#########################################################################################
# Extras #
#########################################################################################
[docs]
@awkward.mixin_class(behavior)
class ObjectID(base.NanoCollection):
"""
Generic Object ID storage, pointing to another collection.
Usually has the <podio::ObjectID> type
"""
_set_repr_name("ObjectID")