plotHME_comparison.py

import ROOT
#from helper import *
import numpy as np
#from localSamplelist import *
import os

treename = "Events"


Plotfolder = "HMEcomparison_MuMu_20210519/"
Plotfolder = "HMEcomparison_MuMu_20210601_METRES/"
#Plotfolder = "HMEcomparison_MuMu_20210601_Rebins/"
Plotfolder = "HMEcomparison_MuMu_20210601_NewWeight/"
#Plotfolder = "HMEcomparison_MuMu_20210609_gravity/"
os.system("mkdir "+Plotfolder)
### plot HME from one file

def findHistWidth(th1):
    nbins = th1.GetNbinsX()
    width = th1.GetBinWidth(1)
    total = th1.Integral()+th1.GetBinContent(0)+th1.GetBinContent(nbins+1)
    if total < 1.0:
        return 0.0,0.0,0.0
    xtotal = th1.GetBinContent(0)
    lowRatio = 0.16; highRatio = 0.84; mRatio = 0.5
    xbin1 = 0; xbin2=0; xbin3=0
    ratio1 = 0.0; ratio2 = 0.0; ratio3=0.0
    findLowEdge = False; findHighEdge = False; findMedian = False
    for i in range(1, nbins+1):
        xtotal += th1.GetBinContent(i)
        #print(i," bincontent ", th1.GetBinContent(i) , " ratio ", xtotal/total," integral ", xtotal, " Total ", total)
        if xtotal > total*lowRatio and not findLowEdge:
            xbin1 = i; findLowEdge = True; ratio1 = xtotal/total
        if xtotal > total*highRatio and not findHighEdge:
            xbin2 = i; findHighEdge = True; ratio2 = xtotal/total
        if xtotal > total*mRatio and not findMedian:
            xbin3 = i; findMedian = True; ratio3 = xtotal/total
        if findLowEdge and findHighEdge:
            break
    print("total ", total, " nbins ", nbins, " lowedge bin ", xbin1, " highedge bin ", xbin2," median bin ", xbin3)
    ratio1_low = (ratio1-th1.GetBinContent(xbin1)/total)
    lowEdge = th1.GetXaxis().GetBinLowEdge(xbin1) + width*(lowRatio-ratio1_low)/(ratio1-ratio1_low)
    ratio2_low = (ratio2-th1.GetBinContent(xbin2)/total)
    highEdge = th1.GetXaxis().GetBinUpEdge(xbin2) - width*(ratio2-highRatio)/(ratio2-ratio2_low)
    ratio3_low = (ratio3-th1.GetBinContent(xbin3)/total)
    median = th1.GetXaxis().GetBinLowEdge(xbin3) + width*(mRatio-ratio3_low)/(ratio3-ratio3_low)

    print("ratio1 ", ratio1, " lowend ", ratio1_low, " low target ", lowRatio, " ratio2 ", ratio2, " lowend ", ratio2_low, " high target ", highRatio )
    print("low Edge ", lowEdge, " high Edge ", highEdge," median ", median)

    return lowEdge, highEdge, median

