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:"  
    if arguments.splitCondorJobs: 
        print "    mergeOutput.py -q    -l " + arguments.localConfig + " -c " + arguments.condorDir
    else: 
        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.16
Committed:	Mon Jun 17 08:49:16 2013 UTC (11 years, 10 months ago) by wulsin
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	V02-03-02, V02-03-01, V02-03-00, HEAD
Changes since 1.15:	+5 -2 lines
Log Message:	When splitting condor jobs, give instruction to use mergeOutput.py
#	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	wulsin	1.16	if arguments.splitCondorJobs:
86			print " mergeOutput.py -q -l " + arguments.localConfig + " -c " + arguments.condorDir
87			else:
88			print " makeBNTreePlot.py -q -l " + arguments.localConfig + " -c " + arguments.condorDir
89			print " makePlots.py -l " + arguments.localConfig + " -c " + arguments.condorDir
90	wulsin	1.7
91			return
92
93
94	wulsin	1.11	watch = TStopwatch()
95			watch1 = TStopwatch()
96
97	jbrinson	1.1
98			parser = OptionParser()
99			parser = set_commandline_arguments(parser)
100	wulsin	1.14
101			### Only used by makeBNTreePlot.py (maybe move to another file?)
102	wulsin	1.15	parser.remove_option("-p")
103	wulsin	1.14	parser.add_option("-D", "--dataset", dest="datasetName",
104			help="Name of dataset (overrides value from local configuration file)")
105			parser.add_option("-C", "--runOnCondor", action="store_true", dest="runOnCondor", default=False,
106			help="Run on condor instead of interactively")
107	wulsin	1.15	parser.add_option("-S", "--splitCondorJobs", action="store_true", dest="splitCondorJobs", default=False,
108			help="Split condor jobs to have one for each file, rather than one for each dataset")
109			parser.add_option("-p", "--condorProcessNum", dest="condorProcessNum", default=-1,
110			help="Specify which condor process to run (default is to run over all).")
111
112	wulsin	1.14
113	jbrinson	1.1	(arguments, args) = parser.parse_args()
114	wulsin	1.14
115	jbrinson	1.1
116			if not arguments.localConfig:
117			sys.exit(" You must specify a localOptions.py file with -l")
118	wulsin	1.6	if arguments.localConfig:
119	jbrinson	1.1	sys.path.append(os.getcwd())
120			exec("from " + arguments.localConfig.rstrip('.py') + " import *")
121			if not arguments.condorDir:
122			sys.exit(" You must specify a condor directory with -c")
123			if arguments.condorDir:
124			condor_dir = "condor/%s" % arguments.condorDir
125
126	wulsin	1.7	if arguments.datasetName: # If datasetName is specified on command line, then override the value from localConfig
127			datasets = [
128			arguments.datasetName,
129			]
130
131	jbrinson	1.2	#save a list of composite datasets
132			composite_datasets = get_composite_datasets(datasets, composite_dataset_definitions)
133			#save a list of datasets with composite datasets split up
134			split_datasets = split_composite_datasets(datasets, composite_dataset_definitions)
135
136	wulsin	1.10	if arguments.quickMerge and arguments.runOnCondor:
137			print "Cannot use -q (--quickMerge) and -C (--runOnCondor) options simultaneously. Please choose one or the other."
138			exit()
139
140	wulsin	1.7	if arguments.runOnCondor:
141			RunOnCondor(arguments, split_datasets)
142	wulsin	1.10	exit()
143	wulsin	1.7
144	jbrinson	1.2	#write new histogram to dataset
145	wulsin	1.10	if not arguments.quickMerge:
146			for dataset in split_datasets:
147	wulsin	1.13	if BNTreeUseScript:
148			chainName = "OSUAnalysis/" + BNTreeChannel + "/BNTree_" + BNTreeChannel
149	wulsin	1.15	command = "root -l -b -q '" + BNTreeScript + "+(\"" + condor_dir + "\",\"" + dataset + "\",\"" + chainName + "\"," + str(arguments.condorProcessNum) + ")'"
150	wulsin	1.13	print "About to execute command: " + command
151			os.system(command)
152			else:
153			for hist in input_histograms:
154			#chain trees together
155			ch = TChain("OSUAnalysis/"+hist['channel']+"/BNTree_"+hist['channel'])
156			ch.Add(condor_dir + "/" + dataset + "/hist_*.root")
157			print ("Looping over chain with # entries = %f; split time = " % ch.GetEntries()),
158			watch1.Stop(); watch1.Print(); watch1.Start()
159
160			outputFile = TFile(condor_dir + "/" + dataset + ".root", "UPDATE")
161			if not outputFile or outputFile.IsZombie(): print "Could not open file: %s/%s.root" % (condor_dir, dataset)
162			outputFile.cd("OSUAnalysis/"+hist['channel'])
163
164			deleteString = hist['histName'] + ";*" # delete all existing instances of the object
165			currentDir = outputFile.GetDirectory("OSUAnalysis/"+hist['channel']);
166			if not currentDir: print "Could not find directory OSUAnalysis/%s in file %s" % (hist['channel'], outputFile.GetName())
167			currentDir.Delete(deleteString);
168			if 'nbinsY' in hist: # only make a 2D histogram if the key "nbinsY" is defined
169			h = TH2D(hist['histName'], hist['histName'],
170			hist['nbins'], hist['xMin'], hist['xMax'],
171			hist['nbinsY'], hist['yMin'], hist['yMax'])
172			else:
173			h = TH1D(hist['histName'], hist['histName'], hist['nbins'], hist['xMin'], hist['xMax'])
174			h.Sumw2() # Needed to get weights correct.
175			cut = TCut(hist['cutString'])
176			ch.Draw(hist['varToPlot']+">>"+hist['histName'], cut)
177			h.Write()
178			outputFile.Close()
179			print "Histogram " + hist['histName'] + " has been added to " + condor_dir + "/"+ dataset + ".root"
180	jbrinson	1.2
181			#merge output if composite dataset
182			for composite_dataset in composite_datasets:
183			component_datasets_list = ""
184			component_dataset_file_path = ""
185			for component_dataset in composite_dataset_definitions[composite_dataset]:
186			component_dataset_dir = "%s/%s" % (condor_dir,component_dataset)
187			component_dataset_file_path = component_dataset_dir + ".root"
188			if os.path.isfile(component_dataset_file_path):
189			component_datasets_list += " " + component_dataset_file_path
190			composite_dataset_dir = "%s/%s" % (condor_dir,composite_dataset)
191			command = "mergeHists -p %s %s" % (composite_dataset_dir, component_datasets_list)
192			print "Merging output for composite dataset: " + composite_dataset
193			os.system(command)
194
195	wulsin	1.15	print ("Total time to run makeBNTreePlot.py: "),
196	jbrinson	1.1	watch.Stop()
197			watch.Print()
198	jbrinson	1.2
199	jbrinson	1.1
200