rv/data/stat/
poisson.rs

1#[cfg(feature = "serde1")]
2use serde::{Deserialize, Serialize};
3
4use crate::data::DataOrSuffStat;
5use crate::dist::Poisson;
6use crate::misc::ln_fact;
7use crate::traits::SuffStat;
8
9/// Poisson sufficient statistic.
10///
11/// Holds the number of observations and their sum.
12#[derive(Debug, Clone, PartialEq)]
13#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
14#[cfg_attr(feature = "serde1", serde(rename_all = "snake_case"))]
15pub struct PoissonSuffStat {
16    /// Number of observations
17    n: usize,
18    /// Sum of observations
19    sum: f64,
20    /// Sum of Log(x!)
21    sum_ln_fact: f64,
22}
23
24impl PoissonSuffStat {
25    /// Create a new empty `SuffStat`
26    #[inline]
27    #[must_use]
28    pub fn new() -> Self {
29        Self {
30            n: 0,
31            sum: 0.0,
32            sum_ln_fact: 0.0,
33        }
34    }
35
36    /// Create a sufficient statistic from components without checking whether
37    /// they are valid.
38    #[inline]
39    #[must_use]
40    pub fn from_parts_unchecked(n: usize, sum: f64, sum_ln_fact: f64) -> Self {
41        Self {
42            n,
43            sum,
44            sum_ln_fact,
45        }
46    }
47
48    /// Get the number of observations
49    #[inline]
50    #[must_use]
51    pub fn n(&self) -> usize {
52        self.n
53    }
54
55    /// Get the sum of all observations
56    #[inline]
57    #[must_use]
58    pub fn sum(&self) -> f64 {
59        self.sum
60    }
61
62    #[inline]
63    #[must_use]
64    pub fn sum_ln_fact(&self) -> f64 {
65        self.sum_ln_fact
66    }
67}
68
69impl Default for PoissonSuffStat {
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75macro_rules! impl_poisson_suffstat {
76    ($kind:ty) => {
77        impl<'a> From<&'a PoissonSuffStat>
78            for DataOrSuffStat<'a, $kind, Poisson>
79        {
80            fn from(stat: &'a PoissonSuffStat) -> Self {
81                DataOrSuffStat::SuffStat(stat)
82            }
83        }
84
85        impl<'a> From<&'a Vec<$kind>> for DataOrSuffStat<'a, $kind, Poisson> {
86            fn from(xs: &'a Vec<$kind>) -> Self {
87                DataOrSuffStat::Data(xs)
88            }
89        }
90
91        impl<'a> From<&'a [$kind]> for DataOrSuffStat<'a, $kind, Poisson> {
92            fn from(xs: &'a [$kind]) -> Self {
93                DataOrSuffStat::Data(xs)
94            }
95        }
96
97        impl SuffStat<$kind> for PoissonSuffStat {
98            fn n(&self) -> usize {
99                self.n
100            }
101
102            fn observe(&mut self, x: &$kind) {
103                let xf = *x as f64;
104                self.n += 1;
105                self.sum += xf;
106                self.sum_ln_fact += ln_fact(*x as usize);
107            }
108
109            fn forget(&mut self, x: &$kind) {
110                if self.n > 1 {
111                    let xf = *x as f64;
112                    self.n -= 1;
113                    self.sum -= xf;
114                    self.sum_ln_fact -= ln_fact(*x as usize);
115                } else {
116                    self.n = 0;
117                    self.sum = 0.0;
118                    self.sum_ln_fact = 0.0;
119                }
120            }
121
122            fn merge(&mut self, other: Self) {
123                self.n += other.n;
124                self.sum += other.sum;
125                self.sum_ln_fact += other.sum_ln_fact;
126            }
127        }
128    };
129}
130
131impl_poisson_suffstat!(u8);
132impl_poisson_suffstat!(u16);
133impl_poisson_suffstat!(u32);
134impl_poisson_suffstat!(usize);