def plotHME(rfiles, legnames, variables, xbins, cut, mass, text, plotsuffix):
    
    nbin = xbins[0]; xmin = xbins[1]; xmax = xbins[2]
    #h1 = ROOT.TH1F("hist1","hist1",xbins[0], xbins[1], xbins[2])

    histlist = []
    tchlist = []

    colors = [800-4, 593, 820, 616-7, 432, 418]
    #colors = [593, 820, 616-7, 432, 418]

    xlegend = 0.53
    if mass > 2000.0:
        xlegend = 0.1
    legend = ROOT.TLegend(xlegend, 0.55, xlegend+0.45, 0.55+len(legnames)*.075); 
    c1 = ROOT.TCanvas("c1","c1",800,600)
    total_env = None
    hs = ROOT.THStack("Shapecomparison", "HME distribution in #mu#mu channel, Signal M=%d GeV"%mass)
    for i,legname in enumerate(legnames):
        variable = variables[i]
        #weight = "total_weight"
        histlist.append(ROOT.TH1F("hist%d"%i,"hist%d"%i,xbins[0], xbins[1], xbins[2]))
        tchlist.append(ROOT.TChain(treename))
        tchlist[i].Add(rfiles[i])
        finalcut = cut+"&& %s >= 250.0"%variable
        tchlist[i].Draw(variable + ">> hist%d"%i, finalcut)
        lowE,highE,median = findHistWidth(histlist[i])

        #integral = histlist[i].Integral()
        mean = histlist[i].GetMean()
        entries = histlist[i].GetEntries()
        histlist[i].SetLineColor(colors[i])
        histlist[i].SetLineWidth(2)
        histlist[i].SetStats(0)
        #legend.AddEntry(histlist[i], legname+": %d events"%histlist[i].Integral()+" Width = [%d, %d]"%(lowE, highE),"l")
        totalevt = histlist[i].Integral()+histlist[i].GetBinContent(0)+histlist[i].GetBinContent(xbins[0]+1)
        legtext = "#splitline{%s: %d events}{Width:[%d, %d]=%d GeV, Median=%d GeV}"%(legname, totalevt, lowE, highE, highE-lowE, median)
        legend.AddEntry(histlist[i], legtext,"l")

        hs.Add(histlist[i])
        #if i == 0:
        #    histlist[i].SetTitle("Drell-Yan distribution, no btagging")
        #    histlist[i].GetXaxis().SetTitle(variable)
        #    histlist[i].Draw()
        #else:
        #    histlist[i].Draw("same")


    #legend.SetHeader("Signal M=%d"%mass)
    #c1.SetLogy()
    c1.SetGridx()
    c1.SetTickx()
    c1.SetGridy()
    c1.SetTicky()
    hs.Draw("nostackhist")
    #hs.GetHistogram().GetXaxis().SetNdivisions(520)
    legend.Draw("same")
    txt = ROOT.TText(0.4, 0.4, text)
    txt.SetNDC()
    txt.SetTextFont(42)
    #txt.SetTextSize(42)
    txt.Draw("same")
    hs.GetHistogram().GetXaxis().SetTitle("Reco HME")
    hs.GetHistogram().GetYaxis().SetTitle("Events")
    c1.SaveAs(Plotfolder+"HMEdistribution_"+plotsuffix+".pdf")
    c1.SaveAs(Plotfolder+"HMEdistribution_"+plotsuffix+".C")



### plot nosolutions from one file
def plotnsolutions(rfile, legnames, variables, xbins, cut, mass, plotsuffix):
    nbin = xbins[0]; xmin = xbins[1]; xmax = xbins[2]
    #h1 = ROOT.TH1F("hist1","hist1",xbins[0], xbins[1], xbins[2])

    histlist = []
    tchlist = []

    colors = [800-4, 593, 820, 616-7, 432, 418]
    #colors = [593, 820, 616-7, 432, 418]

    xlegend = 0.6
    #if mass >= 2000.0:
    #    xlegend = 0.1
    legend = ROOT.TLegend(xlegend,0.65, xlegend+0.3, 0.7+len(legnames)*.05); 
    c1 = ROOT.TCanvas("c1","c1",800,600)
    total_env = None
    #hs = ROOT.THStack("Shapecomparison", "HME distribution in #mu#mu channel, Signal M=%d"%mass)
    hs = ROOT.THStack("Shapecomparison", "Number of iterations with N solutions in HME, #mu#mu, Signal M=%d"%mass)
    for i,legname in enumerate(legnames):
        variable = variables[i]
        #weight = "total_weight"
        histlist.append(ROOT.TH1F("hist%d"%i,"hist%d"%i,xbins[0], xbins[1], xbins[2]))
        tchlist.append(ROOT.TChain(treename))
        tchlist[i].Add(rfile)
        finalcut = cut
        tchlist[i].Draw(variable + ">> hist%d"%i, finalcut)

        #integral = histlist[i].Integral()
        overflow =  histlist[i].GetBinContent(xbins[0]+1)
        lastbinContent = histlist[i].GetBinContent(xbins[0])
        histlist[i].SetBinContent(xbins[0], overflow+lastbinContent)
        mean = histlist[i].GetMean()
        entries = histlist[i].GetEntries()
        #histlist[i].SetLineColor(colors[i])
        histlist[i].SetFillColor(colors[i])
        histlist[i].SetLineWidth(2)
        histlist[i].SetStats(0)
        #legend.AddEntry(histlist[i], legname+": %d events"%histlist[i].Integral(),"l")
        legend.AddEntry(histlist[i], legname,"f")
        hs.Add(histlist[i])
        #if i == 0:
        #    histlist[i].SetTitle("Drell-Yan distribution, no btagging")
        #    histlist[i].GetXaxis().SetTitle(variable)
        #    histlist[i].Draw()
        #else:
        #    histlist[i].Draw("same")


    #legend.SetHeader("Signal M=%d"%mass)
    c1.SetLogy()
    c1.SetGridx()
    c1.SetGridy()
    #hs.Draw("nostackhist")
    hs.Draw()
    legend.Draw("same")
    txt = ROOT.TText(0.4, 0.5, text)
    #txt.Draw("same")
    hs.GetHistogram().GetXaxis().SetTitle("Num of iterations")
    hs.GetHistogram().GetYaxis().SetTitle("Entries")
    c1.SaveAs(Plotfolder+"NumIternations_Nsolutions_"+plotsuffix+".pdf")

