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'")
    parser.add_option("-t", "--no-weights", action="store_true", dest="noWeights", default=False,
                      help="do not apply cross section weights")
    parser.add_option("-q", "--quickMerge", action="store_true", dest="quickMerge", default=False,
                      help="do merge without making cutflow or plots")  


    #### 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):
    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 (
            "keep *",
          ),
          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.28
Committed:	Fri Jul 5 14:04:58 2013 UTC (11 years, 9 months ago) by ahart
Content type:	text/x-python
Branch:	MAIN
Changes since 1.27:	+1 -3 lines
Log Message:	When making BEAN skims, keep all collections in the event.
#	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			parser.add_option("-t", "--no-weights", action="store_true", dest="noWeights", default=False,
61			help="do not apply cross section weights")
62	wulsin	1.18	parser.add_option("-q", "--quickMerge", action="store_true", dest="quickMerge", default=False,
63			help="do merge without making cutflow or plots")
64	lantonel	1.12
65
66			#### Histogram Formatting Options
67	lantonel	1.8	parser.add_option("-n", "--normalize", action="store_true", dest="normalizeToData", default=False,
68	lantonel	1.5	help="normalize total background MC yield to the data")
69	lantonel	1.8	parser.add_option("-u", "--unit-area", action="store_true", dest="normalizeToUnitArea", default=False,
70	lantonel	1.5	help="normalize all samples to unit area (useful to compare shapes)")
71	lantonel	1.8	parser.add_option("-e", "--empty", action="store_true", dest="noStack", default=False,
72	lantonel	1.5	help="don't stack the background samples, draw them as empty histograms instead")
73	lantonel	1.12
74	lantonel	1.8	parser.add_option("-r", "--ratio", action="store_true", dest="makeRatioPlots", default=False,
75	lantonel	1.6	help="draw (data-MC)/MC plots below all 1D histograms")
76	lantonel	1.24	parser.add_option("-R", "--ratioYRange", dest="ratioYRange",
77	wulsin	1.20	help="maximum of range of vertical scale for ratio plots")
78	lantonel	1.12	parser.add_option("-d", "--diff", action="store_true", dest="makeDiffPlots", default=False,
79			help="draw data-MC plots below all 1D histograms")
80	ahart	1.15	parser.add_option("-b", "--rebin", dest="rebinFactor",
81	lantonel	1.12	help="Rebin all the histograms which will have at least 10 bins after rebinning")
82	lantonel	1.13	parser.add_option("--2D", action="store_true", dest="draw2DPlots", default=False,
83			help="draw stacked 2D histograms")
84	ahart	1.15	parser.add_option("-y", "--yields", action="store_true", dest="printYields", default=False,
85			help="Include the yield of each source in the legend")
86	lantonel	1.19	parser.add_option("-p", "--pdfs", action="store_true", dest="savePDFs", default=False,
87			help="Save pdfs files for all plots made")
88	lantonel	1.13
89	lantonel	1.1	return parser
90
91	ahart	1.15	def get_short_name(dataset, dataset_names):
92	lantonel	1.1	for key in dataset_names:
93			if dataset_names[key] == dataset:
94	ahart	1.15	return key
95	lantonel	1.1	return "Unknown"
96	ahart	1.15
97			def stop_ctau (dataset):
98			if not re.match (r"stop[^_]to[^_]_[^_]mm.", dataset):
99			return 0.0
100			return float (re.sub (r"stop[^_]to[^_]_([^_])mm.", r"\1", dataset))
101
102			def source_stop_ctau (ctau):
103			return int (math.pow (10.0, math.ceil (math.log10 (ctau))))
104
105			def add_stops (options, masses, ctaus, bottomBranchingRatios = []):
106			if not bottomBranchingRatios:
107	ahart	1.16	bottomBranchingRatios.append (50.0)
108	ahart	1.15	for mass in masses:
109			for ctau in ctaus:
110			for bottomBranchingRatio in bottomBranchingRatios:
111	ahart	1.16	datasetName = 'stop' + str (mass) + "_" + str (ctau) + "mm_br" + str (int (bottomBranchingRatio))
112	ahart	1.15	bottomDatasetName = 'stop' + str (mass) + "toBl_" + str (ctau) + "mm"
113			sourceBottomDatasetName = 'stop' + str (mass) + "toBl_" + str (source_stop_ctau (ctau)) + "mm"
114			topDatasetName = 'stop' + str (mass) + "toTnu_" + str (ctau) + "mm"
115			sourceTopDatasetName = 'stop' + str (mass) + "toTnu_" + str (source_stop_ctau (ctau)) + "mm"
116			mixedDatasetName = 'stop' + str (mass) + "toBT_" + str (ctau) + "mm"
117			sourceMixedDatasetName = 'stop' + str (mass) + "toBT_" + str (source_stop_ctau (ctau)) + "mm"
118
119			options['datasets'].append (datasetName)
120	ahart	1.16	bottomBranchingRatio /= 100.0
121	ahart	1.15	options['composite_dataset_definitions'][datasetName] = {
122			bottomDatasetName : bottomBranchingRatio * bottomBranchingRatio,
123			topDatasetName : (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio),
124			mixedDatasetName : (1 - bottomBranchingRatio * bottomBranchingRatio - (1 - bottomBranchingRatio) * (1 - bottomBranchingRatio))
125			}
126			options['dataset_names'][bottomDatasetName] = options['dataset_names'][sourceBottomDatasetName]
127			options['dataset_names'][topDatasetName] = options['dataset_names'][sourceTopDatasetName]
128			options['dataset_names'][mixedDatasetName] = options['dataset_names'][sourceMixedDatasetName]
129			options['nJobs'][bottomDatasetName] = 1
130			options['nJobs'][topDatasetName] = 1
131			options['nJobs'][mixedDatasetName] = 1
132			options['maxEvents'][bottomDatasetName] = -1
133			options['maxEvents'][topDatasetName] = -1
134			options['maxEvents'][mixedDatasetName] = -1
135			options['types'][datasetName] = "signalMC"
136			options['types'][bottomDatasetName] = "signalMC"
137			options['types'][topDatasetName] = "signalMC"
138			options['types'][mixedDatasetName] = "signalMC"
139	ahart	1.17	options['labels'][datasetName] = str (mass) + " GeV stop (#LTc#tau#GT = " + str (ctau) + " mm)"
140			options['labels'][bottomDatasetName] = "#tilde{t}#tilde{t}#rightarrowbbll (#LTc#tau#GT = " + str (ctau) + " mm)"
141			options['labels'][topDatasetName] = "#tilde{t}#tilde{t}#rightarrowtt#nu#nu (#LTc#tau#GT = " + str (ctau) + " mm)"
142			options['labels'][mixedDatasetName] = "#tilde{t}#tilde{t}#rightarrowbtl#nu (#LTc#tau#GT = " + str (ctau) + " mm)"
143	ahart	1.23
144			def add_channels (process, channels):
145			suffix = ""
146			if osusub.batchMode:
147			suffix = "_" + str (osusub.jobNumber)
148			i = 0
149			for channel in channels:
150			channelName = channel.name.pythonValue ()
151			channelName = channelName[1:-1]
152	ahart	1.27	try:
153			os.mkdir (channelName)
154			except OSError:
155			pass
156	ahart	1.23	out = cms.OutputModule ("PoolOutputModule",
157	ahart	1.25	fileName = cms.untracked.string (channelName + "/bean" + suffix + ".root"),
158	ahart	1.23	SelectEvents = cms.untracked.PSet (
159			SelectEvents = cms.vstring ("myFilterPath" + str (i))
160			),
161			outputCommands = cms.untracked.vstring (
162	ahart	1.28	"keep *",
163	ahart	1.23	),
164			dropMetaData = cms.untracked.string ("ALL")
165			)
166			myFilter = cms.EDFilter ("FilterOnChannelDecision",
167			src = cms.InputTag ("OSUAnalysis", "channelDecisions"),
168			channelName = cms.string (channelName)
169			)
170			myFilterPath = cms.Path (myFilter)
171			myEndPath = cms.EndPath (out)
172			process.__setattr__ ("out" + str (i), out)
173			process.OSUAnalysis.channels.append (channel)
174			process.OSUAnalysis.useEDMFormat = cms.bool (True)
175			process.__setattr__ ("myFilter" + str (i), myFilter)
176			process.__setattr__ ("myFilterPath" + str (i), myFilterPath)
177			process.__setattr__ ("myEndPath" + str (i), myEndPath)
178			i += 1