NanoAOD/scripts/haddnano.py

0001 #!/bin/env python3
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         if isTree and obj.GetName() == 'Events' and obj.GetEntries() == 0 :
0057             # Zero-events first file. Skip to avoid messing up branches.
0058             print(" 'Events' tree contsins no events; skipping")
0059             obj = fh.GetListOfKeys().FindObject(name).ReadObj()
0060             obj = obj.CloneTree(-1, "fast" if goFast else "")
0061             branchNames = set([x.GetName() for x in obj.GetListOfBranches()])
0062             continue
0063         otherObj = fh.GetListOfKeys().FindObject(name).ReadObj()
0064         if isTree and obj.GetName() == 'Events' and otherObj.GetEntries() == 0 :
0065             # Zero-events file; skip
0066             print(" 'Events' tree contains no events; skipping")
0067             continue
0068         inputs.Add(otherObj)
0069         if isTree and obj.GetName() == 'Events':
0070             otherObj.SetAutoFlush(0)
0071             otherBranches = set([x.GetName()
0072                                  for x in otherObj.GetListOfBranches()])
0073             missingBranches = list(branchNames - otherBranches)
0074             additionalBranches = list(otherBranches - branchNames)
0075             print("missing: " + str(missingBranches) + "\n Additional: " + str(additionalBranches))
0076             for br in missingBranches:
0077                 # fill "Other"
0078                 zeroFill(otherObj, br, obj.GetListOfBranches().FindObject(br))
0079             for br in additionalBranches:
0080                 # fill main
0081                 branchNames.add(br)
0082                 zeroFill(obj, br, otherObj.GetListOfBranches().FindObject(br))
0083             # merge immediately for trees
0084         if isTree and obj.GetName() == 'Runs':
0085             otherObj.SetAutoFlush(0)
0086             otherBranches = set([x.GetName()
0087                                  for x in otherObj.GetListOfBranches()])
0088             missingBranches = list(branchNames - otherBranches)
0089             additionalBranches = list(otherBranches - branchNames)
0090             print("missing: " + str(missingBranches) + "\n Additional: " + str(additionalBranches))
0091             for br in missingBranches:
0092                 # fill "Other"
0093                 zeroFill(otherObj, br, obj.GetListOfBranches(
0094                 ).FindObject(br), allowNonBool=True)
0095             for br in additionalBranches:
0096                 # fill main
0097                 branchNames.add(br)
0098                 zeroFill(obj, br, otherObj.GetListOfBranches(
0099                 ).FindObject(br), allowNonBool=True)
0100             # merge immediately for trees
0101         if isTree:
0102             obj.Merge(inputs, "fast" if goFast else "")
0103             inputs.Clear()
0104
0105     if isTree:
0106         obj.Write()
0107     elif obj.IsA().InheritsFrom(ROOT.TH1.Class()):
0108         obj.Merge(inputs)
0109         obj.Write()
0110     elif obj.IsA().InheritsFrom(ROOT.TObjString.Class()):
0111         for st in inputs:
0112             if st.GetString() != obj.GetString():
0113                 print("Strings are not matching")
0114         obj.Write()
0115     elif obj.IsA().InheritsFrom(ROOT.THnSparse.Class()):
0116         obj.Merge(inputs)
0117         obj.Write()
0118     else:
0119         print("Cannot handle " + str(obj.IsA().GetName()))