analysis_INTT/plot/beamspotReco.py

0001 #! /usr/bin/env python
0002 from optparse import OptionParser
0003 import sys
0004 import os
0005 import datetime
0006 from array import *
0007 from ROOT import TH1F, TH2F, TF1, TFile, TTree, TCanvas, TGraphErrors, TLatex, gROOT, gPad, kRed, kBlue, kGreen, kBlack, gSystem
0008 import numpy as np
0009 import math
0010 import glob
0011
0012 LeftMargin = 0.15
0013 RightMargin = 0.08
0014 TopMargin = 0.08
0015 BottomMargin = 0.13
0016
0017 # gROOT.LoadMacro('./sPHENIXStyle/sPhenixStyle.C')
0018 # gROOT.ProcessLine('SetsPhenixStyle()')
0019 gROOT.SetBatch(True)
0020
0021 # Function to draw TGraphErrors
0022 def Draw_TGraphErrors(gr, xtitle, ytitle, plotname):
0023
0024     data = True if 'BCO' in xtitle else False
0025
0026     # Fit a straight line
0027     f = TF1('f', 'pol1', gr.GetPointX(0)*0.9997, gr.GetPointX(gr.GetN()-1)*1.0003)
0028     gr.Fit('f', 'R')
0029
0030     # Get the maximal of the Y axis
0031     ymax = gr.GetY()[0]
0032     ymin = gr.GetY()[0]
0033     for i in range(gr.GetN()):
0034         if gr.GetY()[i]+gr.GetErrorY(i) > ymax:
0035             ymax = gr.GetY()[i]+gr.GetErrorY(i)
0036         if gr.GetY()[i]-gr.GetErrorY(i) < ymin:
0037             ymin = gr.GetY()[i]-gr.GetErrorY(i)
0038
0039     c = TCanvas('c', 'c', 800, 600)
0040     gPad.SetRightMargin(0.1)
0041     c.cd()
0042     gr.GetXaxis().SetTitle(xtitle)
0043     gr.GetYaxis().SetTitle(ytitle)
0044     # Set the range of the Y axis
0045     w = 0.12 if data else 0.05
0046     gr.GetHistogram().SetMaximum(gr.GetMean(2)+w)
0047     gr.GetHistogram().SetMinimum(gr.GetMean(2)-w)
0048     # gr.GetHistogram().SetMaximum(gr.GetMean(2)+6*gr.GetRMS(2))
0049     # gr.GetHistogram().SetMinimum(gr.GetMean(2)-6*gr.GetRMS(2))
0050     gr.SetMarkerStyle(20)
0051     gr.SetMarkerSize(0.5)
0052     gr.Draw('AP')
0053     f.SetLineColor(kRed)
0054     f.Draw('same')
0055     # print the mean and standard deviation of Y on the plot
0056     text = TLatex()
0057     text.SetNDC()
0058     text.SetTextSize(0.035)
0059     text.DrawLatex(0.2, 0.88, 'Avg. = {:.4g}, stddev. = {:.4g} '.format(gr.GetMean(2), gr.GetRMS(2)))
0060     # print the fit function on the plot
0061     text.DrawLatex(0.2, 0.83, 'Fit function: y = {:.4g} + {:.4g}x'.format(f.GetParameter(0), f.GetParameter(1)))
0062     # print the chi2/ndf on the plot
0063     chi2ndf = f.GetChisquare()/f.GetNDF()
0064     text.DrawLatex(0.2, 0.78, '#chi^{2}/ndf'+' = {:.4g}'.format(chi2ndf))
0065     c.SaveAs('{}.png'.format(plotname))
0066     c.SaveAs('{}.pdf'.format(plotname))
0067     if(c):
0068         c.Close()
0069         gSystem.ProcessEvents()
0070         del c
0071         c = 0
0072
0073
0074 if __name__ == '__main__':
0075     parser = OptionParser(usage="usage: %prog ver [options -h]")
0076     parser.add_option("-d", "--isdata", dest="isdata", action="store_true", default=False, help="Is data")
0077     parser.add_option("-i", "--infiledir", dest="infiledir", default='/sphenix/user/hjheng/sPHENIXRepo/analysis/dNdEta_Run2023/analysis_INTT/minitree/BeamspotMinitree_data_Run20869_HotDead_BCO_ADC_Survey', help="Input file directory")
0078     parser.add_option('-p', '--plotdir', dest='plotdir', type='string', default='Data_Run20869_HotDead_BCO_ADC_Survey', help='Plot directory')
0079
0080     (opt, args) = parser.parse_args()
0081     print('opt: {}'.format(opt))
0082
0083     isdata = opt.isdata
0084     infiledir = opt.infiledir
0085     plotdir = opt.plotdir
0086
0087     if os.path.isfile("{}/minitree_merged.root".format(infiledir)):
0088         os.system("rm {}/minitree_merged.root".format(infiledir))
0089         os.system("hadd -f -j 20 {}/minitree_merged.root {}/minitree_0*.root".format(infiledir, infiledir))
0090     else:
0091         os.system("hadd -f -j 20 {}/minitree_merged.root {}/minitree_0*.root".format(infiledir, infiledir))
0092
0093     os.makedirs('./BeamspotReco/{}'.format(plotdir), exist_ok=True)
0094
0095
0096     bs_x = np.array([])
0097     bs_y = np.array([])
0098     bs_x_err = np.array([])
0099     bs_y_err = np.array([])
0100     bs_bco_med = np.array([])
0101
0102     f = TFile('{}/minitree_merged.root'.format(infiledir), 'r')
0103     tree = f.Get('reco_beamspot')
0104     for idx in range(tree.GetEntries()):
0105         tree.GetEntry(idx)
0106         print ('entry={}, x={}, y={}, intt_bco_median={}'.format(idx, tree.x, tree.y, tree.intt_bco_median))
0107         bs_x = np.append(bs_x, tree.x)
0108         bs_x_err = np.append(bs_x_err, tree.sigma_x)
0109         bs_y = np.append(bs_y, tree.y)
0110         bs_y_err = np.append(bs_y_err, tree.sigma_y)
0111         if isdata:
0112             bs_bco_med = np.append(bs_bco_med, tree.intt_bco_median)
0113         else:
0114             bs_bco_med = np.append(bs_bco_med, idx)
0115
0116     # TGraphErrors for beamspot x
0117     gr_bs_x = TGraphErrors(len(bs_x), bs_bco_med, bs_x, np.zeros(len(bs_x)), bs_x_err)
0118     gr_bs_y = TGraphErrors(len(bs_y), bs_bco_med, bs_y, np.zeros(len(bs_y)), bs_y_err)
0119
0120     Draw_TGraphErrors(gr_bs_x, ('INTT BCO' if isdata else 'Sub-sample index'), 'Beamspot x [cm]', './BeamspotReco/{}/beamspot_x'.format(plotdir))
0121     Draw_TGraphErrors(gr_bs_y, ('INTT BCO' if isdata else 'Sub-sample index'), 'Beamspot y [cm]', './BeamspotReco/{}/beamspot_y'.format(plotdir))
0122
0123     f.Close()
0124
0125     # A TTree to store the beamspot reco results: average x, y
0126     f_out = TFile('{}/beamspot_res.root'.format(infiledir), 'recreate')
0127     t_out = TTree('minitree', 'minitree')
0128     bs_x_mean = np.array([gr_bs_x.GetMean(2)], dtype='float64')
0129     bs_y_mean = np.array([gr_bs_y.GetMean(2)], dtype='float64')
0130     t_out.Branch('bs_x', bs_x_mean, 'bs_x/D')
0131     t_out.Branch('bs_y', bs_y_mean, 'bs_y/D')
0132     t_out.Fill()
0133     t_out.Write()
0134     f_out.Close()
0135
0136     # intermediate plots
0137     f_check = TFile('{}/minitree_00000.root'.format(infiledir), 'r')
0138     d0phi0dist = f_check.Get('d0phi0dist')
0139     d0phi0dist_bkgrm = f_check.Get('d0phi0dist_bkgrm')
0140     d0phi0dist_bkgrm_prof = f_check.Get('d0phi0dist_bkgrm_prof')
0141     maxgraph_rmoutl = f_check.Get('maxgraph_rmoutl')
0142
0143     c = TCanvas('c', 'c', 800, 700)
0144     gPad.SetRightMargin(0.12)
0145     c.cd()
0146     d0phi0dist.GetYaxis().SetTitle('DCA w.r.t origin [cm]')
0147     d0phi0dist.GetXaxis().SetTitle('#phi_{PCA} [rad]')
0148     d0phi0dist.Draw('colz')
0149     c.SaveAs('./BeamspotReco/{}/d0phi0dist.png'.format(plotdir))
0150     c.SaveAs('./BeamspotReco/{}/d0phi0dist.pdf'.format(plotdir))
0151     c.Clear()
0152     d0phi0dist_bkgrm.GetYaxis().SetTitle('DCA w.r.t origin [cm]')
0153     d0phi0dist_bkgrm.GetXaxis().SetTitle('#phi_{PCA} [rad]')
0154     d0phi0dist_bkgrm.Draw('colz')
0155     maxgraph_rmoutl.SetMarkerColorAlpha(kRed, 0.5)
0156     maxgraph_rmoutl.SetMarkerStyle(20)
0157     maxgraph_rmoutl.SetMarkerSize(0.6)
0158     maxgraph_rmoutl.SetLineColorAlpha(kRed, 0.5)
0159     maxgraph_rmoutl.Draw('p same')
0160     c.SaveAs('./BeamspotReco/{}/d0phi0dist_bkgrm_wGraphFit.png'.format(plotdir))
0161     c.SaveAs('./BeamspotReco/{}/d0phi0dist_bkgrm_wGraphFit.pdf'.format(plotdir))
0162     c.Clear()
0163