make_root_from_grid_output.py

Go to the documentation of this file.

#!/usr/bin/env python

from __future__ import print_function
from array import array
from ROOT import *

import glob

def process_one_file(fn):
  with open(fn, 'r') as inf:
    for line in inf:
      line = line.rstrip('\n').rstrip(' ')
      values = line.split(' ')
      run[0] = int(values[0])
      subrun[0] = int(values[1])
      cycle[0] = int(values[2])
      evt[0] = int(values[3])
      subevt[0] = int(values[4])
      multivalues = values[5:]
      ntrks[0] = len(values)/11
      for i in range(ntrks[0]):
        dedxll[0] = float(multivalues[0*ntrks[0]+i])
        scatll[0] = float(multivalues[1*ntrks[0]+i])
        dedx10[0] = float(multivalues[2*ntrks[0]+i])
        dedx20[0] = float(multivalues[3*ntrks[0]+i])
        dedx30[0] = float(multivalues[4*ntrks[0]+i])
        dedx40[0] = float(multivalues[5*ntrks[0]+i])
        # if dedx > 30 => outliers, which affect BDT results dramatically
        # squeeze outliers to some resonable range
        dedx10_squeezed[0] = dedx10[0]
        if dedx10_squeezed[0] > 30.:
          dedx10_squeezed[0] = 30
        dedx20_squeezed[0] = dedx20[0]
        if dedx20_squeezed[0] > 30.:
          dedx20_squeezed[0] = 30
        dedx30_squeezed[0] = dedx30[0]
        if dedx30_squeezed[0] > 30.:
          dedx30_squeezed[0] = 30
        dedx40_squeezed[0] = dedx40[0]
        if dedx40_squeezed[0] > 30.:
          dedx40_squeezed[0] = 30
        pdg[0] = int(float(multivalues[6*ntrks[0]+i]))
        trklen[0] = float(multivalues[7*ntrks[0]+i])
        measfrac[0] = float(multivalues[8*ntrks[0]+i])
        remid[0] = float(multivalues[9*ntrks[0]+i])
        remcont[0] = int(multivalues[10*ntrks[0]+i])
        tr.Fill()
        

file_list = [name for name in glob.glob('grid_output/*.txt')]
#~ for name in glob.glob('grid_output/*.txt'):
#~ for name in glob.glob('grid_output/nd_rhc_numubarcc_selection_stride2_offset0.1_of_200.txt'):
  #~ file_list.append(name)

# make the tree
outf = TFile('training_sample.root', 'recreate')
tr = TTree('tr', 'tree of training sample')

# tree leaves
run = array('i', [0])
subrun = array('i', [0])
cycle = array('i', [0])
evt = array('i', [0])
subevt = array('i', [0])
ntrks = array('i', [0])
dedxll = array('f', [0.])
scatll = array('f', [0.])
dedx10 = array('f', [0.])
dedx10_squeezed = array('f', [0.])
dedx20 = array('f', [0.])
dedx20_squeezed = array('f', [0.])
dedx30 = array('f', [0.])
dedx30_squeezed = array('f', [0.])
dedx40 = array('f', [0.])
dedx40_squeezed = array('f', [0.])
pdg = array('i', [0])
trklen = array('f', [0.])
measfrac = array('f', [0.])
remid = array('f', [0.])
remcont = array('i', [0])
tr.Branch('run', run, 'run/I')
tr.Branch('subrun', subrun, 'subrun/I')
tr.Branch('cycle', cycle, 'cycle/I')
tr.Branch('evt', evt, 'evt/I')
tr.Branch('subevt', subevt, 'subevt/I')
tr.Branch('ntrks', ntrks, 'ntrks/I')
tr.Branch('dedxll', dedxll, 'dedxll/F')
tr.Branch('scatll', scatll, 'scatll/F')
tr.Branch('dedx10', dedx10, 'dedx10/F')
tr.Branch('dedx20', dedx20, 'dedx20/F')
tr.Branch('dedx30', dedx30, 'dedx30/F')
tr.Branch('dedx40', dedx40, 'dedx40/F')
tr.Branch('dedx10_squeezed', dedx10_squeezed, 'dedx10_squeezed/F')
tr.Branch('dedx20_squeezed', dedx20_squeezed, 'dedx20_squeezed/F')
tr.Branch('dedx30_squeezed', dedx30_squeezed, 'dedx30_squeezed/F')
tr.Branch('dedx40_squeezed', dedx40_squeezed, 'dedx40_squeezed/F')
tr.Branch('pdg', pdg, 'pdg/I')
tr.Branch('trklen', trklen, 'trklen/F')
tr.Branch('measfrac', measfrac, 'measfrac/F')
tr.Branch('remid', remid, 'remid/F')
tr.Branch('remcont', remcont, 'remcont/I')

for f in file_list:
  process_one_file(f)
outf.Write()