optirs-core 0.3.2

OptiRS core optimization algorithms and utilities
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! Tests for the Lion optimizer

use approx::assert_abs_diff_eq;
use optirs_core::optimizers::{Lion, Optimizer};
use scirs2_core::ndarray::{Array1, Array2};

#[test]
#[allow(dead_code)]
fn test_lion_new_creates_optimizer_with_default_values() {
    let optimizer: Lion<f64> = Lion::new(0.001);
    assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);
    assert_abs_diff_eq!(optimizer.get_beta1(), 0.9);
    assert_abs_diff_eq!(optimizer.get_beta2(), 0.99);
    assert_abs_diff_eq!(optimizer.get_weight_decay(), 0.0);
}

#[test]
#[allow(dead_code)]
fn test_lion_new_with_config() {
    let optimizer: Lion<f64> = Lion::new_with_config(0.01, 0.85, 0.98, 0.001);
    assert_abs_diff_eq!(optimizer.learning_rate(), 0.01);
    assert_abs_diff_eq!(optimizer.get_beta1(), 0.85);
    assert_abs_diff_eq!(optimizer.get_beta2(), 0.98);
    assert_abs_diff_eq!(optimizer.get_weight_decay(), 0.001);
}

#[test]
#[allow(dead_code)]
fn test_lion_setters() {
    let mut optimizer: Lion<f64> = Lion::new(0.001);

    optimizer.set_lr(0.002);
    assert_abs_diff_eq!(optimizer.learning_rate(), 0.002);

    optimizer.set_beta1(0.8);
    assert_abs_diff_eq!(optimizer.get_beta1(), 0.8);

    optimizer.set_beta2(0.95);
    assert_abs_diff_eq!(optimizer.get_beta2(), 0.95);

    optimizer.set_weight_decay(0.01);
    assert_abs_diff_eq!(optimizer.get_weight_decay(), 0.01);
}

#[test]
#[allow(dead_code)]
fn test_lion_1d_optimization() {
    let mut optimizer: Lion<f64> = Lion::new(0.1);

    // Start from x = 2.0, trying to minimize x^2
    let mut params = Array1::from_vec(vec![2.0]);

    // Run a few optimization steps
    for _ in 0..10 {
        // Gradient of x^2 is 2x
        let gradients = Array1::from_vec(vec![2.0 * params[0]]);
        params = optimizer.step(&params, &gradients).expect("unwrap failed");
    }

    // Should be moving towards 0
    assert!(params[0].abs() < 1.0);
}

#[test]
#[allow(dead_code)]
fn test_lion_2d_optimization() {
    let mut optimizer: Lion<f64> = Lion::new(0.1);

    // Start from (2.0, 3.0), trying to minimize x^2 + y^2
    let mut params = Array1::from_vec(vec![2.0, 3.0]);

    // Run optimization steps
    for _ in 0..20 {
        // Gradient of x^2 + y^2 is (2x, 2y)
        let gradients = Array1::from_vec(vec![2.0 * params[0], 2.0 * params[1]]);
        params = optimizer.step(&params, &gradients).expect("unwrap failed");
    }

    // Should be moving towards (0, 0)
    assert!(params[0].abs() < 1.0);
    assert!(params[1].abs() < 1.0);
}

#[test]
#[allow(dead_code)]
fn test_lion_with_weight_decay() {
    let mut optimizer: Lion<f64> = Lion::new_with_config(0.1, 0.9, 0.99, 0.1);

    // Start from (1.0, 1.0)
    let mut params = Array1::from_vec(vec![1.0, 1.0]);

    // Run optimization with zero gradients (only weight decay effect)
    for _ in 0..10 {
        let gradients = Array1::zeros(2);
        params = optimizer.step(&params, &gradients).expect("unwrap failed");
    }

    // Weight decay should shrink parameters
    assert!(params[0] < 1.0);
    assert!(params[1] < 1.0);
}

#[test]
#[allow(dead_code)]
fn test_lion_reset() {
    let mut optimizer: Lion<f64> = Lion::new(0.1);

    // Perform a step to initialize internal state
    let params = Array1::from_vec(vec![1.0, 2.0]);
    let gradients = Array1::from_vec(vec![0.1, 0.2]);
    let _ = optimizer.step(&params, &gradients).expect("unwrap failed");

    // Reset the optimizer
    optimizer.reset();

    // After reset, the next step should behave like the first step
    let new_params = optimizer.step(&params, &gradients).expect("unwrap failed");

    // This should match the behavior of a fresh optimizer
    let mut fresh_optimizer: Lion<f64> = Lion::new(0.1);
    let fresh_params = fresh_optimizer
        .step(&params, &gradients)
        .expect("unwrap failed");

    assert_abs_diff_eq!(new_params[0], fresh_params[0], epsilon = 1e-6);
    assert_abs_diff_eq!(new_params[1], fresh_params[1], epsilon = 1e-6);
}

#[test]
#[allow(dead_code)]
fn test_lion_multiple_dimensions() {
    let mut optimizer: Lion<f64> = Lion::new(0.1);

    // Test with 2D array
    let params_2d =
        Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("unwrap failed");
    let gradients_2d =
        Array2::from_shape_vec((2, 3), vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6]).expect("unwrap failed");

    let updated_2d = optimizer
        .step(&params_2d, &gradients_2d)
        .expect("unwrap failed");

    // Check dimensions are preserved
    assert_eq!(updated_2d.shape(), params_2d.shape());

    // Check that parameters were updated
    assert_ne!(updated_2d[[0, 0]], params_2d[[0, 0]]);
    assert_ne!(updated_2d[[1, 2]], params_2d[[1, 2]]);
}

#[test]
#[allow(dead_code)]
fn test_lion_trait_methods() {
    let mut optimizer: Lion<f64> = Lion::new(0.001);

    // Test direct methods (not trait methods to avoid dimension inference issues)
    assert_abs_diff_eq!(optimizer.learning_rate(), 0.001);

    optimizer.set_lr(0.005);
    assert_abs_diff_eq!(optimizer.learning_rate(), 0.005);
}