Geometry/python/plot_utils.py

0001 from ROOT import TStyle, kWhite, kTRUE
0002 from ROOT import kGray, kAzure, kMagenta, kOrange, kWhite
0003 from ROOT import kRed, kBlue, kGreen, kPink, kYellow
0004 from ROOT import TLine, TLatex
0005
0006 from collections import namedtuple, OrderedDict
0007 from math import sin, cos, tan, atan, exp, pi
0008
0009 Plot_params = namedtuple('Params',
0010                          ['plotNumber',
0011                           'abscissa', 'ordinate',
0012                           'ymin', 'ymax',
0013                           'xmin', 'xmax',
0014                           'quotaName', 'iDrawEta',
0015                           'histoMin', 'histoMax',
0016                           'zLog', 'iRebin'])
0017 plots = {}
0018 plots.setdefault('x_vs_eta', Plot_params(10, '#eta', 'x/X_{0}', 0.0, 2.575, -4.0, 4.0, '', 0, 0., 0., 0, 1))
0019 plots.setdefault('x_vs_phi', Plot_params(20, '#varphi [rad]', 'x/X_{0}', 0.0, 6.2, -4.0, 4.0, '', 0, 0., 0., 0, 1))
0020 plots.setdefault('x_vs_R',   Plot_params(40, 'R [cm]', 'x/X_{0}', 0.0, 70.0, 0.0, 1200.0, '', 0, 0., 0., 0, 1))
0021 plots.setdefault('l_vs_eta', Plot_params(1010, '#eta', 'x/#lambda_{0}', 0.0, 0.73, -4.0, 4.0, '', 0, 0., 0., 0, 1))
0022 plots.setdefault('l_vs_phi', Plot_params(1020, '#varphi [rad]', 'x/#lambda_{0}', 0.0, 1.2, -4.0, 4.0, '', 0, 0., 0., 0, 1))
0023 plots.setdefault('l_vs_R',   Plot_params(1040, '#R [cm]', 'x/#lambda_{0}', 0.0, 7.5, 0.0, 1200.0, '', 0, 0., 0., 0, 1))
0024 plots.setdefault('x_vs_eta_vs_phi', Plot_params(30, '#eta', '#varphi', -3.2, 3.2, -5.0, 5.0, 'x/X_{0}', 0, -1., -1., 0, 1))
0025 plots.setdefault('l_vs_eta_vs_phi', Plot_params(1030, '#eta', '#varphi', -3.2, 3.2, -5.0, 5.0, 'x/#lambda_{0}', 0, -1, -1, 0, 1))
0026 plots.setdefault('x_vs_z_vs_Rsum', Plot_params(50, 'z [mm]', 'R [mm]', 0., 0., 0., 0., '#Sigmax/X_{0}', 1, 0., 2.5, 0, 0))
0027 plots.setdefault('x_vs_z_vs_R', Plot_params(60, 'z [mm]', 'R [mm]', 0., 0., 0., 0., '1/X_{0}', 1, 0.00001, 0.01, 1, 0))
0028 plots.setdefault('l_vs_z_vs_Rsum', Plot_params(1050, 'z [mm]', 'R [mm]', 0., 0., 0., 0., '#Sigmax/#lambda_{0}', 1, 0., 1., 0, 0))
0029 plots.setdefault('l_vs_z_vs_R', Plot_params(1060, 'z [mm]', 'R [mm]', 0., 1400., -3500., 3500., '1/#lambda_{0}', 1, 0.001, 0.9, 1, 0))
0030 plots.setdefault('l_vs_z_vs_R_geocomp', Plot_params(1060, 'z [mm]', 'R [mm]', 0., 1400., -3500., 3500., '1/#lambda_{0}', 1, 0.001, 0.9, 0, 0))
0031 # x_over_l_vs_eta is not extracted from the ROOT file
0032 # but generated from 10 and 1010 in MaterialBudget.py
0033 plots.setdefault('x_over_l_vs_eta', Plot_params(10, '#eta', '(#frac{x}{X_{0}})/(#frac{x}{#lambda_{0}})', 0., 0., 0., 0., '', 0, -1, -1, 0, 0))
0034 # same goes with x_over_l_vs_phi 20/1020
0035 plots.setdefault('x_over_l_vs_phi', Plot_params(20, '#varphi [rad]', '(#frac{x}{X_{0}})/(#frac{x}{#lambda_{0}})', 0., 0., 0., 0., '', 0, -1, -1, 0, 0))
0036
0037 # Conversion name from the label (key) to the components in CMSSW/Geometry
0038 _LABELS2COMPS = {'BeamPipe': 'BEAM',
0039                  'Tracker': 'Tracker',
0040                  'Pixel':   ['PixelBarrel', 'PixelForwardZplus', 'PixelForwardZminus'],
0041                  'PixBar':  'PixelBarrel',
0042                  'PixFwd':  ['PixelForwardZplus', 'PixelForwardZminus', 'PixelForward'],
0043                  'PixFwdMinus': 'PixelForwardZminus',
0044                  'PixFwdPlus':  'PixelForwardZplus',
0045                  'TIB':         'TIB',
0046                  'TOB':         'TOB',
0047                  'TIDB':        'TIDB',
0048                  'TIDF':        'TIDF',
0049                  'TEC':         'TEC',
0050                  'InnerServices': ['TIBTIDServicesF', 'TIBTIDServicesB'],
0051                  'TkStrct': ['TrackerOuterCylinder', 'TrackerBulkhead'],
0052                  'Phase2PixelBarrel': 'Phase2PixelBarrel',
0053                  'Phase2OTBarrel': 'Phase2OTBarrel',
0054                  'Phase2PixelEndcap': 'Phase2PixelEndcap',
0055                  'Phase2OTForward': 'Phase2OTForward'}
0056
0057 # Compounds are used to stick together different part of the Tracker
0058 # detector, so that their cumulative material description can be
0059 # derived. The key name can be generic, while the names in the
0060 # associated list must be such that an appropriate material
0061 # description file, in ROOT format, is present while producing the
0062 # cumulative plot. A missing element will invalidate the full
0063 # procedure.
0064 COMPOUNDS = OrderedDict()
0065 COMPOUNDS["Tracker"] = ["Tracker"]
0066 COMPOUNDS["TrackerSum"] = ["TIB", "TIDF", "TIDB",
0067                            "BeamPipe", "InnerServices",
0068                            "TOB", "TEC",
0069                            "TkStruct",
0070                            "PixBar", "PixFwdPlus", "PixFwdMinus"]
0071 COMPOUNDS["TrackerSumPhaseII"] = ["BeamPipe",
0072                                   "Phase2PixelBarrel",
0073                                   "Phase2OTBarrel", "Phase2OTForward",
0074                                   "Phase2PixelEndcap"]
0075 COMPOUNDS["Pixel"] = ["PixBar", "PixFwdMinus", "PixFwdPlus"]
0076 COMPOUNDS["Strip"] = ["TIB", "TIDF", "TIDB", "InnerServices", "TOB", "TEC"]
0077 COMPOUNDS["InnerTracker"] = ["TIB", "TIDF", "TIDB", "InnerServices"]
0078
0079 # The DETECTORS must be the single component of the tracker for which
0080 # the user can ask for the corresponding material description.
0081 DETECTORS = OrderedDict()
0082 DETECTORS["BeamPipe"] = kGray+2
0083 DETECTORS["InnerServices"] = kGreen+2
0084 DETECTORS["PixBar"] = kAzure-5
0085 DETECTORS["Phase1PixelBarrel"] = kAzure-5
0086 DETECTORS["Phase2PixelBarrel"] = kAzure-5
0087 DETECTORS["PixFwdPlus"] = kAzure-9
0088 DETECTORS["PixFwdMinus"] = kAzure-9
0089 DETECTORS["Phase2PixelEndcap"] = kAzure-9
0090 DETECTORS["Phase2OTBarrel"] = kMagenta-2
0091 DETECTORS["Phase2OTForward"] = kOrange-2
0092 DETECTORS["TIB"] = kMagenta-6
0093 DETECTORS["TIDF"] = kMagenta+2
0094 DETECTORS["TIDB"] = kMagenta+2
0095 DETECTORS["TOB"] = kOrange+10
0096 DETECTORS["TEC"] = kOrange-2
0097
0098 # sDETS are the label of the Tracker elements in the Reconstruction
0099 # geometry. They are all used to derive the reconstruction material
0100 # profile to be compared to the one obtained directly from the
0101 # simulation. A missing key in the real reconstruction geometry is not
0102 # a problem, since this will imply that the corresponding plotting
0103 # routine will skip that missing part. For this reason this map can be
0104 # made as inclusive as possible with respect to the many
0105 # reconstruction geometries in CMSSW.
0106 sDETS = OrderedDict()
0107 sDETS["PXB"] = kRed
0108 sDETS["PXF"] = kBlue
0109 sDETS["TIB"] = kGreen
0110 sDETS["TID"] = kYellow
0111 sDETS["TOB"] = kOrange
0112 sDETS["TEC"] = kPink
0113
0114 # hist_label_to_num contains the logical names of the Tracker detector
0115 # that holds material. They are therefore not aware of which detector
0116 # they belong to, but they are stored in specific plots in all the
0117 # mat*root files produced. The numbering of the plots is identical
0118 # across all files.
0119 hist_label_to_num = OrderedDict()
0120 hist_label_to_num['SUM'] = [0, kGreen+1, 'Total']
0121 hist_label_to_num['SUP'] = [100, 13, 'Support'] # [Index , color, legend label]
0122 hist_label_to_num['SEN'] = [200, 27, 'Sensitive']
0123 hist_label_to_num['CAB'] = [300, 46, 'Cables']
0124 hist_label_to_num['COL'] = [400, 38, 'Cooling']
0125 hist_label_to_num['ELE'] = [500, 30, 'Electronics']
0126 hist_label_to_num['OTH'] = [600, 42, 'Other']
0127 hist_label_to_num['AIR'] = [700, 29, 'Air']
0128
0129 def setTDRStyle():
0130     """Function to setup a TDR-like style"""
0131
0132     tdrStyle = TStyle("tdrStyle","Style for P-TDR")
0133
0134     # For the canvas:
0135     tdrStyle.SetCanvasBorderMode(0)
0136     tdrStyle.SetCanvasColor(kWhite)
0137     tdrStyle.SetCanvasDefH(600) #Height of canvas
0138     tdrStyle.SetCanvasDefW(600) #Width of canvas
0139     tdrStyle.SetCanvasDefX(0)   #Position on screen
0140     tdrStyle.SetCanvasDefY(0)
0141
0142     # For the Pad:
0143     tdrStyle.SetPadBorderMode(0)
0144     tdrStyle.SetPadColor(kWhite)
0145     tdrStyle.SetPadGridX(False)
0146     tdrStyle.SetPadGridY(False)
0147     tdrStyle.SetGridColor(kWhite)
0148     tdrStyle.SetGridStyle(3)
0149     tdrStyle.SetGridWidth(1)
0150     tdrStyle.SetPadTickX(True)
0151     tdrStyle.SetPadTickY(True)
0152
0153     # For the frame:
0154     tdrStyle.SetFrameBorderMode(0)
0155     tdrStyle.SetFrameBorderSize(1)
0156     tdrStyle.SetFrameFillColor(0)
0157     tdrStyle.SetFrameFillStyle(0)
0158     tdrStyle.SetFrameLineColor(1)
0159     tdrStyle.SetFrameLineStyle(1)
0160     tdrStyle.SetFrameLineWidth(0)
0161
0162     # For the histo:
0163     tdrStyle.SetHistLineColor(1)
0164     tdrStyle.SetHistLineStyle(0)
0165     tdrStyle.SetHistLineWidth(1)
0166     tdrStyle.SetEndErrorSize(1)
0167     #tdrStyle.SetErrorX(0.)
0168     tdrStyle.SetMarkerStyle(20)
0169
0170     #For the fit/function:
0171     tdrStyle.SetOptFit(0)
0172     tdrStyle.SetFitFormat("5.4g")
0173     tdrStyle.SetFuncColor(2)
0174     tdrStyle.SetFuncStyle(1)
0175     tdrStyle.SetFuncWidth(1)
0176
0177     #For the date:
0178     tdrStyle.SetOptDate(0)
0179
0180     # For the statistics box:
0181     tdrStyle.SetOptFile(0)
0182     tdrStyle.SetOptStat(0); # To display the mean and RMS:   SetOptStat("mr")
0183     tdrStyle.SetStatColor(kWhite)
0184     tdrStyle.SetStatFont(42)
0185     tdrStyle.SetStatFontSize(0.025)
0186     tdrStyle.SetStatTextColor(1)
0187     tdrStyle.SetStatFormat("6.4g")
0188     tdrStyle.SetStatBorderSize(1)
0189     tdrStyle.SetStatH(0.1)
0190     tdrStyle.SetStatW(0.15)
0191
0192     # Margins:
0193     tdrStyle.SetPadTopMargin(0.05)
0194     tdrStyle.SetPadBottomMargin(0.1)
0195     tdrStyle.SetPadLeftMargin(0.12)
0196     tdrStyle.SetPadRightMargin(0.05)
0197
0198     # For the Global title:
0199     tdrStyle.SetOptTitle(1)
0200     tdrStyle.SetTitleFont(42)
0201     tdrStyle.SetTitleColor(1)
0202     tdrStyle.SetTitleTextColor(1)
0203     tdrStyle.SetTitleFillColor(10)
0204     tdrStyle.SetTitleFontSize(0.0525)
0205     tdrStyle.SetTitleH(0); # Set the height of the title box
0206     tdrStyle.SetTitleW(0); # Set the width of the title box
0207     tdrStyle.SetTitleX(0.5); # Set the position of the title box
0208     tdrStyle.SetTitleY(1.0); # Set the position of the title box
0209     tdrStyle.SetTitleStyle(1001);
0210     tdrStyle.SetTitleBorderSize(0);
0211     tdrStyle.SetTitleAlign(23)
0212
0213     # For the axis titles:
0214     tdrStyle.SetTitleColor(1, "XYZ")
0215     tdrStyle.SetTitleFont(42, "XYZ")
0216     tdrStyle.SetTitleSize(0.05, "XY")
0217     tdrStyle.SetTitleSize(0.035, "Z")
0218     tdrStyle.SetTitleXOffset(1.0)
0219     tdrStyle.SetTitleYOffset(1.0)
0220
0221     # For the axis labels:
0222     tdrStyle.SetLabelColor(1, "XYZ")
0223     tdrStyle.SetLabelFont(42, "XYZ")
0224     tdrStyle.SetLabelOffset(5e-3, "XYZ")
0225     tdrStyle.SetLabelSize(0.03, "XYZ")
0226
0227     # For the axis:
0228     tdrStyle.SetAxisColor(1, "XYZ")
0229     tdrStyle.SetStripDecimals(kTRUE)
0230     tdrStyle.SetTickLength(0.03, "XYZ")
0231     tdrStyle.SetNdivisions(510, "XYZ")
0232     tdrStyle.SetPadTickX(1)  # To get tick marks on the opposite side of the frame
0233     tdrStyle.SetPadTickY(1)
0234
0235     # Change for log plots:
0236     tdrStyle.SetOptLogx(0)
0237     tdrStyle.SetOptLogy(0)
0238     tdrStyle.SetOptLogz(0)
0239
0240     # Miscellaneous
0241     tdrStyle.SetLegendBorderSize(0)
0242
0243     # Postscript options:
0244     tdrStyle.SetPaperSize(20.,20.)
0245
0246     tdrStyle.cd()
0247
0248 def drawEtaValues():
0249     """Function to draw the eta.
0250
0251     Function to draw the eta references on top of an already existing
0252     TCanvas. The lines and labels drawn are collected inside a list and
0253     the list is returned to the user to extend the live of the objects
0254     contained, otherwise no lines and labels will be drawn, since they
0255     will be garbage-collected as soon as this function returns.
0256     """
0257
0258     # Add eta labels
0259     keep_alive = []
0260     etas = [ 0.2*i for i in range(-17,18) ]
0261
0262     etax = 2850.
0263     etay = 1240.
0264     lineL = 110.
0265     offT = 10.
0266
0267     for ieta in etas:
0268         th = 2*atan(exp(-ieta))
0269         talign = 21
0270
0271         #IP
0272         lineh = TLine(-20.,0.,20.,0.)
0273         lineh.Draw()
0274         linev = TLine(0.,-10.,0.,10.)
0275         linev.Draw()
0276         keep_alive.append(lineh)
0277         keep_alive.append(linev)
0278
0279         x1 = 0
0280         y1 = 0
0281         if ieta>-1.6 and ieta<1.6:
0282             x1 = etay/tan(th)
0283             y1 = etay
0284         elif ieta <=-1.6:
0285             x1 = -etax
0286             y1 = -etax*tan(th)
0287             talign = 11
0288         elif ieta>=1.6:
0289             x1 = etax
0290             y1 = etax*tan(th)
0291             talign = 31
0292         x2 = x1+lineL*cos(th)
0293         y2 = y1+lineL*sin(th)
0294         xt = x2
0295         yt = y2+offT
0296
0297         line1 = TLine(x1,y1,x2,y2)
0298         line1.Draw()
0299         keep_alive.append(line1)
0300
0301         text = "%3.1f" % ieta
0302         t1 = TLatex(xt, yt, '%s' % ('#eta = 0' if ieta == 0 else text))
0303         t1.SetTextSize(0.03)
0304         t1.SetTextAlign(talign)
0305         t1.Draw()
0306         keep_alive.append(t1)
0307     return keep_alive