makeTrainCVSamples.py

Go to the documentation of this file.

import os
import argparse
import subprocess 


parser = argparse.ArgumentParser(description='Dump a tree to text.')
parser.add_argument('--fluxswapdir', '-f', type=str, required=True,
      help="""Directory with fluxswap files.""")
parser.add_argument('--nonswapdir', '-n', type=str, required=True,
      help="""Directory with nonswap files.""")

parser.add_argument('--ntrain', type=int, required=True,
      help="""Number of pairs of fluxswap and nonswap 
              files for training sample.""")

parser.add_argument('--ncv', type=int, required=True,
      help="""Number of pairs of fluxswap and nonswap files for CV smaple .""")

parser.add_argument('--artext', type=str, default=".cvn.root",
      help="""Extension of art-ROOT files.""")

parser.add_argument('--histext', type=str, default=".cvnevents.root",
      help="""Extension of hist-ROOT files with CVN events for training.""")


parser.add_argument('--haddpath', type=str,
      help="""Location to place hadd files for training and lists for cv sample""")


args = parser.parse_args()





class FilePair:
  def __init__(self, artPath, histPath):
    self.artPath = artPath
    self.histPath = histPath 

  def __repr__(self):
    return "File pair:\n %s\n %s\n" %(self.artPath, self.histPath) 

  def __lt__(self, other):
    return self.artPath < other.artPath 

def haddFiles(files, outputName):
  cmdList = ["hadd", "-f", outputName] + files 
  
  subprocess.call(cmdList)
  

dirs = {}
dirs["fluxswap"] = args.fluxswapdir
dirs["nonswap"]  = args.nonswapdir


filePairs = {}

for key, dir in dirs.items():
  fileList = os.listdir(dir)
  artFiles = [f for f in fileList if f.endswith(args.artext)]
  histFiles = [f for f in fileList if f.endswith(args.histext)]
  filePairs[key] = []
  for artFile in artFiles: 
    histFileMatches = [f for f in histFiles \
            if  f.replace(args.histext, "") == artFile.replace(args.artext, "")]
    if len(histFileMatches) > 1:
      raise Exception("too many hist file matches for this file:" + artFile)
    elif len(histFileMatches) == 0:
      print "Warning: failed to find hist match for this art file: ", artFile
      continue
    filePairs[key].append(FilePair(os.path.join(dir, artFile),
                                   os.path.join(dir, histFileMatches[0])))
  filePairs[key].sort()




if args.haddpath:
  trainFilePairs = filePairs["fluxswap"][0:args.ntrain] + \
                   filePairs["nonswap" ][0:args.ntrain]
  haddName = os.path.join(args.haddpath, "train_%s_hadd.root" %(args.ntrain))
  trainListPath = os.path.join(args.haddpath, "train_%s.txt" %(args.ntrain))
  trainList = [f.histPath for f in trainFilePairs]
  trainArtList = [f.artPath for f in trainFilePairs]
  haddFiles(trainList, haddName)
  open(trainListPath, 'w').write("\n".join(trainArtList ))

  cvFilePairs = filePairs["fluxswap"][-args.ncv:] + \
                   filePairs["nonswap" ][-args.ncv:]
  cvArtList = [f.artPath for f in cvFilePairs]
  cvListPath = os.path.join(args.haddpath, "cv_%s_train_%s.txt" %(args.ncv, 
                                                                  args.ntrain))
  open(cvListPath, 'w').write("\n".join(cvArtList))