Merge pull request cbernet#44 from selvaggi/master

added jet substructure example

analyzers/examples/jetsubstructure/TreeProducer.py

@@ -0,0 +1,53 @@
+from heppy.framework.analyzer import Analyzer
+from heppy.statistics.tree import Tree
+from heppy.analyzers.ntuple import *
+from heppy.particles.tlv.resonance import Resonance2 as Resonance
+
+from ROOT import TFile
+
+class TreeProducer(Analyzer):
+
+    def beginLoop(self, setup):
+        super(TreeProducer, self).beginLoop(setup)
+        self.rootfile = TFile('/'.join([self.dirName,
+                                        'tree.root']),
+                              'recreate')
+        self.tree = Tree( 'events', '')
+        self.tree.var('tau1', float)
+        self.tree.var('tau2', float)
+        self.tree.var('tau3', float)
+        self.tree.var('tau32', float)
+        self.tree.var('tau31', float)
+        self.tree.var('tau21', float)
+
+        bookParticle(self.tree, 'Jet')
+        bookParticle(self.tree, 'softDroppedJet')
+        bookParticle(self.tree, 'leadingSoftDroppedSubJet')
+        bookParticle(self.tree, 'trailingSoftDroppedSubJet')
+
+    def process(self, event):
+        self.tree.reset()
+        jets = getattr(event, self.cfg_ana.fatjets)
+
+        # store leading (fat) jet observables
+        if len(jets) > 0 and len(jets[0].subjetsSoftDrop) > 2:
+            self.tree.fill('tau1' , jets[0].tau1 )
+            self.tree.fill('tau2' , jets[0].tau2 )
+            self.tree.fill('tau3' , jets[0].tau3 )
+            self.tree.fill('tau31' , jets[0].tau3/jets[0].tau1 )
+            self.tree.fill('tau32' , jets[0].tau3/jets[0].tau2 )
+            self.tree.fill('tau21' , jets[0].tau2/jets[0].tau1 )
+
+            fillParticle(self.tree, 'Jet', jets[0])
+
+            # first subjet entry is the cleaned jet itself
+            fillParticle(self.tree, 'softDroppedJet', jets[0].subjetsSoftDrop[0])
+            fillParticle(self.tree, 'leadingSoftDroppedSubJet', jets[0].subjetsSoftDrop[1])
+            fillParticle(self.tree, 'trailingSoftDroppedSubJet', jets[0].subjetsSoftDrop[2])
+
+            self.tree.tree.Fill()
+
+    def write(self, setup):
+        self.rootfile.Write()
+        self.rootfile.Close()
+

analyzers/fcc/Reader.py

@@ -98,7 +98,6 @@ def get_collection(class_object, coll_label, sort=True):
             if weightcoll:
                 event.weight = weightcoll[0].value()
 
-
         if hasattr(self.cfg_ana, 'gen_particles'):
             get_collection(Particle, 'gen_particles')
 
@@ -116,7 +115,7 @@ def get_collection(class_object, coll_label, sort=True):
             if hasattr(self.cfg_ana, 'bTags'):
                 for bjet in store.get(self.cfg_ana.bTags):
                     jets[Jet(bjet.jet())].tags['bf'] = bjet.tag()
-        
+
             if hasattr(self.cfg_ana, 'cTags'):
                 for cjet in store.get(self.cfg_ana.cTags):
                     jets[Jet(cjet.jet())].tags['cf'] = cjet.tag()
@@ -125,14 +124,65 @@ def get_collection(class_object, coll_label, sort=True):
                 for taujet in store.get(self.cfg_ana.tauTags):
                     jets[Jet(taujet.jet())].tags['tauf'] = taujet.tag()
 
