Configuration/python/processingUtilities.py

#!/usr/bin/env python
import os
import re
import sys
import math
import datetime
import copy
from optparse import OptionParser
import OSUT3Analysis.DBTools.osusub_cfg as osusub
import FWCore.ParameterSet.Config as cms

def split_composite_datasets(datasets, composite_dataset_definitions):
    for dataset in datasets:
        if dataset in composite_dataset_definitions:
            for component_dataset in composite_dataset_definitions[dataset]:
                datasets.insert(datasets.index(dataset),component_dataset)
            datasets.remove(dataset)
    return datasets

def get_composite_datasets(datasets, composite_dataset_definitions):
    composite_datasets = []
    for dataset in datasets:
        if dataset in composite_dataset_definitions:
            composite_datasets.append(dataset)
    return composite_datasets

def set_condor_submit_dir(arguments):
    if arguments.condorDir:
        condor_dir = "condor/%s" % arguments.condorDir
    else:
        now = datetime.datetime.now()
        date_hash = now.strftime("%Y_%m_%d_%H:%M:%S")
        condor_dir = "condor/condor_%s" % date_hash
    #print "Condor submit directory set to ",condor_dir
    return condor_dir

def set_condor_output_dir(arguments):
    if arguments.condorDir:
        condor_dir = "condor/%s" % arguments.condorDir
    else: #get most recent condor submission directory
        dir_list = []
        for directory in os.listdir("./condor/"):
            if directory.find("condor_") is not -1:
                dir_list.append(directory)
        if len(dir_list) is 0:
            sys.exit("Cannot find last condor working directory")
        dir_list.sort(reverse=True)
        condor_dir = "condor/%s" % dir_list[0]
    #print "Condor output directory set to ",condor_dir
    return condor_dir

def set_commandline_arguments(parser):
    #### Configuration-related Options
    parser.add_option("-l", "--localConfig", dest="localConfig",
                  help="local configuration file")
    parser.add_option("-c", "--condorDir", dest="condorDir",
                      help="condor output directory")
    parser.add_option("-o", "--output-file", dest="outputFileName",
                      help="specify an output file name for the histogram file, default is 'stacked_histograms.root'")


    #### Histogram Formatting Options
    parser.add_option("-n", "--normalize", action="store_true", dest="normalizeToData", default=False,
                      help="normalize total background MC yield to the data")
    parser.add_option("-u", "--unit-area", action="store_true", dest="normalizeToUnitArea", default=False,
                      help="normalize all samples to unit area (useful to compare shapes)")
    parser.add_option("-e", "--empty", action="store_true", dest="noStack", default=False,
                      help="don't stack the background samples, draw them as empty histograms instead")

    parser.add_option("-r", "--ratio", action="store_true", dest="makeRatioPlots", default=False,
                      help="draw (data-MC)/MC plots below all 1D histograms")
    parser.add_option("-R", "--ratioYRange", dest="ratioYRange",
                      help="maximum of range of vertical scale for ratio plots")  
    parser.add_option("-d", "--diff", action="store_true", dest="makeDiffPlots", default=False,
                      help="draw data-MC plots below all 1D histograms")
    parser.add_option("-b", "--rebin", dest="rebinFactor",
                      help="Rebin all the histograms which will have at least 10 bins after rebinning")
    parser.add_option("--2D", action="store_true", dest="draw2DPlots", default=False,
                      help="draw stacked 2D histograms")
    parser.add_option("-y", "--yields", action="store_true", dest="printYields", default=False,
                      help="Include the yield of each source in the legend")
    parser.add_option("-p", "--pdfs",  action="store_true", dest="savePDFs", default=False,
                      help="Save pdfs files for all plots made")

    return parser

def get_short_name(dataset, dataset_names):
    for key in dataset_names:
        if dataset_names[key] == dataset:
            return key
    return "Unknown"