variables = ["hme_mass_peak","hme_mass_peak_divSol","hme_mass_peak_boosted","hme_mass_peak_boosted_divSol"]
legnames =  ["2 subjets, weight case", "2 subjets, average case","fatjet, weight case","fatjet, average case"]

variables = ["hme_mass_peak","hme_mass_peak_divSol"]
legnames =  ["2 subjets, weight case", "2 subjets, average case"]


masslist = [1000, 1250, 1500, 1750, 2000, 2500, 3000, 250, 260, 270, 280, 300, 320, 350, 400, 450, 500, 550, 600, 650, 700,750, 800, 850, 900]
masslist = [1000, 1250, 1500, 1750, 2000, 2500, 3000, 600, 650, 700,750, 800, 850, 900]
#masslist = [1000]
#masslist = [320, 350]

nsolutions_var = ["nsolution0", "nsolution1","nsolution2","nsolution3","nsolution4"]
nsolutions_leg = ["no solution", "1 solution","2 solutions","3 solutions","4 solutions"]

iter_bins = [50, 0, 1000.0]

localdir = "/Users/taohuang/Documents/DiHiggs/20210510_HME/"
rfiledir = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210515/"
rfiledir = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210519/"
rfiledir_rebin2 = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin2_iter10k/"
rfiledir_rebin5 = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin5_iter10k/"
rfiledir_rebin10 = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin10_iter10k/"
rfiledir_rebin20 = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin20_iter10k/"
#rfiledir = "/afs/cern.ch/work/t/tahuang/HHAnalysis/HeavyMassEstimator/scripts/Skim2018ForTao_MuMu_Boosted1B_20210515/"
filelist = []
for mass in masslist:
    rfile = rfiledir+ "GluGluToRadionToHHTo2B2VTo2L2Nu_M-%d_HME.root"%(mass)
    filelist = [rfile, rfile, rfile, rfile]
    #text = "Radion M=%d GeV"%mass
    #cut  = "1"
    cut  = "hme_mass_peak>=250"
    #xbins = [35, 500.0, 4000.0]
    xbins = [90, 400.0, 4000.0]
    plotsuffix = "RadionBoosted_M%d_weight_average_rebin1_xbinsize100"%mass
    text = "HME likelihood hist binsize=1GeV"
    #plotHME(rfile, legnames, variables, xbins, cut, mass, text, plotsuffix)
    #plotnsolutions(rfile, nsolutions_leg, nsolutions_var, iter_bins, cut, mass, plotsuffix)

legnames_rebin = ["likelihood binsize=1GeV", "likelihood binsize=5GeV", "likelihood binsize=10GeV", "likelihood binsize=20GeV"]
variables = ["hme_mass_peak","hme_mass_peak", "hme_mass_peak", "hme_mass_peak" ]
for mass in masslist:
    filename = "GluGluToRadionToHHTo2B2VTo2L2Nu_M-%d_HME.root"%(mass)
    rfilelist_rebin = [rfiledir+filename, rfiledir_rebin5+filename, rfiledir_rebin10+filename, rfiledir_rebin20+filename]
    cut  = "1"
    #xbins = [70, 500.0, 4000.0]
    xbins = [90, 400.0, 4000.0]
    plotsuffix = "RadionBoosted_M%d_weight_average_rebincompare_xbinsize40"%mass
    text = " "
    #plotHME(rfilelist_rebin, legnames_rebin, variables, xbins, cut, mass, text, plotsuffix)


