plot_xsec_2d.py

Go to the documentation of this file.

#!/usr/bin/env python

from __future__ import print_function

from ROOT import * # import pyroot
from bisect import bisect
import argparse
import math
import numpy as np
import os

# define constants used in this script
infp = './'
syst_names = [
    'lightdown','lightup','ckv','calibneg','calibpos','calibshape','MuES_onlyMuKEDw','MuES_onlyMuKEUp',
    '2kA_Dw','2kA_Up','02mmBeamSpotSize_Dw','02mmBeamSpotSize_Up','1mmBeamShiftX_Dw','1mmBeamShiftX_Up',
    '1mmBeamShiftY_Dw','1mmBeamShiftY_Up','3mmHorn1X_Dw','3mmHorn1X_Up','3mmHorn1Y_Dw','3mmHorn1Y_Up',
    '3mmHorn2X_Dw','3mmHorn2X_Up','3mmHorn2Y_Dw','3mmHorn2Y_Up','7mmTargetZ_Dw','7mmTargetZ_Up',
    'MagneticFieldinDecayPipe_Dw','MagneticFieldinDecayPipe_Up','1mmHornWater_Dw','1mmHornWater_Up'
]
Nfinal_fe = 9
final_fe = ["Flux (HP)","GENIE","Light","Calibration","Cherenkov","Calib. Shape","Mu Energy Scale","Normalization","Focusing"]
final_fe_name = ["Flux_HP","GENIE","Light","Calibration","Cherenkov","CalibShape","MuEScale","Normalization","Focusing"]

Nabins2X = 13 # pre calculated N bins X
leg_angle = [
    "0.50<Cos#theta_{#mu}<0.56",
    "0.56<Cos#theta_{#mu}<0.62",
    "0.62<Cos#theta_{#mu}<0.68",
    "0.68<Cos#theta_{#mu}<0.74",
    "0.74<Cos#theta_{#mu}<0.80",
    "0.80<Cos#theta_{#mu}<0.85",
    "0.85<Cos#theta_{#mu}<0.88",
    "0.88<Cos#theta_{#mu}<0.91",
    "0.91<Cos#theta_{#mu}<0.94",
    "0.94<Cos#theta_{#mu}<0.96",
    "0.96<Cos#theta_{#mu}<0.98",
    "0.98<Cos#theta_{#mu}<0.99",
    "0.99<Cos#theta_{#mu}<1.00"
]
col_fe = [2,4,2,4,8,8,kOrange-3,kOrange-3,1]
col_st = [1,1,2,2,1,2,1,2,2]

tit_ke = 'Muon Kinetic Energy (GeV)'
tit_nu = 'Neutrino Energy (GeV)'
tit_fe = 'Fractional Uncertainties'
tit_xs = '#sigma(cm^{2}/nucleon)'


def get_genie_universes(basefn):

    hl = []
    fn = os.path.join(infp, basefn)
    fgenie.append(TFile(fn))
    # find the genie universe IDs from filename
    sIdx = int(basefn.split('_')[1])
    eIdx = int(basefn.split('_')[2])+1
    for i in range(sIdx,eIdx):
        hl.append(fgenie[len(fgenie)-1].Get('xs_nominal/hXS2Dsec_genie{}'.format(i)))

    return hl

def get_ppfx_universes(fppfx_other):

    hl = []
    # by default we have 100 ppfx universes
    for i in range(0,100):
        hl.append(fppfx_other.Get('xs_nominal/hXS2Dsec_ppfx{}'.format(i)))
    
    return hl


def up_down_counts_1sigma(count_list, cv):
    sorted_list = sorted(count_list)
    cv_pos = bisect(sorted_list, cv)
    up_pos = int((len(sorted_list) - cv_pos)*.68 + cv_pos)
    down_pos = int(cv_pos - cv_pos*.68)
    return sorted_list[up_pos]-cv, cv-sorted_list[down_pos]


def process_multiuniverse(isys, huniv):
    
    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            dt_cv = hdt.GetBinContent(i,j)

            if dt_cv <= 0:
                lhup[isys].SetBinContent(i,j,0)
                lhdw[isys].SetBinContent(i,j,0)
                lhfe[isys].SetBinContent(i,j,0)
            else:
                # accumulate bin content from all ppfx universes
                binc_all_uni = [h.GetBinContent(i,j) for h in huniv]
                cup, cdown = up_down_counts_1sigma(binc_all_uni, dt_cv)
                lhup[isys].SetBinContent(i,j,cup)
                lhdw[isys].SetBinContent(i,j,cdown)
                # calculate conservative fractional uncertainty
                conserv_error = max(cup,cdown)
                lhfe[isys].SetBinContent(i,j,conserv_error/dt_cv)


