ipfrs_tensorlogic/
sgd_optimizer.rs

1//! SGD Optimizer variants for tensor parameter optimization.
2//!
3//! Provides stochastic gradient descent with optional momentum, Nesterov
4//! acceleration, weight decay, and dampening.  Designed for training loops
5//! inside the TensorLogic subsystem.
6//!
7//! # Supported variants
8//!
9//! | Variant        | Update rule |
10//! |----------------|-------------|
11//! | **SGD**        | `p -= lr * (g + wd * p)` |
12//! | **SGDMomentum**| `v = m*v + (1-d)*g; p -= lr*(v + wd*p)` |
13//! | **SGDNesterov**| `v = m*v + g; p -= lr*(g + m*v + wd*p)` |
14//!
15//! # Example
16//!
17//! ```
18//! use ipfrs_tensorlogic::sgd_optimizer::{SGDConfig, SGDOptimizer, OptimizerType};
19//! use std::collections::HashMap;
20//!
21//! let config = SGDConfig::default();
22//! let mut opt = SGDOptimizer::new(config);
23//! opt.register_parameter("w", vec![1.0, 2.0, 3.0]);
24//!
25//! let mut grads = HashMap::new();
26//! grads.insert("w".to_string(), vec![0.1, 0.2, 0.3]);
27//! opt.step(&grads).expect("example: should succeed in docs");
28//!
29//! let w = opt.get_parameter("w").expect("example: should succeed in docs");
30//! assert!(w[0] < 1.0); // parameter decreased
31//! ```
32
33use std::collections::HashMap;
34
35// ---------------------------------------------------------------------------
36// Types
37// ---------------------------------------------------------------------------
38
39/// Which SGD variant to use.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum OptimizerType {
42    /// Vanilla SGD (no momentum buffer).
43    SGD,
44    /// Classical momentum SGD.
45    SGDMomentum,
46    /// Nesterov accelerated gradient.
47    SGDNesterov,
48}
49
50impl std::fmt::Display for OptimizerType {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        match self {
53            Self::SGD => write!(f, "SGD"),
54            Self::SGDMomentum => write!(f, "SGDMomentum"),
55            Self::SGDNesterov => write!(f, "SGDNesterov"),
56        }
57    }
58}
59
60/// Configuration for an [`SGDOptimizer`].
61#[derive(Debug, Clone)]
62pub struct SGDConfig {
63    /// Optimizer variant.
64    pub optimizer_type: OptimizerType,
65    /// Step size (default `0.01`).
66    pub learning_rate: f64,
67    /// Momentum coefficient (used by Momentum/Nesterov, default `0.9`).
68    pub momentum: f64,
69    /// L2 weight-decay coefficient (default `0.0`).
70    pub weight_decay: f64,
71    /// Dampening factor for momentum (default `0.0`).
72    pub dampening: f64,
73}
74
75impl Default for SGDConfig {
76    fn default() -> Self {
77        Self {
78            optimizer_type: OptimizerType::SGD,
79            learning_rate: 0.01,
80            momentum: 0.9,
81            weight_decay: 0.0,
82            dampening: 0.0,
83        }
84    }
85}
86
87/// Per-parameter mutable state tracked by the optimizer.
88#[derive(Debug, Clone)]
89pub struct ParameterState {
90    /// Human-readable name.
91    pub name: String,
92    /// Current parameter values.
93    pub values: Vec<f64>,
94    /// Momentum buffer (velocity).
95    pub velocity: Vec<f64>,
96}
97
98/// Summary statistics returned by [`SGDOptimizer::stats`].
99#[derive(Debug, Clone)]
100pub struct SGDOptimizerStats {
101    /// Which variant is active.
102    pub optimizer_type: OptimizerType,
103    /// Current learning rate.
104    pub learning_rate: f64,
105    /// Total number of scalar parameters.
106    pub parameter_count: usize,
107    /// Number of optimiser steps executed so far.
108    pub step_count: u64,
109}
110
111// ---------------------------------------------------------------------------
112// Optimizer
113// ---------------------------------------------------------------------------
114
115/// Stochastic gradient descent optimizer with momentum, Nesterov, and weight
116/// decay support.
117#[derive(Debug, Clone)]
118pub struct SGDOptimizer {
119    config: SGDConfig,
120    parameters: HashMap<String, ParameterState>,
121    step_count: u64,
122}
123
124impl SGDOptimizer {
125    /// Create a new optimizer with the given configuration.
126    pub fn new(config: SGDConfig) -> Self {
127        Self {
128            config,
129            parameters: HashMap::new(),
130            step_count: 0,
131        }
132    }
133
134    /// Register a named parameter vector.  The velocity buffer is initialised
135    /// to zeros with the same length as `initial_values`.
136    pub fn register_parameter(&mut self, name: &str, initial_values: Vec<f64>) {
137        let len = initial_values.len();
138        self.parameters.insert(
139            name.to_string(),
140            ParameterState {
141                name: name.to_string(),
142                values: initial_values,
143                velocity: vec![0.0; len],
144            },
145        );
146    }
147
148    /// Execute one optimiser step.
149    ///
150    /// `gradients` must contain exactly the same keys as the registered
151    /// parameters, and each gradient vector must match the corresponding
152    /// parameter length.
153    pub fn step(&mut self, gradients: &HashMap<String, Vec<f64>>) -> Result<(), String> {
154        // Validate gradient keys match parameters.
155        for key in gradients.keys() {
156            if !self.parameters.contains_key(key) {
157                return Err(format!(
158                    "gradient key '{}' does not match any registered parameter",
159                    key
160                ));
161            }
162        }
163        for key in self.parameters.keys() {
164            if !gradients.contains_key(key) {
165                return Err(format!(
166                    "missing gradient for registered parameter '{}'",
167                    key
168                ));
169            }
170        }
171
172        // Validate sizes.
173        for (key, grad) in gradients {
174            let param = self
175                .parameters
176                .get(key)
177                .ok_or_else(|| format!("parameter '{}' not found", key))?;
178            if grad.len() != param.values.len() {
179                return Err(format!(
180                    "gradient length {} for '{}' does not match parameter length {}",
181                    grad.len(),
182                    key,
183                    param.values.len(),
184                ));
185            }
186        }
187
188        let lr = self.config.learning_rate;
189        let wd = self.config.weight_decay;
190        let mom = self.config.momentum;
191        let damp = self.config.dampening;
192
193        // Collect keys to avoid borrow issues.
194        let keys: Vec<String> = self.parameters.keys().cloned().collect();
195
196        for key in &keys {
197            let grad = gradients
198                .get(key)
199                .ok_or_else(|| format!("missing gradient for '{}'", key))?;
200            let state = self
201                .parameters
202                .get_mut(key)
203                .ok_or_else(|| format!("parameter '{}' disappeared", key))?;
204
205            match self.config.optimizer_type {
206                OptimizerType::SGD => {
207                    for (p, g) in state.values.iter_mut().zip(grad.iter()) {
208                        let effective_grad = g + wd * *p;
209                        *p -= lr * effective_grad;
210                    }
211                }
212                OptimizerType::SGDMomentum => {
213                    for ((p, v), g) in state
214                        .values
215                        .iter_mut()
216                        .zip(state.velocity.iter_mut())
217                        .zip(grad.iter())
218                    {
219                        *v = mom * *v + (1.0 - damp) * g;
220                        let effective = *v + wd * *p;
221                        *p -= lr * effective;
222                    }
223                }
224                OptimizerType::SGDNesterov => {
225                    for ((p, v), g) in state
226                        .values
227                        .iter_mut()
228                        .zip(state.velocity.iter_mut())
229                        .zip(grad.iter())
230                    {
231                        *v = mom * *v + g;
232                        let effective = g + mom * *v + wd * *p;
233                        *p -= lr * effective;
234                    }
235                }
236            }
237        }
238
239        self.step_count += 1;
240        Ok(())
241    }
242
243    /// Return a slice of the current parameter values, if the name exists.
244    pub fn get_parameter(&self, name: &str) -> Option<&[f64]> {
245        self.parameters.get(name).map(|s| s.values.as_slice())
246    }
247
248    /// Return a slice of the velocity buffer, if the name exists.
249    pub fn get_velocity(&self, name: &str) -> Option<&[f64]> {
250        self.parameters.get(name).map(|s| s.velocity.as_slice())
251    }
252
253    /// Dynamically change the learning rate.
254    pub fn set_learning_rate(&mut self, lr: f64) {
255        self.config.learning_rate = lr;
256    }
257
258    /// Total number of registered scalar parameter values.
259    pub fn parameter_count(&self) -> usize {
260        self.parameters.values().map(|s| s.values.len()).sum()
261    }
262
263    /// How many optimiser steps have been performed.
264    pub fn step_count(&self) -> u64 {
265        self.step_count
266    }
267
268    /// Reset all velocity buffers to zero.
269    pub fn zero_velocities(&mut self) {
270        for state in self.parameters.values_mut() {
271            for v in &mut state.velocity {
272                *v = 0.0;
273            }
274        }
275    }
276
277    /// Return a snapshot of the optimizer statistics.
278    pub fn stats(&self) -> SGDOptimizerStats {
279        SGDOptimizerStats {
280            optimizer_type: self.config.optimizer_type,
281            learning_rate: self.config.learning_rate,
282            parameter_count: self.parameter_count(),
283            step_count: self.step_count,
284        }
285    }
286}
287
288// ---------------------------------------------------------------------------
289// Tests
290// ---------------------------------------------------------------------------
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    fn make_grads(name: &str, vals: Vec<f64>) -> HashMap<String, Vec<f64>> {
297        let mut m = HashMap::new();
298        m.insert(name.to_string(), vals);
299        m
300    }
301
302    // ---- SGD basic ----
303
304    #[test]
305    fn sgd_basic_step() {
306        let mut opt = SGDOptimizer::new(SGDConfig::default());
307        opt.register_parameter("w", vec![1.0, 2.0, 3.0]);
308        let grads = make_grads("w", vec![0.1, 0.2, 0.3]);
309        opt.step(&grads).expect("step should succeed");
310        let w = opt.get_parameter("w").expect("param exists");
311        // p -= lr * grad => 1.0 - 0.01*0.1 = 0.999
312        assert!((w[0] - 0.999).abs() < 1e-12);
313        assert!((w[1] - 1.998).abs() < 1e-12);
314        assert!((w[2] - 2.997).abs() < 1e-12);
315    }
316
317    #[test]
318    fn sgd_step_count_increments() {
319        let mut opt = SGDOptimizer::new(SGDConfig::default());
320        opt.register_parameter("w", vec![1.0]);
321        assert_eq!(opt.step_count(), 0);
322        opt.step(&make_grads("w", vec![0.1]))
323            .expect("step should succeed");
324        assert_eq!(opt.step_count(), 1);
325        opt.step(&make_grads("w", vec![0.1]))
326            .expect("step should succeed");
327        assert_eq!(opt.step_count(), 2);
328    }
329
330    #[test]
331    fn sgd_parameter_count() {
332        let mut opt = SGDOptimizer::new(SGDConfig::default());
333        opt.register_parameter("a", vec![1.0, 2.0]);
334        opt.register_parameter("b", vec![3.0]);
335        assert_eq!(opt.parameter_count(), 3);
336    }
337
338    #[test]
339    fn sgd_zero_gradient_no_change() {
340        let mut opt = SGDOptimizer::new(SGDConfig::default());
341        opt.register_parameter("w", vec![5.0, 10.0]);
342        opt.step(&make_grads("w", vec![0.0, 0.0]))
343            .expect("step should succeed");
344        let w = opt.get_parameter("w").expect("param exists");
345        assert!((w[0] - 5.0).abs() < 1e-12);
346        assert!((w[1] - 10.0).abs() < 1e-12);
347    }
348
349    // ---- Weight decay ----
350
351    #[test]
352    fn sgd_weight_decay() {
353        let config = SGDConfig {
354            optimizer_type: OptimizerType::SGD,
355            learning_rate: 0.1,
356            weight_decay: 0.01,
357            ..SGDConfig::default()
358        };
359        let mut opt = SGDOptimizer::new(config);
360        opt.register_parameter("w", vec![10.0]);
361        opt.step(&make_grads("w", vec![0.0]))
362            .expect("step should succeed");
363        let w = opt.get_parameter("w").expect("param exists");
364        // p -= lr * (0 + 0.01 * 10) = 10 - 0.1 * 0.1 = 9.99
365        assert!((w[0] - 9.99).abs() < 1e-12);
366    }
367
368    #[test]
369    fn sgd_weight_decay_with_gradient() {
370        let config = SGDConfig {
371            optimizer_type: OptimizerType::SGD,
372            learning_rate: 0.01,
373            weight_decay: 0.1,
374            ..SGDConfig::default()
375        };
376        let mut opt = SGDOptimizer::new(config);
377        opt.register_parameter("w", vec![2.0]);
378        opt.step(&make_grads("w", vec![1.0]))
379            .expect("step should succeed");
380        let w = opt.get_parameter("w").expect("param exists");
381        // p -= 0.01 * (1.0 + 0.1*2.0) = 2.0 - 0.01*1.2 = 1.988
382        assert!((w[0] - 1.988).abs() < 1e-12);
383    }
384
385    // ---- Momentum ----
386
387    #[test]
388    fn momentum_accumulation() {
389        let config = SGDConfig {
390            optimizer_type: OptimizerType::SGDMomentum,
391            learning_rate: 0.01,
392            momentum: 0.9,
393            dampening: 0.0,
394            weight_decay: 0.0,
395        };
396        let mut opt = SGDOptimizer::new(config);
397        opt.register_parameter("w", vec![1.0]);
398
399        // Step 1: v = 0.9*0 + 1.0*0.5 = 0.5; p = 1.0 - 0.01*0.5 = 0.995
400        opt.step(&make_grads("w", vec![0.5]))
401            .expect("step should succeed");
402        let v1 = opt.get_velocity("w").expect("vel exists")[0];
403        assert!((v1 - 0.5).abs() < 1e-12);
404        let w1 = opt.get_parameter("w").expect("param exists")[0];
405        assert!((w1 - 0.995).abs() < 1e-12);
406
407        // Step 2: v = 0.9*0.5 + 1.0*0.5 = 0.95; p = 0.995 - 0.01*0.95 = 0.9855
408        opt.step(&make_grads("w", vec![0.5]))
409            .expect("step should succeed");
410        let v2 = opt.get_velocity("w").expect("vel exists")[0];
411        assert!((v2 - 0.95).abs() < 1e-12);
412    }
413
414    #[test]
415    fn momentum_with_dampening() {
416        let config = SGDConfig {
417            optimizer_type: OptimizerType::SGDMomentum,
418            learning_rate: 0.1,
419            momentum: 0.9,
420            dampening: 0.5,
421            weight_decay: 0.0,
422        };
423        let mut opt = SGDOptimizer::new(config);
424        opt.register_parameter("w", vec![1.0]);
425
426        // v = 0.9*0 + 0.5*1.0 = 0.5; p = 1.0 - 0.1*0.5 = 0.95
427        opt.step(&make_grads("w", vec![1.0]))
428            .expect("step should succeed");
429        let v = opt.get_velocity("w").expect("vel exists")[0];
430        assert!((v - 0.5).abs() < 1e-12);
431        let w = opt.get_parameter("w").expect("param exists")[0];
432        assert!((w - 0.95).abs() < 1e-12);
433    }
434
435    #[test]
436    fn momentum_with_weight_decay() {
437        let config = SGDConfig {
438            optimizer_type: OptimizerType::SGDMomentum,
439            learning_rate: 0.1,
440            momentum: 0.9,
441            dampening: 0.0,
442            weight_decay: 0.01,
443        };
444        let mut opt = SGDOptimizer::new(config);
445        opt.register_parameter("w", vec![10.0]);
446
447        // v = 0.9*0 + 1.0*0.0 = 0.0 (grad=0)
448        // effective = 0.0 + 0.01*10 = 0.1
449        // p = 10.0 - 0.1*0.1 = 9.99
450        opt.step(&make_grads("w", vec![0.0]))
451            .expect("step should succeed");
452        let w = opt.get_parameter("w").expect("param exists")[0];
453        assert!((w - 9.99).abs() < 1e-12);
454    }
455
456    // ---- Nesterov ----
457
458    #[test]
459    fn nesterov_lookahead() {
460        let config = SGDConfig {
461            optimizer_type: OptimizerType::SGDNesterov,
462            learning_rate: 0.01,
463            momentum: 0.9,
464            dampening: 0.0,
465            weight_decay: 0.0,
466        };
467        let mut opt = SGDOptimizer::new(config);
468        opt.register_parameter("w", vec![1.0]);
469
470        // v = 0.9*0 + 0.5 = 0.5
471        // effective = 0.5 + 0.9*0.5 = 0.95
472        // p = 1.0 - 0.01*0.95 = 0.9905
473        opt.step(&make_grads("w", vec![0.5]))
474            .expect("step should succeed");
475        let w = opt.get_parameter("w").expect("param exists")[0];
476        assert!((w - 0.9905).abs() < 1e-12);
477        let v = opt.get_velocity("w").expect("vel exists")[0];
478        assert!((v - 0.5).abs() < 1e-12);
479    }
480
481    #[test]
482    fn nesterov_two_steps() {
483        let config = SGDConfig {
484            optimizer_type: OptimizerType::SGDNesterov,
485            learning_rate: 0.01,
486            momentum: 0.9,
487            dampening: 0.0,
488            weight_decay: 0.0,
489        };
490        let mut opt = SGDOptimizer::new(config);
491        opt.register_parameter("w", vec![1.0]);
492
493        opt.step(&make_grads("w", vec![1.0]))
494            .expect("step should succeed");
495        // v1 = 0 + 1.0 = 1.0; eff = 1.0 + 0.9*1.0 = 1.9; p = 1.0 - 0.019 = 0.981
496        let w1 = opt.get_parameter("w").expect("param exists")[0];
497        assert!((w1 - 0.981).abs() < 1e-12);
498
499        opt.step(&make_grads("w", vec![1.0]))
500            .expect("step should succeed");
501        // v2 = 0.9*1.0 + 1.0 = 1.9; eff = 1.0 + 0.9*1.9 = 2.71
502        // p = 0.981 - 0.01*2.71 = 0.9539
503        let w2 = opt.get_parameter("w").expect("param exists")[0];
504        assert!((w2 - 0.9539).abs() < 1e-12);
505    }
506
507    #[test]
508    fn nesterov_with_weight_decay() {
509        let config = SGDConfig {
510            optimizer_type: OptimizerType::SGDNesterov,
511            learning_rate: 0.1,
512            momentum: 0.9,
513            dampening: 0.0,
514            weight_decay: 0.01,
515        };
516        let mut opt = SGDOptimizer::new(config);
517        opt.register_parameter("w", vec![10.0]);
518
519        // v = 0 + 0 = 0
520        // effective = 0 + 0.9*0 + 0.01*10 = 0.1
521        // p = 10.0 - 0.1*0.1 = 9.99
522        opt.step(&make_grads("w", vec![0.0]))
523            .expect("step should succeed");
524        let w = opt.get_parameter("w").expect("param exists")[0];
525        assert!((w - 9.99).abs() < 1e-12);
526    }
527
528    // ---- Error handling ----
529
530    #[test]
531    fn gradient_name_mismatch_error() {
532        let mut opt = SGDOptimizer::new(SGDConfig::default());
533        opt.register_parameter("w", vec![1.0]);
534        let grads = make_grads("wrong_name", vec![0.1]);
535        let result = opt.step(&grads);
536        assert!(result.is_err());
537    }
538
539    #[test]
540    fn gradient_size_mismatch_error() {
541        let mut opt = SGDOptimizer::new(SGDConfig::default());
542        opt.register_parameter("w", vec![1.0, 2.0]);
543        let grads = make_grads("w", vec![0.1]);
544        let result = opt.step(&grads);
545        assert!(result.is_err());
546    }
547
548    #[test]
549    fn missing_gradient_error() {
550        let mut opt = SGDOptimizer::new(SGDConfig::default());
551        opt.register_parameter("a", vec![1.0]);
552        opt.register_parameter("b", vec![2.0]);
553        // Only provide gradient for "a"
554        let grads = make_grads("a", vec![0.1]);
555        let result = opt.step(&grads);
556        assert!(result.is_err());
557    }
558
559    // ---- zero_velocities ----
560
561    #[test]
562    fn zero_velocities_resets() {
563        let config = SGDConfig {
564            optimizer_type: OptimizerType::SGDMomentum,
565            learning_rate: 0.01,
566            momentum: 0.9,
567            ..SGDConfig::default()
568        };
569        let mut opt = SGDOptimizer::new(config);
570        opt.register_parameter("w", vec![1.0]);
571        opt.step(&make_grads("w", vec![1.0]))
572            .expect("step should succeed");
573        let v = opt.get_velocity("w").expect("vel exists")[0];
574        assert!(v.abs() > 0.0);
575
576        opt.zero_velocities();
577        let v2 = opt.get_velocity("w").expect("vel exists")[0];
578        assert!((v2).abs() < 1e-15);
579    }
580
581    // ---- Multiple parameters ----
582
583    #[test]
584    fn multiple_parameters() {
585        let mut opt = SGDOptimizer::new(SGDConfig {
586            learning_rate: 0.1,
587            ..SGDConfig::default()
588        });
589        opt.register_parameter("w1", vec![1.0, 2.0]);
590        opt.register_parameter("w2", vec![3.0]);
591
592        let mut grads = HashMap::new();
593        grads.insert("w1".to_string(), vec![0.1, 0.2]);
594        grads.insert("w2".to_string(), vec![0.3]);
595
596        opt.step(&grads).expect("step should succeed");
597        let w1 = opt.get_parameter("w1").expect("param exists");
598        assert!((w1[0] - 0.99).abs() < 1e-12);
599        assert!((w1[1] - 1.98).abs() < 1e-12);
600        let w2 = opt.get_parameter("w2").expect("param exists");
601        assert!((w2[0] - 2.97).abs() < 1e-12);
602    }
603
604    // ---- Learning rate schedule ----
605
606    #[test]
607    fn learning_rate_schedule() {
608        let mut opt = SGDOptimizer::new(SGDConfig {
609            learning_rate: 0.1,
610            ..SGDConfig::default()
611        });
612        opt.register_parameter("w", vec![10.0]);
613
614        opt.step(&make_grads("w", vec![1.0]))
615            .expect("step should succeed");
616        let w1 = opt.get_parameter("w").expect("param exists")[0];
617        assert!((w1 - 9.9).abs() < 1e-12);
618
619        // Halve the LR
620        opt.set_learning_rate(0.05);
621        opt.step(&make_grads("w", vec![1.0]))
622            .expect("step should succeed");
623        let w2 = opt.get_parameter("w").expect("param exists")[0];
624        assert!((w2 - 9.85).abs() < 1e-12);
625    }
626
627    // ---- Convergence test ----
628
629    #[test]
630    fn multiple_steps_convergence() {
631        // Minimise f(x) = 0.5 * x^2  =>  grad = x
632        let mut opt = SGDOptimizer::new(SGDConfig {
633            optimizer_type: OptimizerType::SGD,
634            learning_rate: 0.1,
635            ..SGDConfig::default()
636        });
637        opt.register_parameter("x", vec![10.0]);
638
639        for _ in 0..200 {
640            let x = opt.get_parameter("x").expect("param exists")[0];
641            opt.step(&make_grads("x", vec![x]))
642                .expect("step should succeed");
643        }
644        let x_final = opt.get_parameter("x").expect("param exists")[0];
645        assert!(
646            x_final.abs() < 1e-6,
647            "should converge near zero, got {}",
648            x_final
649        );
650    }
651
652    #[test]
653    fn momentum_convergence_faster() {
654        // Compare SGD vs Momentum on f(x) = 0.5*x^2
655        let steps = 30;
656        let lr = 0.01;
657
658        // Plain SGD
659        let mut sgd = SGDOptimizer::new(SGDConfig {
660            optimizer_type: OptimizerType::SGD,
661            learning_rate: lr,
662            ..SGDConfig::default()
663        });
664        sgd.register_parameter("x", vec![10.0]);
665        for _ in 0..steps {
666            let x = sgd.get_parameter("x").expect("param exists")[0];
667            sgd.step(&make_grads("x", vec![x]))
668                .expect("step should succeed");
669        }
670
671        // Momentum SGD
672        let mut mom = SGDOptimizer::new(SGDConfig {
673            optimizer_type: OptimizerType::SGDMomentum,
674            learning_rate: lr,
675            momentum: 0.9,
676            ..SGDConfig::default()
677        });
678        mom.register_parameter("x", vec![10.0]);
679        for _ in 0..steps {
680            let x = mom.get_parameter("x").expect("param exists")[0];
681            mom.step(&make_grads("x", vec![x]))
682                .expect("step should succeed");
683        }
684
685        let sgd_x = sgd.get_parameter("x").expect("param exists")[0].abs();
686        let mom_x = mom.get_parameter("x").expect("param exists")[0].abs();
687        assert!(
688            mom_x < sgd_x,
689            "momentum should converge faster: sgd={}, mom={}",
690            sgd_x,
691            mom_x
692        );
693    }
694
695    // ---- Stats ----
696
697    #[test]
698    fn stats_accuracy() {
699        let config = SGDConfig {
700            optimizer_type: OptimizerType::SGDNesterov,
701            learning_rate: 0.05,
702            ..SGDConfig::default()
703        };
704        let mut opt = SGDOptimizer::new(config);
705        opt.register_parameter("a", vec![1.0, 2.0, 3.0]);
706        opt.register_parameter("b", vec![4.0, 5.0]);
707
708        let stats = opt.stats();
709        assert_eq!(stats.optimizer_type, OptimizerType::SGDNesterov);
710        assert!((stats.learning_rate - 0.05).abs() < 1e-15);
711        assert_eq!(stats.parameter_count, 5);
712        assert_eq!(stats.step_count, 0);
713    }
714
715    #[test]
716    fn stats_after_steps() {
717        let mut opt = SGDOptimizer::new(SGDConfig::default());
718        opt.register_parameter("w", vec![1.0]);
719        for _ in 0..5 {
720            opt.step(&make_grads("w", vec![0.1]))
721                .expect("step should succeed");
722        }
723        assert_eq!(opt.stats().step_count, 5);
724    }
725
726    // ---- get_parameter / get_velocity on missing ----
727
728    #[test]
729    fn get_missing_parameter_returns_none() {
730        let opt = SGDOptimizer::new(SGDConfig::default());
731        assert!(opt.get_parameter("nonexistent").is_none());
732    }
733
734    #[test]
735    fn get_missing_velocity_returns_none() {
736        let opt = SGDOptimizer::new(SGDConfig::default());
737        assert!(opt.get_velocity("nonexistent").is_none());
738    }
739
740    // ---- Display for OptimizerType ----
741
742    #[test]
743    fn optimizer_type_display() {
744        assert_eq!(format!("{}", OptimizerType::SGD), "SGD");
745        assert_eq!(format!("{}", OptimizerType::SGDMomentum), "SGDMomentum");
746        assert_eq!(format!("{}", OptimizerType::SGDNesterov), "SGDNesterov");
747    }
748
749    // ---- Default config ----
750
751    #[test]
752    fn default_config_values() {
753        let cfg = SGDConfig::default();
754        assert_eq!(cfg.optimizer_type, OptimizerType::SGD);
755        assert!((cfg.learning_rate - 0.01).abs() < 1e-15);
756        assert!((cfg.momentum - 0.9).abs() < 1e-15);
757        assert!((cfg.weight_decay).abs() < 1e-15);
758        assert!((cfg.dampening).abs() < 1e-15);
759    }
760
761    // ---- Register replaces existing ----
762
763    #[test]
764    fn register_replaces_existing_parameter() {
765        let mut opt = SGDOptimizer::new(SGDConfig::default());
766        opt.register_parameter("w", vec![1.0, 2.0]);
767        opt.register_parameter("w", vec![10.0]);
768        assert_eq!(opt.parameter_count(), 1);
769        let w = opt.get_parameter("w").expect("param exists");
770        assert!((w[0] - 10.0).abs() < 1e-12);
771    }
772
773    // ---- Negative gradients ----
774
775    #[test]
776    fn negative_gradient_increases_parameter() {
777        let mut opt = SGDOptimizer::new(SGDConfig {
778            learning_rate: 0.1,
779            ..SGDConfig::default()
780        });
781        opt.register_parameter("w", vec![0.0]);
782        opt.step(&make_grads("w", vec![-1.0]))
783            .expect("step should succeed");
784        let w = opt.get_parameter("w").expect("param exists")[0];
785        // p -= 0.1 * (-1.0) = 0.0 + 0.1 = 0.1
786        assert!((w - 0.1).abs() < 1e-12);
787    }
788
789    // ---- Large gradient ----
790
791    #[test]
792    fn large_gradient_large_step() {
793        let mut opt = SGDOptimizer::new(SGDConfig {
794            learning_rate: 1.0,
795            ..SGDConfig::default()
796        });
797        opt.register_parameter("w", vec![100.0]);
798        opt.step(&make_grads("w", vec![100.0]))
799            .expect("step should succeed");
800        let w = opt.get_parameter("w").expect("param exists")[0];
801        assert!((w - 0.0).abs() < 1e-12);
802    }
803
804    // ---- Nesterov convergence ----
805
806    #[test]
807    fn nesterov_convergence() {
808        let mut opt = SGDOptimizer::new(SGDConfig {
809            optimizer_type: OptimizerType::SGDNesterov,
810            learning_rate: 0.01,
811            momentum: 0.9,
812            ..SGDConfig::default()
813        });
814        opt.register_parameter("x", vec![10.0]);
815
816        for _ in 0..200 {
817            let x = opt.get_parameter("x").expect("param exists")[0];
818            opt.step(&make_grads("x", vec![x]))
819                .expect("step should succeed");
820        }
821        let x_final = opt.get_parameter("x").expect("param exists")[0];
822        assert!(
823            x_final.abs() < 1e-4,
824            "should converge near zero, got {}",
825            x_final
826        );
827    }
828
829    // ---- Empty step ----
830
831    #[test]
832    fn step_with_no_parameters() {
833        let mut opt = SGDOptimizer::new(SGDConfig::default());
834        let grads: HashMap<String, Vec<f64>> = HashMap::new();
835        opt.step(&grads).expect("empty step should succeed");
836        assert_eq!(opt.step_count(), 1);
837    }
838}
ipfrs_tensorlogic/sgd_optimizer.rs

ipfrs_tensorlogic/
sgd_optimizer.rs
