Configuration/scripts/makeBNTreePlot.py

#!/usr/bin/env python

# Must specify options with -l argument, and condor directory with -c argument, e.g.:  
# > makeBNTreePlot.py -l sampleBNTreePlotConfig.py -c myCondorDir

# Additional arugments specify the running mode:
# -D <dataset>:  Run on a single dataset (typically for testing)
# -C:  Submit jobs to run on condor over all datasets
# BNTreeUseScript=True (set in sampleBNTreePlotConfig.py):  run BNTreeScript root macro (also set in sampleBNTreePlotConfig.py),
#    which must take as arguments the condor directory, dataset, and channel 


import sys
import os
import re
from optparse import OptionParser
from array import *
from decimal import *
import subprocess

from OSUT3Analysis.Configuration.configurationOptions import *
from OSUT3Analysis.Configuration.processingUtilities import *

from ROOT import TChain, TCut, TDirectory, TFile, TH1D, TH2D, TStopwatch, TTree, gROOT  

gROOT.SetBatch(True)  # This is to prevent pop-up graphical windows


def MakeCondorSubmitFile(arguments, dataset):
    outdir = "condor/"+arguments.condorDir+"/"+dataset+"/"  
    p = subprocess.Popen(["which", "makeBNTreePlot.py"], stdout=subprocess.PIPE)
    executable = p.communicate()[0]  # read the stdout of the command
    workdir = os.getcwd() + "/" 

    # Count the number of hist*root files in the specified directory 
    p = subprocess.Popen("ls " + workdir + outdir + "hist*root | wc", shell=True, stdout=subprocess.PIPE)  
    out = p.communicate()[0]
    outSplit = out.split()
    totalJobs = int(outSplit[0])  # The first element of the list is the number of files

    ## print "Debug: outSplit = "
    ## print outSplit
    ## out, err = p.communicate()
    ## print "Debug: Size of out:"
    ## print len(out)
    ## print "out = "
    ## print out
    ## print "err = "
    ## print err 
    ## print "Debug:  p.communicate()[0] = " + p.communicate()[0]
    ## totalJobs = len(out)
    out = open(outdir+"/condorBNTree.sub", "w")
    out.write("Executable              = " + executable + " \n")
    out.write("Universe                = vanilla \n")
    out.write("Getenv                  = True \n")
    argCondorProcess = ""  
    if arguments.splitCondorJobs: 
        argCondorProcess = " -p $(Process) " 
    out.write("Arguments               = -D " + dataset + " -l " + arguments.localConfig + " -c " + arguments.condorDir + argCondorProcess + " \n")
    out.write("Output                  = " + workdir + outdir + "condorBNTree_$(Process).out \n")
    out.write("Error                   = " + workdir + outdir + "condorBNTree_$(Process).err \n")
    out.write("Log                     = " + workdir + outdir + "condorBNTree_$(Process).log \n")
    out.write("+IsLocalJob             = true \n")
    out.write("Rank                    = TARGET.IsLocalSlot \n")
    if arguments.splitCondorJobs: 
        out.write("Queue " + str(totalJobs) + " \n")
    else: 
        out.write("Queue 1 \n")
    out.close()  


def RunOnCondor(arguments, split_datasets):  
    print "Running jobs on condor, instead of interactively."
    ## print "Found condorDir = %s" % (condorDir)
    ## print "Found split_datasets = "
    ## print split_datasets

    for dataset in split_datasets:  
        MakeCondorSubmitFile(arguments, dataset)  
        cmd = "condor_submit condor/"+arguments.condorDir+"/"+dataset+"/condorBNTree.sub"
        os.system(cmd)
        print "Submitting job: %s " % cmd

    print "Once condor jobs have finished, merge the composite datasets and then make plots with:"  
    print "    makeBNTreePlot.py -q -l " + arguments.localConfig + " -c " + arguments.condorDir  
    print "    makePlots.py         -l " + arguments.localConfig + " -c " + arguments.condorDir  
        
    return


watch  = TStopwatch()
watch1 = TStopwatch()


