File indexing completed on 2024-04-06 12:23:45
0001
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
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
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
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
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
0078 zeroFill(otherObj, br, obj.GetListOfBranches().FindObject(br))
0079 for br in additionalBranches:
0080
0081 branchNames.add(br)
0082 zeroFill(obj, br, otherObj.GetListOfBranches().FindObject(br))
0083
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
0093 zeroFill(otherObj, br, obj.GetListOfBranches(
0094 ).FindObject(br), allowNonBool=True)
0095 for br in additionalBranches:
0096
0097 branchNames.add(br)
0098 zeroFill(obj, br, otherObj.GetListOfBranches(
0099 ).FindObject(br), allowNonBool=True)
0100
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()))