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).