TagAndProbe/python/cropTnPTrees.py

0001 #!/usr/bin/env python
0002 import ROOT
0003 ROOT.gROOT.SetBatch(True)
0004
0005 import sys
0006
0007 from argparse import ArgumentParser
0008 parser = ArgumentParser()
0009 parser.add_argument("inputFile", type=str)
0010 parser.add_argument("fraction", type=float)
0011 parser.add_argument("outputFile", type=str)
0012 options = parser.parse_args()
0013
0014 frac = options.fraction
0015 input  = ROOT.TFile(options.inputFile)
0016 output = ROOT.TFile(options.outputFile, "RECREATE")
0017 for k in input.GetListOfKeys():
0018     print(k.GetName(), k.GetClassName())
0019     if k.GetClassName() == "TDirectoryFile":
0020         print("  processing directory ",k.GetName())
0021         din  = input.Get(k.GetName())
0022         dout = output.mkdir(k.GetName())
0023         for i in din.GetListOfKeys():
0024             if i.GetClassName() == "TTree":
0025                 src = din.Get(i.GetName()) #i.ReadObj(); # ReadObj doesn't work!!!
0026                 newEntries = int(src.GetEntries()*frac)
0027                 print("    cropped TTree",i.GetName(),", original entries",src.GetEntries(), ", new entries",newEntries)
0028                 cloned = src.CloneTree(newEntries)
0029                 dout.WriteTObject(cloned, i.GetName())
0030             elif i.GetClassName() != "TDirectory":
0031                 dout.WriteTObject(i.ReadObj(), i.GetName())
0032                 print("    copied ",i.GetClassName(),i.GetName())
0033
0034