Configuration/scripts/makeBNTreePlot.py

#!/usr/bin/env python

# Must specify options with -l argument, e.g.:  
# > makeBNTreePlot.py -l sampleBNTreePlotConfig.py  


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, TStopwatch, TTree, gROOT  

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


def MakeCondorSubmitFile(arguments, dataset):
    p = subprocess.Popen(["which", "makeBNTreePlot.py"], stdout=subprocess.PIPE)
    executable = p.communicate()[0]  # read the stdout of the command  
    workdir = os.getcwd() + "/" 
    outdir = "condor/"+arguments.condorDir+"/"+dataset+"/"  
    out = open(outdir+"/condorBNTree.sub", "w")
    out.write("Executable              = " + executable + " \n")
    out.write("Universe                = vanilla \n")
    out.write("Getenv                  = True \n")
    out.write("Arguments               = -D " + dataset + " -l " + arguments.localConfig + " -c " + arguments.condorDir + " \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")
    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)  
        # Now submit the condor job
        # os.system("condor_submit condor/"+arguments.condorDir+"/"+dataset+"condorBNTree.sub")  
        print "Submitted job for: %s/%s" % (arguments.condorDir, dataset)
        
    return


watch = TStopwatch()

parser = OptionParser()
parser = set_commandline_arguments(parser)
(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.runOnCondor:
    RunOnCondor(arguments, split_datasets)
    exit()

print "Running interactively instead"  
    

#write new histogram to dataset
for dataset in split_datasets:
    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" % ch.GetEntries()  

        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);
        
        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)
      
watch.Stop()
watch.Print()
            