parser = OptionParser()
parser = set_commandline_arguments(parser)

###  Only used by makeBNTreePlot.py (maybe move to another file?)        
parser.remove_option("-p")
parser.add_option("-D", "--dataset", dest="datasetName",         
                  help="Name of dataset (overrides value from local configuration file)")        
parser.add_option("-C", "--runOnCondor",  action="store_true", dest="runOnCondor", default=False,        
                  help="Run on condor instead of interactively")         
parser.add_option("-S", "--splitCondorJobs", action="store_true", dest="splitCondorJobs", default=False,                 
                  help="Split condor jobs to have one for each file, rather than one for each dataset")  
parser.add_option("-p", "--condorProcessNum", dest="condorProcessNum", default=-1,       
                  help="Specify which condor process to run (default is to run over all).")      


(arguments, args) = parser.parse_args()
 

if not arguments.localConfig:
    sys.exit(" You must specify a localOptions.py file with -l")
if arguments.localConfig:  
    sys.path.append(os.getcwd())
    exec("from " + arguments.localConfig.rstrip('.py') + " import *")
if not arguments.condorDir:
    sys.exit(" You must specify a condor directory  with -c")
if arguments.condorDir:
    condor_dir = "condor/%s" % arguments.condorDir

if arguments.datasetName:  # If datasetName is specified on command line, then override the value from localConfig  
    datasets = [
        arguments.datasetName,
    ]

#save a list of composite datasets 
composite_datasets = get_composite_datasets(datasets, composite_dataset_definitions)
#save a list of datasets with composite datasets split up 
split_datasets = split_composite_datasets(datasets, composite_dataset_definitions)

if arguments.quickMerge and arguments.runOnCondor:
    print "Cannot use -q (--quickMerge) and -C (--runOnCondor) options simultaneously.  Please choose one or the other."
    exit()  

if arguments.runOnCondor:
    RunOnCondor(arguments, split_datasets)
    exit()    

#write new histogram to dataset
if not arguments.quickMerge: 
    for dataset in split_datasets:
        if BNTreeUseScript:
            chainName = "OSUAnalysis/" + BNTreeChannel + "/BNTree_" + BNTreeChannel
            command = "root -l -b -q '" + BNTreeScript + "+(\"" + condor_dir + "\",\"" + dataset + "\",\"" + chainName + "\"," + str(arguments.condorProcessNum) + ")'"  
            print "About to execute command:  " + command  
            os.system(command)
        else: 
            for hist in input_histograms:
                #chain trees together
                ch = TChain("OSUAnalysis/"+hist['channel']+"/BNTree_"+hist['channel'])
                ch.Add(condor_dir + "/" + dataset + "/hist_*.root")
                print ("Looping over chain with # entries = %f; split time = " % ch.GetEntries()), 
                watch1.Stop(); watch1.Print(); watch1.Start()

                outputFile = TFile(condor_dir + "/" + dataset + ".root", "UPDATE")
                if not outputFile or outputFile.IsZombie():  print "Could not open file: %s/%s.root" % (condor_dir, dataset)  
                outputFile.cd("OSUAnalysis/"+hist['channel'])

                deleteString = hist['histName'] + ";*"  # delete all existing instances of the object
                currentDir = outputFile.GetDirectory("OSUAnalysis/"+hist['channel']);
                if not currentDir: print "Could not find directory OSUAnalysis/%s in file %s" % (hist['channel'], outputFile.GetName())  
                currentDir.Delete(deleteString);
                if 'nbinsY' in hist:  # only make a 2D histogram if the key "nbinsY" is defined  
                    h = TH2D(hist['histName'], hist['histName'],
                             hist['nbins'],  hist['xMin'], hist['xMax'],
                             hist['nbinsY'], hist['yMin'], hist['yMax'])
                else:             
                    h = TH1D(hist['histName'], hist['histName'], hist['nbins'], hist['xMin'], hist['xMax'])
                h.Sumw2()  # Needed to get weights correct.  
                cut = TCut(hist['cutString'])
                ch.Draw(hist['varToPlot']+">>"+hist['histName'], cut)  
                h.Write()
                outputFile.Close()
                print "Histogram " + hist['histName'] + " has been added to " + condor_dir + "/"+ dataset + ".root"