def stop_ctau (dataset):
    if not re.match (r"stop[^_]*to[^_]*_[^_]*mm.*", dataset):
      return 0.0
    return float (re.sub (r"stop[^_]*to[^_]*_([^_]*)mm.*", r"\1", dataset))

def source_stop_ctau (ctau):
    return int (math.pow (10.0, math.ceil (math.log10 (ctau))))

def add_stops (options, masses, ctaus, bottomBranchingRatios = []):
    if not bottomBranchingRatios:
        bottomBranchingRatios.append (50.0)
    for mass in masses:
        for ctau in ctaus:
            for bottomBranchingRatio in bottomBranchingRatios:
                datasetName = 'stop' + str (mass) + "_" + str (ctau) + "mm_br" + str (int (bottomBranchingRatio))
                bottomDatasetName = 'stop' + str (mass) + "toBl_" + str (ctau) + "mm"
                sourceBottomDatasetName = 'stop' + str (mass) + "toBl_" + str (source_stop_ctau (ctau)) + "mm"
                topDatasetName = 'stop' + str (mass) + "toTnu_" + str (ctau) + "mm"
                sourceTopDatasetName = 'stop' + str (mass) + "toTnu_" + str (source_stop_ctau (ctau)) + "mm"
                mixedDatasetName = 'stop' + str (mass) + "toBT_" + str (ctau) + "mm"
                sourceMixedDatasetName = 'stop' + str (mass) + "toBT_" + str (source_stop_ctau (ctau)) + "mm"

                options['datasets'].append (datasetName)
                bottomBranchingRatio /= 100.0
                options['composite_dataset_definitions'][datasetName] = {
                    bottomDatasetName : bottomBranchingRatio * bottomBranchingRatio,
                    topDatasetName : (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio),
                    mixedDatasetName : (1 - bottomBranchingRatio * bottomBranchingRatio - (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio))
                }
                options['dataset_names'][bottomDatasetName] = options['dataset_names'][sourceBottomDatasetName]
                options['dataset_names'][topDatasetName] = options['dataset_names'][sourceTopDatasetName]
                options['dataset_names'][mixedDatasetName] = options['dataset_names'][sourceMixedDatasetName]
                options['nJobs'][bottomDatasetName] = 1
                options['nJobs'][topDatasetName] = 1
                options['nJobs'][mixedDatasetName] = 1
                options['maxEvents'][bottomDatasetName] = -1
                options['maxEvents'][topDatasetName] = -1
                options['maxEvents'][mixedDatasetName] = -1
                options['types'][datasetName] = "signalMC"
                options['types'][bottomDatasetName] = "signalMC"
                options['types'][topDatasetName] = "signalMC"
                options['types'][mixedDatasetName] = "signalMC"
                options['labels'][datasetName] = str (mass) + " GeV stop (#LTc#tau#GT = " + str (ctau) + " mm)"
                options['labels'][bottomDatasetName] = "#tilde{t}#tilde{t}#rightarrowbbll (#LTc#tau#GT = " + str (ctau) + " mm)"
                options['labels'][topDatasetName] = "#tilde{t}#tilde{t}#rightarrowtt#nu#nu (#LTc#tau#GT = " + str (ctau) + " mm)"
                options['labels'][mixedDatasetName] = "#tilde{t}#tilde{t}#rightarrowbtl#nu (#LTc#tau#GT = " + str (ctau) + " mm)"

