diff --git a/graphicle/base.py b/graphicle/base.py index 6e8db52..8318413 100644 --- a/graphicle/base.py +++ b/graphicle/base.py @@ -45,6 +45,35 @@ DataType = ty.TypeVar("DataType") +class LheEventInterface(ty.Protocol): + """Interface for a generic Les Houches event. + + :group: base + + .. versionadded:: 0.4.0 + """ + + @property + def pdg(self) -> IntVector: + ... + + @property + def pmu(self) -> AnyVector: + ... + + @property + def color(self) -> AnyVector: + ... + + @property + def helicity(self) -> HalfIntVector: + ... + + @property + def status(self) -> HalfIntVector: + ... + + class EventInterface(ty.Protocol): """Defines the interface for a generic event object expected by graphicle's routines. Attributes are stored as numpy arrays, with diff --git a/graphicle/data.py b/graphicle/data.py index eb7399f..492a466 100644 --- a/graphicle/data.py +++ b/graphicle/data.py @@ -2296,6 +2296,36 @@ def copy(self) -> "ParticleSet": """Copies the underlying data into a new ParticleSet instance.""" return _composite_copy(self) + @classmethod + def from_lhe_event(cls, event: base.LheEventInterface) -> "ParticleSet": + """Creates a ParticleSet instance directly from a data structure + holding LHE data. This is useful when you want to study the hard + process, without needing to shower or hadronize. + + .. versionadded:: 0.4.0 + + Parameters + ---------- + event : LheEventInterface + A data structure with attributes "pdg", "pmu", "color", + "helicity", and "status". These attributes provide access to + numpy arrays, which hold the underlying Les Houches event + data. + + Returns + ------- + ParticleSet + A composite object, wrapping the data provided in Graphicle + objects, and providing a unified interface to them. + """ + props = ("pdg", "pmu", "color", "helicity", "status") + pcls = cls.from_numpy(**{name: getattr(event, name) for name in props}) + pcls.final = MaskArray(np.zeros(len(pcls), dtype=np.bool_)) + status_map = {-1: -21, 1: -23, -2: -22, 2: -22, 3: -22, -9: -12} + for old_code, new_code in status_map.items(): + pcls.status.data[pcls.status.data == old_code] = new_code + return pcls + @classmethod def from_numpy( cls,