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#[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 n: usize,
18 sum: f64,
20 sum_ln_fact: f64,
22}
23
24impl PoissonSuffStat {
25 #[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 #[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 #[inline]
50 #[must_use]
51 pub fn n(&self) -> usize {
52 self.n
53 }
54
55 #[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);