minidx_core/
misc.rs

1use serde::{Deserialize, Serialize};
2
3/// Computes a rolling exponential average, based on the given alpha value.
4///
5/// An alpha of 0.5 to 1.0 gives higher weight to the most recent sample, whereas
6/// an alpha of 0.1 to 0.3 gives higher weight to past samples.
7pub struct ExpAvg {
8    alpha: f32,       // Smoothing factor (0 < alpha ≤ 1)
9    avg: Option<f32>, // Optional to handle the first sample
10}
11
12impl ExpAvg {
13    /// Creates a new `ExpAvg` with the given smoothing factor.
14    pub fn new(alpha: f32) -> Self {
15        assert!(
16            (0.0..=1.0).contains(&alpha),
17            "Alpha must be in the range (0,1]"
18        );
19        Self { alpha, avg: None }
20    }
21
22    /// Updates the moving average with a new sample.
23    pub fn update(&mut self, sample: f32) -> f32 {
24        self.avg = Some(match self.avg {
25            Some(current_avg) => self.alpha * sample + (1.0 - self.alpha) * current_avg,
26            None => sample, // Initialize with the first sample
27        });
28        self.avg.unwrap()
29    }
30
31    /// Returns the current average.
32    pub fn get(&self) -> Option<f32> {
33        self.avg
34    }
35}
36
37fn cosine_decay(current_timestep: usize, end_timestep: usize, start_val: f32, end_val: f32) -> f32 {
38    if current_timestep >= end_timestep {
39        return end_val;
40    }
41
42    use std::f32::consts::PI;
43    let decay_ratio = (current_timestep as f32) / (end_timestep as f32);
44    let cosine_decay = 0.5 * (1.0 + (PI * decay_ratio).cos());
45    end_val + (start_val - end_val) * cosine_decay
46}
47
48/// A value which can decay according to some schedule over the progression of timesteps.
49#[derive(Clone, Serialize, Deserialize, Debug)]
50pub enum Decay {
51    /// No decay: The specified constant is used regardless of timestep.
52    None(f32),
53    /// Linear decay: The value decays from `start` by `decay` each timestep.
54    Linear { start: f32, decay: f32 },
55    /// Cosine decay: The value smoothly decays from `start` to `end` over `num_steps` timesteps,
56    /// with the decay matching the shape of the cosine function.
57    Cosine {
58        start: f32,
59        end: f32,
60        num_steps: usize,
61    },
62}
63
64impl Decay {
65    /// The output value at timestep 0.
66    pub fn start_value(&self) -> f32 {
67        use Decay::*;
68        match self {
69            None(v) => *v,
70            Linear { start, .. } => *start,
71            Cosine { start, .. } => *start,
72        }
73    }
74
75    /// The output value at the specified timestep.
76    pub fn at_timestep(&self, timestep: usize) -> f32 {
77        use Decay::*;
78        match self {
79            None(v) => *v,
80            Linear { start, decay } => {
81                let decayed = *start - (*decay * timestep as f32);
82                if decayed < *decay {
83                    *decay
84                } else {
85                    decayed
86                }
87            }
88            Cosine {
89                start,
90                end,
91                num_steps,
92            } => cosine_decay(timestep, *num_steps, *start, *end),
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_decay_linear() {
103        let decay = Decay::Linear {
104            start: 2.0,
105            decay: 0.25,
106        };
107        assert_eq!(decay.at_timestep(0), 2.0);
108        assert_eq!(decay.at_timestep(1), 1.75);
109        assert_eq!(decay.at_timestep(999999), 0.25);
110    }
111
112    #[test]
113    fn test_decay_cosine() {
114        let decay = Decay::Cosine {
115            start: 2.0,
116            end: 1.0,
117            num_steps: 1000,
118        };
119        assert_eq!(decay.at_timestep(0), 2.0);
120
121        let mid = decay.at_timestep(500);
122        assert!(mid > 1.49);
123        assert!(mid < 1.51);
124
125        assert_eq!(decay.at_timestep(999999), 1.0);
126    }
127}
minidx_core/misc.rs

minidx_core/
misc.rs
