clawft_kernel/sensor_graph/
train.rs

1//! Train / eval scaffolding for [`KstemitModel`] on synthetic arrays.
2
3use super::fixture::SyntheticDataset;
4use super::model::{KstemitConfig, KstemitModel};
5
6/// SGD hyper-parameters.
7#[derive(Debug, Clone)]
8pub struct TrainConfig {
9    pub epochs: usize,
10    pub learning_rate: f32,
11    pub weight_decay: f32,
12    /// Shuffle samples each epoch (LCG).
13    pub shuffle: bool,
14    pub seed: u64,
15}
16
17impl Default for TrainConfig {
18    fn default() -> Self {
19        Self {
20            epochs: 30,
21            learning_rate: 0.05,
22            weight_decay: 1e-4,
23            shuffle: true,
24            seed: 7,
25        }
26    }
27}
28
29/// Result of a training run.
30#[derive(Debug, Clone)]
31pub struct TrainResult {
32    pub model: KstemitModel,
33    /// Mean CE loss per epoch.
34    pub epoch_losses: Vec<f32>,
35    /// Final train accuracy.
36    pub train_accuracy: f32,
37}
38
39/// Classification metrics.
40#[derive(Debug, Clone)]
41pub struct EvalMetrics {
42    pub accuracy: f32,
43    pub mean_loss: f32,
44    pub n: usize,
45    /// Confusion matrix flattened row-major `[num_classes * num_classes]`.
46    pub confusion: Vec<usize>,
47    pub num_classes: usize,
48}
49
50/// Train a new model on a synthetic dataset.
51pub fn train_sgd(
52    dataset: &SyntheticDataset,
53    model_cfg: KstemitConfig,
54    train_cfg: TrainConfig,
55) -> TrainResult {
56    let mut model = KstemitModel::new(model_cfg);
57    let mut rng = train_cfg.seed;
58    let mut epoch_losses = Vec::with_capacity(train_cfg.epochs);
59    let n = dataset.samples.len();
60    assert!(n > 0, "empty dataset");
61
62    for _epoch in 0..train_cfg.epochs {
63        let mut order: Vec<usize> = (0..n).collect();
64        if train_cfg.shuffle {
65            // Fisher–Yates with LCG.
66            for i in (1..n).rev() {
67                rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
68                let j = (rng as usize) % (i + 1);
69                order.swap(i, j);
70            }
71        }
72        let mut loss_sum = 0.0f32;
73        for &si in &order {
74            let s = &dataset.samples[si];
75            let loss = model.train_step(
76                &dataset.graph,
77                s.node_idx,
78                s.label,
79                train_cfg.learning_rate,
80                train_cfg.weight_decay,
81            );
82            loss_sum += loss;
83        }
84        epoch_losses.push(loss_sum / n as f32);
85    }
86
87    let eval = evaluate(&model, dataset);
88    TrainResult {
89        model,
90        epoch_losses,
91        train_accuracy: eval.accuracy,
92    }
93}
94
95/// Evaluate mean CE and accuracy on all samples.
96pub fn evaluate(model: &KstemitModel, dataset: &SyntheticDataset) -> EvalMetrics {
97    let num_classes = model.config.num_classes;
98    let mut confusion = vec![0usize; num_classes * num_classes];
99    let mut correct = 0usize;
100    let mut loss_sum = 0.0f32;
101    let n = dataset.samples.len();
102
103    for s in &dataset.samples {
104        let out = model.forward_node(&dataset.graph, s.node_idx);
105        loss_sum += super::layers::cross_entropy(&out.probs, s.label);
106        let pred = argmax(&out.logits);
107        if pred == s.label {
108            correct += 1;
109        }
110        if s.label < num_classes && pred < num_classes {
111            confusion[s.label * num_classes + pred] += 1;
112        }
113    }
114
115    EvalMetrics {
116        accuracy: if n == 0 {
117            0.0
118        } else {
119            correct as f32 / n as f32
120        },
121        mean_loss: if n == 0 { 0.0 } else { loss_sum / n as f32 },
122        n,
123        confusion,
124        num_classes,
125    }
126}
127
128fn argmax(v: &[f32]) -> usize {
129    let mut best_i = 0;
130    let mut best_v = f32::NEG_INFINITY;
131    for (i, &x) in v.iter().enumerate() {
132        if x > best_v {
133            best_v = x;
134            best_i = i;
135        }
136    }
137    best_i
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::sensor_graph::fixture::{generate_synthetic_array, SyntheticConfig};
144    use crate::sensor_graph::graph::PHYSICS_DIM;
145
146    #[test]
147    fn train_improves_over_chance_on_synthetic() {
148        let ds = generate_synthetic_array(SyntheticConfig {
149            n_buoys: 12,
150            series_len: 32,
151            connect_radius: 3.0,
152            noise: 0.04,
153            seed: 1234,
154        });
155        let model_cfg = KstemitConfig {
156            feature_dim: PHYSICS_DIM,
157            temporal_window: 16,
158            hidden: 12,
159            num_classes: 3,
160            physics_scale: 0.25,
161            seed: 3,
162        };
163        let train_cfg = TrainConfig {
164            epochs: 40,
165            learning_rate: 0.08,
166            weight_decay: 1e-4,
167            shuffle: true,
168            seed: 9,
169        };
170        let result = train_sgd(&ds, model_cfg, train_cfg);
171        assert!(result.epoch_losses.len() == 40);
172        // Loss should trend down.
173        let first = result.epoch_losses[0];
174        let last = *result.epoch_losses.last().unwrap();
175        assert!(
176            last < first,
177            "loss should drop: first={first:.4} last={last:.4}"
178        );
179        // Better than random (1/3) with margin — synthetic is separable.
180        assert!(
181            result.train_accuracy > 0.45,
182            "acc={:.3} expected > 0.45",
183            result.train_accuracy
184        );
185    }
186
187    #[test]
188    fn evaluate_shapes() {
189        let ds = generate_synthetic_array(SyntheticConfig::default());
190        let model = KstemitModel::new(KstemitConfig::default());
191        let m = evaluate(&model, &ds);
192        assert_eq!(m.n, ds.samples.len());
193        assert_eq!(m.confusion.len(), 9);
194    }
195}
clawft_kernel/sensor_graph/train.rs

clawft_kernel/sensor_graph/
train.rs
