From b2086f74e621d282dfd8d891a8b934358557e075 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 26 Feb 2026 22:49:44 +0100
Subject: [PATCH] fix(ml): correct checkpoint path in evaluate_supervised +
 persist NormStats

evaluate_supervised looked for checkpoints at {models_dir}/{model}_fold{N}_best
but training saves to {models_dir}/{model}/{model}_fold{N}_best (model subdirectory).
Also saves NormStats JSON per fold during training so evaluation uses
training-time normalization instead of computing from test data (data leakage).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/examples/evaluate_supervised.rs       |  2 ++
 crates/ml/examples/train_baseline_supervised.rs | 17 ++++++++++++++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/crates/ml/examples/evaluate_supervised.rs b/crates/ml/examples/evaluate_supervised.rs
index 12ecb615f..b0b760922 100644
--- a/crates/ml/examples/evaluate_supervised.rs
+++ b/crates/ml/examples/evaluate_supervised.rs
@@ -402,6 +402,7 @@ fn evaluate_fold(
 ) -> Result<(Vec<f64>, [usize; 3], f64)> {
     let ckpt_path = args
         .models_dir
+        .join(model_name)
         .join(format!("{}_fold{}_best", model_name, fold));
 
     // Check if checkpoint metadata exists
@@ -606,6 +607,7 @@ fn main() -> Result<()> {
         // Load NormStats from training
         let norm_path = args
             .models_dir
+            .join(&args.model)
             .join(format!("norm_stats_fold{}.json", window.fold));
         let norm_stats: NormStats = if norm_path.exists() {
             let norm_json = std::fs::read_to_string(&norm_path)
diff --git a/crates/ml/examples/train_baseline_supervised.rs b/crates/ml/examples/train_baseline_supervised.rs
index fe7a9487c..3a973ee4a 100644
--- a/crates/ml/examples/train_baseline_supervised.rs
+++ b/crates/ml/examples/train_baseline_supervised.rs
@@ -364,7 +364,7 @@ fn prepare_fold_data(
     val_bars: &[OHLCVBar],
     args: &Args,
     device: &Device,
-) -> Result<(Vec<(Tensor, Tensor)>, Vec<(Tensor, Tensor)>)> {
+) -> Result<(Vec<(Tensor, Tensor)>, Vec<(Tensor, Tensor)>, NormStats)> {
     let train_features = extract_ml_features(train_bars)
         .context("Failed to extract training features")?;
     let val_features = extract_ml_features(val_bars)
@@ -393,7 +393,7 @@ fn prepare_fold_data(
     let train_pairs = build_tensor_pairs(&norm_train, train_bars, train_bar_offset, args, device)?;
     let val_pairs = build_tensor_pairs(&norm_val, val_bars, val_bar_offset, args, device)?;
 
-    Ok((train_pairs, val_pairs))
+    Ok((train_pairs, val_pairs, norm_stats))
 }
 
 /// Convert normalized features + bars into (input, target) tensor pairs.
@@ -687,7 +687,7 @@ fn run_training(args: &Args) -> Result<Vec<TrainingResult>> {
                 window.test.len()
             );
 
-            let (train_pairs, val_pairs) =
+            let (train_pairs, val_pairs, norm_stats) =
                 match prepare_fold_data(&window.train, &window.val, args, &device) {
                     Ok(data) => data,
                     Err(e) => {
@@ -696,6 +696,17 @@ fn run_training(args: &Args) -> Result<Vec<TrainingResult>> {
                     }
                 };
 
+            // Persist NormStats for evaluation
+            let norm_path = model_output.join(format!("norm_stats_fold{}.json", fold));
+            match serde_json::to_string_pretty(&norm_stats) {
+                Ok(json) => {
+                    if let Err(e) = std::fs::write(&norm_path, json) {
+                        warn!("Failed to save NormStats to {}: {}", norm_path.display(), e);
+                    }
+                }
+                Err(e) => warn!("Failed to serialize NormStats: {}", e),
+            }
+
             if train_pairs.is_empty() || val_pairs.is_empty() {
                 warn!(
                     "Skipping fold {} -- empty data after feature extraction",