File indexing completed on 2026-07-16 08:11:23
0001
0002 """
0003 find_sv_smear_fracs.py
0004 ----------------------
0005 Finds per-pT-bin Gaussian smearing fractions for the SV-extrapolated simulation
0006 files so that the smeared SV-sim invariant mass from daughters matches the data
0007 KFP mass width.
0008
0009 Why this is the correct comparison
0010 -----------------------------------
0011 • SV-sim stores momenta at the secondary vertex, which are essentially the
0012 KFP-refitted values → inv mass from SV daughters ≈ KFP mass in sim.
0013 • Data KFP mass is the definitive measured width.
0014 • Both sides use the same (post-KFP / SV) momenta, so smearing the SV-sim
0015 daughters to reach the data KFP sigma is a clean, apples-to-apples match.
0016
0017 Two-step procedure
0018 ------------------
0019 Step 1 K0S → π⁺π⁻
0020 Both daughters are pions; smear each by the same fraction f per pT bin.
0021 Target: σ_data_KFP (data K_S0_mass width).
0022
0023 Step 2 Λ → pπ⁻
0024 Fix pion smearing from Step 1 (per pion-pT lookup).
0025 Vary proton fraction f_p per proton-pT bin.
0026 Target: σ_data_KFP (data Lambda0_mass width).
0027
0028 Outputs
0029 -------
0030 optimal_sv_smear_fracs.txt K0S pion fracs (CSV)
0031 optimal_sv_xi_fracs.txt Lambda proton fracs (CSV)
0032 plot_sv_ks_result.C ROOT macro — K0S sigma vs pT
0033 plot_sv_xi_result.C ROOT macro — Lambda sigma vs pT
0034 """
0035
0036 import argparse, sys
0037 import numpy as np
0038 import uproot
0039 from scipy.optimize import brentq, curve_fit
0040
0041
0042 KS_DATA_FILE = "output_Kshort_run3pp_ana530_2025p009_v001.root"
0043 KS_SV_FILE = "outputKFParticleKShortRecoSV.root"
0044 XI_DATA_FILE = "Ximinus_fullDataset_finalCuts_0p2pTCut_rapidity1p0Cut_BCOcut_charge_.root"
0045 XI_SV_FILE = "outputKFParticleXiminusSV.root"
0046 TREE_NAME = "DecayTree"
0047
0048
0049 KS_PT_BINS = [0.2, 0.25, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0, 1e9]
0050 XI_PT_BINS = [0.3, 0.5, 0.75, 1.0, 1.5, 1e9]
0051
0052
0053 KS_LO, KS_HI, KS_BINS = 460.0, 540.0, 80
0054 XI_LO, XI_HI, XI_BINS = 1300.0, 1340.0, 80
0055
0056
0057 PION_MASS = 0.13957
0058 PROTON_MASS = 0.938272
0059 RNG_SEED = 42
0060
0061
0062 KS_CUT_BR = [
0063 "track_1_MVTX_nStates", "track_2_MVTX_nStates",
0064 "track_1_INTT_nStates", "track_2_INTT_nStates",
0065 "track_1_TPC_nStates", "track_2_TPC_nStates",
0066 "track_1_chi2", "track_1_nDoF", "track_2_chi2", "track_2_nDoF",
0067 "track_1_IP_xy", "track_2_IP_xy",
0068 "track_1_track_2_DCA", "track_1_track_2_DCA_xy",
0069 "K_S0_DIRA", "K_S0_chi2", "K_S0_nDoF",
0070 ]
0071
0072 XI_CUT_BR = [
0073 "Lambda0_track_1_MVTX_nStates", "Lambda0_track_2_MVTX_nStates",
0074 "Lambda0_track_1_INTT_nStates", "Lambda0_track_2_INTT_nStates",
0075 "Lambda0_track_1_TPC_nStates", "Lambda0_track_2_TPC_nStates",
0076 "track_3_MVTX_nStates", "track_3_INTT_nStates", "track_3_TPC_nStates",
0077 "Lambda0_track_1_chi2", "Lambda0_track_1_nDoF", "Lambda0_track_2_chi2", "Lambda0_track_2_nDoF",
0078 "track_3_chi2", "track_3_nDoF",
0079 "track_1_track_2_DCA", "track_1_track_3_DCA", "track_2_track_3_DCA",
0080 "Ximinus_decayLength_xy", "Ximinus_chi2", "Ximinus_nDoF",
0081 "Lambda0_mass", "Lambda0_decayLength_xy", "Lambda0_chi2", "Lambda0_nDoF",
0082 ]
0083
0084
0085
0086 def ks_mass_mev(px1, py1, pz1, px2, py2, pz2):
0087 E1 = np.sqrt(px1**2 + py1**2 + pz1**2 + PION_MASS**2)
0088 E2 = np.sqrt(px2**2 + py2**2 + pz2**2 + PION_MASS**2)
0089 m2 = (E1+E2)**2 - (px1+px2)**2 - (py1+py2)**2 - (pz1+pz2)**2
0090 return np.sqrt(np.maximum(m2, 0.0)) * 1000.0
0091
0092 def xi_mass_mev(px_p, py_p, pz_p, px_pi1, py_pi1, pz_pi1, px_pi2, py_pi2, pz_pi2):
0093 Ep = np.sqrt(px_p**2 + py_p**2 + pz_p**2 + PROTON_MASS**2)
0094 Epi1 = np.sqrt(px_pi1**2 + py_pi1**2 + pz_pi1**2 + PION_MASS**2)
0095 Epi2 = np.sqrt(px_pi2**2 + py_pi2**2 + pz_pi2**2 + PION_MASS**2)
0096 m2 = (Ep+Epi1+Epi2)**2 - (px_p+px_pi1+px_pi2)**2 - (py_p+py_pi1+px_pi2)**2 - (pz_p+pz_pi1+pz_pi2)**2
0097 return np.sqrt(np.maximum(m2, 0.0)) * 1000.0
0098
0099 def apply_smear(px, py, pz, frac, z):
0100 """Smear pT by frac × pT × z (same z draws → smooth σ(frac))."""
0101 pt = np.sqrt(px**2 + py**2)
0102 phi = np.arctan2(py, px)
0103 ptot = np.sqrt(pz**2 + pt**2)
0104 eta = np.arctanh(np.clip(pz / ptot, -1+1e-9, 1-1e-9))
0105 spt = pt + frac * pt * z
0106 return spt*np.cos(phi), spt*np.sin(phi), spt*np.sinh(eta)
0107
0108 def apply_smear_pervec(px, py, pz, fracs, z):
0109 """Same but fracs is a per-event array."""
0110 pt = np.sqrt(px**2 + py**2)
0111 phi = np.arctan2(py, px)
0112 ptot = np.sqrt(pz**2 + pt**2)
0113 eta = np.arctanh(np.clip(pz / ptot, -1+1e-9, 1-1e-9))
0114 spt = pt + fracs * pt * z
0115 return spt*np.cos(phi), spt*np.sin(phi), spt*np.sinh(eta)
0116
0117
0118
0119
0120 def ks_cut_mask(d):
0121 return (
0122 (np.minimum(d["track_1_MVTX_nStates"], d["track_2_MVTX_nStates"]) > 0) &
0123 (np.minimum(d["track_1_INTT_nStates"], d["track_2_INTT_nStates"]) > 0) &
0124 (np.minimum(d["track_1_TPC_nStates"], d["track_2_TPC_nStates"]) >= 20) &
0125 (np.maximum(d["track_1_chi2"]/d["track_1_nDoF"],
0126 d["track_2_chi2"]/d["track_2_nDoF"]) <= 300) &
0127 (np.minimum(np.abs(d["track_1_IP_xy"]),
0128 np.abs(d["track_2_IP_xy"])) >= 0.05) &
0129 (d["track_1_track_2_DCA"] <= 0.5) &
0130 (d["track_1_track_2_DCA_xy"] <= 1.0) &
0131 (d["K_S0_DIRA"] >= 0.99) &
0132 (d["K_S0_chi2"] / d["K_S0_nDoF"] <= 20)
0133 )
0134
0135 def xi_cut_mask(d):
0136 return (
0137 (np.minimum(d["Lambda0_track_1_MVTX_nStates"], d["Lambda0_track_2_MVTX_nStates"], d["track_3_MVTX_nStates"]) > 0) &
0138 (np.minimum(d["Lambda0_track_1_INTT_nStates"], d["Lambda0_track_2_INTT_nStates"], d["track_3_INTT_nStates"]) > 0) &
0139 (np.minimum(d["Lambda0_track_1_TPC_nStates"], d["Lambda0_track_2_TPC_nStates"], d["track_3_TPC_nStates"]) >= 20) &
0140 (np.maximum(d["Lambda0_track_1_chi2"]/d["Lambda0_track_1_nDoF"],
0141 d["Lambda0_track_2_chi2"]/d["Lambda0_track_2_nDoF"],
0142 d["track_3_chi2"]/d["track_3_nDoF"]) <= 400) &
0143 (np.maximum(d["track_1_track_2_DCA"],
0144 d["track_1_track_3_DCA"],
0145 d["track_2_track_3_DCA"]) <= 0.5) &
0146 (d["Ximinus_chi2"]/d["Ximinus_nDoF"] <= 50) &
0147 (d["Lambda0_chi2"]/d["Lambda0_nDoF"] <= 50) &
0148 (np.abs(d["Lambda0_mass"] - 1.1157) <= 0.01) &
0149 (d["Ximinus_decayLength_xy"] >= 0.15) &
0150 (d["Lambda0_decayLength_xy"] >= 0.01)
0151 )
0152
0153
0154
0155 def _gauss(x, amp, mu, sig):
0156 return amp * np.exp(-0.5 * ((x - mu) / sig)**2)
0157
0158 def fit_sigma(masses_mev, lo, hi, nbins):
0159 counts, edges = np.histogram(masses_mev, bins=nbins, range=(lo, hi))
0160 if counts.sum() < 20:
0161 return np.nan, np.nan
0162 ctrs = 0.5*(edges[:-1]+edges[1:])
0163 pk = np.argmax(counts)
0164 try:
0165 popt, _ = curve_fit(_gauss, ctrs, counts.astype(float),
0166 p0=[counts[pk], ctrs[pk], 5.0],
0167 bounds=([0, lo, 0.5], [1e9, hi, 60.0]))
0168 return popt[1], abs(popt[2])
0169 except Exception:
0170 return np.nan, np.nan
0171
0172 def fit_ks(m): return fit_sigma(m, KS_LO, KS_HI, KS_BINS)
0173 def fit_xi(m): return fit_sigma(m, XI_LO, XI_HI, XI_BINS)
0174
0175
0176
0177
0178 def make_pion_lookup(edges, sigmas):
0179 """Return a function pt_arr → per-event frac array."""
0180 edges_arr = np.array(edges)
0181 sigmas_arr = np.array(sigmas)
0182 def lookup(pt_arr):
0183 idx = np.searchsorted(edges_arr, pt_arr, side='right') - 1
0184 idx = np.clip(idx, 0, len(sigmas_arr)-1)
0185 return sigmas_arr[idx]
0186 return lookup
0187
0188
0189
0190
0191
0192
0193 def run_ks(frac_max):
0194 print("=" * 68)
0195 print("STEP 1 — K0S → π⁺π⁻ (SV sim inv mass → data KFP mass)")
0196 print("=" * 68)
0197
0198
0199 print(f"\nLoading KS DATA : {KS_DATA_FILE}")
0200 ks_phys = ["track_1_px","track_1_py","track_1_pz","track_1_pT",
0201 "track_2_px","track_2_py","track_2_pz","track_2_pT","K_S0_mass"]
0202 with uproot.open(KS_DATA_FILE) as f:
0203 d = f[TREE_NAME].arrays(ks_phys + KS_CUT_BR, library="np")
0204 dm = ks_cut_mask(d)
0205 print(f" {dm.sum():,} / {len(dm):,} pass quality cuts")
0206 d = {k: v[dm] for k, v in d.items()}
0207 d_pt1 = d["track_1_pT"].astype(np.float64)
0208 d_pt2 = d["track_2_pT"].astype(np.float64)
0209 d_kfp = d["K_S0_mass"].astype(np.float64) * 1000.0
0210
0211
0212 print(f"Loading KS SV sim: {KS_SV_FILE}")
0213 ks_sim_phys = ks_phys + ["track_1_PDG_ID","track_2_PDG_ID"]
0214 with uproot.open(KS_SV_FILE) as f:
0215 s = f[TREE_NAME].arrays(ks_sim_phys + KS_CUT_BR, library="np")
0216 sm = (ks_cut_mask(s) &
0217 (np.abs(s["track_1_PDG_ID"])==211) &
0218 (np.abs(s["track_2_PDG_ID"])==211))
0219 print(f" {sm.sum():,} / {len(sm):,} pass quality + pion-PDG cuts\n")
0220 s = {k: v[sm] for k, v in s.items()}
0221 s_pt1 = s["track_1_pT"].astype(np.float64)
0222 s_pt2 = s["track_2_pT"].astype(np.float64)
0223 s_px1 = s["track_1_px"].astype(np.float64)
0224 s_py1 = s["track_1_py"].astype(np.float64)
0225 s_pz1 = s["track_1_pz"].astype(np.float64)
0226 s_px2 = s["track_2_px"].astype(np.float64)
0227 s_py2 = s["track_2_py"].astype(np.float64)
0228 s_pz2 = s["track_2_pz"].astype(np.float64)
0229
0230 nBins = len(KS_PT_BINS) - 1
0231 results = []
0232
0233 HDR = (f"{'Bin':>3} {'pT range':>18} "
0234 f"{'σ_data_KFP':>11} {'σ_sv_inv0':>10} "
0235 f"{'best%':>7} {'σ_matched':>10}")
0236 print(HDR)
0237 print("─" * len(HDR.rstrip()))
0238
0239 for i in range(nBins):
0240 lo, hi = KS_PT_BINS[i], KS_PT_BINS[i+1]
0241 hi_s = f"{hi:.3f}" if hi < 1e8 else "∞ "
0242
0243 dm_b = (np.minimum(d_pt1,d_pt2)>=lo) & (np.maximum(d_pt1,d_pt2)<hi)
0244 sm_b = (np.minimum(s_pt1,s_pt2)>=lo) & (np.maximum(s_pt1,s_pt2)<hi)
0245 n_d, n_s = dm_b.sum(), sm_b.sum()
0246 avg_pt = float((d_pt1[dm_b].mean()+d_pt2[dm_b].mean())/2) if n_d>0 else (lo+hi)/2
0247
0248 _, sig_data_kfp = fit_ks(d_kfp[dm_b])
0249 px1_b = s_px1[sm_b]; py1_b = s_py1[sm_b]; pz1_b = s_pz1[sm_b]
0250 px2_b = s_px2[sm_b]; py2_b = s_py2[sm_b]; pz2_b = s_pz2[sm_b]
0251 sv_inv0 = ks_mass_mev(px1_b,py1_b,pz1_b, px2_b,py2_b,pz2_b)
0252 _, sig_sv_inv0 = fit_ks(sv_inv0)
0253
0254 def _bad(note):
0255 results.append(dict(lo=lo, hi=hi, frac=np.nan, avg_pt=avg_pt,
0256 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0257 sig_match=np.nan, note=note, n_d=n_d, n_s=n_s))
0258
0259 if n_d < 20 or np.isnan(sig_data_kfp):
0260 _bad("too few data"); print(f"{i:>3} [{lo:.3f},{hi_s}) too few data"); continue
0261 if n_s < 20 or np.isnan(sig_sv_inv0):
0262 _bad("too few sim"); print(f"{i:>3} [{lo:.3f},{hi_s}) too few sim"); continue
0263
0264 rng = np.random.default_rng(RNG_SEED + i)
0265 z1 = rng.standard_normal(n_s)
0266 z2 = rng.standard_normal(n_s)
0267
0268 def sim_sigma(frac):
0269 spx1,spy1,spz1 = apply_smear(px1_b,py1_b,pz1_b,frac,z1)
0270 spx2,spy2,spz2 = apply_smear(px2_b,py2_b,pz2_b,frac,z2)
0271 _, sig = fit_ks(ks_mass_mev(spx1,spy1,spz1, spx2,spy2,spz2))
0272 return sig
0273
0274 def residual(frac): return sim_sigma(frac) - sig_data_kfp
0275
0276 r0 = residual(0.0)
0277 if r0 >= 0:
0278 note = "sim≥data at 0%"
0279 results.append(dict(lo=lo, hi=hi, frac=0.0, avg_pt=avg_pt,
0280 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0281 sig_match=sig_sv_inv0, note=note, n_d=n_d, n_s=n_s))
0282 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0283 f"{sig_sv_inv0:>10.3f} {'0.00%':>7} {note}")
0284 continue
0285
0286 r_max = residual(frac_max)
0287 if r_max < 0:
0288 sig_at_max = sim_sigma(frac_max)
0289 note = f">={frac_max*100:.0f}%"
0290 results.append(dict(lo=lo, hi=hi, frac=frac_max, avg_pt=avg_pt,
0291 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0292 sig_match=sig_at_max, note=note, n_d=n_d, n_s=n_s))
0293 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0294 f"{sig_sv_inv0:>10.3f} {note:>7} {sig_at_max:>10.3f}")
0295 continue
0296
0297 best = brentq(residual, 0.0, frac_max, xtol=5e-5, maxiter=80)
0298 sig_m = sim_sigma(best)
0299 results.append(dict(lo=lo, hi=hi, frac=best, avg_pt=avg_pt,
0300 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0301 sig_match=sig_m, note="ok", n_d=n_d, n_s=n_s))
0302 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0303 f"{sig_sv_inv0:>10.3f} {best*100:>6.2f}% {sig_m:>10.3f}")
0304
0305 valid = [r for r in results if not np.isnan(r["frac"])]
0306 edges_cpp = ", ".join(f"{r['lo']:.4f}" for r in valid)
0307 sigmas_cpp = ", ".join(f"{r['frac']:.5f}" for r in valid)
0308 print(f"\n{'─'*68}")
0309 print("K0S pion fracs (C++ for smear_KS_decaytree_ptbin.C):\n")
0310 print(f" pt_bin_edges = {{{edges_cpp}}}")
0311 print(f" pt_bin_sigmas = {{{sigmas_cpp}}}")
0312
0313 with open("optimal_sv_smear_fracs.txt","w") as fh:
0314 fh.write("# Auto-generated by find_sv_smear_fracs.py\n")
0315 fh.write("# SV sim inv mass smeared to match data KFP mass.\n")
0316 fh.write(f"# pt_bin_edges = {{{edges_cpp}}}\n")
0317 fh.write(f"# pt_bin_sigmas = {{{sigmas_cpp}}}\n\n")
0318 fh.write("bin_lo,bin_hi,avg_pt,best_frac_pct,sig_data_kfp,sig_sv_inv0,sig_match,n_data,n_sim,note\n")
0319 for r in results:
0320 hi_w = f"{r['hi']:.4f}" if r["hi"]<1e8 else "inf"
0321 _f = lambda v: f"{v:.4f}" if not np.isnan(v) else "nan"
0322 fp = f"{r['frac']*100:.4f}" if not np.isnan(r["frac"]) else "nan"
0323 fh.write(f"{r['lo']:.4f},{hi_w},{r['avg_pt']:.4f},{fp},{_f(r['sig_data_kfp'])},"
0324 f"{_f(r['sig_sv_inv0'])},{_f(r['sig_match'])},"
0325 f"{r['n_d']},{r['n_s']},{r['note']}\n")
0326 print(f"\nCSV saved → optimal_sv_smear_fracs.txt")
0327
0328 return results
0329
0330
0331
0332
0333
0334 def run_xi(ks_results, frac_max):
0335 print("\n" + "=" * 68)
0336 print("STEP 2 — Xi → ppipi (SV sim inv mass → data KFP mass)")
0337 print("=" * 68)
0338
0339
0340 valid_ks = [r for r in ks_results if not np.isnan(r["frac"])]
0341 pi_edges = [r["lo"] for r in valid_ks]
0342 pi_sigmas = [r["frac"] for r in valid_ks]
0343 pion_frac = make_pion_lookup(pi_edges, pi_sigmas)
0344 print(f"\nPion smearing calibration from K0S SV results:")
0345 for lo, frac in zip(pi_edges, pi_sigmas):
0346 print(f" pT >= {lo:.3f} GeV → {frac*100:.3f}%")
0347 print(f" (below {pi_edges[0]:.2f} GeV → {pi_sigmas[0]*100:.3f}% [lowest bin])")
0348
0349
0350 print(f"\nLoading Xi DATA : {XI_DATA_FILE}")
0351 xi_phys = ["Lambda0_track_1_px","Lambda0_track_1_py","Lambda0_track_1_pz","Lambda0_track_1_pT",
0352 "Lambda0_track_2_px","Lambda0_track_2_py","Lambda0_track_2_pz","Lambda0_track_2_pT",
0353 "track_3_px","track_3_py","track_3_pz","track_3_pT",
0354 "Ximinus_mass","Lambda0_track_1_PDG_ID","Lambda0_track_2_PDG_ID","track_3_PDG_ID"]
0355 with uproot.open(XI_DATA_FILE) as f:
0356 d = f[TREE_NAME].arrays(xi_phys + XI_CUT_BR, library="np")
0357 dm = (xi_cut_mask(d) &
0358 (np.abs(d["Lambda0_track_1_PDG_ID"])==211) &
0359 (np.abs(d["Lambda0_track_2_PDG_ID"])==2212) &
0360 (np.abs(d["track_3_PDG_ID"])==211))
0361 print(f" {dm.sum():,} / {len(dm):,} pass quality + PDG cuts")
0362 d = {k: v[dm] for k, v in d.items()}
0363 d_p_pt = d["Lambda0_track_2_pT"].astype(np.float64)
0364 d_kfp = d["Ximinus_mass"].astype(np.float64) * 1000.0
0365
0366
0367 print(f"Loading Xi SV sim: {XI_SV_FILE}")
0368 with uproot.open(XI_SV_FILE) as f:
0369 s = f[TREE_NAME].arrays(xi_phys + XI_CUT_BR, library="np")
0370 sm = (xi_cut_mask(s) &
0371 (np.abs(s["Lambda0_track_1_PDG_ID"])==211) &
0372 (np.abs(s["Lambda0_track_2_PDG_ID"])==2212) &
0373 (np.abs(s["track_3_PDG_ID"])==211))
0374 print(f" {sm.sum():,} / {len(sm):,} pass quality + PDG cuts\n")
0375 s = {k: v[sm] for k, v in s.items()}
0376 s_pi1_pt = s["Lambda0_track_1_pT"].astype(np.float64)
0377 s_pi1_px = s["Lambda0_track_1_px"].astype(np.float64)
0378 s_pi1_py = s["Lambda0_track_1_py"].astype(np.float64)
0379 s_pi1_pz = s["Lambda0_track_1_pz"].astype(np.float64)
0380 s_p_pt = s["Lambda0_track_2_pT"].astype(np.float64)
0381 s_p_px = s["Lambda0_track_2_px"].astype(np.float64)
0382 s_p_py = s["Lambda0_track_2_py"].astype(np.float64)
0383 s_p_pz = s["Lambda0_track_2_pz"].astype(np.float64)
0384 s_pi2_pt = s["track_3_pT"].astype(np.float64)
0385 s_pi2_px = s["track_3_px"].astype(np.float64)
0386 s_pi2_py = s["track_3_py"].astype(np.float64)
0387 s_pi2_pz = s["track_3_pz"].astype(np.float64)
0388
0389 nBins = len(XI_PT_BINS) - 1
0390 results = []
0391
0392 HDR = (f"{'Bin':>3} {'proton pT':>16} "
0393 f"{'σ_data_KFP':>11} {'σ_sv_inv0':>10} "
0394 f"{'σ_sv_pi-smr':>12} {'best%':>7} {'σ_matched':>10}")
0395 print(HDR)
0396 print("─" * len(HDR.rstrip()))
0397
0398 for i in range(nBins):
0399 lo, hi = XI_PT_BINS[i], XI_PT_BINS[i+1]
0400 hi_s = f"{hi:.3f}" if hi<1e8 else "∞ "
0401
0402 dm_b = (d_p_pt>=lo) & (d_p_pt<hi)
0403 sm_b = (s_p_pt>=lo) & (s_p_pt<hi)
0404 n_d, n_s = int(dm_b.sum()), int(sm_b.sum())
0405 avg_pt = float(d_p_pt[dm_b].mean()) if n_d>0 else (lo+hi)/2
0406
0407 _, sig_data_kfp = fit_xi(d_kfp[dm_b])
0408
0409 pi1_px_b = s_pi1_px[sm_b]; pi1_py_b = s_pi1_py[sm_b]; pi1_pz_b = s_pi1_pz[sm_b]
0410 pi1_pt_b = s_pi1_pt[sm_b]
0411 pi2_px_b = s_pi2_px[sm_b]; pi2_py_b = s_pi2_py[sm_b]; pi2_pz_b = s_pi2_pz[sm_b]
0412 pi2_pt_b = s_pi2_pt[sm_b]
0413 p_px_b = s_p_px[sm_b]; p_py_b = s_p_py[sm_b]; p_pz_b = s_p_pz[sm_b]
0414 sv_inv0 = xi_mass_mev(p_px_b,p_py_b,p_pz_b,pi1_px_b,pi1_py_b,pi1_pz_b,pi2_px_b,pi2_py_b,pi2_pz_b)
0415 _, sig_sv_inv0 = fit_xi(sv_inv0)
0416
0417 def _bad(note):
0418 results.append(dict(lo=lo, hi=hi, frac=np.nan, avg_pt=avg_pt,
0419 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0420 sig_pi_only=np.nan, sig_match=np.nan,
0421 note=note, n_d=n_d, n_s=n_s))
0422
0423 if n_d < 20 or np.isnan(sig_data_kfp):
0424 _bad("too few data"); print(f"{i:>3} [{lo:.3f},{hi_s}) too few data"); continue
0425 if n_s < 20 or np.isnan(sig_sv_inv0):
0426 _bad("too few sim"); print(f"{i:>3} [{lo:.3f},{hi_s}) too few sim"); continue
0427
0428 pi1_fracs = pion_frac(pi1_pt_b)
0429 pi2_fracs = pion_frac(pi2_pt_b)
0430 rng = np.random.default_rng(RNG_SEED + i)
0431 z_pi1 = rng.standard_normal(n_s)
0432 z_pi2 = rng.standard_normal(n_s)
0433 z_p = rng.standard_normal(n_s)
0434
0435 def sim_sigma(frac_p):
0436 spx_pi1,spy_pi1,spz_pi1 = apply_smear_pervec(pi1_px_b,pi1_py_b,pi1_pz_b,pi1_fracs,z_pi1)
0437 spx_pi2,spy_pi2,spz_pi2 = apply_smear_pervec(pi2_px_b,pi2_py_b,pi2_pz_b,pi2_fracs,z_pi2)
0438 spx_p, spy_p, spz_p = apply_smear(p_px_b,p_py_b,p_pz_b,frac_p,z_p)
0439 _, sig = fit_xi(xi_mass_mev(spx_p,spy_p,spz_p,spx_pi1,spy_pi1,spz_pi1,spx_pi2,spy_pi2,spz_pi2))
0440 return sig
0441
0442 def residual(frac_p): return sim_sigma(frac_p) - sig_data_kfp
0443
0444
0445 r0 = residual(0.0)
0446 sig_pi_only = sig_data_kfp + r0
0447
0448 if r0 >= 0:
0449 note = "sim≥data at 0%"
0450 results.append(dict(lo=lo, hi=hi, frac=0.0, avg_pt=avg_pt,
0451 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0452 sig_pi_only=sig_pi_only, sig_match=sig_pi_only,
0453 note=note, n_d=n_d, n_s=n_s))
0454 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0455 f"{sig_sv_inv0:>10.3f} {sig_pi_only:>12.3f} {'0.00%':>7} "
0456 f"{sig_pi_only:>10.3f} ({note})")
0457 continue
0458
0459 r_max = residual(frac_max)
0460 if r_max < 0:
0461 sig_at_max = sim_sigma(frac_max)
0462 note = f">={frac_max*100:.0f}%"
0463 results.append(dict(lo=lo, hi=hi, frac=frac_max, avg_pt=avg_pt,
0464 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0465 sig_pi_only=sig_pi_only, sig_match=sig_at_max,
0466 note=note, n_d=n_d, n_s=n_s))
0467 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0468 f"{sig_sv_inv0:>10.3f} {sig_pi_only:>12.3f} {note:>7} {sig_at_max:>10.3f}")
0469 continue
0470
0471 best = brentq(residual, 0.0, frac_max, xtol=5e-5, maxiter=80)
0472 sig_m = sim_sigma(best)
0473 results.append(dict(lo=lo, hi=hi, frac=best, avg_pt=avg_pt,
0474 sig_data_kfp=sig_data_kfp, sig_sv_inv0=sig_sv_inv0,
0475 sig_pi_only=sig_pi_only, sig_match=sig_m,
0476 note="ok", n_d=n_d, n_s=n_s))
0477 print(f"{i:>3} [{lo:.3f},{hi_s}) {sig_data_kfp:>11.3f} "
0478 f"{sig_sv_inv0:>10.3f} {sig_pi_only:>12.3f} {best*100:>6.2f}% {sig_m:>10.3f}")
0479
0480 valid = [r for r in results if not np.isnan(r["frac"])]
0481 edges_cpp = ", ".join(f"{r['lo']:.4f}" for r in valid)
0482 sigmas_cpp = ", ".join(f"{r['frac']:.5f}" for r in valid)
0483 print(f"\n{'─'*68}")
0484 print("Xi proton fracs (C++ for smear_Xi_ptbin.C):\n")
0485 print(f" pt_bin_edges = {{{edges_cpp}}}")
0486 print(f" pt_bin_sigmas = {{{sigmas_cpp}}}")
0487
0488 with open("optimal_sv_xi_fracs.txt","w") as fh:
0489 fh.write("# Auto-generated by find_sv_smear_fracs.py\n")
0490 fh.write("# SV sim inv mass smeared to match data KFP mass.\n")
0491 fh.write("# Pion fracs fixed from K0S SV analysis (see optimal_sv_smear_fracs.txt).\n")
0492 fh.write(f"# pt_bin_edges = {{{edges_cpp}}}\n")
0493 fh.write(f"# pt_bin_sigmas = {{{sigmas_cpp}}}\n\n")
0494 fh.write("bin_lo,bin_hi,avg_pt,best_frac_pct,sig_data_kfp,sig_sv_inv0,"
0495 "sig_pi_only,sig_match,n_data,n_sim,note\n")
0496 for r in results:
0497 hi_w = f"{r['hi']:.4f}" if r["hi"]<1e8 else "inf"
0498 _f = lambda v: f"{v:.4f}" if not np.isnan(v) else "nan"
0499 fp = f"{r['frac']*100:.4f}" if not np.isnan(r["frac"]) else "nan"
0500 fh.write(f"{r['lo']:.4f},{hi_w},{r['avg_pt']:.4f},{fp},{_f(r['sig_data_kfp'])},"
0501 f"{_f(r['sig_sv_inv0'])},{_f(r['sig_pi_only'])},"
0502 f"{_f(r['sig_match'])},{r['n_d']},{r['n_s']},{r['note']}\n")
0503 print(f"CSV saved → optimal_sv_xi_fracs.txt")
0504
0505 return results
0506
0507
0508
0509
0510
0511
0512 def _fmt(lst, fmt=".4f"):
0513 return "{" + ", ".join(format(v if not np.isnan(v) else 0, fmt) for v in lst) + "}"
0514
0515 def write_ks_macro(results):
0516 n = len(results)
0517 pts = [r["avg_pt"] for r in results]
0518 sig_dk = [r["sig_data_kfp"] if not np.isnan(r["sig_data_kfp"]) else 0 for r in results]
0519 sig_sv0 = [r["sig_sv_inv0"] if not np.isnan(r["sig_sv_inv0"]) else 0 for r in results]
0520 sig_m = [r["sig_match"] if not np.isnan(r["sig_match"]) else 0 for r in results]
0521 fracs = [r["frac"]*100 if not np.isnan(r["frac"]) else 0 for r in results]
0522 ymax = max(max(sig_dk), max(sig_m)) * 1.35
0523
0524 macro = f"""// Auto-generated by find_sv_smear_fracs.py
0525 // K0S: SV sim inv mass smeared to match data KFP mass.
0526 // Run with: root -l -b -q plot_sv_ks_result.C
0527
0528 void plot_sv_ks_result() {{
0529 gStyle->SetOptStat(0); gStyle->SetOptTitle(0);
0530
0531 const int N = {n};
0532 double avg_pt[] = {_fmt(pts)};
0533 double sig_dk[] = {_fmt(sig_dk)}; // data KFP mass sigma ← TARGET
0534 double sig_sv0[] = {_fmt(sig_sv0)}; // SV sim inv mass (no smear)
0535 double sig_m[] = {_fmt(sig_m)}; // SV sim inv mass (matched smear)
0536 double fracs[] = {_fmt(fracs)}; // optimal pion smearing %
0537
0538 TCanvas *c1 = new TCanvas("c1","K0S SV smearing",950,620);
0539 c1->SetLeftMargin(0.12); c1->SetBottomMargin(0.13); c1->SetRightMargin(0.05);
0540
0541 auto sty = [](TGraph* g, Color_t col, Style_t mrk, Style_t ls=1){{
0542 g->SetMarkerColor(col); g->SetLineColor(col);
0543 g->SetMarkerStyle(mrk); g->SetMarkerSize(1.6);
0544 g->SetLineWidth(2); g->SetLineStyle(ls);
0545 }};
0546
0547 TGraph *gDK = new TGraph(N,avg_pt,sig_dk);
0548 TGraph *gS0 = new TGraph(N,avg_pt,sig_sv0);
0549 TGraph *gM = new TGraph(N,avg_pt,sig_m);
0550
0551 sty(gDK, kBlack, 20, 1);
0552 sty(gS0, kBlue+1, 24, 2);
0553 sty(gM, kRed+1, 22, 1);
0554
0555 TMultiGraph *mg = new TMultiGraph();
0556 mg->Add(gS0,"PL"); mg->Add(gM,"PL"); mg->Add(gDK,"PL");
0557 mg->Draw("A");
0558 mg->GetXaxis()->SetTitle("p_{{T}} (GeV/#it{{c}})");
0559 mg->GetYaxis()->SetTitle("#sigma_{{K^{{0}}_{{S}}}} (MeV/#it{{c}}^{{2}})");
0560 mg->GetYaxis()->SetRangeUser(0,{ymax:.1f});
0561 mg->GetXaxis()->SetTitleSize(0.05);
0562 mg->GetYaxis()->SetTitleSize(0.05);
0563
0564 TLegend *leg = new TLegend(0.14,0.68,0.72,0.93);
0565 leg->SetBorderSize(0); leg->SetFillStyle(0); leg->SetTextSize(0.032);
0566 leg->AddEntry(gDK,"Data (K_{{S}}^{{0}} mass, KFP) — target","lp");
0567 leg->AddEntry(gS0,"SV sim (inv mass, no smear)","lp");
0568 leg->AddEntry(gM, "SV sim (inv mass, matched smear)","lp");
0569 leg->Draw();
0570
0571 c1->SaveAs("plot_sv_ks_result.pdf");
0572 c1->SaveAs("plot_sv_ks_result.png");
0573 std::cout << "Saved plot_sv_ks_result.pdf/.png\\n";
0574
0575 TCanvas *c2 = new TCanvas("c2","K0S pion smearing frac",900,550);
0576 c2->SetLeftMargin(0.13); c2->SetBottomMargin(0.13);
0577 TGraph *gF = new TGraph(N,avg_pt,fracs);
0578 sty(gF, kGreen+2, 23);
0579 gF->SetTitle(";p_{{T}} (GeV/#it{{c}});Optimal pion smearing (% of p_{{T}})");
0580 gF->GetYaxis()->SetRangeUser(0,15);
0581 gF->GetXaxis()->SetTitleSize(0.05); gF->GetYaxis()->SetTitleSize(0.05);
0582 gF->Draw("APL");
0583 c2->SaveAs("plot_sv_ks_frac.pdf");
0584 c2->SaveAs("plot_sv_ks_frac.png");
0585 std::cout << "Saved plot_sv_ks_frac.pdf/.png\\n";
0586 }}
0587 """
0588 with open("plot_sv_ks_result.C","w") as fh:
0589 fh.write(macro)
0590
0591
0592 def write_xi_macro(results):
0593 n = len(results)
0594 pts = [r["avg_pt"] for r in results]
0595 sig_dk = [r["sig_data_kfp"] if not np.isnan(r["sig_data_kfp"]) else 0 for r in results]
0596 sig_sv0 = [r["sig_sv_inv0"] if not np.isnan(r["sig_sv_inv0"]) else 0 for r in results]
0597 sig_pi = [r["sig_pi_only"] if not np.isnan(r["sig_pi_only"]) else 0 for r in results]
0598 sig_m = [r["sig_match"] if not np.isnan(r["sig_match"]) else 0 for r in results]
0599 fracs = [r["frac"]*100 if not np.isnan(r["frac"]) else 0 for r in results]
0600 ymax = max(max(sig_dk), max(sig_m), max(sig_pi)) * 1.35
0601
0602 macro = f"""// Auto-generated by find_sv_smear_fracs.py
0603 // Xi: SV sim inv mass smeared to match data KFP mass.
0604 // Run with: root -l -b -q plot_sv_xi_result.C
0605
0606 void plot_sv_xi_result() {{
0607 gStyle->SetOptStat(0); gStyle->SetOptTitle(0);
0608
0609 const int N = {n};
0610 double avg_pt[] = {_fmt(pts)};
0611 double sig_dk[] = {_fmt(sig_dk)}; // data KFP mass sigma ← TARGET
0612 double sig_sv0[] = {_fmt(sig_sv0)}; // SV sim inv mass (no smear)
0613 double sig_pi[] = {_fmt(sig_pi)}; // SV sim inv mass (pion smeared, proton 0%)
0614 double sig_m[] = {_fmt(sig_m)}; // SV sim inv mass (pion + proton matched)
0615 double fracs[] = {_fmt(fracs)}; // optimal proton smearing %
0616
0617 TCanvas *c1 = new TCanvas("c1","Xi SV smearing",950,620);
0618 c1->SetLeftMargin(0.12); c1->SetBottomMargin(0.13); c1->SetRightMargin(0.05);
0619
0620 auto sty = [](TGraph* g, Color_t col, Style_t mrk, Style_t ls=1){{
0621 g->SetMarkerColor(col); g->SetLineColor(col);
0622 g->SetMarkerStyle(mrk); g->SetMarkerSize(1.6);
0623 g->SetLineWidth(2); g->SetLineStyle(ls);
0624 }};
0625
0626 TGraph *gDK = new TGraph(N,avg_pt,sig_dk);
0627 TGraph *gS0 = new TGraph(N,avg_pt,sig_sv0);
0628 TGraph *gPI = new TGraph(N,avg_pt,sig_pi);
0629 TGraph *gM = new TGraph(N,avg_pt,sig_m);
0630
0631 sty(gDK, kBlack, 20, 1);
0632 sty(gS0, kBlue+1, 24, 2);
0633 sty(gPI, kOrange+1, 25, 2);
0634 sty(gM, kRed+1, 22, 1);
0635
0636 TMultiGraph *mg = new TMultiGraph();
0637 mg->Add(gS0,"PL"); mg->Add(gPI,"PL"); mg->Add(gM,"PL"); mg->Add(gDK,"PL");
0638 mg->Draw("A");
0639 mg->GetXaxis()->SetTitle("Proton p_{{T}} (GeV/#it{{c}})");
0640 mg->GetYaxis()->SetTitle("#sigma_{{#Xi}} (MeV/#it{{c}}^{{2}})");
0641 mg->GetYaxis()->SetRangeUser(0,{ymax:.1f});
0642 mg->GetXaxis()->SetTitleSize(0.05);
0643 mg->GetYaxis()->SetTitleSize(0.05);
0644
0645 TLegend *leg = new TLegend(0.14,0.60,0.76,0.93);
0646 leg->SetBorderSize(0); leg->SetFillStyle(0); leg->SetTextSize(0.030);
0647 leg->AddEntry(gDK,"Data (#Xi mass, KFP) — target","lp");
0648 leg->AddEntry(gS0,"SV sim (inv mass, no smear)","lp");
0649 leg->AddEntry(gPI,"SV sim (inv mass, #pi smeared, proton 0%)","lp");
0650 leg->AddEntry(gM, "SV sim (inv mass, #pi + proton matched)","lp");
0651 leg->Draw();
0652
0653 c1->SaveAs("plot_sv_xi_result.pdf");
0654 c1->SaveAs("plot_sv_xi_result.png");
0655 std::cout << "Saved plot_sv_xi_result.pdf/.png\\n";
0656
0657 TCanvas *c2 = new TCanvas("c2","Xi proton smearing frac",900,550);
0658 c2->SetLeftMargin(0.13); c2->SetBottomMargin(0.13);
0659 TGraph *gF = new TGraph(N,avg_pt,fracs);
0660 sty(gF, kMagenta+1, 23);
0661 gF->SetTitle(";Proton p_{{T}} (GeV/#it{{c}});Optimal proton smearing (% of p_{{T}})");
0662 gF->GetYaxis()->SetRangeUser(0,15);
0663 gF->GetXaxis()->SetTitleSize(0.05); gF->GetYaxis()->SetTitleSize(0.05);
0664 gF->Draw("APL");
0665 c2->SaveAs("plot_sv_xi_frac.pdf");
0666 c2->SaveAs("plot_sv_xi_frac.png");
0667 std::cout << "Saved plot_sv_xi_frac.pdf/.png\\n";
0668 }}
0669 """
0670 with open("plot_sv_xi_result.C","w") as fh:
0671 fh.write(macro)
0672
0673
0674
0675
0676
0677
0678 def main():
0679 ap = argparse.ArgumentParser(description=__doc__,
0680 formatter_class=argparse.RawDescriptionHelpFormatter)
0681 ap.add_argument("--frac-max", type=float, default=0.15,
0682 help="Upper smearing fraction bound (default 0.15 = 15%%)")
0683 args = ap.parse_args()
0684
0685 ks_results = run_ks(args.frac_max)
0686 xi_results = run_xi(ks_results, args.frac_max)
0687
0688 write_ks_macro(ks_results)
0689 write_xi_macro(xi_results)
0690
0691 print("\nROOT macros written:")
0692 print(" plot_sv_ks_result.C plot_sv_xi_result.C")
0693 print("Run with:")
0694 print(" root -l -b -q plot_sv_ks_result.C")
0695 print(" root -l -b -q plot_sv_xi_result.C")
0696
0697
0698 if __name__ == "__main__":
0699 main()