Revision:	1.7
Committed:	Thu Jun 6 23:34:29 2013 UTC (11 years, 10 months ago) by wulsin
Content type:	text/x-python
Branch:	MAIN
Changes since 1.6:	+54 -6 lines
Log Message:	Add option to submit to condor
#	User	Rev	Content
1	jbrinson	1.1	#!/usr/bin/env python
2	wulsin	1.3
3			# Must specify options with -l argument, e.g.:
4			# > makeBNTreePlot.py -l sampleBNTreePlotConfig.py
5
6
7	jbrinson	1.1	import sys
8			import os
9			import re
10			from optparse import OptionParser
11			from array import *
12			from decimal import *
13	wulsin	1.7	import subprocess
14	jbrinson	1.1
15			from OSUT3Analysis.Configuration.configurationOptions import *
16			from OSUT3Analysis.Configuration.processingUtilities import *
17
18	wulsin	1.5	from ROOT import TChain, TCut, TDirectory, TFile, TH1D, TStopwatch, TTree, gROOT
19	wulsin	1.4
20			gROOT.SetBatch(True) # This is to prevent pop-up graphical windows
21	jbrinson	1.1
22	wulsin	1.7
23			def MakeCondorSubmitFile(arguments, dataset):
24			p = subprocess.Popen(["which", "makeBNTreePlot.py"], stdout=subprocess.PIPE)
25			executable = p.communicate()[0] # read the stdout of the command
26			workdir = os.getcwd() + "/"
27			outdir = "condor/"+arguments.condorDir+"/"+dataset+"/"
28			out = open(outdir+"/condorBNTree.sub", "w")
29			out.write("Executable = " + executable + " \n")
30			out.write("Universe = vanilla \n")
31			out.write("Getenv = True \n")
32			out.write("Arguments = -D " + dataset + " -l " + arguments.localConfig + " -c " + arguments.condorDir + " \n")
33			out.write("Output = " + workdir + outdir + "condorBNTree_$(Process).out \n")
34			out.write("Error = " + workdir + outdir + "condorBNTree_$(Process).err \n")
35			out.write("Log = " + workdir + outdir + "condorBNTree_$(Process).log \n")
36			out.write("+IsLocalJob = true \n")
37			out.write("Rank = TARGET.IsLocalSlot \n")
38			out.write("Queue 1 \n")
39			out.close()
40
41
42			def RunOnCondor(arguments, split_datasets):
43			print "Running jobs on condor, instead of interactively."
44			## print "Found condorDir = %s" % (condorDir)
45			## print "Found split_datasets = "
46			## print split_datasets
47
48			for dataset in split_datasets:
49			MakeCondorSubmitFile(arguments, dataset)
50			# Now submit the condor job
51			# os.system("condor_submit condor/"+arguments.condorDir+"/"+dataset+"condorBNTree.sub")
52			print "Submitted job for: %s/%s" % (arguments.condorDir, dataset)
53
54			return
55
56
57	jbrinson	1.1	watch = TStopwatch()
58
59			parser = OptionParser()
60			parser = set_commandline_arguments(parser)
61			(arguments, args) = parser.parse_args()
62
63			if not arguments.localConfig:
64			sys.exit(" You must specify a localOptions.py file with -l")
65	wulsin	1.6	if arguments.localConfig:
66	jbrinson	1.1	sys.path.append(os.getcwd())
67			exec("from " + arguments.localConfig.rstrip('.py') + " import *")
68			if not arguments.condorDir:
69			sys.exit(" You must specify a condor directory with -c")
70			if arguments.condorDir:
71			condor_dir = "condor/%s" % arguments.condorDir
72
73	wulsin	1.7	if arguments.datasetName: # If datasetName is specified on command line, then override the value from localConfig
74			datasets = [
75			arguments.datasetName,
76			]
77
78	jbrinson	1.2	#save a list of composite datasets
79			composite_datasets = get_composite_datasets(datasets, composite_dataset_definitions)
80			#save a list of datasets with composite datasets split up
81			split_datasets = split_composite_datasets(datasets, composite_dataset_definitions)
82
83	wulsin	1.7	if arguments.runOnCondor:
84			RunOnCondor(arguments, split_datasets)
85			exit()
86
87			print "Running interactively instead"
88
89
90	jbrinson	1.2	#write new histogram to dataset
91			for dataset in split_datasets:
92	jbrinson	1.1	for hist in input_histograms:
93	jbrinson	1.2	#chain trees together
94	jbrinson	1.1	ch = TChain("OSUAnalysis/"+hist['channel']+"/BNTree_"+hist['channel'])
95			ch.Add(condor_dir + "/" + dataset + "/hist_*.root")
96	wulsin	1.6	print "Looping over chain with # entries = %f" % ch.GetEntries()
97	jbrinson	1.2
98	wulsin	1.7	outputFile = TFile(condor_dir + "/" + dataset + ".root", "UPDATE")
99			if not outputFile or outputFile.IsZombie(): print "Could not open file: %s/%s.root" % (condor_dir, dataset)
100			outputFile.cd("OSUAnalysis/"+hist['channel'])
101	wulsin	1.5
102			deleteString = hist['histName'] + ";*" # delete all existing instances of the object
103	wulsin	1.7	currentDir = outputFile.GetDirectory("OSUAnalysis/"+hist['channel']);
104			if not currentDir: print "Could not find directory OSUAnalysis/%s in file %s" % (hist['channel'], outputFile.GetName())
105	wulsin	1.5	currentDir.Delete(deleteString);
106
107	jbrinson	1.1	h = TH1D(hist['histName'], hist['histName'], hist['nbins'], hist['xMin'], hist['xMax'])
108	wulsin	1.6	h.Sumw2() # Needed to get weights correct.
109	jbrinson	1.1	cut = TCut(hist['cutString'])
110	wulsin	1.5	ch.Draw(hist['varToPlot']+">>"+hist['histName'], cut)
111	jbrinson	1.1	h.Write()
112	wulsin	1.7	outputFile.Close()
113	jbrinson	1.1	print "Histogram " + hist['histName'] + " has been added to " + condor_dir + "/"+ dataset + ".root"
114	jbrinson	1.2
115			#merge output if composite dataset
116			for composite_dataset in composite_datasets:
117			component_datasets_list = ""
118			component_dataset_file_path = ""
119			for component_dataset in composite_dataset_definitions[composite_dataset]:
120			component_dataset_dir = "%s/%s" % (condor_dir,component_dataset)
121			component_dataset_file_path = component_dataset_dir + ".root"
122			if os.path.isfile(component_dataset_file_path):
123			component_datasets_list += " " + component_dataset_file_path
124			composite_dataset_dir = "%s/%s" % (condor_dir,composite_dataset)
125			command = "mergeHists -p %s %s" % (composite_dataset_dir, component_datasets_list)
126			print "Merging output for composite dataset: " + composite_dataset
127			os.system(command)
128
129	jbrinson	1.1	watch.Stop()
130			watch.Print()
131	jbrinson	1.2
132	jbrinson	1.1
133