Update TrkPID to use track-based observables

[michaelmackenzie]

This update includes:

• Moving to flat electron and muon samples for training inputs
• Replacing E - P with E/P
• Removing R(cluster) and track p dot cluster x observables
• Adding track p(chi^2) and tracker hit dt/dz over the expected slope from the track fit

To add the tracker hit dt/dz, I'm adding a branch to the EventNtuple tree for each track, as it's not yet in EventNtuple. I think this observable should ideally be added to the KalSeed data product and then naturally added to the TrkInfo struct in EventNtuple.

As the ONNX features aren't yet available in pyenv rootana the model has not been exported to ONNX or added to Offline processing yet.

Input features and output predictions:

[oksuzian]

PR Review Summary — "Update TrkPID to use track-based observables" (#9)

Reviewed at head 3935ad52 (2026-08-04). First review of this PR; no prior reviews or comments to carry forward. Scope: retrain TrkPID on MDC2025 flat e−/μ− samples with track-based observables (E/P, p(χ²), hit-time-vs-track-time slope), plus a new calo-free TrkOnlyPID package, ntuple-skimming tools that add the fitted dt/dz and dt/dt branches, and a slope-method comparison script. Note: MLTrain is outside the Offline C++ coding-standards scope; findings below are about workflow correctness, not style conformance.

Decision

• 🔴 Request changes — two S1s, each a one/two-line fix: the new TrkOnlyPID trainer still uses the mixed_float16 policy that commit e202463 identified as incompatible with TMVA SOFIE and fixed in TrkPID only, and the version-tagged output renaming left plot_history reading a filename that is no longer written, so the default TrkPID run now ends in FileNotFoundError.

Findings

1. 🟠 [S1] TrkOnlyPID keeps the mixed_float16 policy your own SOFIE fix removed from TrkPID

   • Evidence: TrkOnlyPID/TrkOnlyPIDTrain.py:116 sets set_global_policy('mixed_float16'). Commit e202463 ("Update to float32 for TMVA SOFIE") changed exactly this line to 'float32' — but only in TrkPID/TrackPIDTrain.py:126. TrkOnlyPID was added before that fix (1665f05) and never picked it up.
   • Impact: the trained/saved TrkOnlyPID model reproduces the configuration you already found breaks the ONNX→SOFIE conversion, so it will fail at export time once the ONNX tooling is available (the PR body notes export is deferred), likely after the training work is considered done.
   • Suggested fix: 'mixed_float16' → 'float32' in TrkOnlyPIDTrain.py, matching e202463.

2. 🟠 [S1] Default TrkPID run crashes at the end: plot_history reads a filename the versioned save no longer writes

   • Evidence: train_model now saves PID_model_v{version}.keras and train_history_v{version}.json (TrkPID/TrackPIDTrain.py:149-151, renamed in e7dc3a8), but the plotting block still calls plot_history("train_history.json", ...) (TrkPID/TrackPIDTrain.py:439). On main, save and read were both unversioned, so this is a regression introduced by the rename. *.json is gitignored, so in a fresh area the unversioned file never exists.
   • Impact: a full default invocation (import → train → export → plots) dies with FileNotFoundError after the training time has been spent; plot_history and plot_model output is lost.
   • Suggested fix: plot_history(f"train_history_v{version}.json", ...).

3. 🟡 [S2] TrkOnlyPID save/load names are asymmetric, so --skip-train can never load what a training run wrote

   • Evidence: 3935ad5 version-tagged the load path (TrkOnlyPID_model_v{version}.keras, train_history_v{version}.json, TrkOnlyPIDTrain.py:374-375) and the ONNX output, but the save path still writes unversioned TrkOnlyPID_model.keras / train_history.json (TrkOnlyPIDTrain.py:139-141); plot_history at line 408 also reads the unversioned name.
   • Impact: train-then---skip-train round trip fails unless files are renamed by hand; same class of drift as finding 2, mirrored.
   • Suggested fix: version-tag the save path (and the plot_history argument) to match the load path, as TrkPID does.

4. 🟡 [S2] The new background-efficiency bisection in make_results has no termination guarantee

   • Evidence: TrkPID/TrackPIDTrain.py:169-174 and TrkOnlyPID/TrkOnlyPIDTrain.py:159-164 (new in this PR) loop while abs(eff - bkg_eff)/bkg_eff > tolerance with no iteration cap. eff is a step function of the threshold (quantized in units of 1/nbkg, with plateaus wherever predictions tie).
   • Impact: potential risk, not observed — if no achievable efficiency lands within 1% of the target (small test set, or tied/saturated sigmoid outputs, e.g. a degenerate training), the interval collapses onto a plateau and the script hangs forever.
   • Suggested fix: replace the loop with threshold = np.quantile(dataset.loc[dataset['label']==0, 'prediction'], 1 - bkg_eff) — exact, one line, and removes the failure mode; otherwise add an iteration cap.

5. 🟡 [S2] perform_tz_fit has no guard on degenerate hit lists

   • Evidence: both make_inputs.py files fit every track unconditionally; with ≤2 hits curve_fit raises, and with exactly 3 points minus 2 parameters the code is fine but dof = len(z) - 2 reaches 0 at 2 hits → ZeroDivisionError; curve_fit can also raise RuntimeError on non-convergence.
   • Impact: potential risk — one pathological track aborts the whole multi-file skim, losing the run. Tracks in these ntuples normally have ≫3 hits, so this is defensive, but the skim is the expensive step.
   • Suggested fix: guard len(hits) < 3 (and wrap the fit in try/except), pushing a sentinel (e.g. NaN) for that track so downstream selection can drop it.

6. 🟡 [S2] Duplication across TrkPID/TrkOnlyPID is already biting (simplification lens — does not gate approval)

   • Evidence: the two make_inputs.py differ by one blank line and one comment; the two pairs of 150-line .files lists are byte-identical; TrkOnlyPIDTrain.py is ~85% identical to TrackPIDTrain.py (import_evtntuple, apply_cut, make_results, all plotting). Finding 1 is the concrete cost: a fix applied to one copy silently missed the other.
   • Suggested fix: not asking for a refactor in this PR — but consider one shared make_inputs.py (the branch additions are identical) and one copy of the dataset file lists, with the trainers importing shared helpers as a follow-up.

7. 🟡 [S2] TrkOnlyPID/README.md documents a feature the model no longer uses

   • Evidence: README lists "tracker hit dt/dz slope divided by the expected slope from the track fit assuming an electron mass" as a v0 input, but v0 features are ['nActiveFrac', 'nNullFrac', 'fitcon', 'dtdt_slope'] (TrkOnlyPIDTrain.py:346) — the dt/dz ratio was replaced by the dt/dt slope in 3935ad5, and the TrkPID README was updated while this one wasn't. The README also names the output "TrkOnlyPID.onnx" vs the actual TrkOnlyPID_v{version}.onnx. (The PR description likewise still says "dt/dz over the expected slope".)
   • Suggested fix: update the v0 feature list and output name; one line in the PR body noting the dt/dz→dt/dt switch would help future archaeology.

8. ⚪ [S3] ROOT.TMath.GausI does not exist

   • Evidence: TrkPID/compare_hit_slopes.py:36 calls TMath.GausI in the equal-sigma branch of analytic_gaussian_overlap. TMath has no GausI in any ROOT version (checked v5.34 and v6.16 headers on cvmfs and current master docs); the normal CDF is TMath::Freq.
   • Impact: nearly dead code — the branch runs only when the two fitted sigmas agree to math.isclose default 1e-9 — but it would AttributeError if it ever fired.
   • Suggested fix: return 2 * ROOT.TMath.Freq(z).

9. ⚪ [S3] Housekeeping batch (no action required for approval)

   • Inherited "cosmic muon" wording is now wrong: apply_cut docstrings and the confusion-matrix/rate printouts in both trainers label the background "cosmic muons", but the samples are flat μ− from target stops.
   • Dead/unused: time_mod (mod 1695) is computed and never read; branches_mc imports trkmc/trksegsmc but only trkmcsim is used (extra I/O); nMatActiveFrac is derived and unused in both feature sets; pconv → pcov.
   • Typos: "Unknown training verion" (both trainers), "unnessary" (both READMEs).
   • Pre-existing, noting only since the lines are nearby: plot_ROC prints "Accuracy at this threshold" but the value is the TPR; the TrkPID README link [TrackPIDTrain.py](TrkPIDTrain.py) targets a nonexistent filename.

Questions (non-gating)

• The t–z / t–t fits in make_inputs.py use every trkhits entry unconditionally; TrkStrawHitInfo.state encodes activity. Is including fit-deactivated hits intended (they may carry discrimination for the wrong-mass hypothesis) or should the fit filter on state?
• tf.keras.layers.Input(..., batch_size=32) bakes a fixed batch of 32 into the ONNX signature (pre-existing on main, copied into TrkOnlyPID). Does the SOFIE-generated inference in Offline evaluate with batch 32, or does per-track batch-1 inference need batch_size=None/1 at export? Worth settling before the first real export since that step is being redone anyway.

Verified 🟢 (checked, no action needed)

• 🟢 make_inputs.py field usage verified against EventNtuple/inc/TrkStrawHitInfo.hh: earlyend is an int (valid array index), etime is TDCTimes, and tottdrift, ptoca, udt, poca all exist with the semantics used; the dt/dt construction (fit of hit time vs track TOCA; slope ≈ 1 under the correct mass hypothesis) is coherent.
• 🟢 The SetBranchStatus sequence is correct: heavy branches dropped, CloneTree(0) excludes them from the output, trkhits re-enabled after cloning for reading only — the skimmed tree carries the six new fit branches but not trkhits.
• 🟢 TrkPID keeps the trkcalohit.active requirement; TrkOnlyPID correctly drops it (tracker-only).
• 🟢 MC-truth selection (trkmcsim[...,0].pdg == 11/13, highest-rank match) matches the FlateMinus/FlatMuMinus samples; the electron-hypothesis dtdz_exp (m = 0.511 MeV) matches the documented intent.
• 🟢 TrkPID README v1 feature list matches the implemented version == 1 features.
• 🟢 .gitignore additions appropriately keep models, CSVs, ROOT files, and figures out of the repo.

Validation check

• Build/tests run: none — MLTrain has no CI, and this review did not execute the scripts (no rootana env in session). Verification is static: three-dot diff against merge-base 2e5b8c3, per-commit archaeology, and field-level checks against EventNtuple headers and ROOT (cvmfs + docs).
• Config contract check: n/a (no FHiCL).
• Cross-repo consistency: the PR body's plan to move the dt/dz-type observable into KalSeed/TrkInfo (EventNtuple) is the right long-term home; nothing in Offline/EventNtuple is required for this PR to merge.

Residual risk

• The exported-model path (ONNX → SOFIE → Offline) is untested by construction (tooling unavailable); findings 1 and the batch-size question both live there and will surface at export time.

Author follow-ups

1. 'mixed_float16' → 'float32' in TrkOnlyPIDTrain.py (finding 1).
2. Version-tag plot_history's filename in TrackPIDTrain.py:439 and the TrkOnlyPID save path (findings 2–3).
3. Consider the np.quantile replacement for the threshold bisection and a hit-count guard in perform_tz_fit (findings 4–5).
4. Refresh TrkOnlyPID/README.md for the dt/dt feature (finding 7).

[michaelmackenzie]

Responding to the AI comments:

• 1: Fixed
• 2: Fixed
• 3: Fixed
• 4: I switched to np.quantile
• 5: I added a guard, but I don't think we have < 3 hit tracks
• 6: There's a lot of duplication, but I expect these to MVAs will diverge in the near future
• 7: Fixed
• 8: Fixed
• 9: I fixed some printouts/typos
  I also switched from a batch size of 32 --> 1 so the memory overhead will be smaller in TMVA::SOFIE

[oksuzian]

PR Review Summary — re-review of "Update TrkPID to use track-based observables" (#9)

Reviewed at head 4c7926c4 (2026-08-04). Delta since the previously reviewed 3935ad52: one commit ("Address comments") responding point-by-point to the request-changes review. Every fix was verified in the diff, not the reply; the one behavioral question (fixed batch size vs Keras enforcement) was settled empirically in the rootana/current environment.

Decision

• 🟢 Approve. Both S1s and all actionable S2s from the request-changes review are fixed and verified; the remaining items are ⚪ S3 polish and one non-gating physics question. This approval supersedes my earlier request-changes.

---

Carry-forward accounting (vs the review posted at 3935ad52)

1. 🟢 [was S1] TrkOnlyPID mixed_float16 policy — FIXED in 4c7926c4, verified.
TrkOnlyPIDTrain.py now sets set_global_policy('float32'), matching the e202463 fix in TrkPID. The SOFIE-incompatible configuration is gone from both trainers.

2. 🟢 [was S1] plot_history filename regression — FIXED, verified.
Both trainers now call plot_history(f"train_history_v{version}.json", ...), matching what train_model writes. The default import→train→plot path no longer ends in FileNotFoundError.

3. 🟢 [was S2] TrkOnlyPID save/load asymmetry — FIXED, verified.
The save path is now version-tagged (TrkOnlyPID_model_v{version}.keras, train_history_v{version}.json), symmetric with the --skip-train load path.

4. 🟢 [was S2] Threshold bisection non-termination — FIXED, verified.
The loop is replaced with the exact np.quantile(dataset.loc[dataset['label']==0, 'prediction'], 1 - bkg_eff) one-liner in both trainers, plus a threshold printout. The hang mode is gone by construction.

5. 🟢 [was S2] perform_tz_fit degenerate-track guard — FIXED, verified.
if len(hits) < 3: return 0., 0., 0. in both make_inputs.py. A zero slope is far from the physical value (~1 for the dt/dt observable), so such tracks land as obvious outliers rather than crashing the skim. The author notes <3-hit tracks shouldn't exist; agreed — this is defensive, and sufficient.

6. ⚪ [was S2] Duplication across TrkPID/TrkOnlyPID — WITHDRAWN as an action item.
Author declines with rationale: the two MVAs are expected to diverge in the near future. That is a judgment call the author is entitled to, and the finding never gated. The observation stands as context (the mixed_float16 miss was its concrete cost) — no further action requested.

7. 🟢 [was S2] TrkOnlyPID/README.md stale feature list — FIXED, verified (minor residue below).
The v0 feature list now reads "straw hit time vs track time at the hit slope" and matches the implemented features in order. Residue: both READMEs still name the unversioned ONNX outputs ("TrkOnlyPID.onnx" / "TrackPID.onnx") while the code writes *_v{version}.onnx — now ⚪ S3.

8. 🟢 [was S3] ROOT.TMath.GausI — FIXED, verified. Now 2 * ROOT.TMath.Freq(z), as prescribed.

9. 🟡 [was S3] Housekeeping batch — PARTIAL, remaining bits are take-or-leave.
Fixed: the misleading "cosmic muon" confusion-matrix print is deleted outright, the rate printouts are relabeled — and the new FPR/FNR lines are actually more correct than the old ones (the old wording had FPR/FNR semantics crossed; the new FPR = FP/total_bkg, FNR = FN/total_sig is standard); "verion" and "unnessary" typos fixed; the pre-existing broken README link TrkPIDTrain.py → TrackPIDTrain.py fixed too. Still open (all ⚪, no action needed for approval): unused time_mod computation, trkmc/trksegsmc imported but unread (extra I/O), unused nMatActiveFrac, pconv → pcov, and the pre-existing "Accuracy at this threshold" print that reports the TPR.

---

New in the delta (3935ad52..4c7926c4)

1. 🟢 Fixed batch_size=1 verified safe end-to-end.
The Input-layer batch switch 32 → 1 (both trainers) raised the question of whether Keras enforces the static batch against fit/predict, which use their default batch of 32. Tested empirically in the Mu2e rootana/current environment (TF 2.20.0): a model declared with batch_size=1 trains and predicts cleanly at batch 32. So training is unaffected, and the ONNX signature now carries batch 1 — the right shape for per-track SOFIE inference in Offline. This also settles the second question from the previous review.

2. ⚪ [S3] The relabeled rate printouts drop the class names entirely.
"True Positive Rate: ...%" no longer says which class is positive. One header line ("positive = electron") would keep the log self-describing. Cosmetic.

Open question (non-gating, carried)

• The t–z / t–t fits in make_inputs.py still use every trkhits entry unconditionally; TrkStrawHitInfo.state encodes activity. Is including fit-deactivated hits intended (they may carry discrimination for the wrong-mass hypothesis) or should the fit filter on state? Worth a one-line answer for the record; either answer is fine for this PR.

Validation check

• Build/tests run: none — MLTrain has no CI. Verification: full diff 3935ad52..4c7926c4 read line-by-line against the prior findings; the batch-size behavior tested by running TF in rootana/current.
• Config contract check: n/a.
• Cross-repo consistency: unchanged — the planned KalSeed/TrkInfo (EventNtuple) home for the dt observable remains future work; nothing required for this merge.

Residual risk

• Low. The ONNX→SOFIE export step remains untested by construction (tooling unavailable), but the two known landmines on that path (mixed_float16, fixed batch 32) are both cleared.

Author follow-ups (non-blocking)

1. README ONNX names → *_v{version}.onnx (finding 7 residue).
2. Optional S3 sweep: time_mod, unused MC branch imports, nMatActiveFrac, pconv→pcov, positive-class label in the rate printouts.
3. A one-line answer on the hit-state question for the record.