rfiledir_res30_cosdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes30_cosdphi/"
rfiledir_res40_cosdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes40_cosdphi/"
rfiledir_res50_cosdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes50_cosdphi/"
rfiledir_res25_coshdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes25_wrongcosh/"
rfiledir_res30_coshdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes30_wrongcosh/"
rfiledir_res40_coshdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes40_wrongcosh/"
rfiledir_res50_coshdphi = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210525_rebin1_iter10k_weight1_metRes50_wrongcosh/"
rfiledir_res25_cosdphi_rebin1_neww = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_Newweight_metRes25_cosdphi/"
rfiledir_res25_cosdphi_rebin20_neww = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210601_rebin20_iter10k_Newweight_metRes25_cosdphi/"
rfiledir_res30_cosdphi_rebin20_neww = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210601_rebin20_iter10k_Newweight_metRes30_cosdphi/"
rfiledir_res40_cosdphi_rebin20_neww = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210601_rebin20_iter10k_Newweight_metRes40_cosdphi/"
rfiledir_res50_cosdphi_rebin20_neww = "/Users/taohuang/Documents/DiHiggs/20210510_HME/Skim2018ForTao_MuMu_Boosted1B_20210601_rebin20_iter10k_Newweight_metRes50_cosdphi/"
legnames_res = ["MET Res=25.2GeV", "MET Res=30.0GeV", "MET Res=40.0GeV", "MET Res=50.0GeV"]
variables = ["hme_mass_peak","hme_mass_peak", "hme_mass_peak", "hme_mass_peak" ]
for mass in masslist:
    filename = "GluGluToRadionToHHTo2B2VTo2L2Nu_M-%d_HME.root"%(mass)
    filelist_cosdphi = [rfiledir+filename, rfiledir_res30_cosdphi+filename, rfiledir_res40_cosdphi+filename, rfiledir_res50_cosdphi+filename]
    filelist_coshdphi = [rfiledir_res25_coshdphi+filename, rfiledir_res30_coshdphi+filename, rfiledir_res40_coshdphi+filename, rfiledir_res50_coshdphi+filename]
    filelist_cosdphi_rebin20 = [rfiledir_res25_cosdphi_rebin20_neww+filename, rfiledir_res30_cosdphi_rebin20_neww+filename, rfiledir_res40_cosdphi_rebin20_neww+filename, rfiledir_res50_cosdphi_rebin20_neww+filename]
    
    #cut  = "1"
    cut  = "hme_mass_peak>=250"
    #xbins = [35, 500.0, 4000.0]
    #xbins = [70, 500.0, 4000.0]
    xbins = [90, 400.0, 4000.0]
    plotsuffix = "RadionBoosted_M%d_weight_cosdphi_rebin1_xbinsize40_METRes"%mass
    text = "HME likelihood hist binsize=1GeV"
    #plotHME(filelist_cosdphi, legnames_res, variables, xbins, cut, mass, text, plotsuffix)

    plotsuffix2 = "RadionBoosted_M%d_weight_coshdphi_rebin1_xbinsize40_METRes"%mass
    text2 = "HME likelihood hist binsize=1GeV"
    #plotHME(filelist_coshdphi, legnames_res, variables, xbins, cut, mass, text2, plotsuffix2)

    plotsuffix3 = "RadionBoosted_M%d_Newweight_cosphi_rebin20_xbinsize40_METRes"%mass
    text3 = "HME likelihood hist binsize=20GeV"
    #plotHME(filelist_cosdphi_rebin20, legnames_res, variables, xbins, cut, mass, text3, plotsuffix3)

