extract_stat (#37)

* progress

* typos

* cargo fmt

* progress

* cleanup

* cleanup

* cargo fmt

* inline

Author: cscherrer

src/data/mod.rs

@@ -216,38 +216,36 @@ where
     }
 }
 
-/// Convert a `DataOrSuffStat` into a `Stat`
 #[inline]
 pub fn extract_stat<X, Fx, Pr>(pr: &Pr, x: &DataOrSuffStat<X, Fx>) -> Fx::Stat
 where
     Fx: HasSuffStat<X> + HasDensity<X>,
     Fx::Stat: Clone,
     Pr: ConjugatePrior<X, Fx>,
 {
-    match x {
-        DataOrSuffStat::SuffStat(s) => (*s).clone(),
-        DataOrSuffStat::Data(xs) => {
-            let mut stat = pr.empty_stat();
-            stat.observe_many(xs);
-            stat
-        }
-    }
+    extract_stat_then(pr, x, |s| s.clone())
 }
 
 /// Convert a `DataOrSuffStat` into a `Stat` then do something with it
+#[inline]
 pub fn extract_stat_then<X, Fx, Pr, Fnx, Y>(
     pr: &Pr,
     x: &DataOrSuffStat<X, Fx>,
     f_stat: Fnx,
 ) -> Y
 where
     Fx: HasSuffStat<X> + HasDensity<X>,
-    Fx::Stat: Clone,
     Pr: ConjugatePrior<X, Fx>,
-    Fnx: Fn(Fx::Stat) -> Y,
+    Fnx: Fn(&Fx::Stat) -> Y,
 {
-    let stat = extract_stat(pr, x);
-    f_stat(stat)
+    match x {
+        DataOrSuffStat::SuffStat(s) => f_stat(s),
+        DataOrSuffStat::Data(xs) => {
+            let mut stat = pr.empty_stat();
+            stat.observe_many(xs);
+            f_stat(&stat)
+        }
+    }
 }
 
 #[cfg(test)]
@@ -649,9 +647,10 @@ mod tests {
             let data: DataOrSuffStat<f64, Gaussian> =
                 DataOrSuffStat::Data(&data_vec);
 
-            let result = extract_stat_then(&pr, &data, |stat| {
-                stat.n() * 10 + (stat.sum_x() as usize)
-            });
+            let result =
+                extract_stat_then(&pr, &data, |stat: &GaussianSuffStat| {
+                    stat.n() * 10 + (stat.sum_x() as usize)
+                });
 
             assert_eq!(result, 36); // 3 * 10 + 6
         }

src/dist/dirichlet/categorical_prior.rs

@@ -31,7 +31,7 @@ impl<X: CategoricalDatum> ConjugatePrior<X, Categorical>
     }
 
     fn posterior(&self, x: &CategoricalData<X>) -> Self::Posterior {
-        extract_stat_then(self, x, |stat: CategoricalSuffStat| {
+        extract_stat_then(self, x, |stat: &CategoricalSuffStat| {
             let alphas: Vec<f64> =
                 stat.counts().iter().map(|&ct| self.alpha() + ct).collect();
 
@@ -54,7 +54,7 @@ impl<X: CategoricalDatum> ConjugatePrior<X, Categorical>
     ) -> f64 {
         let sum_alpha = self.alpha() * self.k() as f64;
 
-        extract_stat_then(self, x, |stat: CategoricalSuffStat| {
+        extract_stat_then(self, x, |stat: &CategoricalSuffStat| {
             let b = ln_gammafn(sum_alpha + stat.n() as f64);
             let c = stat
                 .counts()
@@ -101,7 +101,7 @@ impl<X: CategoricalDatum> ConjugatePrior<X, Categorical> for Dirichlet {
     }
 
     fn posterior(&self, x: &CategoricalData<X>) -> Self::Posterior {
-        extract_stat_then(self, x, |stat: CategoricalSuffStat| {
+        extract_stat_then(self, x, |stat: &CategoricalSuffStat| {
             let alphas: Vec<f64> = self
                 .alphas()
                 .iter()
@@ -130,7 +130,7 @@ impl<X: CategoricalDatum> ConjugatePrior<X, Categorical> for Dirichlet {
         x: &CategoricalData<X>,
     ) -> f64 {
         let (sum_alpha, ln_norm) = cache;
-        extract_stat_then(self, x, |stat: CategoricalSuffStat| {
+        extract_stat_then(self, x, |stat: &CategoricalSuffStat| {
             let b = ln_gammafn(sum_alpha + stat.n() as f64);
             let c = self
                 .alphas()

src/dist/niw/mvg_prior.rs

@@ -1,5 +1,6 @@
 use crate::consts::LN_2PI;
-use crate::data::{extract_stat_then, DataOrSuffStat, MvGaussianSuffStat};
+use crate::data::extract_stat_then;
+use crate::data::{DataOrSuffStat, MvGaussianSuffStat};
 use crate::dist::{MvGaussian, NormalInvWishart};
 use crate::misc::lnmv_gamma;
 use crate::traits::ConjugatePrior;
@@ -36,7 +37,7 @@ impl ConjugatePrior<DVector<f64>, MvGaussian> for NormalInvWishart {
         }
 
         let nf = x.n() as f64;
-        extract_stat_then(self, x, |stat: MvGaussianSuffStat| {
+        extract_stat_then(self, x, |stat: &MvGaussianSuffStat| {
             let xbar = stat.sum_x() / stat.n() as f64;
             let diff = &xbar - self.mu();
             // s = \sum_{i=1}^N (x_i - \bar{x}) (x_i - \bar{x})^T

src/dist/normal_inv_chi_squared/gaussian_prior.rs

@@ -70,7 +70,7 @@ impl ConjugatePrior<f64, Gaussian> for NormalInvChiSquared {
     }
 
     fn posterior(&self, x: &DataOrSuffStat<f64, Gaussian>) -> Self {
-        extract_stat_then(self, x, |stat: GaussianSuffStat| {
+        extract_stat_then(self, x, |stat: &GaussianSuffStat| {
             posterior_from_stat(self, &stat).into()
         })
     }
@@ -85,7 +85,7 @@ impl ConjugatePrior<f64, Gaussian> for NormalInvChiSquared {
         cache: &Self::MCache,
         x: &DataOrSuffStat<f64, Gaussian>,
     ) -> f64 {
-        extract_stat_then(self, x, |stat: GaussianSuffStat| {
+        extract_stat_then(self, x, |stat: &GaussianSuffStat| {
             let n = stat.n() as f64;
             let post: Self = posterior_from_stat(self, &stat).into();
             let lnz_n = post.ln_z();