def process_light_calibration_mue(isys, hidxdown, hidxup):

    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            dt_cv = hdt.GetBinContent(i,j)

            lhup[isys].SetBinContent(i,j,0)
            lhdw[isys].SetBinContent(i,j,0)
            lhfe[isys].SetBinContent(i,j,0)

            sigma_dw = dt_cv - hother[hidxdown].GetBinContent(i,j)
            sigma_up = dt_cv - hother[hidxup].GetBinContent(i,j)
            sigma = max(abs(sigma_dw), abs(sigma_up))

            if sigma_dw < 0 and sigma_up < 0:
                lhup[isys].SetBinContent(i,j,sigma)
            elif sigma_dw > 0 and sigma_up > 0:
                lhdw[isys].SetBinContent(i,j,sigma)
            else:
                lhup[isys].SetBinContent(i,j,abs(sigma_up))
                lhdw[isys].SetBinContent(i,j,abs(sigma_dw))
            
            if dt_cv > 0:
                lhfe[isys].SetBinContent(i,j,sigma/dt_cv)
            else:
                lhfe[isys].SetBinContent(i,j,0)


def process_ckv_calibshape(isys, hidx):

    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            dt_cv = hdt.GetBinContent(i,j)

            lhup[isys].SetBinContent(i,j,0)
            lhdw[isys].SetBinContent(i,j,0)
            lhfe[isys].SetBinContent(i,j,0)

            sigma = dt_cv - hother[hidx].GetBinContent(i,j)

            # if sigma > 0:
            #     lhdw[isys].SetBinContent(i,j,abs(sigma))
            # if sigma <= 0:
            #     lhup[isys].SetBinContent(i,j,abs(sigma))
            lhdw[isys].SetBinContent(i,j,abs(sigma))
            lhup[isys].SetBinContent(i,j,abs(sigma))
            
            if dt_cv > 0:
                lhfe[isys].SetBinContent(i,j,sigma/dt_cv)
            else:
                lhfe[isys].SetBinContent(i,j,0)


def process_normalization(isys):

    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            dt_cv = hdt.GetBinContent(i,j)

            sigma = dt_cv * fe_norm
            # lhup[isys].SetBinContent(i,j,dt_cv+sigma)
            # lhdw[isys].SetBinContent(i,j,dt_cv-sigma)
            lhup[isys].SetBinContent(i,j,sigma)
            lhdw[isys].SetBinContent(i,j,sigma)
            if dt_cv > 0:
                lhfe[isys].SetBinContent(i,j,fe_norm)
            else:
                lhfe[isys].SetBinContent(i,j,0)


def process_focusing(isys, sstart, nsys):

    # debug information
    print('Focusing systematics:')
    for k in range(nsys):
        print(syst_names[sstart+k*2], syst_names[sstart+k*2+1])
    
    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            dt_cv = hdt.GetBinContent(i,j)

            lhup[isys].SetBinContent(i,j,0)
            lhdw[isys].SetBinContent(i,j,0)
            lhfe[isys].SetBinContent(i,j,0)

            larger_sigma = []
            for k in range(nsys):
                sigma_dw = dt_cv - hother[sstart+k*2].GetBinContent(i,j)
                sigma_up = dt_cv - hother[sstart+k*2+1].GetBinContent(i,j)
                larger_sigma.append(max(abs(sigma_dw), abs(sigma_up)))

            qsum = math.sqrt(sum([x**2 for x in larger_sigma]))
            lhup[isys].SetBinContent(i,j,qsum)
            lhdw[isys].SetBinContent(i,j,qsum)
            
            if dt_cv > 0:
                lhfe[isys].SetBinContent(i,j,qsum/dt_cv)
            else:
                lhfe[isys].SetBinContent(i,j,0)


def databin(xbin, ybin):
    answer = False
    if xbin==1 and ybin<=4: answer = True
    if (xbin==2 or xbin==3) and ybin<=5: answer = True
    if xbin==4 and ybin<=7: answer = True
    if xbin==5 and ybin<=9: answer = True
    if xbin==6 and ybin<=10: answer = True
    if xbin==7 and ybin<=12: answer = True
    if xbin==8 and ybin<=15: answer = True
    if xbin==9 and ybin<=19: answer = True
    if (xbin==10 or xbin==11 or xbin==12 or xbin==13) and ybin<=20: answer = True

    return answer


