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

#[cfg(feature = "serde1")]
use serde_derive::{Deserialize, Serialize};

use crate::data::DataOrSuffStat;
use crate::dist::Poisson;
use crate::misc::ln_fact;
use crate::traits::SuffStat;

/// Poisson sufficient statistic.
///
/// Holds the number of observations and their sum.
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct PoissonSuffStat {
    /// Number of observations
    n: usize,
    /// Sum of observations
    sum: f64,
    /// Sum of Log(x!)
    sum_ln_fact: f64,
}

impl PoissonSuffStat {
    /// Create a new empty SuffStat
    #[inline]
    pub fn new() -> Self {
        Self {
            n: 0,
            sum: 0.0,
            sum_ln_fact: 0.0,
        }
    }

    /// Get the number of observations
    #[inline]
    pub fn n(&self) -> usize {
        self.n
    }

    /// Get the sum of all observations
    #[inline]
    pub fn sum(&self) -> f64 {
        self.sum
    }

    #[inline]
    pub fn sum_ln_fact(&self) -> f64 {
        self.sum_ln_fact
    }
}

impl Default for PoissonSuffStat {
    fn default() -> Self {
        Self::new()
    }
}

macro_rules! impl_poisson_suffstat {
    ($kind:ty) => {
        impl<'a> Into<DataOrSuffStat<'a, $kind, Poisson>>
            for &'a PoissonSuffStat
        {
            fn into(self) -> DataOrSuffStat<'a, $kind, Poisson> {
                DataOrSuffStat::SuffStat(self)
            }
        }

        impl<'a> Into<DataOrSuffStat<'a, $kind, Poisson>> for &'a Vec<$kind> {
            fn into(self) -> DataOrSuffStat<'a, $kind, Poisson> {
                DataOrSuffStat::Data(self)
            }
        }

        impl SuffStat<$kind> for PoissonSuffStat {
            fn n(&self) -> usize {
                self.n
            }

            fn observe(&mut self, x: &$kind) {
                let xf = f64::from(*x);
                self.n += 1;
                self.sum += xf;
                self.sum_ln_fact += ln_fact(*x as usize);
            }

            fn forget(&mut self, x: &$kind) {
                if self.n > 1 {
                    let xf = f64::from(*x);
                    self.n -= 1;
                    self.sum -= xf;
                    self.sum_ln_fact -= ln_fact(*x as usize);
                } else {
                    self.n = 0;
                    self.sum = 0.0;
                    self.sum_ln_fact = 0.0;
                }
            }
        }
    };
}

impl_poisson_suffstat!(u8);
impl_poisson_suffstat!(u16);
impl_poisson_suffstat!(u32);