diff --git a/src/R2_R3_cent_comparison.py b/src/R2_R3_cent_comparison.py new file mode 100644 index 0000000..193ea1f --- /dev/null +++ b/src/R2_R3_cent_comparison.py @@ -0,0 +1,296 @@ +import ROOT +import os +import array + +# Set global style for professional plots +def set_global_style(): + ROOT.gStyle.SetOptTitle(0) + ROOT.gStyle.SetOptStat(0) + ROOT.gStyle.SetPadTickX(1) + ROOT.gStyle.SetPadTickY(1) + ROOT.gStyle.SetLegendBorderSize(0) + ROOT.gStyle.SetLegendFillColor(0) + ROOT.gStyle.SetLegendFont(42) + ROOT.gStyle.SetLegendTextSize(0.035) + ROOT.gStyle.SetLabelSize(0.04, "XYZ") + ROOT.gStyle.SetTitleSize(0.045, "XYZ") + ROOT.gStyle.SetTitleOffset(1.1, "Y") + ROOT.gStyle.SetTitleOffset(1.0, "X") + ROOT.gStyle.SetPadLeftMargin(0.15) + ROOT.gStyle.SetPadBottomMargin(0.15) + ROOT.gStyle.SetPadRightMargin(0.05) + ROOT.gStyle.SetPadTopMargin(0.05) + ROOT.gStyle.SetFrameLineWidth(2) + +set_global_style() + +# Open the files +f1 = ROOT.TFile.Open("resosp0_100_v3.root") # R3: 0-100% +f2 = ROOT.TFile.Open("resospk0100_inte.root") # R2: 0-100% + +if not f1 or not f1.IsOpen(): + print("Error: Could not open resosp0_100_v3.root") + exit() +if not f2 or not f2.IsOpen(): + print("Error: Could not open resospk0100_inte.root") + exit() + +# List of detector combinations to compare +det_combinations = [ + "FT0c_FT0a_FV0a", + "FT0c_FT0a_TPCpos", + "FT0c_FT0a_TPCneg", + "FT0c_FT0a_TPCtot", + "FT0c_FV0a_TPCpos", + "FT0c_FV0a_TPCneg", + "FT0c_FV0a_TPCtot", + "FT0c_TPCpos_TPCneg", + "FT0a_FV0a_TPCpos", + "FT0a_FV0a_TPCneg", + "FT0a_FV0a_TPCtot", + "FT0a_TPCpos_TPCneg", + "FT0m_FV0a_TPCpos", + "FT0m_FV0a_TPCneg", + "FT0m_FV0a_TPCtot", + "FT0m_TPCpos_TPCneg", + "FV0a_TPCpos_TPCneg" +] + +# Colors for different resolution types +resolution_colors = { + "R3": ROOT.kRed+1, + "R2": ROOT.kBlue+1 +} + +markers = { + "R3": 20, + "R2": 21 +} + +# Create output directory if it doesn't exist +if not os.path.exists("plots"): + os.makedirs("plots") + +# Function to get resolution histogram from a directory +def get_resolution_hist(directory): + if not directory: + return None + keys = directory.GetListOfKeys() + for key in keys: + obj_name = key.GetName() + # Look for resolution histogram + if "histo_reso" in obj_name and not "delta_cent" in obj_name: + obj = key.ReadObj() + if "TH1" in obj.ClassName(): + return obj + return None + +# Create output file for results +output_file = ROOT.TFile("plots/resolution_comparison_results.root", "RECREATE") + +# Store results for summary +results_summary = [] + +# Create individual plots for each detector combination +for det in det_combinations: + print(f"\nAnalyzing {det}...") + + # Get directories + d1 = f1.Get(det) # R3 + d2 = f2.Get(det) # R2 + + if not d1 or not d2: + print(f"Skipping {det}, not found in one of the files") + continue + + # Get resolution histograms + h_r3 = d1.Get("histo_reso") + h_r2 = d2.Get("histo_reso") + + if not h_r3 or not h_r2: + print(f"Skipping {det}, resolution histogram not found") + continue + + # Create comparison canvas + c = ROOT.TCanvas(f"c_{det}", f"Resolution Comparison - {det}", 800, 600) + + # Style the histograms + h_r3.SetLineColor(resolution_colors["R3"]) + h_r3.SetMarkerColor(resolution_colors["R3"]) + h_r3.SetMarkerStyle(markers["R3"]) + h_r3.SetLineWidth(2) + h_r3.SetMarkerSize(1.5) + h_r3.SetTitle("") + h_r3.GetXaxis().SetTitle("Centrality (%)") + h_r3.GetYaxis().SetTitle("Resolution") + h_r3.GetYaxis().SetTitleOffset(1.4) + + h_r2.SetLineColor(resolution_colors["R2"]) + h_r2.SetMarkerColor(resolution_colors["R2"]) + h_r2.SetMarkerStyle(markers["R2"]) + h_r2.SetLineWidth(2) + h_r2.SetMarkerSize(1.5) + + # Determine y-axis range + y_min = min(h_r3.GetMinimum(), h_r2.GetMinimum()) * 0.9 + y_max = max(h_r3.GetMaximum(), h_r2.GetMaximum()) * 1.1 + h_r3.GetYaxis().SetRangeUser(y_min, y_max) + + # Draw without uncertainties + h_r3.Draw("P") + h_r2.Draw("P SAME") + + # Add legend + leg = ROOT.TLegend(0.65, 0.75, 0.88, 0.88) + leg.AddEntry(h_r3, "R3 (v3 method)", "lp") + leg.AddEntry(h_r2, "R2 (standard)", "lp") + leg.Draw() + + # Add detector label + label = ROOT.TLatex() + label.SetNDC() + label.SetTextSize(0.04) + label.DrawLatex(0.2, 0.85, det.replace("_", " & ")) + + # Add ALICE label + alice_label = ROOT.TLatex() + alice_label.SetNDC() + alice_label.SetTextFont(42) + alice_label.SetTextSize(0.04) + # alice_label.DrawLatex(0.2, 0.92, "") + + c.Update() + c.SaveAs(f"plots/resolution_comparison_{det}.pdf") + c.SaveAs(f"plots/resolution_comparison_{det}.png") + + # Save histograms to output file + output_file.cd() + h_r3.Write(f"h_r3_{det}") + h_r2.Write(f"h_r2_{det}") + + # Calculate average values for summary + n_bins = h_r3.GetNbinsX() + avg_r3 = 0 + avg_r2 = 0 + count = 0 + + for bin in range(1, n_bins + 1): + if h_r3.GetBinContent(bin) > 0 and h_r2.GetBinContent(bin) > 0: + avg_r3 += h_r3.GetBinContent(bin) + avg_r2 += h_r2.GetBinContent(bin) + count += 1 + + if count > 0: + avg_r3 /= count + avg_r2 /= count + + # Store results for summary + results_summary.append({ + 'detector': det, + 'avg_r3': avg_r3, + 'avg_r2': avg_r2, + 'count': count + }) + + print(f"Created comparison plot for {det}") + +# Create summary comparison plot +c_summary = ROOT.TCanvas("c_summary", "R3 vs R2 Summary Comparison", 1200, 800) + +# Create graphs for average values +n_det = len(results_summary) +if n_det > 0: + x = array.array('d') + y_r3 = array.array('d') + y_r2 = array.array('d') + labels = [] + + for i, result in enumerate(results_summary): + x.append(i + 1) + y_r3.append(result['avg_r3']) + y_r2.append(result['avg_r2']) + labels.append(result['detector'].replace('_', '\n')) + + # Create graphs + g_r3 = ROOT.TGraph(n_det, x, y_r3) + g_r2 = ROOT.TGraph(n_det, x, y_r2) + + # Style the graphs + g_r3.SetTitle("Average Resolution by Detector Combination") + g_r3.GetXaxis().SetTitle("Detector Combination") + g_r3.GetYaxis().SetTitle("Average Resolution") + g_r3.GetYaxis().SetTitleOffset(1.4) + g_r3.SetMarkerColor(resolution_colors["R3"]) + g_r3.SetLineColor(resolution_colors["R3"]) + g_r3.SetMarkerStyle(markers["R3"]) + g_r3.SetMarkerSize(1.5) + g_r3.SetLineWidth(2) + + g_r2.SetMarkerColor(resolution_colors["R2"]) + g_r2.SetLineColor(resolution_colors["R2"]) + g_r2.SetMarkerStyle(markers["R2"]) + g_r2.SetMarkerSize(1.5) + g_r2.SetLineWidth(2) + + # Draw graphs + g_r3.Draw("APL") + g_r2.Draw("PL SAME") + + # Set x-axis labels + for i in range(n_det): + label = ROOT.TLatex() + label.SetTextSize(0.02) + label.SetTextAngle(45) + label.DrawLatex(x[i], min(y_r3 + y_r2) * 0.9, labels[i]) + + # Add legend + leg_summary = ROOT.TLegend(0.7, 0.8, 0.9, 0.9) + leg_summary.AddEntry(g_r3, "R3 (v3 method)", "lp") + leg_summary.AddEntry(g_r2, "R2 (standard)", "lp") + leg_summary.Draw() + + # Add ALICE label + alice_label = ROOT.TLatex() + alice_label.SetNDC() + alice_label.SetTextFont(42) + alice_label.SetTextSize(0.04) + alice_label.DrawLatex(0.2, 0.95, " R3 vs R2 Resolution Comparison") + + c_summary.Update() + c_summary.SaveAs("plots/r3_r2_summary_comparison.pdf") + c_summary.SaveAs("plots/r3_r2_summary_comparison.png") + + # Save summary results + output_file.cd() + g_r3.Write("summary_avg_r3") + g_r2.Write("summary_avg_r2") + +# Create text file with results summary +with open("plots/results_summary.txt", "w") as f: + f.write("R3 vs R2 Resolution Comparison Summary\n") + f.write("=" * 60 + "\n\n") + f.write(f"{'Detector Combination':<25} {'Avg R3':<10} {'Avg R2':<10} {'Bins':<5}\n") + f.write("-" * 60 + "\n") + + for result in results_summary: + f.write(f"{result['detector']:<25} {result['avg_r3']:<10.3f} {result['avg_r2']:<10.3f} {result['count']:<5}\n") + + f.write("\nLegend:\n") + f.write("- R3: Resolution from v3 method (0-100% centrality)\n") + f.write("- R2: Standard resolution (0-100% centrality)\n") + +# Close files +output_file.Close() +f1.Close() +f2.Close() + +print("\n" + "="*60) +print("COMPARISON COMPLETED SUCCESSFULLY") +print("="*60) +print(f"Compared {len(results_summary)} detector combinations") +print("Results saved in 'plots/' directory:") +print("- Individual comparison plots (PDF/PNG)") +print("- Summary comparison plot") +print("- ROOT file with all histograms") +print("- Text file with results summary") +print("="*60) \ No newline at end of file diff --git a/src/compute_Resolution.py b/src/compute_Resolution.py new file mode 100644 index 0000000..fa4fee5 --- /dev/null +++ b/src/compute_Resolution.py @@ -0,0 +1,147 @@ +import sys +import argparse +import ROOT +import os + +# Append the utils/ directory to Python path +utils_path = os.path.join(os.path.dirname(__file__), 'utils') +sys.path.append(utils_path) + +# Import your custom functions from utils +from utils import get_resolution, get_centrality_bins, getListOfHisots +from StyleFormatter import SetObjectStyle, SetGlobalStyle + +SetGlobalStyle(padleftmargin=0.15, padbottommargin=0.15, + padrightmargin=0.15, titleoffsety=1.1, maxdigits=3, titlesizex=0.03, + labelsizey=0.04, setoptstat=0, setopttitle=0, palette=ROOT.kGreyScale) + +ROOT.gROOT.SetBatch(False) + +# TODO: move this to the StyleFormatter +def SetFrameStyle(hFrame, xtitle, ytitle, ytitleoffset, ytitlesize, ylabelsize, + ylabeloffset, xticklength, yticklength, xtitlesize, xlabelsize, + xtitleoffset, xlabeloffset, ydivisions, xmoreloglabels, ycentertitle, ymaxdigits): + hFrame.GetXaxis().SetTitle(xtitle) + hFrame.GetYaxis().SetTitle(ytitle) + hFrame.GetYaxis().SetTitleOffset(ytitleoffset) + hFrame.GetYaxis().SetTitleSize(ytitlesize) + hFrame.GetYaxis().SetLabelSize(ylabelsize) + hFrame.GetYaxis().SetLabelOffset(ylabeloffset) + hFrame.GetXaxis().SetTickLength(xticklength) + hFrame.GetYaxis().SetTickLength(yticklength) + hFrame.GetXaxis().SetTitleSize(xtitlesize) + hFrame.GetXaxis().SetLabelSize(xlabelsize) + hFrame.GetXaxis().SetTitleOffset(xtitleoffset) + hFrame.GetXaxis().SetLabelOffset(xlabeloffset) + hFrame.GetYaxis().SetNdivisions(ydivisions) + hFrame.GetXaxis().SetMoreLogLabels(xmoreloglabels) + hFrame.GetYaxis().CenterTitle(ycentertitle) + hFrame.GetYaxis().SetMaxDigits(ymaxdigits) + +def compute_reso(an_res_file, centClass, wagon_id, outputdir, suffix): + + _, cent_min_max = get_centrality_bins(centClass) + + # only SP method + histos_triplets, histos_triplets_lables = getListOfHisots(an_res_file, wagon_id) + + # prepare output file + ytitle = 'Q^{A} Q^{B}' + outfile_name = f'{outputdir}resosp{suffix}.root' + outfile = ROOT.TFile(outfile_name, 'RECREATE') + + # loop over all possible combinations of detectors + latex = ROOT.TLatex() + latex.SetNDC() + latex.SetTextSize(0.05) + for i, (histo_triplet, histo_triplet_label) in enumerate(zip(histos_triplets, histos_triplets_lables)): + histos_mean, histos_mean_deltacent, histo_reso, histo_reso_deltacent = get_resolution(histo_triplet, + histo_triplet_label, + cent_min_max) + detA_label = histo_triplet_label[0] + detB_label = histo_triplet_label[1] + detC_label = histo_triplet_label[2] + outfile.cd() + outfile.mkdir(f'{detA_label}_{detB_label}_{detC_label}') + outfile.cd(f'{detA_label}_{detB_label}_{detC_label}') + canvas = ROOT.TCanvas(f'canvas_{detA_label}_{detB_label}_{detC_label}', + f'canvas_{detA_label}_{detB_label}_{detC_label}', + 2400, 800) + canvas.Divide(3, 1) + leg = ROOT.TLegend(0.2, 0.2, 0.5, 0.3) + leg.SetBorderSize(0) + leg.SetFillStyle(0) + leg.SetTextSize(0.03) + for i, (hist_det, hist_mean, histo_mean_deltacent) in enumerate(zip(histo_triplet, + histos_mean, + histos_mean_deltacent)): + SetObjectStyle(hist_mean, color=ROOT.kRed, markerstyle=ROOT.kFullCircle, + markersize=1, fillstyle=0, linewidth=2) + SetObjectStyle(histo_mean_deltacent, color=ROOT.kBlue, markerstyle=ROOT.kOpenCircle, + markersize=1, fillstyle=0, linestyle=2, linewidth=3) + canvas.cd(i+1) + canvas.cd(i+1).SetLogz() + hFrame = canvas.cd(i+1).DrawFrame(0, -2, 100, 2) + SetFrameStyle(hFrame, + xtitle='Cent. FT0c (%)', + ytitle=ytitle, + ytitleoffset=1.15, + ytitlesize=0.05, + ylabelsize=0.04, + ylabeloffset=0.01, + xticklength=0.04, + yticklength=0.03, + xtitlesize=0.05, + xlabelsize=0.04, + xtitleoffset=1.1, + xlabeloffset=0.020, + ydivisions=406, + xmoreloglabels=True, + ycentertitle=True, + ymaxdigits=5) + hist_det.Draw('same colz') + histo_mean_deltacent.Draw('same pl') + hist_mean.Draw('same pl') + if i == 0: + leg.AddEntry(hist_mean, 'Average 1% centrality', 'lp') + leg.AddEntry(histo_mean_deltacent, + f'Average {cent_min_max[1]-cent_min_max[0]}% centrality', 'lp') + leg.Draw() + latex.DrawLatex(0.2, 0.85, f'A: {detA_label}, B: {detB_label}') + elif i == 1: + latex.DrawLatex(0.2, 0.85, f'A: {detA_label}, B: {detC_label}') + else: + latex.DrawLatex(0.2, 0.85, f'A: {detB_label}, B: {detC_label}') + histo_mean_deltacent.Write() + hist_mean.Write() + hist_det.Write() + canvas.Update() + canvas.Write() + histo_reso.SetDirectory(outfile) + histo_reso_deltacent.SetDirectory(outfile) + histo_reso.Write() + histo_reso_deltacent.Write() + outfile.cd('..') + + input('Resolutions computed. Press any key to continue') + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Arguments") + parser.add_argument("an_res_file", metavar="text", + default="an_res.root", help="input ROOT file with anres") + parser.add_argument('--centClass', '-c', metavar='text', default='k0100') + parser.add_argument("--wagon_id", "-w", metavar="text", + default="", help="wagon ID", required=False) + parser.add_argument("--outputdir", "-o", metavar="text", + default=".", help="output directory") + parser.add_argument("--suffix", "-s", metavar="text", + default="", help="suffix for output files") + args = parser.parse_args() + + compute_reso( + an_res_file=args.an_res_file, + centClass=args.centClass, + wagon_id=args.wagon_id, + outputdir=args.outputdir, + suffix=args.suffix + ) \ No newline at end of file diff --git a/utils/utils.py b/utils/utils.py index 0c6efdc..fb62ba7 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -3,7 +3,12 @@ import sys import ROOT import ctypes -from ROOT import TH1, TH2, TH3, TFile +from itertools import combinations +from PIL import Image +import math +import glob +import re +from ROOT import TH1, TH2, TH3 import numpy as np def check_dir(dir): @@ -42,6 +47,10 @@ def logger(message, level='INFO'): else: print(f"\033[37m{message}\033[0m") # Default to white for unknown levels +def make_dir_root_file(dir, file): + if not file.GetDirectory(dir): + file.mkdir(dir) + def make_dir_root_file(directory, file): if not file.GetDirectory(directory): file.mkdir(directory) @@ -356,65 +365,6 @@ def check_histo_exists(file, histo_name): if file.Get(histo_name): histo_exists = True return histo_exists - -def get_refl_histo(reflFile, ptMins, ptMaxs): - ''' - Method that loads MC histograms for the reflections of D0 - - Input: - - reflFile: - TFile, ROOT file, include reflections of D0 - - centMinMax: - list, min and max centrality - - ptMins: - list, min pt bins - - ptMaxs: - list, max pt bins - - Output: - - useRefl: - bool, if True, MC histograms for the reflections of D0 exists - - hMCSgn: - lsit, signal histograms of D0 - - hMCRefl: - list, reflection histograms of D0 - ''' - hMCSgn, hMCRefl = [], [] - if not check_file_exists(reflFile): - logger(f'Reflections file {reflFile} does not exist! Turning off reflections usage', level='ERROR') - return False - - reflFile = TFile(reflFile, 'READ') - for iPt, (ptMin, ptMax) in enumerate(zip(ptMins, ptMaxs)): - ptMinSuf, ptMaxSuf = int(ptMin*10), int(ptMax*10) - dirName = f'pt_{ptMinSuf}_{ptMaxSuf}' - if not reflFile.GetDirectory(dirName): - logger(f'No directory {dirName} found! Turning off reflections usage', level='ERROR') - return False - - hMCSgn.append(reflFile.Get(f'{dirName}/hFDMass')) - hMCSgn[iPt].Add(reflFile.Get(f'{dirName}/hPromptMass'), 1) - if not isinstance(hMCSgn[iPt], TH1) or hMCSgn[iPt] == None: - logger(f'In directory {dirName}, hFDMass/hPromptMass_{ptMinSuf}_{ptMaxSuf} not found! Turning off reflections usage', level='ERROR') - return False - hMCSgn[iPt].SetName(f'histSgn_{iPt}') - hMCSgn[iPt].SetDirectory(0) - - hMCRefl.append(reflFile.Get(f'{dirName}/hRecoReflMass')) - if not isinstance(hMCRefl[iPt], TH1) or hMCRefl[iPt] == None: - logger(f'In directory {dirName}, hRecoReflMass not found! Turning off reflections usage', level='ERROR') - return False - hMCRefl[iPt].SetName(f'histRfl_{iPt}') - hMCRefl[iPt].SetDirectory(0) - - if hMCRefl[iPt].Integral() <= 0: - logger(f'Error: Empty reflection template for pt bin {ptMin}-{ptMax}! Turning off reflections usage', level='ERROR') - return False - - reflFile.Close() - - return True, hMCSgn, hMCRefl - def get_centrality_bins(centrality): ''' Get centrality bins @@ -477,40 +427,291 @@ def get_centrality_bins(centrality): print(f"ERROR: cent class \'{centrality}\' is not supported! Exit") sys.exit() -def reweight_histo_1D(histo, weights, binned=False): - for iBin in range(1, histo.GetNbinsX()+1): - ptCent = histo.GetBinCenter(iBin) - weight = weights[iBin-1] if binned else weights(ptCent) if weights(ptCent) > 0 else 0 - histo.SetBinContent(iBin, histo.GetBinContent(iBin) * weight) - histo.SetBinError(iBin, histo.GetBinError(iBin) * weight) - proj_hist = histo.Clone(histo.GetName()) - return proj_hist - -def reweight_histo_2D(histo, weights, binned=False): - for iBinX in range(1, histo.GetXaxis().GetNbins()+1): - for iBinY in range(1, histo.GetYaxis().GetNbins()+1): - if binned: - weight = weights[iBinY-1] - else: - binCentVal = histo.GetYaxis().GetBinCenter(iBinY) - weight = weights(binCentVal) if weights(binCentVal) > 0 else 0 - weighted_content = histo.GetBinContent(iBinX, iBinY) * weight - weighted_error = histo.GetBinError(iBinX, iBinY) * weight if weight > 0 else 0 - histo.SetBinContent(iBinX, iBinY, weighted_content) - histo.SetBinError(iBinX, iBinY, weighted_error) - proj_hist = histo.ProjectionX(histo.GetName(), 0, histo.GetYaxis().GetNbins()+1, 'e') - return proj_hist - -def reweight_histo_3D(histo, weightsY, weightsZ): - for iBinX in range(1, histo.GetXaxis().GetNbins()+1): - for iBinY in range(1, histo.GetYaxis().GetNbins()+1): - for iBinZ in range(1, histo.GetZaxis().GetNbins()+1): - binCenterY = histo.GetYaxis().GetBinCenter(iBinY) - weight = weightsZ[iBinZ-1]*weightsY(binCenterY) if weightsY(binCenterY) > 0 else weightsZ[iBinZ-1] - weighted_content = histo.GetBinContent(iBinX, iBinY, iBinZ) * weight - weighted_error = histo.GetBinError(iBinX, iBinY, iBinZ) * weight if weight > 0 else 0 - histo.SetBinContent(iBinX, iBinY, iBinZ, weighted_content) - histo.SetBinError(iBinX, iBinY, iBinZ, weighted_error) - proj_hist = histo.ProjectionX(histo.GetName(), 0, histo.GetYaxis().GetNbins()+1, - 0, histo.GetZaxis().GetNbins()+1, 'e') - return proj_hist \ No newline at end of file + + + +def get_resolution(dets, det_lables, cent_min_max): + ''' + Compute resolution for SP method + ''' + histo_projs, histo_means, histo_means_deltacent = [], [], [] + + # collect the qvecs and prepare histo for mean and resolution + for _, (det, det_label) in enumerate(zip(dets, det_lables)): + print(f'Processing {det_label}') + # th1 for mean 1% centrality bins + histo_means.append(ROOT.TH1F('', '', cent_min_max[1]-cent_min_max[0], cent_min_max[0], cent_min_max[1])) + histo_means[-1].SetDirectory(0) + histo_means[-1].SetName(f'proj_{det_label}_mean') + # th1 for mean CentMin-CentMax + histo_projs.append([]) + hist_proj_dummy = det.ProjectionY(f'proj_{det.GetName()}_mean_deltacent', + det.GetXaxis().FindBin(cent_min_max[0]), + det.GetXaxis().FindBin(cent_min_max[1])-1) + histo_means_deltacent.append(ROOT.TH1F('', '', 1, cent_min_max[0], cent_min_max[1])) + histo_means_deltacent[-1].SetDirectory(0) + histo_means_deltacent[-1].SetName(f'proj_{det_label}_mean_deltacent') + + # Set mean values for CentMin-CentMax + histo_means_deltacent[-1].SetBinContent(1, hist_proj_dummy.GetMean()) + histo_means_deltacent[-1].SetBinError(1, hist_proj_dummy.GetMeanError()) + del hist_proj_dummy + + # collect projections 1% centrality bins + for cent in range(cent_min_max[0], cent_min_max[1]): + bin_cent = det.GetXaxis().FindBin(cent) # common binning + histo_projs[-1].append(det.ProjectionY(f'proj_{det_label}_{cent}', + bin_cent, bin_cent)) + # Set mean values for 1% centrality bins + for ihist, _ in enumerate(histo_projs[-1]): + histo_means[-1].SetBinContent(ihist+1, histo_projs[-1][ihist].GetMean()) + + # Compute resolution for 1% centrality bins + histo_reso = ROOT.TH1F('histo_reso', 'histo_reso', + cent_min_max[1]-cent_min_max[0], + cent_min_max[0], cent_min_max[1]) + histo_reso.SetDirectory(0) + for icent in range(cent_min_max[0], cent_min_max[1]): + reso = compute_resolution([histo_means[i].GetBinContent(icent-cent_min_max[0]+1) for i in range(len(dets))]) + centbin = histo_reso.GetXaxis().FindBin(icent) + histo_reso.SetBinContent(centbin, reso) + + # Compute resolution for CentMin-CentMax + histo_reso_delta_cent = ROOT.TH1F('histo_reso_delta_cent', 'histo_reso_delta_cent', + 1, cent_min_max[0], cent_min_max[1]) + res_deltacent = compute_resolution([histo_means_deltacent[i].GetBinContent(1) for i in range(len(dets))]) + histo_reso_delta_cent.SetBinContent(1, res_deltacent) + histo_reso_delta_cent.SetDirectory(0) + + return histo_means, histo_means_deltacent, histo_reso, histo_reso_delta_cent + +def getListOfHisots(an_res_file, wagon_id): + ''' + Get list of histograms for SP resolution + Now only SP method + ''' + # Build the path - only SP method now + infile_path = f'hf-task-flow-charm-hadrons' + if wagon_id: + infile_path = f'{infile_path}_id{wagon_id}' + infile_path = f'{infile_path}/spReso' + prefix = 'hSpReso' + + print(f"🔍 Looking for directory in ROOT file: {infile_path}") + + infile = ROOT.TFile(an_res_file, 'READ') + directory = infile.GetDirectory(infile_path) + + if not directory: + print(f"❌ ERROR: Directory '{infile_path}' not found in file {an_res_file}") + print("📁 Available top-level keys:") + infile.ls() + return [], [] + + histos = [key.ReadObj() for key in directory.GetListOfKeys()] + for histo in histos: + histo.SetDirectory(0) + pairs = [key.GetName() for key in directory.GetListOfKeys()] + + triplets = list(combinations(pairs, 3)) + histo_triplets = list(combinations(histos, 3)) + correct_histo_triplets = [] + correct_histo_labels = [] + detsA = ['FT0c', 'FT0a', 'FV0a', 'TPCpos', 'FT0m', 'TPCneg'] + + for i, triplet in enumerate(triplets): + for detA in detsA: + detB = triplet[0].replace(prefix, '').replace(detA, '') + detC = triplet[1].replace(prefix, '').replace(detA, '') + if (detA in triplet[0] and detA in triplet[1]) and \ + (detB in triplet[0] and detB in triplet[2]) and \ + (detC in triplet[1] and detC in triplet[2]): + correct_histo_triplets.append(histo_triplets[i]) + correct_histo_labels.append((detA, detB, detC)) + + return correct_histo_triplets, correct_histo_labels + +def compute_resolution(subMean): + ''' + Compute resolution for SP method + ''' + print(subMean) + if len(subMean) == 1: + resolution = subMean[0] + if resolution <= 0: + return 0 + else: + return np.sqrt(resolution) + elif len(subMean) == 3: + print('3 subsystems') + resolution = (subMean[0] * subMean[1]) / subMean[2] if subMean[2] != 0 else 0 + if resolution <= 0: + return 0 + else: + print(resolution, np.sqrt(resolution)) + return np.sqrt(resolution) + else: + print('ERROR: dets must be a list of 2 or 3 subsystems') + sys.exit(1) + +# FIXED: Removed compute_r2 function as it was EP-specific + +# FIXED: Removed get_invmass_vs_deltaphi function as it was EP-specific + +# FIXED: Removed get_ep_vn function as it was EP-specific + +def get_cut_sets(npt_bins, sig_cut, bkg_cut_maxs, correlated_cuts=True): + ''' + Get cut sets + ''' + nCutSets = [] + sig_cuts_lower, sig_cuts_upper, bkg_cuts_lower, bkg_cuts_upper = {}, {}, {}, {} + if correlated_cuts: + sig_cut_mins = sig_cut['min'] + sig_cut_maxs = sig_cut['max'] + sig_cut_steps = sig_cut['step'] + + # compute the signal cutsets for each pt bin + sig_cuts_lower = [list(np.arange(sig_cut_mins[iPt], sig_cut_maxs[iPt], sig_cut_steps[iPt])) for iPt in range(npt_bins)] + sig_cuts_upper = [[1.0 for _ in range(len(sig_cuts_lower[iPt]))] for iPt in range(npt_bins)] + + # compute the ncutsets by signal cut for each pt bin + nCutSets = [len(sig_cuts_lower[iPt]) for iPt in range(npt_bins)] + + # bkg cuts lower edge should always be 0 + bkg_cuts_lower = [[0. for _ in range(nCutSets[iPt])] for iPt in range(npt_bins)] + bkg_cuts_upper = [[bkg_cut_maxs[iPt] for _ in range(nCutSets[iPt])] for iPt in range(npt_bins)] + + else: + # load the signal cut + sig_cuts_lower = [sig_cut[iPt]['min'] for iPt in range(npt_bins)] + sig_cuts_upper = [sig_cut[iPt]['max'] for iPt in range(npt_bins)] + + # compute the ncutsets by the signal cut for each pt bin + nCutSets = [len(sig_cuts_lower[iPt]) for iPt in range(npt_bins)] + + # load the background cut + bkg_cuts_lower = [[0. for _ in range(nCutSets[iPt])] for iPt in range(npt_bins)] + # different max bkg cuts for different cutsets, list of the max bkg cuts given + bkg_cuts_upper = [bkg_cut_maxs[iPt] for iPt in range(npt_bins)] + + # safety check + + for iPt in range(npt_bins): + assert len(sig_cuts_lower[iPt]) == len(sig_cuts_upper[iPt]) == len(bkg_cuts_lower[iPt]) == len(bkg_cuts_upper[iPt]) == nCutSets[iPt], ( + f"Mismatch in lengths for pt bin {iPt}: \n" + f"sig_low:{len(sig_cuts_lower[iPt])}, \n" + f"sig_up: {len(sig_cuts_upper[iPt])}, \n" + f"bkg_low: {len(bkg_cuts_lower[iPt])}, \n" + f"bkg_up: {len(bkg_cuts_upper[iPt])}, \n" + f"nCutSets: {nCutSets[iPt]}" + ) + + return nCutSets, sig_cuts_lower, sig_cuts_upper, bkg_cuts_lower, bkg_cuts_upper + +def get_cut_sets_config(config): + ''' + Get cut sets from configuration file + ''' + import yaml + with open(config, 'r') as ymlCfgFile: + config = yaml.load(ymlCfgFile, yaml.FullLoader) + + ptmins = config['ptmins'] + ptmaxs = config['ptmaxs'] + correlated_cuts = config['minimisation']['correlated'] + if correlated_cuts: + sig_cut = config['cut_variation']['corr_bdt_cut']['sig'] + bkg_cut_maxs = config['cut_variation']['corr_bdt_cut']['bkg_max'] + else: + sig_cut = config['cut_variation']['uncorr_bdt_cut']['sig'] + bkg_cut_maxs = config['cut_variation']['uncorr_bdt_cut']['bkg_max'] + + return get_cut_sets(len(ptmins), sig_cut, bkg_cut_maxs, correlated_cuts) + +def cut_var_image_merger(config, cut_var_dir, suffix): + ''' + Merge cut variation images + ''' + def pdf_to_images(pdf_path, dpi=100): + """Extract high-quality images from a PDF.""" + import fitz + doc = fitz.open(pdf_path) + images = [] + + for page in doc: + pix = page.get_pixmap(dpi=dpi) + img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) + images.append(img) + + return images + + def create_multipanel(images, iImage, images_per_row=None, images_per_col=None, bg_color="white"): + """Combine multiple images into a grid layout with customizable rows and columns.""" + if not images: + raise ValueError("At least one image is required.") + + num_images = len(images) + # Auto-calculate rows and columns if not specified + if images_per_row is None and images_per_col is None: + images_per_row = math.ceil(math.sqrt(num_images)) + if images_per_row is None: + images_per_row = math.ceil(num_images / images_per_col) + if images_per_col is None: + images_per_col = math.ceil(num_images / images_per_row) + + # Resize images to match the smallest width and height + min_width = min(img[iImage].width for img in images) + min_height = min(img[iImage].height for img in images) + resized_images = [img[iImage].resize((min_width, min_height), Image.LANCZOS) for img in images] + + # Create a blank canvas + combined_width = min_width * images_per_row + combined_height = min_height * images_per_col + combined = Image.new("RGB", (combined_width, combined_height), bg_color) + + # Paste images into the grid + for index, img in enumerate(resized_images): + row, col = divmod(index, images_per_row) + x_offset = col * min_width + y_offset = row * min_height + combined.paste(img, (x_offset, y_offset)) + + return combined + + def process_pdfs(config, pdf_list, output_folder, images_names, images_per_row=None, images_per_col=None): + """Processes PDFs and saves multipanel images with high-quality settings.""" + if len(pdf_list) == 0: + raise ValueError("No PDF provided!") + + images = [pdf_to_images(pdf, dpi=100) for pdf in pdf_list] + num_pages = min(len(imgs) for imgs in images) + + os.makedirs(output_folder, exist_ok=True) + + for i in range(num_pages): + panel = create_multipanel(images, i, images_per_row, images_per_col) + output_path = os.path.join(output_folder, f"{images_names}_pt_{int(config['ptmins'][i]*10)}_{int(config['ptmaxs'][i]*10)}.png") + panel.save(output_path, format="PNG", compression_level=1) + + print(f"Saved {num_pages} high-quality multipanel images in '{output_folder}'.") + + cutvar_files = [ + f"{cut_var_dir}/CutVarFrac/CutVarFrac_{suffix}_CorrMatrix.pdf", + f"{cut_var_dir}/CutVarFrac/CutVarFrac_{suffix}_Distr.pdf", + f"{cut_var_dir}/CutVarFrac/CutVarFrac_{suffix}_Eff.pdf", + f"{cut_var_dir}/CutVarFrac/CutVarFrac_{suffix}_Frac.pdf", + f"{cut_var_dir}/V2VsFrac/FracV2_{suffix}.pdf" + ] + + try: + process_pdfs(config, cutvar_files, f"{cut_var_dir}/merged_images/cutvar", 'cutvar_summary') + except: + print("Error in merging cut variation files") + + try: + fit_files = glob.glob(f"{cut_var_dir}/ry/*.pdf") + fit_files_sorted = sorted(fit_files, key=lambda x: int(re.search(r"_(\d+)_D\w*.pdf$", x).group(1))) + process_pdfs(config, fit_files_sorted, f"{cut_var_dir}/merged_images/fits/", 'fit_summary', 5) + except: + print("Error in merging fit files") \ No newline at end of file