def process_data():
    
    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            kin   = hdt.GetYaxis().GetBinCenter(j)
            binw  = hdt.GetYaxis().GetBinWidth(j)
            cont  = hdt.GetBinContent(i,j)
            errUp = 0
            errDw = 0

            for p in range(Nfinal_fe):
                up = lhup[p].GetBinContent(i,j)
                dw = lhdw[p].GetBinContent(i,j)
                errUp += up**2
                errDw += dw**2
            
            errUp = math.sqrt(errUp)
            errDw = math.sqrt(errDw)

            if databin(i,j):
                count[i-1] += 1
                gdt[i-1].SetPoint(count[i-1], kin, cont)
                gdt[i-1].SetPointError(count[i-1],binw/2, binw/2, errDw, errUp)


def process_mc():
    
    for i in range(1, NbinsX+1):
        for j in range(1, NbinsY+1):
            if databin(i,j):
                binw  = hdt.GetYaxis().GetBinWidth(j)
                cont = hmc.GetBinContent(i,j)
                h1MC[i-1].SetBinContent(j,cont)
                for p in range(Nfinal_fe):
                    ferr = lhfe[p].GetBinContent(i,j)
                    h1fe[i-1][p].SetBinContent(j,ferr)

def setMC1D(hIn):
    hIn.SetLineColor(2)
    hIn.SetLineWidth(2)
    hIn.SetStats(0)

def setData1D(gIn):
    gIn.SetMarkerColor(kBlack)
    gIn.SetMarkerStyle(20)
    gIn.SetLineColor(kBlack)
    gIn.SetLineWidth(3)

def setUncer1D(hIn, ihist):
    hIn.SetLineColor(col_fe[ihist])
    hIn.SetLineWidth(3)
    hIn.SetLineStyle(col_st[ihist])
    hIn.SetStats(0)


