rill_ml/bandit/
stats.rs

1//! Per-arm diagnostic statistics.
2//!
3//! [`ArmStats`] is a lightweight snapshot of a single arm's state, returned by
4//! [`Bandit::arm_stats`](crate::bandit::Bandit::arm_stats). It does not own any
5//! data from the bandit — it is a copy for diagnostics.
6
7use crate::error::RillError;
8
9/// A snapshot of a single arm's statistics.
10///
11/// Produced by [`Bandit::arm_stats`](crate::bandit::Bandit::arm_stats) for
12/// diagnostics. All fields are plain values — this struct does not borrow from
13/// the bandit.
14#[derive(Debug, Clone, Copy, PartialEq)]
15#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
16pub struct ArmStats {
17    /// How many times this arm has been pulled.
18    pub pulls: u64,
19    /// The total reward accumulated by this arm.
20    pub total_reward: f64,
21    /// The mean reward of this arm (`total_reward / pulls`), or `0.0` if the
22    /// arm has never been pulled.
23    pub mean_reward: f64,
24}
25
26impl ArmStats {
27    /// Create a new `ArmStats` for an arm that has never been pulled.
28    pub const fn empty() -> Self {
29        Self {
30            pulls: 0,
31            total_reward: 0.0,
32            mean_reward: 0.0,
33        }
34    }
35
36    /// Create a new `ArmStats` from raw per-arm data.
37    ///
38    /// If `pulls` is zero, `mean_reward` is set to `0.0`.
39    pub fn new(pulls: u64, total_reward: f64) -> Result<Self, RillError> {
40        if !total_reward.is_finite() {
41            return Err(RillError::NonFiniteValue {
42                field: "total_reward",
43                value: total_reward,
44            });
45        }
46        let mean_reward = if pulls > 0 {
47            total_reward / pulls as f64
48        } else {
49            0.0
50        };
51        Ok(Self {
52            pulls,
53            total_reward,
54            mean_reward,
55        })
56    }
57}
58
59impl Default for ArmStats {
60    fn default() -> Self {
61        Self::empty()
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn empty_stats() {
71        let s = ArmStats::empty();
72        assert_eq!(s.pulls, 0);
73        assert!((s.total_reward - 0.0).abs() < 1e-12);
74        assert!((s.mean_reward - 0.0).abs() < 1e-12);
75    }
76
77    #[test]
78    fn default_is_empty() {
79        assert_eq!(ArmStats::default(), ArmStats::empty());
80    }
81
82    #[test]
83    fn new_with_pulls() {
84        let s = ArmStats::new(10, 5.0).unwrap();
85        assert_eq!(s.pulls, 10);
86        assert!((s.total_reward - 5.0).abs() < 1e-12);
87        assert!((s.mean_reward - 0.5).abs() < 1e-12);
88    }
89
90    #[test]
91    fn new_with_zero_pulls() {
92        let s = ArmStats::new(0, 0.0).unwrap();
93        assert_eq!(s.pulls, 0);
94        assert!((s.mean_reward - 0.0).abs() < 1e-12);
95    }
96
97    #[test]
98    fn new_rejects_non_finite() {
99        assert!(ArmStats::new(5, f64::NAN).is_err());
100        assert!(ArmStats::new(5, f64::INFINITY).is_err());
101    }
102
103    #[test]
104    fn mean_reward_for_various_pulls() {
105        let s = ArmStats::new(3, 9.0).unwrap();
106        assert!((s.mean_reward - 3.0).abs() < 1e-12);
107
108        let s = ArmStats::new(100, 250.0).unwrap();
109        assert!((s.mean_reward - 2.5).abs() < 1e-12);
110    }
111
112    #[cfg(feature = "serde")]
113    #[test]
114    fn serde_roundtrip() {
115        let s = ArmStats::new(42, 123.5).unwrap();
116        let json = serde_json::to_string(&s).unwrap();
117        let restored: ArmStats = serde_json::from_str(&json).unwrap();
118        assert_eq!(restored.pulls, 42);
119        assert!((restored.total_reward - 123.5).abs() < 1e-12);
120        assert!((restored.mean_reward - (123.5 / 42.0)).abs() < 1e-12);
121    }
122}
rill_ml/bandit/stats.rs

rill_ml/bandit/
stats.rs