+        ############################
+        #  jet substructure stuff  #
+        ############################
+
+        fatjetcoll = get_collection(Jet, 'fatjets')
+        if fatjetcoll:
+            fatjets = dict()
+            for jet in fatjetcoll:
+                fatjets[jet] = jet
+            # store N-subjettiness up to 3
+            if hasattr(self.cfg_ana, 'jetsOneSubJettiness'):
+                for tjet in store.get(self.cfg_ana.jetsOneSubJettiness):
+                    fatjets[Jet(tjet.jet())].tau1 = tjet.tag()
+
+            if hasattr(self.cfg_ana, 'jetsTwoSubJettiness'):
+                for tjet in store.get(self.cfg_ana.jetsTwoSubJettiness):
+                    fatjets[Jet(tjet.jet())].tau2 = tjet.tag()
+            if hasattr(self.cfg_ana, 'jetsThreeSubJettiness'):
+                for tjet in store.get(self.cfg_ana.jetsThreeSubJettiness):
+                    fatjets[Jet(tjet.jet())].tau3 = tjet.tag()
+
+            from collections import defaultdict
+            # store subjets according to various algorithms
+            # the first entry of subjets list is the "cleaned" fastjet itself
+
+            # trimming
+            if hasattr(self.cfg_ana, 'subjetsTrimming') and hasattr(self.cfg_ana, 'subjetsTrimmingTagged'):
+                 relations = defaultdict(list)
+                 for tjet in store.get(self.cfg_ana.subjetsTrimmingTagged):
+                      for i in range(tjet.subjets_size()):
+                          relations[Jet(tjet.jet())].append(Jet(tjet.subjets(i)))
+                 for fatjet, subjets in relations.items():
+                     fatjets[fatjet].subjetsTrimming = subjets
+
+            # pruning
+            if hasattr(self.cfg_ana, 'subjetsPruning') and hasattr(self.cfg_ana, 'subjetsPruningTagged'):
+                 relations = defaultdict(list)
+                 for tjet in store.get(self.cfg_ana.subjetsPruningTagged):
+                      for i in range(tjet.subjets_size()):
+                          relations[Jet(tjet.jet())].append(Jet(tjet.subjets(i)))
+                 for fatjet, subjets in relations.items():
+                     fatjets[fatjet].subjetsPruning = subjets
+
+            # soft drop
+            if hasattr(self.cfg_ana, 'subjetsSoftDrop') and hasattr(self.cfg_ana, 'subjetsSoftDropTagged'):
+                 relations = defaultdict(list)
+                 for tjet in store.get(self.cfg_ana.subjetsSoftDropTagged):
+                      for i in range(tjet.subjets_size()):
+                          relations[Jet(tjet.jet())].append(Jet(tjet.subjets(i)))
+                 for fatjet, subjets in relations.items():
+                     fatjets[fatjet].subjetsSoftDrop = subjets
+
 
         class Iso(object):
             def __init__(self):
                 self.sumpt=-9999
                 self.sume=-9999
                 self.num=-9999
 
-
         electrons = dict()
         if hasattr(self.cfg_ana, 'electrons'):
             event.electrons = map(Particle, store.get(self.cfg_ana.electrons))
@@ -149,7 +199,6 @@ def __init__(self):
                     if ele.sim() and ele.rec():
                         electrons[Particle(ele.rec())].gen = Particle(ele.sim())
 
-
         muons = dict()
         if hasattr(self.cfg_ana, 'muons'):
             event.muons = map(Particle, store.get(self.cfg_ana.muons))

doc/Heppy_-_Parallel_Processing.md

@@ -178,4 +178,10 @@ or
 
 ```
 nohup ./batchScript.sh &
-```
+```
+
+Alternatively, if you want to resubmit all failed jobs on the batch queue, do something like:
+
+```
+heppy_check.py Outdir/*Chunk* -b 'bsub -q 1nh'
+```

test/test_substructure_cfg.py

@@ -0,0 +1,89 @@
+import os
+import copy
+import heppy.framework.config as cfg
+
+import logging
+# next 2 lines necessary to deal with reimports from ipython
+logging.shutdown()
+reload(logging)
+logging.basicConfig(level=logging.WARNING)
+
+comp = cfg.Component(
+    'example',
+    files = ['FCCDelphesOutput.root']
+   )
+selectedComponents = [comp]
+
+from heppy.analyzers.fcc.Reader import Reader
+source = cfg.Analyzer(
+    Reader,
+
+    fatjets = 'fatjets',
+    jetsOneSubJettiness = 'jetsOneSubJettiness', 
+    jetsTwoSubJettiness = 'jetsTwoSubJettiness', 
+    jetsThreeSubJettiness = 'jetsThreeSubJettiness', 
+    subjetsTrimmingTagged = 'subjetsTrimmingTagged', 
+    subjetsTrimming = 'subjetsTrimming', 
+    subjetsPruningTagged = 'subjetsPruningTagged', 
+    subjetsPruning = 'subjetsPruning', 
+    subjetsSoftDropTagged = 'subjetsSoftDropTagged', 
+    subjetsSoftDrop = 'subjetsSoftDrop', 
+
+)
+
+
+from ROOT import gSystem
+gSystem.Load("libdatamodelDict")
+from EventStore import EventStore as Events
+
+#############################
+##   Reco Level Analysis   ##
+#############################
+
+from heppy.analyzers.Selector import Selector
+# select jet above 400 GeV
+fatjets_400 = cfg.Analyzer(
+    Selector,
+    'fatjets_400',
+    output = 'fatjets_400',
+    input_objects = 'fatjets',
+    filter_func = lambda jet: jet.pt()>400.
+)
+
+# produce flat root tree containing jet substructure information
+from heppy.analyzers.examples.jetsubstructure.TreeProducer import TreeProducer
+tree = cfg.Analyzer(
+    TreeProducer,
+    fatjets = 'fatjets_400',
+)
+
+
+# definition of a sequence of analyzers,
+# the analyzers will process each event in this order
+sequence = cfg.Sequence( [
+    source,
+    fatjets_400,
+    tree,
+    ] )
+
+
+config = cfg.Config(
+    components = selectedComponents,
+    sequence = sequence,
+    services = [],
+    events_class = Events
+)
+
+if __name__ == '__main__':
+    import sys
+    from heppy.framework.looper import Looper
+
+    def next():
+        loop.process(loop.iEvent+1)
+
+    loop = Looper( 'looper', config,
+                   nEvents=100,
+                   nPrint=0,
+                   timeReport=True)
+    loop.process(6)
+    print loop.event