NanoAOD/scripts/haddnano.py

0001 #!/bin/env python
0002 import ROOT
0003 import numpy
0004 import sys
0005
0006 if len(sys.argv) < 3:
0007     print("Syntax: haddnano.py out.root input1.root input2.root ...")
0008 ofname = sys.argv[1]
0009 files = sys.argv[2:]
0010
0011
0012 def zeroFill(tree, brName, brObj, allowNonBool=False):
0013     # typename: (numpy type code, root type code)
0014     branch_type_dict = {'Bool_t': ('?', 'O'), 'Float_t': ('f4', 'F'), 'UInt_t': (
0015         'u4', 'i'), 'Long64_t': ('i8', 'L'), 'Double_t': ('f8', 'D')}
0016     brType = brObj.GetLeaf(brName).GetTypeName()
0017     if (not allowNonBool) and (brType != "Bool_t"):
0018         print(("Did not expect to back fill non-boolean branches", tree, brName, brObj.GetLeaf(br).GetTypeName()))
0019     else:
0020         if brType not in branch_type_dict:
0021             raise RuntimeError('Impossible to backfill branch of type %s' % brType)
0022         buff = numpy.zeros(1, dtype=numpy.dtype(branch_type_dict[brType][0]))
0023         b = tree.Branch(brName, buff, brName + "/" +
0024                         branch_type_dict[brType][1])
0025         # be sure we do not trigger flushing
0026         b.SetBasketSize(tree.GetEntries() * 2)
0027         for x in range(0, tree.GetEntries()):
0028             b.Fill()
0029         b.ResetAddress()
0030
0031
0032 fileHandles = []
0033 goFast = True
0034 for fn in files:
0035     print("Adding file" + str(fn))
0036     fileHandles.append(ROOT.TFile.Open(fn))
0037     if fileHandles[-1].GetCompressionSettings() != fileHandles[0].GetCompressionSettings():
0038         goFast = False
0039         print("Disabling fast merging as inputs have different compressions")
0040 of = ROOT.TFile(ofname, "recreate")
0041 if goFast:
0042     of.SetCompressionSettings(fileHandles[0].GetCompressionSettings())
0043 of.cd()
0044
0045 for e in fileHandles[0].GetListOfKeys():
0046     name = e.GetName()
0047     print("Merging" + str(name))
0048     obj = e.ReadObj()
0049     cl = ROOT.TClass.GetClass(e.GetClassName())
0050     inputs = ROOT.TList()
0051     isTree = obj.IsA().InheritsFrom(ROOT.TTree.Class())
0052     if isTree:
0053         obj = obj.CloneTree(-1, "fast" if goFast else "")
0054         branchNames = set([x.GetName() for x in obj.GetListOfBranches()])
0055     for fh in fileHandles[1:]:
0056         otherObj = fh.GetListOfKeys().FindObject(name).ReadObj()
0057         inputs.Add(otherObj)
0058         if isTree and obj.GetName() == 'Events':
0059             otherObj.SetAutoFlush(0)
0060             otherBranches = set([x.GetName()
0061                                  for x in otherObj.GetListOfBranches()])
0062             missingBranches = list(branchNames - otherBranches)
0063             additionalBranches = list(otherBranches - branchNames)
0064             print("missing: " + str(missingBranches) + "\n Additional:" + str(additionalBranches))
0065             for br in missingBranches:
0066                 # fill "Other"
0067                 zeroFill(otherObj, br, obj.GetListOfBranches().FindObject(br))
0068             for br in additionalBranches:
0069                 # fill main
0070                 branchNames.add(br)
0071                 zeroFill(obj, br, otherObj.GetListOfBranches().FindObject(br))
0072             # merge immediately for trees
0073         if isTree and obj.GetName() == 'Runs':
0074             otherObj.SetAutoFlush(0)
0075             otherBranches = set([x.GetName()
0076                                  for x in otherObj.GetListOfBranches()])
0077             missingBranches = list(branchNames - otherBranches)
0078             additionalBranches = list(otherBranches - branchNames)
0079             print("missing: " + str(missingBranches) + "\n Additional:" + str(additionalBranches))
0080             for br in missingBranches:
0081                 # fill "Other"
0082                 zeroFill(otherObj, br, obj.GetListOfBranches(
0083                 ).FindObject(br), allowNonBool=True)
0084             for br in additionalBranches:
0085                 # fill main
0086                 branchNames.add(br)
0087                 zeroFill(obj, br, otherObj.GetListOfBranches(
0088                 ).FindObject(br), allowNonBool=True)
0089             # merge immediately for trees
0090         if isTree:
0091             obj.Merge(inputs, "fast" if goFast else "")
0092             inputs.Clear()
0093
0094     if isTree:
0095         obj.Write()
0096     elif obj.IsA().InheritsFrom(ROOT.TH1.Class()):
0097         obj.Merge(inputs)
0098         obj.Write()
0099     elif obj.IsA().InheritsFrom(ROOT.TObjString.Class()):
0100         for st in inputs:
0101             if st.GetString() != obj.GetString():
0102                 print("Strings are not matching")
0103         obj.Write()
0104     else:
0105         print("Cannot handle " + str(obj.IsA().GetName()))