EWP-BplusToKstMuMu-AngAna/Code/Scripts/LaTeX/putIntoSlides.py

from __future__ import print_function
import numpy as np 
from pathlib import Path

from putIntoSlides_Utils import *

HD_FOLDER = "/home/renata/B2KMuMu/sigma0/B2KstarMuMu/code/ewp-Bplus2Kstmumu-AngAna/FCNCfitter/"
FIG_FOLDER = HD_FOLDER + "plots/"
MASS_FIG_FOLDER = FIG_FOLDER + "MassFit/"
MC_FIG_FOLDER = FIG_FOLDER + "MCfit/"
TOYS_FIG_FOLDER = FIG_FOLDER + "Toys/"
MAIN_FIG_FOLDER = FIG_FOLDER + "MainFit/"
BKG_FIG_FOLDER = FIG_FOLDER + "bkgFit/"
GENLVL_FIG_FOLDER = FIG_FOLDER + "GenLvlFit/"

TEX_FILE_FOLDER = "includes/" #Don't use ./ as it efs up the print of includes

#FIG_FOLDER  = "./figures"

OUTPUT_TAG = ""
TITLE_TAG = ""

decName = "KplusPi0Resolved"
ANGLES = ["ctk", "ctl", "phi"]
MASS_PLOTS = ["m_log","m", "mkpi"]
ANG_PARAMS  = ["Fl", "S3", "S4", "S5", "Afb", "S7", "S8", "S9"]
MASS_PARAMS = ["m_b", "m_sigma_1", "alpha_1", "alpha_2", "n_1", "n_2"]
MASS_BKG = ["m_lambda","f_sig"]
ANGLE_BKG = ["cbkgctk1","cbkgctk2","cbkgctk3","cbkgctk4","cbkgctl1","cbkgctl2"] 
SWAVE_PARAMS = ["FS","SS1","SS2","SS3","SS4","SS5"]

ANGLE_FOLDING = [
        ["Fl", "S3", "Afb", "S9"],
        ["Fl", "S3", "S4"],
        ["Fl", "S3", "S5"],
        ["Fl", "S3", "S7"],
        ["Fl", "S3", "S8"]
]

ANGLE_FOLDING_S = [ #FS is set extra, so to not complicate it, it is not listed here
        ["SS1"],
        ["SS1","SS2"],
        ["SS1","SS3"],
        ["SS1","SS4"],
        ["SS4"],
]


def projections(angle,Run): #Plot all angle projections in 18 bins
    listOfEntries = list(map(str,np.arange(0,18,1)))
    defaultTitle = angle + " projections, " + RunTag(Run)
    outputName = angle +  "proj_" + RunTag(Run)
    figsPerSlide=6
    figsPerRow=3
    make_figure_slides(FIG_FOLDER+"/angular/",angle+"eff__LOOP___" + decName + "_" + RunTag(Run), listOfEntries, defaultTitle, figsPerSlide, figsPerRow, ".eps", outputName+".tex")


def massFit(nBins, observable, Run): #Plot mass-only fits
    massFitName = "_massfit" + rareJpsiTag(False) 
    path = observable + massFitName + "Fit" + binTag(nBins) + "_bin__LOOP___"+ RunTag(Run) + ""
    figsPerSlide=6
    figsPerRow=3
    
    make_figure_slides(MASS_FIG_FOLDER,path,
                       list(map(str,np.arange(0,nBins,1))),"Signal mass fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MassFit_Signal_"+observable+"_"+RunTag(Run)+".tex")
    return 

def massFitRef(Run): 
    massFitName = "_massfit" + rareJpsiTag(True) 
    path = "__LOOP__" + massFitName + "Fit" + binTag(1) + "_bin0_" + RunTag(Run) + ""
    figsPerSlide=2
    figsPerRow=2
    
    make_figure_slides(MASS_FIG_FOLDER,path,
                       MASS_PLOTS,"Reference mass fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MassFit_Jpsi_"+RunTag(Run)+".tex")
    return 

def massFitMCRef(Run): #Plot only mass fit of reference MC 
    MCfitName = "_MC" + rareJpsiTag(True) + "Fit"
    path = "__LOOP__" + MCfitName + binTag(1) + "_bin0_SimultaneousFit_" + RunTag(Run) + "_OnlyMass"
    figsPerSlide=2
    figsPerRow=2
    
    make_figure_slides(MC_FIG_FOLDER + "Jpsi/",path,
                       MASS_PLOTS,"MC Reference mass fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_MassFit_Jpsi_"+RunTag(Run)+".tex")
    return 
    
def massFitMC(nBins, observable, Run): 
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    path = observable + MCfitName + binTag(nBins) + "_bin__LOOP___SimultaneousFit_" + RunTag(Run) 
    figsPerSlide=5
    figsPerRow=3
    
    make_figure_slides(MC_FIG_FOLDER + "Signal/",path,
                       list(map(str,np.arange(0,nBins,1))),"MC Signal mass fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_MassFit_Signal_"+observable+"_"+RunTag(Run)+".tex")
    return 