legnames_neww = ["One iteration with weight=1", "Weight depends on Reco mass"]
variables = ["hme_mass_peak","hme_mass_peak"]
for mass in masslist:
    filename = "GluGluToRadionToHHTo2B2VTo2L2Nu_M-%d_HME.root"%(mass)
    #filelist_neww = [rfiledir_rebin20+filename, rfiledir_res25_cosdphi_rebin20_neww+filename]
    filelist_neww = [rfiledir+filename, rfiledir_res25_cosdphi_rebin1_neww+filename]

    cut  = "hme_mass_peak>=250"
    #xbins = [70, 500.0, 4000.0]
    xbins = [90, 400.0, 4000.0]
    plotsuffix = "RadionBoosted_M%d_cosdphi_rebin1_xbinsize40_METRes25_NewwCompare"%mass
    text = "HME likelihood hist binsize=1GeV"
    plotHME(filelist_neww, legnames_neww, variables, xbins, cut, mass, text, plotsuffix)

rfiledir_rebin1_gravity18 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_weightgravity_metRes25_cosdphi/"
rfiledir_rebin5_gravity18 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin5_iter10k_weightgravity_metRes25_cosdphi/"
rfiledir_rebin10_gravity18 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin10_iter10k_weightgravity_metRes25_cosdphi/"
rfiledir_rebin20_gravity18 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin20_iter10k_weightgravity_metRes25_cosdphi/"

rfiledir_rebin1_gravity12 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_weightgravity_metRes25_cosdphi_gravityp12/"
rfiledir_rebin1_gravity24 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_weightgravity_metRes25_cosdphi_gravityp24/"
rfiledir_rebin1_gravity30 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_weightgravity_metRes25_cosdphi_gravityp30/"
rfiledir_rebin1_gravity40 = localdir+"Skim2018ForTao_MuMu_Boosted1B_20210601_rebin1_iter10k_weightgravity_metRes25_cosdphi_gravityp40/"
legnames_gravity = ["most probable", "Most probable with gravity center,12%", "Most probable with gravity center,18%", "Most probable with gravity center,30%","Most probable with gravity center,40%"]
legnames_gravity_rebin=["Most probable with gravity center,18%,binsize=1GeV","Most probable with gravity center,18%,binsize=5GeV","Most probable with gravity center,18%,binsize=10GeV","Most probable with gravity center,18%,binsize=20GeV"]

#legnames_gravity = ["most probable", "Most probable with gravity center,18%"]
variables_gravity = ["hme_mass_peak","hme_mass_peak_gravity", "hme_mass_peak_gravity", "hme_mass_peak_gravity","hme_mass_peak_gravity"]
variables_gravity_rebin = ["hme_mass_peak_gravity", "hme_mass_peak_gravity", "hme_mass_peak_gravity","hme_mass_peak_gravity"]
for mass in masslist:
    filename = "GluGluToRadionToHHTo2B2VTo2L2Nu_M-%d_HME.root"%(mass)
    filelist_gravity = [rfiledir+filename, rfiledir_rebin1_gravity12+filename, rfiledir_rebin1_gravity18+filename, rfiledir_rebin1_gravity30+filename,rfiledir_rebin1_gravity40+filename]
    #filelist_gravity = [rfiledir+filename, rfiledir_rebin1_gravity18+filename]
    filelist_gravity_rebin = [rfiledir_rebin1_gravity18+filename, rfiledir_rebin5_gravity18+filename, rfiledir_rebin10_gravity18+filename,rfiledir_rebin20_gravity18+filename]

    #cut  = "hme_mass_peak_divSol>=250"
    cut = "1"
    #xbins = [70, 500.0, 4000.0]
    xbins = [90, 400.0, 4000.0]
    plotsuffix = "RadionBoosted_M%d_cosdphi_rebin1_xbinsize40_METRes25_GravityCompare"%mass
    text = "HME likelihood hist binsize=1GeV"
    #plotHME(filelist_gravity, legnames_gravity, variables_gravity, xbins, cut, mass, text, plotsuffix)

    plotsuffix = "RadionBoosted_M%d_cosdphi_rebincompare_xbinsize40_METRes25_Gravityp18"%mass
    text = " "
    #plotHME(filelist_gravity_rebin, legnames_gravity_rebin, variables_gravity_rebin, xbins, cut, mass, text, plotsuffix)