#merge output if composite dataset
for composite_dataset in composite_datasets:
    component_datasets_list = ""
    component_dataset_file_path = ""
    for component_dataset in composite_dataset_definitions[composite_dataset]:
        component_dataset_dir = "%s/%s" % (condor_dir,component_dataset)
        component_dataset_file_path = component_dataset_dir + ".root"
        if os.path.isfile(component_dataset_file_path):
            component_datasets_list += " " + component_dataset_file_path
    composite_dataset_dir = "%s/%s" % (condor_dir,composite_dataset)
    command = "mergeHists -p %s %s" % (composite_dataset_dir, component_datasets_list)
    print "Merging output for composite dataset: " + composite_dataset
    os.system(command)
      
print ("Total time to run makeBNTreePlot.py: "), 
watch.Stop()
watch.Print()
            

Revision:	1.15
Committed:	Fri Jun 14 16:08:18 2013 UTC (11 years, 10 months ago) by wulsin
Content type:	text/x-python
Branch:	MAIN
Changes since 1.14:	+36 -6 lines
Log Message:	Allow to split condor jobs for each dataset with -S option
#	User	Rev	Content
1	jbrinson	1.1	#!/usr/bin/env python
2	wulsin	1.3
3	wulsin	1.13	# Must specify options with -l argument, and condor directory with -c argument, e.g.:
4			# > makeBNTreePlot.py -l sampleBNTreePlotConfig.py -c myCondorDir
5
6			# Additional arugments specify the running mode:
7			# -D <dataset>: Run on a single dataset (typically for testing)
8			# -C: Submit jobs to run on condor over all datasets
9			# BNTreeUseScript=True (set in sampleBNTreePlotConfig.py): run BNTreeScript root macro (also set in sampleBNTreePlotConfig.py),
10			# which must take as arguments the condor directory, dataset, and channel
11	wulsin	1.3
12
13	jbrinson	1.1	import sys
14			import os
15			import re
16			from optparse import OptionParser
17			from array import *
18			from decimal import *
19	wulsin	1.7	import subprocess
20	jbrinson	1.1
21			from OSUT3Analysis.Configuration.configurationOptions import *
22			from OSUT3Analysis.Configuration.processingUtilities import *
23
24	wulsin	1.9	from ROOT import TChain, TCut, TDirectory, TFile, TH1D, TH2D, TStopwatch, TTree, gROOT
25	wulsin	1.4
26			gROOT.SetBatch(True) # This is to prevent pop-up graphical windows
27	jbrinson	1.1
28	wulsin	1.7
29			def MakeCondorSubmitFile(arguments, dataset):
30	wulsin	1.15	outdir = "condor/"+arguments.condorDir+"/"+dataset+"/"
31	wulsin	1.7	p = subprocess.Popen(["which", "makeBNTreePlot.py"], stdout=subprocess.PIPE)
32	wulsin	1.15	executable = p.communicate()[0] # read the stdout of the command
33	wulsin	1.7	workdir = os.getcwd() + "/"
34	wulsin	1.15
35			# Count the number of hist*root files in the specified directory
36			p = subprocess.Popen("ls " + workdir + outdir + "hist*root \| wc", shell=True, stdout=subprocess.PIPE)
37			out = p.communicate()[0]
38			outSplit = out.split()
39			totalJobs = int(outSplit[0]) # The first element of the list is the number of files
40
41			## print "Debug: outSplit = "
42			## print outSplit
43			## out, err = p.communicate()
44			## print "Debug: Size of out:"
45			## print len(out)
46			## print "out = "
47			## print out
48			## print "err = "
49			## print err
50			## print "Debug: p.communicate()[0] = " + p.communicate()[0]
51			## totalJobs = len(out)
52	wulsin	1.7	out = open(outdir+"/condorBNTree.sub", "w")
53			out.write("Executable = " + executable + " \n")
54			out.write("Universe = vanilla \n")
55			out.write("Getenv = True \n")
56	wulsin	1.15	argCondorProcess = ""
57			if arguments.splitCondorJobs:
58			argCondorProcess = " -p $(Process) "
59			out.write("Arguments = -D " + dataset + " -l " + arguments.localConfig + " -c " + arguments.condorDir + argCondorProcess + " \n")
60	wulsin	1.7	out.write("Output = " + workdir + outdir + "condorBNTree_$(Process).out \n")
61			out.write("Error = " + workdir + outdir + "condorBNTree_$(Process).err \n")
62			out.write("Log = " + workdir + outdir + "condorBNTree_$(Process).log \n")
63			out.write("+IsLocalJob = true \n")
64			out.write("Rank = TARGET.IsLocalSlot \n")
65	wulsin	1.15	if arguments.splitCondorJobs:
66			out.write("Queue " + str(totalJobs) + " \n")
67			else:
68			out.write("Queue 1 \n")
69	wulsin	1.7	out.close()
70
71
72			def RunOnCondor(arguments, split_datasets):
73			print "Running jobs on condor, instead of interactively."
74			## print "Found condorDir = %s" % (condorDir)
75			## print "Found split_datasets = "
76			## print split_datasets
77
78			for dataset in split_datasets:
79			MakeCondorSubmitFile(arguments, dataset)
80	wulsin	1.8	cmd = "condor_submit condor/"+arguments.condorDir+"/"+dataset+"/condorBNTree.sub"
81			os.system(cmd)
82			print "Submitting job: %s " % cmd
83	wulsin	1.10
84			print "Once condor jobs have finished, merge the composite datasets and then make plots with:"
85			print " makeBNTreePlot.py -q -l " + arguments.localConfig + " -c " + arguments.condorDir
86			print " makePlots.py -l " + arguments.localConfig + " -c " + arguments.condorDir
87	wulsin	1.7
88			return
89
90
91	wulsin	1.11	watch = TStopwatch()
92			watch1 = TStopwatch()
93
94	jbrinson	1.1
95			parser = OptionParser()
96			parser = set_commandline_arguments(parser)
97	wulsin	1.14
98			### Only used by makeBNTreePlot.py (maybe move to another file?)
99	wulsin	1.15	parser.remove_option("-p")
100	wulsin	1.14	parser.add_option("-D", "--dataset", dest="datasetName",
101			help="Name of dataset (overrides value from local configuration file)")
102			parser.add_option("-C", "--runOnCondor", action="store_true", dest="runOnCondor", default=False,
103			help="Run on condor instead of interactively")
104	wulsin	1.15	parser.add_option("-S", "--splitCondorJobs", action="store_true", dest="splitCondorJobs", default=False,
105			help="Split condor jobs to have one for each file, rather than one for each dataset")
106			parser.add_option("-p", "--condorProcessNum", dest="condorProcessNum", default=-1,
107			help="Specify which condor process to run (default is to run over all).")
108
109	wulsin	1.14
110	jbrinson	1.1	(arguments, args) = parser.parse_args()
111	wulsin	1.14
112	jbrinson	1.1
113			if not arguments.localConfig:
114			sys.exit(" You must specify a localOptions.py file with -l")
115	wulsin	1.6	if arguments.localConfig:
116	jbrinson	1.1	sys.path.append(os.getcwd())
117			exec("from " + arguments.localConfig.rstrip('.py') + " import *")
118			if not arguments.condorDir:
119			sys.exit(" You must specify a condor directory with -c")
120			if arguments.condorDir:
121			condor_dir = "condor/%s" % arguments.condorDir
122
123	wulsin	1.7	if arguments.datasetName: # If datasetName is specified on command line, then override the value from localConfig
124			datasets = [
125			arguments.datasetName,
126			]
127
128	jbrinson	1.2	#save a list of composite datasets
129			composite_datasets = get_composite_datasets(datasets, composite_dataset_definitions)
130			#save a list of datasets with composite datasets split up
131			split_datasets = split_composite_datasets(datasets, composite_dataset_definitions)
132
133	wulsin	1.10	if arguments.quickMerge and arguments.runOnCondor:
134			print "Cannot use -q (--quickMerge) and -C (--runOnCondor) options simultaneously. Please choose one or the other."
135			exit()
136
137	wulsin	1.7	if arguments.runOnCondor:
138			RunOnCondor(arguments, split_datasets)
139	wulsin	1.10	exit()
140	wulsin	1.7
141	jbrinson	1.2	#write new histogram to dataset
142	wulsin	1.10	if not arguments.quickMerge:
143			for dataset in split_datasets:
144	wulsin	1.13	if BNTreeUseScript:
145			chainName = "OSUAnalysis/" + BNTreeChannel + "/BNTree_" + BNTreeChannel
146	wulsin	1.15	command = "root -l -b -q '" + BNTreeScript + "+(\"" + condor_dir + "\",\"" + dataset + "\",\"" + chainName + "\"," + str(arguments.condorProcessNum) + ")'"
147	wulsin	1.13	print "About to execute command: " + command
148			os.system(command)
149			else:
150			for hist in input_histograms:
151			#chain trees together
152			ch = TChain("OSUAnalysis/"+hist['channel']+"/BNTree_"+hist['channel'])
153			ch.Add(condor_dir + "/" + dataset + "/hist_*.root")
154			print ("Looping over chain with # entries = %f; split time = " % ch.GetEntries()),
155			watch1.Stop(); watch1.Print(); watch1.Start()
156
157			outputFile = TFile(condor_dir + "/" + dataset + ".root", "UPDATE")
158			if not outputFile or outputFile.IsZombie(): print "Could not open file: %s/%s.root" % (condor_dir, dataset)
159			outputFile.cd("OSUAnalysis/"+hist['channel'])
160
161			deleteString = hist['histName'] + ";*" # delete all existing instances of the object
162			currentDir = outputFile.GetDirectory("OSUAnalysis/"+hist['channel']);
163			if not currentDir: print "Could not find directory OSUAnalysis/%s in file %s" % (hist['channel'], outputFile.GetName())
164			currentDir.Delete(deleteString);
165			if 'nbinsY' in hist: # only make a 2D histogram if the key "nbinsY" is defined
166			h = TH2D(hist['histName'], hist['histName'],
167			hist['nbins'], hist['xMin'], hist['xMax'],
168			hist['nbinsY'], hist['yMin'], hist['yMax'])
169			else:
170			h = TH1D(hist['histName'], hist['histName'], hist['nbins'], hist['xMin'], hist['xMax'])
171			h.Sumw2() # Needed to get weights correct.
172			cut = TCut(hist['cutString'])
173			ch.Draw(hist['varToPlot']+">>"+hist['histName'], cut)
174			h.Write()
175			outputFile.Close()
176			print "Histogram " + hist['histName'] + " has been added to " + condor_dir + "/"+ dataset + ".root"
177	jbrinson	1.2
178			#merge output if composite dataset
179			for composite_dataset in composite_datasets:
180			component_datasets_list = ""
181			component_dataset_file_path = ""
182			for component_dataset in composite_dataset_definitions[composite_dataset]:
183			component_dataset_dir = "%s/%s" % (condor_dir,component_dataset)
184			component_dataset_file_path = component_dataset_dir + ".root"
185			if os.path.isfile(component_dataset_file_path):
186			component_datasets_list += " " + component_dataset_file_path
187			composite_dataset_dir = "%s/%s" % (condor_dir,composite_dataset)
188			command = "mergeHists -p %s %s" % (composite_dataset_dir, component_datasets_list)
189			print "Merging output for composite dataset: " + composite_dataset
190			os.system(command)
191
192	wulsin	1.15	print ("Total time to run makeBNTreePlot.py: "),
193	jbrinson	1.1	watch.Stop()
194			watch.Print()
195	jbrinson	1.2
196	jbrinson	1.1
197