systematics/trackMomScale/smear_sv_xi.py

0001 #!/usr/bin/env python3
0002 """
0003 smear_sv_xi.py
0004 ------------------
0005 Reads outputKFParticleXiSV_filtered.root, applies per-track Gaussian pT smearing,
0006 and writes outputKFParticleXiSV_smeared.root with updated kinematics.
0007
0008 Two independent smearing calibrations are applied:
0009   Track 1 (pion, ±211)  — fracs from optimal_sv_smear_fracs.txt  (K0S calibration)
0010   Track 2 (proton, 2212)— fracs from optimal_sv_xi_fracs.txt  (Xi calibration)
0011
0012 Both frac tables use linear interpolation between calibrated avg_pt points,
0013 with flat extrapolation outside the range.  eta is preserved for each track.
0014
0015 Branches updated
0016 ----------------
0017   track_{1,2}_px, _py, _pz, _pT, _p, _pE
0018   track_{1,2}_phi, _pseudorapidity, _rapidity, _theta
0019   Lambda0_px, _py, _pz, _pT, _p, _pE
0020   Lambda0_phi, _pseudorapidity, _rapidity, _theta
0021   Lambda0_mass                  (pion+proton inv mass from smeared daughters)
0022   secondary_vertex_mass_pionPID (pion+pion inv mass from smeared daughters)
0023
0024 Branches NOT updated
0025 --------------------
0026   *_chi2, *_nDoF, *_Covariance, *Err  KFP fitter outputs
0027   *_IP*, *_DCA, *_DIRA                require track re-propagation
0028   *_decayLength*, *_decayTime*        vertex geometry
0029   *_x, *_y, *_z                       vertex positions
0030   true_* branches                     MC truth
0031
0032 Usage
0033 -----
0034   python3 smear_sv_lambda.py
0035   python3 smear_sv_lambda.py --input foo.root --output foo_smeared.root
0036 """
0037
0038 import argparse
0039 import numpy as np
0040 import uproot
0041
0042 PION_MASS   = 0.13957018   # GeV
0043 PROTON_MASS = 0.93827209   # GeV
0044 TREE_NAME   = "DecayTree"
0045 RNG_SEED    = 42
0046
0047
0048 # ── helpers (shared with smear_sv_ks.py) ──────────────────────────────────
0049
0050 def load_fracs(csv_path):
0051     """Parse a fracs CSV, return (avg_pt, frac) arrays for np.interp."""
0052     pts, fracs = [], []
0053     avg_col = frc_col = None
0054     with open(csv_path) as fh:
0055         for line in fh:
0056             line = line.strip()
0057             if not line or line.startswith("#"):
0058                 continue
0059             if line.startswith("bin_lo"):
0060                 cols = line.split(",")
0061                 avg_col = cols.index("avg_pt")
0062                 frc_col = cols.index("best_frac_pct")
0063                 continue
0064             parts = line.split(",")
0065             frc = parts[frc_col]
0066             if frc == "nan":
0067                 continue
0068             pts.append(float(parts[avg_col]))
0069             fracs.append(float(frc) / 100.0)
0070     return np.array(pts, dtype=np.float64), np.array(fracs, dtype=np.float64)
0071
0072
0073 def interp_frac(pt_arr, pts, fracs):
0074     """Linear interpolation with flat extrapolation."""
0075     return np.interp(pt_arr.astype(np.float64), pts, fracs)
0076
0077
0078 def smear_track(px, py, pz, frac_arr, z):
0079     """Smear pT = pT*(1 + frac*z), preserve eta. Returns float64."""
0080     px, py, pz = (np.asarray(a, dtype=np.float64) for a in (px, py, pz))
0081     pT   = np.sqrt(px**2 + py**2)
0082     phi  = np.arctan2(py, px)
0083     ptot = np.sqrt(pT**2 + pz**2)
0084     eta  = np.arctanh(np.clip(pz / np.where(ptot > 0, ptot, 1e-12),
0085                                -1+1e-9, 1-1e-9))
0086     pT_s = np.maximum(pT * (1.0 + frac_arr * z), 0.0)
0087     return pT_s*np.cos(phi), pT_s*np.sin(phi), pT_s*np.sinh(eta)
0088
0089
0090 def derived(px, py, pz, mass_gev):
0091     """(pT, p, pE, phi, eta, rapidity, theta) from 3-momentum + mass."""
0092     px, py, pz = (np.asarray(a, dtype=np.float64) for a in (px, py, pz))
0093     pT    = np.sqrt(px**2 + py**2)
0094     p     = np.sqrt(px**2 + py**2 + pz**2)
0095     pE    = np.sqrt(p**2 + mass_gev**2)
0096     phi   = np.arctan2(py, px)
0097     p_s   = np.where(p > 0, p, 1e-12)
0098     eta   = np.arctanh(np.clip(pz / p_s, -1+1e-9, 1-1e-9))
0099     denom = np.abs(pE - np.abs(pz))
0100     rap   = np.where(denom > 1e-12,
0101                      0.5*np.log((pE+pz) / np.where(denom>1e-12, np.abs(pE-pz), 1e-12)),
0102                      np.sign(pz)*1e6)
0103     theta = np.arctan2(pT, pz)
0104     return pT, p, pE, phi, eta, rap, theta
0105
0106
0107 def inv_mass(px1, py1, pz1, m1, px2, py2, pz2, m2, px3, py3, pz3, m3):
0108     """Invariant mass in GeV."""
0109     E1 = np.sqrt(px1**2+py1**2+pz1**2 + m1**2)
0110     E2 = np.sqrt(px2**2+py2**2+pz2**2 + m2**2)
0111     E3 = np.sqrt(px3**2+py3**2+pz3**2 + m3**2)
0112     m2v = (E1+E2+E3)**2 - (px1+px2+px3)**2 - (py1+py2+py3)**2 - (pz1+pz2+pz3)**2
0113     return np.sqrt(np.maximum(m2v, 0.0))
0114
0115
0116 # ── main ───────────────────────────────────────────────────────────────────
0117
0118 def main():
0119     parser = argparse.ArgumentParser()
0120     parser.add_argument("--input",       default="outputKFParticleXiminusSV_filtered.root")
0121     parser.add_argument("--pion_fracs",  default="mass_smear_fracs_ks.txt",
0122                         help="K0S pion smearing fracs")
0123     parser.add_argument("--proton_fracs",default="mass_smear_fracs_xi_proton.txt",
0124                         help="Xi proton smearing fracs")
0125     parser.add_argument("--output",      default="outputKFParticleXiSV_smeared.root")
0126     parser.add_argument("--seed",        type=int, default=RNG_SEED)
0127     args = parser.parse_args()
0128
0129     # ── calibration ────────────────────────────────────────────────────────
0130     pi_pts,  pi_fracs  = load_fracs(args.pion_fracs)
0131     pro_pts, pro_fracs = load_fracs(args.proton_fracs)
0132
0133     print(f"Pion fracs ({args.pion_fracs}):")
0134     for p,f in zip(pi_pts, pi_fracs): print(f"  {p:.4f} GeV → {f*100:.4f}%")
0135     print(f"\nProton fracs ({args.proton_fracs}):")
0136     for p,f in zip(pro_pts, pro_fracs): print(f"  {p:.4f} GeV → {f*100:.4f}%")
0137     print()
0138
0139     # ── branch list ────────────────────────────────────────────────────────
0140     with uproot.open(args.input) as fin:
0141         tree = fin[TREE_NAME]
0142         skip = [b for b in tree.keys()
0143                 if "vector" in tree[b].typename.lower()
0144                 or "std::" in tree[b].typename.lower()]
0145         load_br = [b for b in tree.keys() if b not in skip]
0146
0147     print(f"Reading {args.input} ...")
0148     print(f"  Loading {len(load_br)} branches, skipping {len(skip)} vector branches")
0149
0150     with uproot.open(args.input) as fin:
0151         d = fin[TREE_NAME].arrays(load_br, library="np")
0152     n = len(d[load_br[0]])
0153     print(f"  {n:,} events read\n")
0154
0155     # ── smear ──────────────────────────────────────────────────────────────
0156     rng = np.random.default_rng(args.seed)
0157     z1  = rng.standard_normal(n)   # lam pion
0158     z2  = rng.standard_normal(n)   # proton
0159     z3  = rng.standard_normal(n)   # bach pion
0160
0161     pT1 = d["Lambda0_track_1_pT"].astype(np.float64)   # lam pion pT
0162     pT2 = d["Lambda0_track_2_pT"].astype(np.float64)   # proton pT
0163     pT3 = d["track_3_pT"].astype(np.float64)   # bach pion pT
0164
0165     f_pi1  = interp_frac(pT1, pi_pts,  pi_fracs)
0166     f_pro = interp_frac(pT2, pro_pts, pro_fracs)
0167     f_pi3  = interp_frac(pT3, pi_pts,  pi_fracs)
0168
0169     print("Smearing tracks ...")
0170     px1s, py1s, pz1s = smear_track(d["Lambda0_track_1_px"], d["Lambda0_track_1_py"], d["Lambda0_track_1_pz"], f_pi1, z1)
0171     px2s, py2s, pz2s = smear_track(d["Lambda0_track_2_px"], d["Lambda0_track_2_py"], d["Lambda0_track_2_pz"], f_pro, z2)
0172     px3s, py3s, pz3s = smear_track(d["track_3_px"], d["track_3_py"], d["track_3_pz"], f_pi3, z3)
0173
0174     # ── derived kinematics ─────────────────────────────────────────────────
0175     pT1s, p1s, pE1s, phi1s, eta1s, rap1s, th1s = derived(px1s, py1s, pz1s, PION_MASS)
0176     pT2s, p2s, pE2s, phi2s, eta2s, rap2s, th2s = derived(px2s, py2s, pz2s, PROTON_MASS)
0177     pT3s, p3s, pE3s, phi3s, eta3s, rap3s, th3s = derived(px3s, py3s, pz3s, PION_MASS)
0178
0179     xi_mass  = inv_mass(px1s,py1s,pz1s, PION_MASS,  px2s,py2s,pz2s, PROTON_MASS, px3s,py3s,pz3s, PION_MASS)
0180     pion_mass = inv_mass(px1s,py1s,pz1s, PION_MASS,  px2s,py2s,pz2s, PION_MASS, px3s,py3s,pz3s, PION_MASS)
0181
0182     xi_px = px1s+px2s+px3s;  xi_py = py1s+py2s+py3s;  xi_pz = pz1s+pz2s+pz3s
0183     xi_pTs, xi_ps, xi_pEs, xi_phis, xi_etas, xi_raps, xi_ths = derived(
0184         xi_px, xi_py, xi_pz, xi_mass)
0185
0186     # sanity
0187     win = (d["Ximinus_mass"]*1000 > 1300) & (d["Ximinus_mass"]*1000 < 1340)
0188     before = d["Ximinus_mass"][win].astype(np.float64)*1000
0189     after  = xi_mass[win]*1000
0190     print(f"Signal-window mass (MeV) [1300,1340]:")
0191     print(f"  Before smearing: mean={before.mean():.2f}  σ_rms={before.std():.2f}")
0192     print(f"  After  smearing: mean={after.mean():.2f}   σ_rms={after.std():.2f}\n")
0193
0194     # ── replacement map ────────────────────────────────────────────────────
0195     replacements = {
0196         "Lambda0_track_1_px":              px1s,
0197         "Lambda0_track_1_py":              py1s,
0198         "Lambda0_track_1_pz":              pz1s,
0199         "Lambda0_track_1_pT":              pT1s,
0200         "Lambda0_track_1_p":               p1s,
0201         "Lambda0_track_1_pE":              pE1s,
0202         "Lambda0_track_1_phi":             phi1s,
0203         "Lambda0_track_1_pseudorapidity":  eta1s,
0204         "Lambda0_track_1_rapidity":        rap1s,
0205         "Lambda0_track_1_theta":           th1s,
0206         "Lambda0_track_2_px":              px2s,
0207         "Lambda0_track_2_py":              py2s,
0208         "Lambda0_track_2_pz":              pz2s,
0209         "Lambda0_track_2_pT":              pT2s,
0210         "Lambda0_track_2_p":               p2s,
0211         "Lambda0_track_2_pE":              pE2s,
0212         "Lambda0_track_2_phi":             phi2s,
0213         "Lambda0_track_2_pseudorapidity":  eta2s,
0214         "Lambda0_track_2_rapidity":        rap2s,
0215         "Lambda0_track_2_theta":           th2s,
0216         "track_3_px":              px3s,
0217         "track_3_py":              py3s,
0218         "track_3_pz":              pz3s,
0219         "track_3_pT":              pT3s,
0220         "track_3_p":               p3s,
0221         "track_3_pE":              pE3s,
0222         "track_3_phi":             phi3s,
0223         "track_3_pseudorapidity":  eta3s,
0224         "track_3_rapidity":        rap3s,
0225         "track_3_theta":           th3s,
0226         "Ximinus_px":              xi_px,
0227         "Ximinus_py":              xi_py,
0228         "Ximinus_pz":              xi_pz,
0229         "Ximinus_pT":              xi_pTs,
0230         "Ximinus_p":               xi_ps,
0231         "Ximinus_pE":              xi_pEs,
0232         "Ximinus_phi":             xi_phis,
0233         "Ximinus_pseudorapidity":  xi_etas,
0234         "Ximinus_rapidity":        xi_raps,
0235         "Ximinus_theta":           xi_ths,
0236         "Ximinus_mass":            xi_mass,
0237         "secondary_vertex_mass_pionPID": pion_mass,
0238     }
0239
0240     # ── assemble & write ───────────────────────────────────────────────────
0241     out = {}
0242     for b in load_br:
0243         out[b] = replacements[b].astype(d[b].dtype) if b in replacements else d[b]
0244
0245     not_found = [b for b in replacements if b not in load_br]
0246     if not_found:
0247         print(f"WARNING: branches not found in file: {not_found}")
0248
0249     print(f"Writing {args.output} ...")
0250     with uproot.recreate(args.output) as fout:
0251         fout[TREE_NAME] = out
0252
0253     import os
0254     print(f"Done — {n:,} events, {len(out)} branches → {args.output} "
0255           f"({os.path.getsize(args.output)/1e6:.1f} MB)")
0256     if skip:
0257         print(f"\nNote: {len(skip)} vector branches omitted (hit IDs, residuals, etc.)")
0258
0259
0260 if __name__ == "__main__":
0261     main()