From 711b09f2ad01349cf4c6cc30b14f46ef0fc200ff Mon Sep 17 00:00:00 2001 From: RaresAli Date: Wed, 1 Jul 2026 15:16:57 +0200 Subject: [PATCH] Analyze multitrial parameters --- syst/multitrial/multitrial_params.py | 326 +++++++++++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 syst/multitrial/multitrial_params.py diff --git a/syst/multitrial/multitrial_params.py b/syst/multitrial/multitrial_params.py new file mode 100644 index 0000000..45a1f98 --- /dev/null +++ b/syst/multitrial/multitrial_params.py @@ -0,0 +1,326 @@ +#!/usr/bin/env python3 +""" +To run the script you need to have previously run "run_multitrial_fit.sh" to generate the trials_results.parquet file with the trial results. + +For Rebin, Chi2, Significance, MassMin, MassMax, Sigma and Mean the script produces a plot of variable vs trial +In addition by using "--plot A:B A:C B:C" you can produce a plot of whichever variable you like + as long as it has the name of one of the .parquet columns +""" + +import argparse +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from matplotlib.ticker import FormatStrFormatter, MaxNLocator +from matplotlib.backends.backend_pdf import PdfPages +from pathlib import Path +from collections import defaultdict + +try: + import ROOT +except ImportError: + ROOT = None + +# ---------------------------------------------------------------------- +def parse_mass_fit_ranges(ranges_series): + """Parse mass fit ranges from string or list format.""" + mass_min, mass_max = [], [] + for val in ranges_series: + if isinstance(val, str): + try: + s = val.strip("[]") + parts = [float(x.strip()) for x in s.split(",")] + if len(parts) == 2: + mass_min.append(parts[0]) + mass_max.append(parts[1]) + else: + mass_min.append(np.nan) + mass_max.append(np.nan) + except: + mass_min.append(np.nan) + mass_max.append(np.nan) + elif isinstance(val, (list, tuple)) and len(val) == 2: + mass_min.append(float(val[0])) + mass_max.append(float(val[1])) + else: + mass_min.append(np.nan) + mass_max.append(np.nan) + return np.array(mass_min, dtype=float), np.array(mass_max, dtype=float) + +# ---------------------------------------------------------------------- +def compute_ylimits(var_name, y_vals): + """Calculate y_min and y_max based on the variable name (without suffixes).""" + valid = y_vals[~np.isnan(y_vals)] + if len(valid) == 0: + return 0, 1 + + if var_name in ["Rebin"]: + y_min = np.nanmin(y_vals) - 1 + y_max = np.nanmax(y_vals) + 1 + elif var_name == "Chi2": + y_min = 0.0 + y_max = np.ceil(np.nanmax(y_vals)) + if y_max == 0: + y_max = 1 + elif var_name == "Significance": + y_min = np.nanmin(y_vals) - 5 + y_max = np.nanmax(y_vals) + 5 + elif var_name in ["MassMin", "MassMax"]: + y_min = np.nanmin(y_vals) - 0.01 + y_max = np.nanmax(y_vals) + 0.01 + elif var_name == "Sigma": + y_min = np.nanmin(y_vals) - 0.00005 + y_max = np.nanmax(y_vals) + 0.00005 + else: # Mean and anything else + y_min = np.nanmin(y_vals) * 0.9995 + y_max = np.nanmax(y_vals) * 1.0005 + + if y_min == y_max: + y_min -= 0.001 + y_max += 0.001 + return y_min, y_max + +# ---------------------------------------------------------------------- +def clean_name(name): + """Remove _BkgFunc or _BkgFuncVn suffixes from the task name.""" + for suffix in ["_BkgFuncVn", "_BkgFunc"]: + if name.endswith(suffix): + return name[:-len(suffix)] + return name # fallback (should not happen) + +def main(): + parser = argparse.ArgumentParser( + description="Multi-Variable vs Trial plots with optional custom plots" + ) + parser.add_argument("--centrality", required=True, + choices=["0-20", "20-50", "50-100"], + help="Centrality class") + parser.add_argument("--output", required=True, + help="Output directory") + parser.add_argument("--plot", nargs='+', type=str, default=None, + help='Custom scatter plot in the form "Y:X" (e.g. "Mean:MassMin"). ' + 'Multiple pairs can be added separated by spaces.') + args = parser.parse_args() + + # Process custom pairs + custom_plots = [] # list of tuples (y_col, x_col) + if args.plot: + for pair in args.plot: + parts = pair.split(':') + if len(parts) != 2: + parser.error(f"Invalid format '{pair}'. Use 'Y:X'.") + y_col, x_col = parts[0], parts[1] + custom_plots.append((y_col, x_col)) + + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + parquet_path = Path( + f"/home/rdrosu/alice/analisi/misure_v2_inclusive/crystal_ball/" + f"{args.centrality}/misure/multitrial/syst/multitrial/fit/summary/trials_results.parquet" + ) + if not parquet_path.exists(): + print(f"ERROR: file not found: {parquet_path}") + return + + try: + df = pd.read_parquet(parquet_path) + except Exception as e: + print(f"ERROR reading Parquet file: {e}") + return + + df_filtered = df[df["V2Type"] == "Cutset_0"].copy() + if df_filtered.empty: + print("No rows with V2Type == 'Cutset_0'.") + return + + required = ["TrialIdx", "PtLabel", "BkgFunc", "BkgFuncVn"] + for c in required: + if c not in df_filtered.columns: + print(f"ERROR: column '{c}' missing.") + return + + if "MassFitRanges" in df_filtered.columns: + mass_min, mass_max = parse_mass_fit_ranges(df_filtered["MassFitRanges"]) + df_filtered["MassMin"] = mass_min + df_filtered["MassMax"] = mass_max + else: + df_filtered["MassMin"] = np.nan + df_filtered["MassMax"] = np.nan + + for y_col, x_col in custom_plots: + for col in [y_col, x_col]: + if col not in df_filtered.columns: + print(f"ERROR: column '{col}' required for custom plot not found.") + return + + color_map_mpl = {"kLin": "red", "kPol2": "green", "kExpo": "blue"} + default_color_mpl = "black" + if ROOT is not None: + color_map_root = {"kLin": ROOT.kRed, "kPol2": ROOT.kGreen + 2, "kExpo": ROOT.kBlue} + default_color_root = ROOT.kBlack + + for pt_label, group in df_filtered.groupby("PtLabel", observed=True): + group = group.sort_values("TrialIdx") + x_trial = group["TrialIdx"].astype(int).values + bkgfunc = group["BkgFunc"].values + bkgfuncvn = group["BkgFuncVn"].values + + standard_vars = [] + for var in ["Mean", + "Sigma", + "Significance", + "Chi2", + #"Rebin", + #"MassMin", + #"MassMax", + ]: + if var in group.columns: + y_vals = pd.to_numeric(group[var], errors='coerce').astype(float).values + if not np.all(np.isnan(y_vals)): + standard_vars.append((var, y_vals)) + + plot_tasks = [] + for var_name, y_vals in standard_vars: + plot_tasks.append((f"{var_name}_BkgFunc", x_trial, y_vals, bkgfunc)) + plot_tasks.append((f"{var_name}_BkgFuncVn", x_trial, y_vals, bkgfuncvn)) + + for y_col, x_col in custom_plots: + y_c = pd.to_numeric(group[y_col], errors='coerce').astype(float).values + x_c = pd.to_numeric(group[x_col], errors='coerce').astype(float).values + valid = ~np.isnan(y_c) & ~np.isnan(x_c) + if np.any(valid): + x_cut = x_c[valid] + y_cut = y_c[valid] + bkgfunc_cut = bkgfunc[valid] + bkgfuncvn_cut = bkgfuncvn[valid] + plot_tasks.append((f"{y_col}_vs_{x_col}_BkgFunc", x_cut, y_cut, bkgfunc_cut)) + plot_tasks.append((f"{y_col}_vs_{x_col}_BkgFuncVn", x_cut, y_cut, bkgfuncvn_cut)) + else: + print(f"Warning: no valid data for {y_col} vs {x_col} in {pt_label}") + + if not plot_tasks: + print(f"No plot tasks for {pt_label}") + continue + + # ----- Calculate shared Y limits for each variable ----- + groups = defaultdict(list) + for name, x_data, y_data, color_var in plot_tasks: + base_name = clean_name(name) + groups[base_name].append((name, x_data, y_data, color_var)) + + group_ylims = {} + for base, tasks in groups.items(): + all_y = np.concatenate([t[2] for t in tasks]) + y_min, y_max = compute_ylimits(base, all_y) + group_ylims[base] = (y_min, y_max) + + safe_label = str(pt_label).replace("/", "_").replace(" ", "_") + + + pdf_path = output_dir / f"summary_{safe_label}_{args.centrality}.pdf" + with PdfPages(pdf_path) as pdf: + n_tasks = len(plot_tasks) + for start in range(0, n_tasks, 4): + subset = plot_tasks[start:start+4] + n_sub = len(subset) + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + axes = axes.flatten() + + for idx, (name, x_data, y_data, color_var) in enumerate(subset): + ax = axes[idx] + colors = [color_map_mpl.get(c, default_color_mpl) for c in color_var] + ax.scatter(x_data, y_data, c=colors, s=8, edgecolors='none', alpha=0.7) + + # X limits + x_min, x_max = np.nanmin(x_data), np.nanmax(x_data) + if x_min == x_max: + x_min -= 0.5 + x_max += 0.5 + ax.set_xlim(x_min - 0.05*(x_max-x_min), x_max + 0.05*(x_max-x_min)) + + # Y limits from group (shared) + base_name = clean_name(name) + y_min, y_max = group_ylims[base_name] + ax.set_ylim(y_min, y_max) + + # Ticks + ax.xaxis.set_major_locator(MaxNLocator(nbins=6)) + ax.yaxis.set_major_locator(MaxNLocator(nbins=5)) + ax.yaxis.set_major_formatter(FormatStrFormatter('%.4f')) + + # Legend + unique_vals = sorted(set(color_var)) + legend_elements = [ + plt.Line2D([0], [0], marker='o', color='w', label=val, + markerfacecolor=color_map_mpl.get(val, default_color_mpl), + markersize=8) + for val in unique_vals + ] + ax.legend(handles=legend_elements, loc='best', fontsize='x-small') + + # Axis labels + if "_vs_" in base_name: + parts = base_name.split("_vs_") + y_label = parts[0] + x_label = parts[1] if len(parts) > 1 else "X" + else: + x_label = "TrialIdx" + y_label = base_name + ax.set_xlabel(x_label) + ax.set_ylabel(y_label) + ax.set_title(f"{name} ({args.centrality}%)") + ax.grid(True, alpha=0.3) + + for j in range(n_sub, 4): + axes[j].set_visible(False) + + fig.suptitle(f"{pt_label} (Centrality {args.centrality}%)", fontsize=14, fontweight='bold') + fig.tight_layout(rect=[0, 0, 1, 0.97]) + pdf.savefig(fig) + plt.close(fig) + print(f"PDF saved: {pdf_path}") + + # ----- ROOT ----- + if ROOT is not None: + root_path = output_dir / f"summary_{safe_label}_{args.centrality}.root" + f_root = ROOT.TFile(str(root_path), "RECREATE") + + for name, x_data, y_data, color_var in plot_tasks: + mg = ROOT.TMultiGraph() + base_name = clean_name(name) + if "_vs_" in base_name: + parts = base_name.split("_vs_") + y_ax_label = parts[0] + x_ax_label = parts[1] if len(parts) > 1 else "X" + else: + x_ax_label = "TrialIdx" + y_ax_label = base_name + mg.SetTitle(f"{name} - {pt_label} ({args.centrality}%);{x_ax_label};{y_ax_label}") + + for cat in sorted(set(color_var)): + mask = color_var == cat + x_b = x_data[mask] + y_b = y_data[mask] + n = len(x_b) + if n == 0: + continue + + gr = ROOT.TGraph(n, x_b.astype(np.float64), y_b.astype(np.float64)) + gr.SetName(f"gr_{cat.replace('k', '')}") + gr.SetTitle(cat) + gr.SetLineWidth(0) + gr.SetMarkerStyle(20) + gr.SetMarkerSize(0.5) + col = color_map_root.get(cat, default_color_root) + gr.SetMarkerColor(col) + gr.SetLineColor(col) + mg.Add(gr) + + mg.Write(name) + + f_root.Close() + print(f"ROOT saved: {root_path}") + +if __name__ == "__main__": + main() \ No newline at end of file