def angleFitMC(nBins, bin, Run): #Plot only ANGLES fit of signal MC
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + MCfitName + binTagFull(nBins,bin)+"_SimultaneousFit_" + RunTag(Run) + "_OnlyANGLES_AllPDFs"
    figsPerSlide=3
    figsPerRow=3
    
    make_figure_slides(MC_FIG_FOLDER + "Signal/",path,
                       ANGLES,"MC Signal fit, " + binTitle(nBins,bin),
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_Signal"+binTagFull(nBins,bin)+"_"+RunTag(Run)+"_onlyANGLES.tex")
    return 

def MainFitRef(Run, Toy, fold, fullBkg): 
    fitName = rareJpsiTag(True) + "Fit"
    path = "__LOOP__" + fitName + toyTag(Toy)+ binTag(1) + "_bin"+str(0)+"_SimultaneousFit" + foldTag(fold) + "_" + RunTag(Run) + plotBkgTag(fullBkg)+"_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    title = "Main fit, Ref"+ (", toy" if Toy else "")
    if (fold != -1): title += ", fld " + str(fold)
    
    make_figure_slides(MAIN_FIG_FOLDER, path,
                       ANGLES+MASS_PLOTS,title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFit_Ref"+toyTag(Toy)+ foldTag(fold)+ "_"+RunTag(Run)+plotBkgTag(fullBkg)+".tex")
    return 

def MainFit(Run, nBins, bin, fold, fullBkg): 
    fitName =  rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + fitName + binTagFull(nBins,bin) +"_SimultaneousFit" + foldTag(fold) +"_" + RunTag(Run) + plotBkgTag(fullBkg)+ "_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    title = "Main fit, toy, "+str(bin+1)+"/" + str(nBins) + " bin(s)"
    if (fold != -1): title += ", fld " + str(fold)
    
    make_figure_slides(MAIN_FIG_FOLDER, path,
                       ANGLES+MASS_PLOTS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFit"+ binTagFull(nBins,bin) + foldTag(fold)+"_"+RunTag(Run)+plotBkgTag(fullBkg)+".tex")
    return 

def ToyFitCheckRef():
    path = "Init_vs_Fit___LOOP___Ref"
    plot_list = ANG_PARAMS + ANGLE_BKG +MASS_PARAMS + MASS_BKG 
    figsPerSlide=6
    figsPerRow=3

    make_figure_slides("./figures/FullFitToy/",path,
                       plot_list,"Toy init vs fit, Reference-like",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"ToyInitFit.tex")
    return 

def MainFitToyCheck(isRef):
    refTag = "_Ref" if isRef else ""
    path = "FinalToy_Init_vs_Fit___LOOP__" + refTag
    plot_list = ANG_PARAMS + ANGLE_BKG +MASS_PARAMS + MASS_BKG 
    figsPerSlide=6
    figsPerRow=3

    make_figure_slides(TOYS_FIG_FOLDER,path,
                       plot_list,"Main fit vs init," + " Reference-like" if isRef else " Signal-like",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFitToy"+ refTag+".tex")
    return 

def localToy(folder, jobID, isRef, bin, nBins, fold, incBkg, onlyBkg, onlySig):
    refTag = "Ref" if isRef else "Signal"
    path = str(jobID)+"/__LOOP___toyfit__"+str(jobID)
    path = path + "_"+rareJpsiTag(isRef)+"Fit" + sigBkgTag(onlySig,onlyBkg)
    path = path + binTag(nBins)+"_bin"+str(bin)
    path = path +"_Run12_SimultaneousFit"+foldTag(fold)+"_HighBmass"+plotBkgTag(incBkg) + "_AllPDFs"
    #Split the path definition this stupid cause easier debuging


    plot_list = ANGLES + MASS_PLOTS
    if (onlyBkg): plot_list = ["ctk", "ctl", "phi", "m"]
    figsPerSlide= 4 if onlyBkg else 6
    figsPerRow= 4 if onlyBkg else 3

    title = "Toys" + rareJpsiTag(isRef) + sigBkgTag(onlySig,onlyBkg) +binTag(nBins)+"_bin"+str(bin) +foldTag(fold) +plotBkgTag(incBkg)
    title = title.replace("_"," ") 


    fileTag = str(jobID)+rareJpsiTag(isRef)+str(bin)+str(nBins)+foldTag(fold)+plotBkgTag(incBkg)
    make_figure_slides(folder,path,
                       plot_list,title, #TODO: check the title
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"Toy"+ fileTag+".tex")
    return 

def MCfit(nBins, bin, Run, fold): #Plot MC fit everything
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    
    path = "__LOOP__" + MCfitName + binTagFull(nBins,bin)+"_SimultaneousFit" + foldTag(fold) +"_" + RunTag(Run) + "_AllPDFs"
    figsPerSlide=6
    figsPerRow=3
    
    title = "MC Signal fit, "+binTitle(nBins,bin)
    if (fold != -1): title += ", fld " + str(fold)
    make_figure_slides(MC_FIG_FOLDER + "Signal/",path,
                       ANGLES+MASS_PLOTS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_Signal"+binTagFull(nBins,bin)+ foldTag(fold)+"_"+RunTag(Run)+".tex")
    return 

def BkgFitRef(Run, sideband): 

    var_list = ["ctk", "ctl", "phi", "mkpi"]

    massFitName = "_Bckgnd" + rareJpsiTag(True) 
    sidebandTag = ""
    if (sideband ==-1):  sidebandTag = "_LowBmass"
    elif (sideband ==0): sidebandTag = "_HighBmass"

    sidebandTitle = ""
    if (sideband ==-1):  sidebandTitle = ", lower sideband"
    elif (sideband ==0): sidebandTitle = ", upper sideband"


    path = "__LOOP__" + massFitName + "Fit" + binTag(1) + "_bin0_" + RunTag(Run)+ sidebandTag
    figsPerSlide=4
    figsPerRow=4
    
    make_figure_slides(BKG_FIG_FOLDER,path,
                       var_list,"BKG mass fit"+sidebandTitle,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"BkgFit_Jpsi_"+RunTag(Run)+sidebandTag+".tex")
    return 

def genLvlMC(nBins, bin, Run, fold):
    genLvlName = "_genLvl_MC" + rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + genLvlName +  binTagFull(nBins,bin)+ "_OnlyANGLES_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    title = "GenLvl MC fit, "+binTitle(nBins,bin)
    make_figure_slides(GENLVL_FIG_FOLDER,path,
                       ANGLES, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"GenLvlMC"+binTagFull(nBins,bin)+ foldTag(fold)+"_"+RunTag(Run)+".tex")

    return

def compareGenLvlAndMC(toDavid):
    plotName = "GenLvl_vs_MC_"
    path = plotName + "__LOOP__" + ("_MyVsDavids" if toDavid else "")
    figsPerSlide=8
    figsPerRow=4

    title = "GenLvl vs "
    if (toDavid): title += "David's GenLvl"
    else:         title += "signal MC"

    make_figure_slides(MC_FIG_FOLDER,path,
                       ANG_PARAMS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"GenLvl" + ("_Vs_Davids" if toDavid else "_Vs_SigMC")+".tex")
    return

def compareFolding():
    path = "Folding___LOOP__" 
    figsPerSlide=8
    figsPerRow=4

    title = "Folding comparisons"

    make_figure_slides(MC_FIG_FOLDER,path,
                       ANG_PARAMS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"Folding_comparison.tex")
    return

def KstarMass(Ref,nBins, Run):
    MCfitName = "mkpi_MC" + rareJpsiTag(Ref) + "Fit"
    
    path = MCfitName + binTag(nBins) +"_bin__LOOP__" +"_SimultaneousFit_" + RunTag(Run) + ("_AllPDFs" if (Run == 12) else "") + ("_OnlyMass" if (Ref) else "") 
    figsPerSlide=6
    figsPerRow=3
    
    binsList = list(map(str,np.arange(0,nBins,1)))
    title = "MC "+("Jpsi" if Ref else "Signal") +" fit, "+str(nBins) + " bin" + (", " if nBins==1 else "s, ") + RunTag(Run)
    make_figure_slides(MC_FIG_FOLDER + ("Jpsi/" if Ref else "Signal/"),path,
                       binsList, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC"+ rareJpsiTag(Ref)+"_Kpi"+binTag(nBins)+ "_"+RunTag(Run)+".tex")
    return 

#######################

def toyPulls(bin,jobID,sWave=False):
    all_vars_list  = ANG_PARAMS + ANGLE_BKG + MASS_BKG
    if(sWave): all_vars_list = all_vars_list + SWAVE_PARAMS

    figName = "/__LOOP___bin" + str(bin) +"_Signal_HighBmass_Pulls"
    figsPerSlide=8
    figsPerRow=4
    
    outFolder = "outputs/"+str(jobID)+"/"
    Path(outFolder).mkdir(parents=True, exist_ok=True)

    title = str(jobID) + ", bin " + str(bin)
    make_figure_slides(TOYS_FIG_FOLDER + str(jobID),figName,
                       all_vars_list, title,
                       figsPerSlide, figsPerRow,
                       ".eps",outFolder+"bin"+str(bin)+".tex")



#######################

totBins = 5
totFld = 4

# for b in range(0,5):
#     toyPulls(b,353)

#BkgFitRef(12,0)

#for x in range(325,370):
#     toyPulls(0,x,True)

toys = False
for fld in range(-1,totFld+1):
    MainFitRef(12, toys, fld, False)
    MainFitRef(12, toys, fld, True)


for fld in range(-1,totFld+1):
    #localToy("Toys/Bkg/",7,True,0,1,fld,False,True,False)
    #localToy("Toys/Ref/",5,True,0,1,fld,False,False,False)
    for bin in range(0,totBins):
        localToy("Toys/Sig/",10,False,bin,totBins,fld,False, False, False)
        localToy("Toys/Sig/",10,False,bin,totBins,fld,True, False, False)