def add_channels (process, channels, outputCommands = ["keep *"]):
    suffix = ""
    if osusub.batchMode:
        suffix = "_" + str (osusub.jobNumber)
    i = 0
    for channel in channels:
        channelName = channel.name.pythonValue ()
        channelName = channelName[1:-1]
        try:
            os.mkdir (channelName)
        except OSError:
            pass
        out = cms.OutputModule ("PoolOutputModule",
          fileName = cms.untracked.string (channelName + "/bean" + suffix + ".root"),
          SelectEvents = cms.untracked.PSet (
            SelectEvents = cms.vstring ("myFilterPath" + str (i))
          ),
          outputCommands = cms.untracked.vstring (
            outputCommands
          ),
          dropMetaData = cms.untracked.string ("ALL")
        )
        myFilter = cms.EDFilter ("FilterOnChannelDecision",
            src = cms.InputTag ("OSUAnalysis", "channelDecisions"),
            channelName = cms.string (channelName)
        )
        myFilterPath = cms.Path (myFilter)
        myEndPath = cms.EndPath (out)
        process.__setattr__ ("out" + str (i), out)
        process.OSUAnalysis.channels.append (channel)
        process.OSUAnalysis.useEDMFormat = cms.bool (True)
        process.__setattr__ ("myFilter" + str (i), myFilter)
        process.__setattr__ ("myFilterPath" + str (i), myFilterPath)
        process.__setattr__ ("myEndPath" + str (i), myEndPath)
        i += 1
