File indexing completed on 2026-07-16 08:11:23
0001
0002 """
0003 smear_sv_ks.py
0004 --------------
0005 Reads outputKFParticleKShortRecoSV.root (or any KFParticle K0S SV file),
0006 applies per-pion Gaussian pT smearing using the calibrated fracs from
0007 optimal_sv_smear_fracs.txt, and writes a new ROOT file with updated
0008 track and K0S kinematics.
0009
0010 Smearing model
0011 --------------
0012 δpT = frac(pT) × pT × z, z ~ N(0,1)
0013
0014 frac(pT) is **linearly interpolated** between the calibrated (avg_pt, frac)
0015 control points, with flat extrapolation outside the range. Each pion is
0016 smeared independently using its own pT. eta is preserved: only pT is
0017 scaled, pz follows via pz_new = pT_new × sinh(eta).
0018
0019 Branches updated
0020 ----------------
0021 track_{1,2}_px, _py, _pz smeared momenta
0022 track_{1,2}_pT, _p, _pE recomputed from smeared 3-momentum
0023 track_{1,2}_phi recomputed
0024 track_{1,2}_pseudorapidity recomputed (= eta)
0025 track_{1,2}_rapidity recomputed from smeared E, pz
0026 track_{1,2}_theta recomputed
0027 K_S0_px, _py, _pz sum of smeared daughter momenta
0028 K_S0_pT, _p, _pE recomputed
0029 K_S0_phi, _pseudorapidity, recomputed
0030 _rapidity, _theta
0031 K_S0_mass invariant mass from smeared daughters (GeV)
0032 secondary_vertex_mass_pionPID same as K_S0_mass (pion-PID assumption)
0033
0034 Branches NOT updated (KFP-fitter outputs / vertex geometry)
0035 -----------------------------------------------------------
0036 *_chi2, *_nDoF, *_Covariance KFP fitter quantities
0037 *_massErr, *_pTErr, *_pErr KFP uncertainties
0038 *_IP*, *_DCA, *_DIRA require re-propagating track to vertex
0039 *_decayLength*, *_decayTime* vertex geometry
0040 *_x, *_y, *_z (positions) vertex positions unchanged
0041 true_* branches MC truth unchanged
0042
0043 Note on vector branches
0044 -----------------------
0045 Variable-length branches (std::vector<T>: hit IDs, residuals, IP_allPV,
0046 track history, etc.) are NOT included in the output. None of these are
0047 used by the standard quality cuts or kinematic analysis. If you need them,
0048 use ROOT's hadd or TTree::AddFriend to merge with the original file.
0049
0050 Usage
0051 -----
0052 python3 smear_sv_ks.py # defaults
0053 python3 smear_sv_ks.py --input foo.root --output foo_smeared.root
0054 python3 smear_sv_ks.py --seed 123
0055 """
0056
0057 import argparse, sys
0058 import numpy as np
0059 import uproot
0060
0061
0062 PION_MASS = 0.13957018
0063 TREE_NAME = "DecayTree"
0064 RNG_SEED = 42
0065
0066
0067
0068
0069 def load_fracs(csv_path):
0070 """
0071 Parse optimal_sv_smear_fracs.txt.
0072 Returns (avg_pt_array, frac_array) suitable for np.interp.
0073 Skips rows with nan frac.
0074 """
0075 pts, fracs = [], []
0076 avg_col = frc_col = None
0077 with open(csv_path) as fh:
0078 for line in fh:
0079 line = line.strip()
0080 if not line or line.startswith("#"):
0081 continue
0082 if line.startswith("bin_lo"):
0083 cols = line.split(",")
0084 avg_col = cols.index("avg_pt")
0085 frc_col = cols.index("best_frac_pct")
0086 continue
0087 parts = line.split(",")
0088 frc = parts[frc_col]
0089 if frc == "nan":
0090 continue
0091 pts.append(float(parts[avg_col]))
0092 fracs.append(float(frc) / 100.0)
0093 return np.array(pts, dtype=np.float64), np.array(fracs, dtype=np.float64)
0094
0095
0096 def interp_frac(pt_arr, pts, fracs):
0097 """
0098 Linearly interpolate smearing frac at each pT value.
0099 Flat extrapolation outside [pts[0], pts[-1]].
0100 """
0101 return np.interp(pt_arr.astype(np.float64), pts, fracs)
0102
0103
0104 def smear_track(px, py, pz, frac_arr, z):
0105 """
0106 Smear pT of a track array; preserve eta.
0107 pT_new = pT * (1 + frac * z)
0108 Returns (px_new, py_new, pz_new) as float64.
0109 """
0110 px, py, pz = (np.asarray(a, dtype=np.float64) for a in (px, py, pz))
0111 pT = np.sqrt(px**2 + py**2)
0112 phi = np.arctan2(py, px)
0113 ptot = np.sqrt(pT**2 + pz**2)
0114 eta = np.arctanh(np.clip(pz / np.where(ptot > 0, ptot, 1e-12),
0115 -1 + 1e-9, 1 - 1e-9))
0116 pT_s = np.maximum(pT * (1.0 + frac_arr * z), 0.0)
0117 return pT_s * np.cos(phi), pT_s * np.sin(phi), pT_s * np.sinh(eta)
0118
0119
0120 def derived(px, py, pz, mass_gev):
0121 """
0122 Compute (pT, p, pE, phi, eta, rapidity, theta) from 3-momentum + mass.
0123 All inputs/outputs float64.
0124 """
0125 px, py, pz = (np.asarray(a, dtype=np.float64) for a in (px, py, pz))
0126 pT = np.sqrt(px**2 + py**2)
0127 p = np.sqrt(px**2 + py**2 + pz**2)
0128 pE = np.sqrt(p**2 + mass_gev**2)
0129 phi = np.arctan2(py, px)
0130 p_s = np.where(p > 0, p, 1e-12)
0131 eta = np.arctanh(np.clip(pz / p_s, -1 + 1e-9, 1 - 1e-9))
0132
0133 denom = np.abs(pE - np.abs(pz))
0134 rap = np.where(denom > 1e-12,
0135 0.5 * np.log((pE + pz) / np.where(denom > 1e-12, np.abs(pE - pz), 1e-12)),
0136 np.sign(pz) * 1e6)
0137 theta = np.arctan2(pT, pz)
0138 return pT, p, pE, phi, eta, rap, theta
0139
0140
0141 def inv_mass(px1, py1, pz1, px2, py2, pz2, m1=PION_MASS, m2=PION_MASS):
0142 """K0S invariant mass in GeV from two daughters."""
0143 E1 = np.sqrt(px1**2 + py1**2 + pz1**2 + m1**2)
0144 E2 = np.sqrt(px2**2 + py2**2 + pz2**2 + m2**2)
0145 m2v = (E1+E2)**2 - (px1+px2)**2 - (py1+py2)**2 - (pz1+pz2)**2
0146 return np.sqrt(np.maximum(m2v, 0.0))
0147
0148
0149
0150
0151 def main():
0152 parser = argparse.ArgumentParser(
0153 description="Apply SV K0S pion smearing and write updated ROOT file.")
0154 parser.add_argument("--input", default="outputKFParticleKShortRecoSV_filtered.root")
0155 parser.add_argument("--fracs", default="mass_smear_fracs_ks.txt")
0156 parser.add_argument("--output", default="outputKFParticleKShortRecoSV_smeared.root")
0157 parser.add_argument("--seed", type=int, default=RNG_SEED)
0158 args = parser.parse_args()
0159
0160
0161 pts, fracs = load_fracs(args.fracs)
0162 print(f"Loaded {len(fracs)} calibration points from {args.fracs}:")
0163 for p, f in zip(pts, fracs):
0164 print(f" avg_pT = {p:.4f} GeV → frac = {f*100:.4f}%")
0165 print(f" (linear interp; flat extrapolation outside [{pts[0]:.3f}, {pts[-1]:.3f}] GeV)\n")
0166
0167
0168
0169
0170 print(f"Reading branch list from {args.input} ...")
0171 with uproot.open(args.input) as fin:
0172 tree = fin[TREE_NAME]
0173 all_branches = tree.keys()
0174 skip_vector = [b for b in all_branches
0175 if "vector" in tree[b].typename.lower()
0176 or "std::" in tree[b].typename.lower()]
0177 load_branches = [b for b in all_branches if b not in skip_vector]
0178
0179 print(f" Loading {len(load_branches)} scalar/fixed-array branches "
0180 f"(skipping {len(skip_vector)} vector branches)")
0181 if skip_vector:
0182 print(f" Skipped: {', '.join(skip_vector[:6])}"
0183 + (" ..." if len(skip_vector) > 6 else ""))
0184
0185
0186 print(f"\nReading events ...")
0187 with uproot.open(args.input) as fin:
0188 d = fin[TREE_NAME].arrays(load_branches, library="np")
0189 n = len(d[load_branches[0]])
0190 print(f" {n:,} events read\n")
0191
0192
0193 rng = np.random.default_rng(args.seed)
0194 z1 = rng.standard_normal(n)
0195 z2 = rng.standard_normal(n)
0196
0197 pT1 = d["track_1_pT"].astype(np.float64)
0198 pT2 = d["track_2_pT"].astype(np.float64)
0199
0200 f1 = interp_frac(pT1, pts, fracs)
0201 f2 = interp_frac(pT2, pts, fracs)
0202
0203 print("Smearing tracks ...")
0204 px1s, py1s, pz1s = smear_track(d["track_1_px"], d["track_1_py"], d["track_1_pz"], f1, z1)
0205 px2s, py2s, pz2s = smear_track(d["track_2_px"], d["track_2_py"], d["track_2_pz"], f2, z2)
0206
0207
0208 pT1s, p1s, pE1s, phi1s, eta1s, rap1s, th1s = derived(px1s, py1s, pz1s, PION_MASS)
0209 pT2s, p2s, pE2s, phi2s, eta2s, rap2s, th2s = derived(px2s, py2s, pz2s, PION_MASS)
0210
0211 ks_mass = inv_mass(px1s, py1s, pz1s, px2s, py2s, pz2s)
0212 ks_px, ks_py, ks_pz = px1s+px2s, py1s+py2s, pz1s+pz2s
0213 ks_pTs, ks_ps, ks_pEs, ks_phis, ks_etas, ks_raps, ks_ths = derived(
0214 ks_px, ks_py, ks_pz, ks_mass)
0215
0216
0217 win = (d["K_S0_mass"]*1000 > 480) & (d["K_S0_mass"]*1000 < 516)
0218 before = d["K_S0_mass"][win].astype(np.float64)*1000
0219 after = ks_mass[win]*1000
0220 print(f"Signal-window mass (MeV):")
0221 print(f" Before smearing: mean={before.mean():.2f} σ_rms={before.std():.2f}")
0222 print(f" After smearing: mean={after.mean():.2f} σ_rms={after.std():.2f}\n")
0223
0224
0225 replacements = {
0226 "track_1_px": px1s,
0227 "track_1_py": py1s,
0228 "track_1_pz": pz1s,
0229 "track_1_pT": pT1s,
0230 "track_1_p": p1s,
0231 "track_1_pE": pE1s,
0232 "track_1_phi": phi1s,
0233 "track_1_pseudorapidity": eta1s,
0234 "track_1_rapidity": rap1s,
0235 "track_1_theta": th1s,
0236 "track_2_px": px2s,
0237 "track_2_py": py2s,
0238 "track_2_pz": pz2s,
0239 "track_2_pT": pT2s,
0240 "track_2_p": p2s,
0241 "track_2_pE": pE2s,
0242 "track_2_phi": phi2s,
0243 "track_2_pseudorapidity": eta2s,
0244 "track_2_rapidity": rap2s,
0245 "track_2_theta": th2s,
0246 "K_S0_px": ks_px,
0247 "K_S0_py": ks_py,
0248 "K_S0_pz": ks_pz,
0249 "K_S0_pT": ks_pTs,
0250 "K_S0_p": ks_ps,
0251 "K_S0_pE": ks_pEs,
0252 "K_S0_phi": ks_phis,
0253 "K_S0_pseudorapidity": ks_etas,
0254 "K_S0_rapidity": ks_raps,
0255 "K_S0_theta": ks_ths,
0256 "K_S0_mass": ks_mass,
0257 "secondary_vertex_mass_pionPID": ks_mass,
0258 }
0259
0260
0261 out = {}
0262 for b in load_branches:
0263 if b in replacements:
0264 out[b] = replacements[b].astype(d[b].dtype)
0265 else:
0266 out[b] = d[b]
0267
0268 not_found = [b for b in replacements if b not in load_branches]
0269 if not_found:
0270 print(f"WARNING: replacement branches not in file: {not_found}")
0271
0272
0273 print(f"Writing {args.output} ...")
0274 with uproot.recreate(args.output) as fout:
0275 fout[TREE_NAME] = out
0276
0277 import os
0278 sz = os.path.getsize(args.output) / 1e6
0279 print(f"Done — {n:,} events, {len(out)} branches → {args.output} ({sz:.1f} MB)")
0280 if skip_vector:
0281 print(f"\nNote: {len(skip_vector)} vector branches omitted (hit IDs, "
0282 f"residuals, IP_allPV, track history).")
0283 print(f"To recover them use ROOT:\n"
0284 f" new_tree->AddFriend(original_tree) # or hadd after friend setup")
0285
0286
0287 if __name__ == "__main__":
0288 main()
