diff --git a/scribblers/lhetop.py b/scribblers/lhetop.py index 74ee0cc..e9db19f 100755 --- a/scribblers/lhetop.py +++ b/scribblers/lhetop.py @@ -51,3 +51,28 @@ def end(self): self.out = None ##__________________________________________________________________|| +class NLheLeptons(object): + + def __repr__(self): + name_value_pairs = () + return '{}({})'.format( + self.__class__.__name__, + ', '.join(['{}={!r}'.format(n, v) for n, v in name_value_pairs]), + ) + + def begin(self, event): + self.out = [ ] + self._attach_to_event(event) + + def _attach_to_event(self, event): + event.nLheLeptons = self.out + + def event(self, event): + self._attach_to_event(event) + + status = np.array(event.GenPart_status) + abspdgid = np.absolute(np.array(event.GenPart_pdgId)) + b = (status == 23) & ( (11 == abspdgid) | (13 == abspdgid) | (15 == abspdgid)) + self.out[:] = [abspdgid[b].size] + +##__________________________________________________________________||