Revision:	1.30
Committed:	Thu Jul 11 14:24:20 2013 UTC (11 years, 10 months ago) by lantonel
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	V02-03-02, HEAD
Changes since 1.29:	+0 -4 lines
Log Message:	moved specific options to mergeOutput.py
#	User	Rev	Content
1	lantonel	1.1	#!/usr/bin/env python
2			import os
3	ahart	1.15	import re
4	lantonel	1.1	import sys
5	ahart	1.15	import math
6	lantonel	1.1	import datetime
7	ahart	1.23	import copy
8	lantonel	1.8	from optparse import OptionParser
9	ahart	1.23	import OSUT3Analysis.DBTools.osusub_cfg as osusub
10			import FWCore.ParameterSet.Config as cms
11	lantonel	1.2
12	ahart	1.15	def split_composite_datasets(datasets, composite_dataset_definitions):
13	lantonel	1.1	for dataset in datasets:
14			if dataset in composite_dataset_definitions:
15			for component_dataset in composite_dataset_definitions[dataset]:
16			datasets.insert(datasets.index(dataset),component_dataset)
17			datasets.remove(dataset)
18			return datasets
19
20	ahart	1.15	def get_composite_datasets(datasets, composite_dataset_definitions):
21	lantonel	1.1	composite_datasets = []
22			for dataset in datasets:
23			if dataset in composite_dataset_definitions:
24			composite_datasets.append(dataset)
25			return composite_datasets
26
27	lantonel	1.7	def set_condor_submit_dir(arguments):
28			if arguments.condorDir:
29			condor_dir = "condor/%s" % arguments.condorDir
30	lantonel	1.1	else:
31			now = datetime.datetime.now()
32			date_hash = now.strftime("%Y_%m_%d_%H:%M:%S")
33			condor_dir = "condor/condor_%s" % date_hash
34	lantonel	1.3	#print "Condor submit directory set to ",condor_dir
35	lantonel	1.1	return condor_dir
36
37	lantonel	1.7	def set_condor_output_dir(arguments):
38			if arguments.condorDir:
39			condor_dir = "condor/%s" % arguments.condorDir
40	lantonel	1.1	else: #get most recent condor submission directory
41			dir_list = []
42			for directory in os.listdir("./condor/"):
43			if directory.find("condor_") is not -1:
44			dir_list.append(directory)
45			if len(dir_list) is 0:
46			sys.exit("Cannot find last condor working directory")
47			dir_list.sort(reverse=True)
48			condor_dir = "condor/%s" % dir_list[0]
49	wulsin	1.21	#print "Condor output directory set to ",condor_dir
50	lantonel	1.1	return condor_dir
51
52			def set_commandline_arguments(parser):
53	lantonel	1.12	#### Configuration-related Options
54	lantonel	1.8	parser.add_option("-l", "--localConfig", dest="localConfig",
55	lantonel	1.1	help="local configuration file")
56	lantonel	1.8	parser.add_option("-c", "--condorDir", dest="condorDir",
57	lantonel	1.5	help="condor output directory")
58	lantonel	1.12	parser.add_option("-o", "--output-file", dest="outputFileName",
59			help="specify an output file name for the histogram file, default is 'stacked_histograms.root'")
60
61
62			#### Histogram Formatting Options
63	lantonel	1.8	parser.add_option("-n", "--normalize", action="store_true", dest="normalizeToData", default=False,
64	lantonel	1.5	help="normalize total background MC yield to the data")
65	lantonel	1.8	parser.add_option("-u", "--unit-area", action="store_true", dest="normalizeToUnitArea", default=False,
66	lantonel	1.5	help="normalize all samples to unit area (useful to compare shapes)")
67	lantonel	1.8	parser.add_option("-e", "--empty", action="store_true", dest="noStack", default=False,
68	lantonel	1.5	help="don't stack the background samples, draw them as empty histograms instead")
69	lantonel	1.12
70	lantonel	1.8	parser.add_option("-r", "--ratio", action="store_true", dest="makeRatioPlots", default=False,
71	lantonel	1.6	help="draw (data-MC)/MC plots below all 1D histograms")
72	lantonel	1.24	parser.add_option("-R", "--ratioYRange", dest="ratioYRange",
73	wulsin	1.20	help="maximum of range of vertical scale for ratio plots")
74	lantonel	1.12	parser.add_option("-d", "--diff", action="store_true", dest="makeDiffPlots", default=False,
75			help="draw data-MC plots below all 1D histograms")
76	ahart	1.15	parser.add_option("-b", "--rebin", dest="rebinFactor",
77	lantonel	1.12	help="Rebin all the histograms which will have at least 10 bins after rebinning")
78	lantonel	1.13	parser.add_option("--2D", action="store_true", dest="draw2DPlots", default=False,
79			help="draw stacked 2D histograms")
80	ahart	1.15	parser.add_option("-y", "--yields", action="store_true", dest="printYields", default=False,
81			help="Include the yield of each source in the legend")
82	lantonel	1.19	parser.add_option("-p", "--pdfs", action="store_true", dest="savePDFs", default=False,
83			help="Save pdfs files for all plots made")
84	lantonel	1.13
85	lantonel	1.1	return parser
86
87	ahart	1.15	def get_short_name(dataset, dataset_names):
88	lantonel	1.1	for key in dataset_names:
89			if dataset_names[key] == dataset:
90	ahart	1.15	return key
91	lantonel	1.1	return "Unknown"
92	ahart	1.15
93			def stop_ctau (dataset):
94			if not re.match (r"stop[^_]to[^_]_[^_]mm.", dataset):
95			return 0.0
96			return float (re.sub (r"stop[^_]to[^_]_([^_])mm.", r"\1", dataset))
97
98			def source_stop_ctau (ctau):
99			return int (math.pow (10.0, math.ceil (math.log10 (ctau))))
100
101			def add_stops (options, masses, ctaus, bottomBranchingRatios = []):
102			if not bottomBranchingRatios:
103	ahart	1.16	bottomBranchingRatios.append (50.0)
104	ahart	1.15	for mass in masses:
105			for ctau in ctaus:
106			for bottomBranchingRatio in bottomBranchingRatios:
107	ahart	1.16	datasetName = 'stop' + str (mass) + "_" + str (ctau) + "mm_br" + str (int (bottomBranchingRatio))
108	ahart	1.15	bottomDatasetName = 'stop' + str (mass) + "toBl_" + str (ctau) + "mm"
109			sourceBottomDatasetName = 'stop' + str (mass) + "toBl_" + str (source_stop_ctau (ctau)) + "mm"
110			topDatasetName = 'stop' + str (mass) + "toTnu_" + str (ctau) + "mm"
111			sourceTopDatasetName = 'stop' + str (mass) + "toTnu_" + str (source_stop_ctau (ctau)) + "mm"
112			mixedDatasetName = 'stop' + str (mass) + "toBT_" + str (ctau) + "mm"
113			sourceMixedDatasetName = 'stop' + str (mass) + "toBT_" + str (source_stop_ctau (ctau)) + "mm"
114
115			options['datasets'].append (datasetName)
116	ahart	1.16	bottomBranchingRatio /= 100.0
117	ahart	1.15	options['composite_dataset_definitions'][datasetName] = {
118			bottomDatasetName : bottomBranchingRatio * bottomBranchingRatio,
119			topDatasetName : (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio),
120			mixedDatasetName : (1 - bottomBranchingRatio * bottomBranchingRatio - (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio))
121			}
122			options['dataset_names'][bottomDatasetName] = options['dataset_names'][sourceBottomDatasetName]
123			options['dataset_names'][topDatasetName] = options['dataset_names'][sourceTopDatasetName]
124			options['dataset_names'][mixedDatasetName] = options['dataset_names'][sourceMixedDatasetName]
125			options['nJobs'][bottomDatasetName] = 1
126			options['nJobs'][topDatasetName] = 1
127			options['nJobs'][mixedDatasetName] = 1
128			options['maxEvents'][bottomDatasetName] = -1
129			options['maxEvents'][topDatasetName] = -1
130			options['maxEvents'][mixedDatasetName] = -1
131			options['types'][datasetName] = "signalMC"
132			options['types'][bottomDatasetName] = "signalMC"
133			options['types'][topDatasetName] = "signalMC"
134			options['types'][mixedDatasetName] = "signalMC"
135	ahart	1.17	options['labels'][datasetName] = str (mass) + " GeV stop (#LTc#tau#GT = " + str (ctau) + " mm)"
136			options['labels'][bottomDatasetName] = "#tilde{t}#tilde{t}#rightarrowbbll (#LTc#tau#GT = " + str (ctau) + " mm)"
137			options['labels'][topDatasetName] = "#tilde{t}#tilde{t}#rightarrowtt#nu#nu (#LTc#tau#GT = " + str (ctau) + " mm)"
138			options['labels'][mixedDatasetName] = "#tilde{t}#tilde{t}#rightarrowbtl#nu (#LTc#tau#GT = " + str (ctau) + " mm)"
139	ahart	1.23
140	ahart	1.29	def add_channels (process, channels, outputCommands = ["keep *"]):
141	ahart	1.23	suffix = ""
142			if osusub.batchMode:
143			suffix = "_" + str (osusub.jobNumber)
144			i = 0
145			for channel in channels:
146			channelName = channel.name.pythonValue ()
147			channelName = channelName[1:-1]
148	ahart	1.27	try:
149			os.mkdir (channelName)
150			except OSError:
151			pass
152	ahart	1.23	out = cms.OutputModule ("PoolOutputModule",
153	ahart	1.25	fileName = cms.untracked.string (channelName + "/bean" + suffix + ".root"),
154	ahart	1.23	SelectEvents = cms.untracked.PSet (
155			SelectEvents = cms.vstring ("myFilterPath" + str (i))
156			),
157			outputCommands = cms.untracked.vstring (
158	ahart	1.29	outputCommands
159	ahart	1.23	),
160			dropMetaData = cms.untracked.string ("ALL")
161			)
162			myFilter = cms.EDFilter ("FilterOnChannelDecision",
163			src = cms.InputTag ("OSUAnalysis", "channelDecisions"),
164			channelName = cms.string (channelName)
165			)
166			myFilterPath = cms.Path (myFilter)
167			myEndPath = cms.EndPath (out)
168			process.__setattr__ ("out" + str (i), out)
169			process.OSUAnalysis.channels.append (channel)
170			process.OSUAnalysis.useEDMFormat = cms.bool (True)
171			process.__setattr__ ("myFilter" + str (i), myFilter)
172			process.__setattr__ ("myFilterPath" + str (i), myFilterPath)
173			process.__setattr__ ("myEndPath" + str (i), myEndPath)
174			i += 1