Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-08-05 08:11:19

0001 #! /usr/bin/env python
0002 from optparse import OptionParser
0003 import sys
0004 import os
0005 import datetime
0006 from array import *
0007 import ROOT
0008 import numpy as np
0009 import pandas as pd
0010 import math
0011 import glob
0012 import ctypes
0013 from plotUtil import Draw_2Dhist
0014 from sigmaEff import minimum_size_range
0015 
0016 ROOT.gROOT.LoadMacro('./sPHENIXStyle/sPhenixStyle.C')
0017 ROOT.gROOT.ProcessLine('SetsPhenixStyle()')
0018 ROOT.gROOT.SetBatch(True)
0019 
0020 TickSize = 0.03
0021 AxisTitleSize = 0.05
0022 AxisLabelSize = 0.04
0023 LeftMargin = 0.15
0024 RightMargin = 0.08
0025 TopMargin = 0.08
0026 BottomMargin = 0.13
0027 
0028 global nbinsdphi, nbinsdca 
0029 global dphiinterval, dcainterval
0030 
0031 def Draw_1Dhist_wTF1(hist, histdata, norm1, logy, ymaxscale, XaxisName, Ytitle_unit, outname):
0032     xmin, xmax = minimum_size_range(histdata, 68.5)
0033     effsigma = xmax - xmin
0034 
0035     hist.Sumw2()
0036     binwidth = hist.GetXaxis().GetBinWidth(1)
0037     histmax = hist.GetMaximum()
0038     c = ROOT.TCanvas('c', 'c', 800, 700)
0039     if norm1:
0040         hist.Scale(1. / hist.Integral(-1, -1))
0041     if logy:
0042         c.SetLogy()
0043     c.cd()
0044     ROOT.gPad.SetRightMargin(RightMargin)
0045     ROOT.gPad.SetTopMargin(TopMargin)
0046     ROOT.gPad.SetLeftMargin(LeftMargin)
0047     ROOT.gPad.SetBottomMargin(BottomMargin)
0048     if norm1:
0049         if Ytitle_unit == '':
0050             hist.GetYaxis().SetTitle(
0051                 'Normalized entries / ({:g})'.format(binwidth))
0052         else:
0053             hist.GetYaxis().SetTitle(
0054                 'Normalized entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
0055     else:
0056         if Ytitle_unit == '':
0057             hist.GetYaxis().SetTitle('Entries / ({:g})'.format(binwidth))
0058         else:
0059             hist.GetYaxis().SetTitle(
0060                 'Entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
0061 
0062     # hist.GetXaxis().SetRangeUser(hist.GetBinLowEdge(1)-binwidth, hist.GetBinLowEdge(hist.GetNbinsX())+2*binwidth)
0063     if logy:
0064         hist.GetYaxis().SetRangeUser(hist.GetMinimum(0)*0.5, histmax * ymaxscale)
0065     else:
0066         hist.GetYaxis().SetRangeUser(0., histmax * ymaxscale)
0067     hist.GetXaxis().SetTitle(XaxisName)
0068     hist.GetXaxis().SetTickSize(TickSize)
0069     hist.GetXaxis().SetTitleSize(AxisTitleSize)
0070     hist.GetXaxis().SetLabelSize(AxisLabelSize)
0071     hist.GetYaxis().SetTickSize(TickSize)
0072     hist.GetYaxis().SetTitleSize(AxisTitleSize)
0073     hist.GetYaxis().SetLabelSize(AxisLabelSize)
0074     hist.GetXaxis().SetTitleOffset(1.1)
0075     hist.GetYaxis().SetTitleOffset(1.35)
0076     hist.SetLineColor(1)
0077     hist.SetLineWidth(2)
0078     hist.Draw('hist')
0079     # Gaussian fit
0080     # maxbincenter = hist.GetBinCenter(hist.GetMaximumBin())
0081     # f1 = ROOT.TF1('f1', 'gaus', maxbincenter - hist.GetStdDev(), maxbincenter + hist.GetStdDev())
0082     # f1.SetParameter(1,hist.GetMean())
0083     # f1.SetParLimits(1,-1,1)
0084     # f1.SetParameter(2,hist.GetStdDev())
0085     # f1.SetParLimits(2,0,5)
0086     # f1.SetLineColorAlpha(ROOT.TColor.GetColor('#F54748'), 0.9)
0087     # hist.Fit(f1,'NQ')
0088     # f1.Draw('same')
0089     # Draw lines
0090     l1 = ROOT.TLine(xmin, 0, xmin, histmax)
0091     l1.SetLineColor(ROOT.kRed)
0092     l1.SetLineStyle(ROOT.kDashed)
0093     l1.SetLineWidth(2)
0094     l1.Draw()
0095     l2 = ROOT.TLine(xmax, 0, xmax, histmax)
0096     l2.SetLineColor(ROOT.kRed)
0097     l2.SetLineStyle(ROOT.kDashed)
0098     l2.SetLineWidth(2)
0099     l2.Draw()
0100     leg = ROOT.TLegend((1-RightMargin)-0.45, (1-TopMargin)-0.1,
0101                   (1-RightMargin)-0.1, (1-TopMargin)-0.03)
0102     leg.SetTextSize(0.045)
0103     leg.SetFillStyle(0)
0104     leg.AddEntry('', '#it{#bf{sPHENIX}} Simulation', '')
0105     # leg.AddEntry('','Au+Au #sqrt{s_{NN}}=200 GeV','')
0106     leg.Draw()
0107     leg2 = ROOT.TLegend(0.45, 0.67, 0.88, 0.8)
0108     leg2.SetTextSize(0.033)
0109     leg2.SetFillStyle(0)
0110     # leg2.AddEntry(f1,'Gaussian fit','l')
0111     # leg2.AddEntry('', '#mu={0:.2e}#pm{1:.2e}'.format(f1.GetParameter(1), f1.GetParError(1)), '')
0112     # leg2.AddEntry('', '#sigma={0:.2e}#pm{1:.2e}'.format(f1.GetParameter(2), f1.GetParError(2)), '')
0113     leg2.AddEntry('', '#LT\Deltaz#GT={0:.2e}#pm{1:.2e} {2}'.format(hist.GetMean(), hist.GetMeanError(), Ytitle_unit), '')
0114     # leg2.AddEntry('', '#sigma={0:.2e}#pm{1:.2e}'.format(hist.GetStdDev(), hist.GetStdDevError()), '')
0115     leg2.AddEntry('', '#sigma_{{eff}} (68.5%) ={0:.4g} {1}'.format(effsigma, Ytitle_unit), '')
0116     leg2.Draw()
0117     c.RedrawAxis()
0118     c.Draw()
0119     c.SaveAs(outname+'.pdf')
0120     c.SaveAs(outname+'.png')
0121     if(c):
0122         c.Close()
0123         ROOT.gSystem.ProcessEvents()
0124         del c
0125         c = 0
0126     # return f1.GetParameter(1), f1.GetParError(1), f1.GetParameter(2), f1.GetParError(2)
0127     # return hist.GetMean(), hist.GetMeanError(), hist.GetStdDev(), hist.GetStdDevError()
0128         return hist.GetMean(), hist.GetMeanError(), effsigma, 0.0
0129 
0130 
0131 def Draw2Dhisttable(hist, XaxisName, YaxisName, ZaxisName, DrawOpt, outname):
0132     c = ROOT.TCanvas('c', 'c', 4000, 3700)
0133     c.cd()
0134     ROOT.gPad.SetRightMargin(0.185)
0135     ROOT.gPad.SetTopMargin(TopMargin)
0136     ROOT.gPad.SetLeftMargin(LeftMargin)
0137     ROOT.gPad.SetBottomMargin(BottomMargin)
0138     ROOT.gStyle.SetPaintTextFormat('1.5f')
0139     hist.SetMarkerSize(0.4)
0140     hist.GetXaxis().SetTitle(XaxisName)
0141     hist.GetYaxis().SetTitle(YaxisName)
0142     hist.GetZaxis().SetTitle(ZaxisName)
0143     hist.GetXaxis().SetTickSize(TickSize)
0144     hist.GetYaxis().SetTickSize(TickSize)
0145     hist.GetXaxis().SetTitleSize(AxisTitleSize)
0146     hist.GetYaxis().SetTitleSize(AxisTitleSize)
0147     hist.GetZaxis().SetTitleSize(AxisTitleSize)
0148     hist.GetXaxis().SetLabelSize(AxisLabelSize)
0149     hist.GetYaxis().SetLabelSize(AxisLabelSize)
0150     hist.GetXaxis().SetTitleOffset(1.1)
0151     hist.GetYaxis().SetTitleOffset(1.3)
0152     hist.GetZaxis().SetTitleOffset(1.3)
0153     hist.GetZaxis().SetLabelSize(AxisLabelSize)
0154     hist.GetZaxis().SetRangeUser(hist.GetMinimum(), hist.GetMaximum())
0155     hist.LabelsOption("v")
0156     hist.SetContour(10000)
0157     hist.Draw(DrawOpt)
0158     # bx,by,bz = (ctypes.c_int(),ctypes.c_int(),ctypes.c_int()) # Ref: https://github.com/svenkreiss/decouple/blob/master/Decouple/plot_utils.py#L44-L52
0159     # hist.GetBinXYZ(hist.GetMinimumBin(),bx,by,bz)
0160     # binxCenter = hist.GetXaxis().GetBinCenter(bx.value)
0161     # binyCenter = hist.GetYaxis().GetBinCenter(by.value)
0162     # el = ROOT.TEllipse(binxCenter,binyCenter,.1,.1)
0163     # el.SetFillColor(2)
0164     # el.Draw()
0165     c.RedrawAxis()
0166     c.Draw()
0167     c.SaveAs(outname+'.pdf')
0168     c.SaveAs(outname+'.png')
0169     if(c):
0170         c.Close()
0171         ROOT.gSystem.ProcessEvents()
0172         del c
0173         c = 0
0174 
0175 def doFitSaveFitresult(reshist, resdata, fitresname, histwTF1name):
0176     print ('In doFitSaveFitresult()')
0177     out = ROOT.TFile(fitresname, 'recreate') 
0178     tree = ROOT.TTree('tree', 'tree')
0179     arrmean = array('f', [0])
0180     arrmeanerr = array('f', [0])
0181     arrsigma = array('f', [0])
0182     arrsigmaerr = array('f', [0])
0183     tree.Branch('mean', arrmean, 'mean/F') 
0184     tree.Branch('mean_err', arrmeanerr, 'mean_err/F')
0185     tree.Branch('sigma', arrsigma, 'sigma/F') 
0186     tree.Branch('sigma_err', arrsigmaerr, 'sigma_err/F')
0187     arrmean[0], arrmeanerr[0], arrsigma[0], arrsigmaerr[0] = Draw_1Dhist_wTF1(reshist, resdata, False, True, 50, '#Deltaz(PV_{Truth}, PV_{Reco}) [cm]', 'cm', histwTF1name)
0188     tree.Fill()
0189     tree.Write('', ROOT.TObject.kOverwrite) 
0190     out.Close()
0191     del reshist, out, tree, arrmean, arrmeanerr, arrsigma, arrsigmaerr
0192 
0193 def main(dphitxt, dcatxt, infiledir, outfiledir):
0194     print ('In main()')
0195 
0196     hM_DiffVtxZ = ROOT.TH1F('hM_DiffVtxZ', 'hM_DiffVtxZ_altrange', 200, -5, 5)
0197     hM_DiffVtxZ_NClusLayer1 = ROOT.TH2F('hM_DiffVtxZ_NClusLayer1', 'hM_DiffVtxZ_NClusLayer1', 200, -5, 5, 200, 0, 4000)
0198 
0199     df = ROOT.RDataFrame('minitree', '{}/dphi{}deg_dca{}cm/INTTVtxZ.root'.format(infiledir, dphitxt, dcatxt))
0200     data = df.AsNumpy(columns=['NTruthVtx','PV_z','TruthPV_z','NClusLayer1'])
0201 
0202     for i, ntruthvtx in enumerate(data['NTruthVtx']):
0203         if ntruthvtx == 1:
0204             hM_DiffVtxZ.Fill((data['PV_z'][i] - data['TruthPV_z'][i]))
0205             hM_DiffVtxZ_NClusLayer1.Fill((data['PV_z'][i] - data['TruthPV_z'][i]), data['NClusLayer1'][i])
0206 
0207     diffarray = data['PV_z'] - data['TruthPV_z']
0208     doFitSaveFitresult(hM_DiffVtxZ, diffarray, outfiledir+'fitresult_dPhi{}deg_dca{}cm.root'.format(dphitxt, dcatxt), outfiledir+'DiffVtxZ_wTF1')
0209     Draw_2Dhist(hM_DiffVtxZ_NClusLayer1, False, False, False, 0.13, '#Deltaz(PV_{Truth}, PV_{Reco}) [cm]', 'Number of clusters (inner)', 'colz', outfiledir+'DiffVtxZ_NClusLayer1_dphi{}deg_dca{}cm'.format(dphitxt, dcatxt))
0210 
0211     del df, data, hM_DiffVtxZ, hM_DiffVtxZ_NClusLayer1
0212     
0213 def get2DTable(fitresdir):
0214     print('In get2DTable')
0215     hM_mean_dPhidcaCut = ROOT.TH2F('hM_mean_dPhidcaCut', 'hM_mean_dPhidcaCut', nbinsdphi, 0, nbinsdphi, nbinsdca, 0, nbinsdca)
0216     hM_sigma_dPhidcaCut = ROOT.TH2F('hM_sigma_dPhidcaCut', 'hM_sigma_dPhidcaCut', nbinsdphi, 0, nbinsdphi, nbinsdca, 0, nbinsdca)
0217 
0218     l_sigma = []
0219     l_mean = []
0220     l_dPhi = []
0221     l_dca = []
0222     for i in range(0, nbinsdphi):
0223         for j in range(0, nbinsdca):
0224             dphicut_deg = dphiinterval * (i+1)
0225             dcacut = dcainterval * (j+1)
0226             dphitext = '{:.3f}'.format(dphicut_deg).replace('.', 'p')
0227             dcadtext = '{:.3f}'.format(dcacut).replace('.', 'p')
0228             try:
0229                 fitrs = ROOT.RDataFrame('tree', '{}/dPhi{}deg_dca{}cm/fitresult_dPhi{}deg_dca{}cm.root'.format(fitresdir, dphitext, dcadtext, dphitext, dcadtext))
0230                 res = fitrs.AsNumpy(columns=['mean','sigma'])
0231                 l_dPhi.append(dphicut_deg)
0232                 l_dca.append(dcacut)
0233                 l_mean.append(res['mean'][0])
0234                 l_sigma.append(res['sigma'][0])
0235                 hM_mean_dPhidcaCut.SetBinContent(i + 1, j + 1, res['mean'][0])
0236                 hM_sigma_dPhidcaCut.SetBinContent(i + 1, j + 1, res['sigma'][0])
0237                 hM_mean_dPhidcaCut.GetXaxis().SetBinLabel(i+1, '{:.2f}'.format((i+1)*dphiinterval))
0238                 hM_mean_dPhidcaCut.GetYaxis().SetBinLabel(j+1, '{:.3f}'.format((j+1)*dcainterval))
0239                 hM_sigma_dPhidcaCut.GetXaxis().SetBinLabel(i+1, '{:.2f}'.format((i+1)*dphiinterval))
0240                 hM_sigma_dPhidcaCut.GetYaxis().SetBinLabel(j+1, '{:.3f}'.format((j+1)*dcainterval))
0241                 del fitrs, res
0242 
0243             except Exception as e:
0244                 print(e)
0245 
0246     df_opt = pd.DataFrame(list(zip(l_dPhi, l_dca, l_mean, l_sigma)), columns =['dPhiBin', 'dZBin', 'mean', 'sigma'])
0247     df_opt = df_opt.sort_values(by='sigma', ascending=True, ignore_index=True)        
0248 
0249     del l_dPhi, l_dca, l_mean, l_sigma
0250     return hM_mean_dPhidcaCut, hM_sigma_dPhidcaCut, df_opt
0251 
0252 if __name__ == '__main__':
0253     parser = OptionParser(usage="usage: %prog ver [options -h]")
0254     parser.add_option("-i", "--infiledir", dest="infiledir", default='/sphenix/user/hjheng/TrackletAna/minitree/INTT/VtxEvtMap_ana382_zvtx-20cm_dummyAlignParams', help="Input file directory")
0255     parser.add_option("-o", "--outdirprefix", dest="outdirprefix", default='ana382_zvtx-20cm_dummyAlignParams', help="Output file directory")
0256     parser.add_option("-f", "--dofit", action="store_true", dest="dofit", default=False, help="Do fit")
0257     parser.add_option("-p", "--nbinsdphi", dest="nbinsdphi", default=20, help="Number of bins in dPhi")
0258     parser.add_option("-d", "--nbinsdca", dest="nbinsdca", default=20, help="Number of bins in dca")
0259 
0260     (opt, args) = parser.parse_args()
0261     print('opt: {}'.format(opt)) 
0262 
0263     infiledir = opt.infiledir
0264     outdirprefix = opt.outdirprefix
0265     dofit = opt.dofit
0266     nbinsdphi = int(opt.nbinsdphi)
0267     nbinsdca = int(opt.nbinsdca)
0268 
0269     plotpath = './RecoPV_optimization/'
0270     os.makedirs('{}/{}'.format(plotpath,outdirprefix), exist_ok=True)
0271 
0272     dphiinterval = 0.1
0273     dcainterval = 0.025
0274 
0275     for i in range(0, nbinsdphi):
0276         for j in range(0, nbinsdca):
0277             dphicut_deg = dphiinterval * (i+1)
0278             dphicut_rad = dphicut_deg * (np.pi / 180.)
0279             dcacut = dcainterval * (j+1)
0280             dphitext = '{:.3f}'.format(dphicut_deg).replace('.', 'p')
0281             dcadtext = '{:.3f}'.format(dcacut).replace('.', 'p')
0282 
0283             outpath = './RecoPV_optimization/{}/dPhi{}deg_dca{}cm/'.format(outdirprefix,dphitext,dcadtext)
0284             os.makedirs(outpath, exist_ok=True)
0285             try:
0286                 if dofit == True:
0287                     main(dphitext,dcadtext,infiledir,outpath)
0288             except Exception as e:
0289                 print(e)
0290 
0291     hM_mean_dPhidcaCut, hM_sigma_dPhidcaCut, df_opt = get2DTable('./RecoPV_optimization/{}'.format(outdirprefix))
0292 
0293     ROOT.gStyle.SetPalette(ROOT.kThermometer)
0294     ROOT.TGaxis.SetMaxDigits(3)
0295     # set the number of divisions to show the bin labels 
0296     # hM_mean_dPhidcaCut.GetYaxis().SetMaxDigits(4)
0297     # hM_sigma_dPhidcaCut.GetYaxis().SetMaxDigits(4)
0298     # hM_mean_dPhidcaCut.GetYaxis().SetDecimals()
0299     # hM_sigma_dPhidcaCut.GetYaxis().SetDecimals()
0300     
0301     Draw2Dhisttable(hM_mean_dPhidcaCut, '#Delta#phi [degree]', 'DCA [cm]', '#LT\Deltaz#GT [cm]', 'colztext45', '{}/{}/Optimization_GaussianMean_Table'.format(plotpath,outdirprefix))
0302     Draw2Dhisttable(hM_sigma_dPhidcaCut, '#Delta#phi [degree]', 'DCA [cm]', '#sigma_{eff} (68.5%) [cm]', 'colztext45', '{}/{}/Optimization_GaussianSigma_Table'.format(plotpath,outdirprefix))
0303 
0304     for i in range(5):
0305         sigmainfo = df_opt.iloc[i]
0306         print(sigmainfo)
0307 
0308