if __name__ == '__main__':
    '''
    Remember:
    The unit of the output plots from this script has to be
    rescaled by bin size when showing the final results!
    '''

    # fake data / real data switch
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--data', action='store_true', help='Switch for using data instead of fake data.')
    args = parser.parse_args()
    use_data = args.data

    # read histograms into containers

    # file containers
    fgenie = []
    fmc = TFile(os.path.join(infp, 'nominal_xsec.root'))
    if not use_data:
        fppfx_other = TFile(os.path.join(infp, 'ppfx_others_v2.root'))
    else:
        fppfx_other = TFile(os.path.join(infp, 'ppfx_others_data.root'))

    # historam containers
    hgenie = []
    hppfx = []
    hother = []

    # get nominal cross section
    hmc = fmc.Get('hXS2Dsec_nominal')
    # get (fake) data cross section
    hdt = fppfx_other.Get('xs_nominal/hXS2Dsec_cv')

    # store genie histograms
    if not use_data:
        hgenie += get_genie_universes('genie_0_99_v2.root')
        hgenie += get_genie_universes('genie_100_299_v2.root')
        hgenie += get_genie_universes('genie_300_499_v2.root')
    else:
        hgenie += get_genie_universes('genie_0_99_data.root')
        hgenie += get_genie_universes('genie_100_299_data.root')
        hgenie += get_genie_universes('genie_300_499_data.root')

    # store ppfx histograms
    hppfx += get_ppfx_universes(fppfx_other)
    
    # load all other systematics
    for syst_name in syst_names:
        hother.append(fppfx_other.Get('xs_nominal/hXS2Dsec_{}'.format(syst_name)))
    
    # below starts to calculate up and down histograms
    lhup = []
    lhdw = []
    lhfe = []
    for i in range(Nfinal_fe):
        lhup.append(hmc.Clone('hUp_{}'.format(i)))
        lhdw.append(hmc.Clone('hDw_{}'.format(i)))
        lhfe.append(hmc.Clone('hFE_{}'.format(i)))
    
    # some global variables common to all systematics
    # number of bins
    NbinsX = hdt.GetNbinsX()
    NbinsY = hdt.GetNbinsY()

    # ppfx
    process_multiuniverse(0, hppfx)

    # genie
    process_multiuniverse(1, hgenie)
    
    # light
    process_light_calibration_mue(2, 0, 1)
    
    # calibration
    process_light_calibration_mue(3, 3, 4)

    # cherenkov
    # this does not look right
    process_ckv_calibshape(4, 2)
    
    # calibshape
    # this does not look right
    process_ckv_calibshape(5, 5)
    
    # muon energy scale
    process_light_calibration_mue(6, 6, 7)
    
    # normalization
    # Normalization (TN, doc-):
    # POT: 0.28% (Section 9.1)
    # Intensity: 2% (Section 9.2)
    # Particle propagation: 1-2%, so we take 2% (Section 9.5).
    # fe_norm = sqrt(pow(0.0028,2)+pow(0.02,2)+pow(0.02,2));
    fe_norm = 0.028422526
    process_normalization(7)

    # focusing
    # process_light_calibration_mue_foc(8, 8, 9)
    process_focusing(8, 8, 11)

    # data
    gdt = [TGraphAsymmErrors() for i in range(Nabins2X)]
    count = [0 for i in range(Nabins2X)]
    process_data()

    # MC
    h1MC = [TH1D('h1MC_{}'.format(i),'',20,0.5,2.5) for i in range(Nabins2X)]
    h1fe = [[TH1D('h1fe_{}_{}'.format(i,j),'',20,0.5,2.5) for j in range(Nfinal_fe)] for i in range(Nabins2X)]
    process_mc()

    # make sure the output directory exists
    out_subdir = 'real_data' if use_data else 'fake_data'
    dest_path = os.path.join('output', out_subdir)
    if not os.path.exists(dest_path):
        os.makedirs(dest_path)

    # make cross section plots
    cxs = []
    leg = [TLegend(0.67,0.68,0.87,0.88) for i in range(Nabins2X)]
    lat = [TLatex() for i in range(Nabins2X)]
    maxY = -1.
    for i in range(NbinsX):
        maxY = h1MC[i].GetMaximum()
        h1MC[i].SetMaximum(maxY*2)
        cxs.append(TCanvas())
        lat[i].SetNDC()
        leg[i].SetBorderSize(0)
        setMC1D(h1MC[i])
        setData1D(gdt[i])
        h1MC[i].GetXaxis().SetTitle('Muon Kinetic Energy (GeV)')
        h1MC[i].GetXaxis().CenterTitle()
        h1MC[i].GetYaxis().SetTitle('d^{2}#sigma/dT_{#mu}dcos#theta_{#mu} (cm^{2}/GeV/nucleon)')
        h1MC[i].GetYaxis().CenterTitle()
        h1MC[i].Draw("hist")
        gdt[i].Draw("p1 same")
        leg[i].SetTextSize(0.05)
        leg[i].AddEntry(h1MC[i],  "MC", "fl")
        leg[i].AddEntry(gdt[i],"Data (Syst.)", "lp")
        leg[i].Draw()
        lat[i].DrawLatex(0.2,0.83,'{}'.format(leg_angle[i]))
        cxs[i].SaveAs('{}/xs2D_anglebin{}.png'.format(dest_path,i))
        cxs[i].SaveAs('{}/xs2D_anglebin{}.pdf'.format(dest_path,i))
    
    # make fractional uncertainty plots
    cfe = []
    latfe = [TLatex() for i in range(Nabins2X)]
    lfe = []
    for i in range(NbinsX):
        if i > 6:
            lfe = TLegend(0.2,0.5,0.45,0.88)
        else:
            lfe = TLegend(0.6,0.4,0.85,0.78)
        lfe.SetBorderSize(0)
        cfe.append(TCanvas())
        latfe[i].SetNDC()
        for j in range(Nfinal_fe):
            setUncer1D(h1fe[i][j],j)
            h1fe[i][j].GetXaxis().SetTitle('Muon Kinetic Energy (GeV)')
            h1fe[i][j].GetXaxis().CenterTitle()
            h1fe[i][j].GetYaxis().SetTitle('Fractional Uncertainty')
            h1fe[i][j].GetYaxis().CenterTitle()
            if j == 0:
                h1fe[i][j].Draw('hist')
            else:
                h1fe[i][j].Draw('histsame')
            h1fe[i][j].SetMaximum(.4)
            h1fe[i][j].SetMinimum(0)
            lfe.AddEntry(h1fe[i][j],final_fe[j],'l')
        latfe[i].DrawLatex(0.63,0.83,'{}'.format(leg_angle[i]))
        lfe.Draw()
        cfe[i].SaveAs('{}/fe2D_anglebin{}.png'.format(dest_path,i))
        cfe[i].SaveAs('{}/fe2D_anglebin{}.pdf'.format(dest_path,i))

    # storing the final histograms
    fileOut = TFile('{}/numucc_inc_{}.root'.format(dest_path,out_subdir), 'recreate')
    fileOut.cd()
    fileOut.mkdir('CrossSections')
    for i in range(NbinsX):
        fileOut.cd('CrossSections')
        gdt[i].GetXaxis().SetTitle(tit_ke)
        gdt[i].GetYaxis().SetTitle(tit_xs)
        gdt[i].Write('mock_data_bin{}'.format(i))
        h1MC[i].Write('mc_bin{}'.format(i))
        fileOut.mkdir('FractUncer_bin{}'.format(i))
        fileOut.cd('FractUncer_bin{}'.format(i))
        for j in range(Nfinal_fe):
            h1fe[i][j].Write('FractUncer_bin{}_{}'.format(i,final_fe_name[j]))

    raw_input('Press enter to continue.')