Feature/faster nix postpred

CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [0.19.0] - 2024-12-25
+
+### Changed
+- Merry Christmas
+- `NormalInvChiSquared`, `NormalGamma`, and `NormalInvGamme` `PpCache` for Gaussian conjugate analysis changed. `ln_pp_with_cache` is much faster.
+- `Gamma` `PpCache` for Poisson conjugate analysis has been optimized. `ln_pp_with_cache` is faster.
+
 ## [0.18.0] - 2024-06-24
 
 ### Added
@@ -215,6 +222,7 @@
 - Remove dependency on `quadrature` crate in favor of hand-rolled adaptive
     Simpson's rule, which handles multimodal distributions better.
 
+[0.19.0]: https://github.com/promise-ai/rv/compare/v0.18.0...v0.19.0
 [0.18.0]: https://github.com/promise-ai/rv/compare/v0.17.0...v0.18.0
 [0.17.0]: https://github.com/promise-ai/rv/compare/v0.16.5...v0.17.0
 [0.16.5]: https://github.com/promise-ai/rv/compare/v0.16.4...v0.16.5

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rv"
-version = "0.18.0"
+version = "0.19.0"
 authors = ["Baxter Eaves", "Michael Schmidt", "Chad Scherrer"]
 description = "Random variables"
 repository = "https://github.com/promised-ai/rv"
@@ -13,7 +13,7 @@ include = ["README.md", "src/**/*", "benches/*", "Cargo.toml"]
 rust-version = "1.72"
 
 [badges]
-github = { repository = "promised-ai/rv", tag = "v0.17.0" }
+github = { repository = "promised-ai/rv", tag = "v0.19.0" }
 maintenance = { status = "actively-developed" }
 
 [dependencies]
@@ -90,3 +90,11 @@ required-features = ["arraydist"]
 [[bench]]
 name = "mixture_entropy"
 harness = false
+
+[[bench]]
+name = "nix"
+harness = false
+
+[[bench]]
+name = "ng"
+harness = false

benches/ng.rs

@@ -0,0 +1,57 @@
+use criterion::black_box;
+use criterion::BatchSize;
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main};
+use rv::data::GaussianSuffStat;
+use rv::dist::Gaussian;
+use rv::dist::NormalGamma;
+use rv::traits::*;
+
+fn bench_ng_postpred(c: &mut Criterion) {
+    let mut group = c.benchmark_group("NG ln pp(x)");
+    let ng = NormalGamma::new_unchecked(0.1, 1.2, 2.3, 3.4);
+    let mut rng = rand::thread_rng();
+    let g = Gaussian::standard();
+
+    group.bench_function(format!("No cache"), |b| {
+        b.iter_batched(
+            || {
+                let stat = {
+                    let mut stat = GaussianSuffStat::new();
+                    g.sample_stream(&mut rng).take(10).for_each(|x: f64| {
+                        stat.observe(&x);
+                    });
+                    stat
+                };
+                let y: f64 = g.draw(&mut rng);
+                (y, stat)
+            },
+            |(y, stat)| {
+                black_box(ng.ln_pp(&y, &DataOrSuffStat::SuffStat(&stat)))
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function(format!("With cache"), |b| {
+        b.iter_batched(
+            || {
+                let stat = {
+                    let mut stat = GaussianSuffStat::new();
+                    g.sample_stream(&mut rng).take(10).for_each(|x: f64| {
+                        stat.observe(&x);
+                    });
+                    stat
+                };
+                let y: f64 = g.draw(&mut rng);
+                let cache = ng.ln_pp_cache(&DataOrSuffStat::SuffStat(&stat));
+                (y, cache)
+            },
+            |(y, cache)| black_box(ng.ln_pp_with_cache(&cache, &y)),
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group!(ng_benches, bench_ng_postpred);
+criterion_main!(ng_benches);

benches/nix.rs

@@ -0,0 +1,99 @@
+use criterion::black_box;
+use criterion::BatchSize;
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main};
+use rv::data::GaussianSuffStat;
+use rv::dist::Gaussian;
+use rv::dist::NormalInvChiSquared;
+use rv::traits::*;
+
+fn bench_nix_postpred(c: &mut Criterion) {
+    let mut group = c.benchmark_group("NIX ln pp(x)");
+    let nix = NormalInvChiSquared::new_unchecked(0.1, 1.2, 2.3, 3.4);
+    let mut rng = rand::thread_rng();
+    let g = Gaussian::standard();
+
+    group.bench_function(format!("No cache"), |b| {
+        b.iter_batched(
+            || {
+                let stat = {
+                    let mut stat = GaussianSuffStat::new();
+                    g.sample_stream(&mut rng).take(10).for_each(|x: f64| {
+                        stat.observe(&x);
+                    });
+                    stat
+                };
+                let y: f64 = g.draw(&mut rng);
+                (y, stat)
+            },
+            |(y, stat)| {
+                black_box(nix.ln_pp(&y, &DataOrSuffStat::SuffStat(&stat)))
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function(format!("With cache"), |b| {
+        b.iter_batched(
+            || {
+                let stat = {
+                    let mut stat = GaussianSuffStat::new();
+                    g.sample_stream(&mut rng).take(10).for_each(|x: f64| {
+                        stat.observe(&x);
+                    });
+                    stat
+                };
+                let y: f64 = g.draw(&mut rng);
+                let cache = nix.ln_pp_cache(&DataOrSuffStat::SuffStat(&stat));
+                (y, cache)
+            },
+            |(y, cache)| black_box(nix.ln_pp_with_cache(&cache, &y)),
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+fn bench_gauss_stat(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Gaussian Suffstat");
+
+    let mut rng = rand::thread_rng();
+    let g = Gaussian::standard();
+
+    group.bench_function(format!("Forget"), |b| {
+        b.iter_batched(
+            || {
+                let mut stat = GaussianSuffStat::new();
+                for _ in 0..3 {
+                    let x: f64 = g.draw(&mut rng);
+                    stat.observe(&x);
+                }
+                let x: f64 = g.draw(&mut rng);
+                stat.observe(&x);
+                (x, stat)
+            },
+            |(x, mut stat)| {
+                black_box(stat.forget(&x));
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function(format!("Observe"), |b| {
+        b.iter_batched(
+            || {
+                let mut stat = GaussianSuffStat::new();
+                let x: f64 = g.draw(&mut rng);
+                stat.observe(&x);
+                let x: f64 = g.draw(&mut rng);
+                (x, stat)
+            },
+            |(x, mut stat)| {
+                black_box(stat.observe(&x));
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group!(nix_benches, bench_nix_postpred, bench_gauss_stat);
+criterion_main!(nix_benches);

src/data/mod.rs

@@ -17,6 +17,8 @@ pub use stat::UnitPowerLawSuffStat;
 use crate::dist::{
     Bernoulli, Categorical, Gaussian, InvGamma, InvGaussian, Poisson,
 };
+use crate::traits::ConjugatePrior;
+use crate::traits::HasDensity;
 use crate::traits::{HasSuffStat, SuffStat};
 
 pub type BernoulliData<'a, X> = DataOrSuffStat<'a, X, Bernoulli>;
@@ -159,7 +161,7 @@ where
     pub fn n(&self) -> usize {
         match &self {
             DataOrSuffStat::Data(data) => data.len(),
-            DataOrSuffStat::SuffStat(s) => s.n(),
+            DataOrSuffStat::SuffStat(s) => <Fx::Stat as SuffStat<X>>::n(s),
         }
     }
 
@@ -212,39 +214,35 @@ where
 
 /// Convert a `DataOrSuffStat` into a `Stat`
 #[inline]
-pub fn extract_stat<Fx, X, Ctor>(
-    x: &DataOrSuffStat<X, Fx>,
-    stat_ctor: Ctor,
-) -> Fx::Stat
+pub fn extract_stat<X, Fx, Pr>(pr: &Pr, x: &DataOrSuffStat<X, Fx>) -> Fx::Stat
 where
-    Fx: HasSuffStat<X>,
+    Fx: HasSuffStat<X> + HasDensity<X>,
     Fx::Stat: Clone,
-    Ctor: Fn() -> Fx::Stat,
+    Pr: ConjugatePrior<X, Fx>,
 {
     match x {
         DataOrSuffStat::SuffStat(s) => (*s).clone(),
         DataOrSuffStat::Data(xs) => {
-            let mut stat = stat_ctor();
-            xs.iter().for_each(|y| stat.observe(y));
+            let mut stat = pr.empty_stat();
+            stat.observe_many(xs);
             stat
         }
     }
 }
 
 /// Convert a `DataOrSuffStat` into a `Stat` then do something with it
-#[inline]
-pub fn extract_stat_then<Fx, X, Ctor, Fnx, Y>(
+pub fn extract_stat_then<X, Fx, Pr, Fnx, Y>(
+    pr: &Pr,
     x: &DataOrSuffStat<X, Fx>,
-    stat_ctor: Ctor,
     f_stat: Fnx,
 ) -> Y
 where
-    Fx: HasSuffStat<X>,
+    Fx: HasSuffStat<X> + HasDensity<X>,
     Fx::Stat: Clone,
-    Ctor: Fn() -> Fx::Stat,
+    Pr: ConjugatePrior<X, Fx>,
     Fnx: Fn(Fx::Stat) -> Y,
 {
-    let stat = extract_stat(x, stat_ctor);
+    let stat = extract_stat(pr, x);
     f_stat(stat)
 }
 

src/data/stat/bernoulli.rs

@@ -121,6 +121,11 @@ impl<X: Booleable> SuffStat<X> for BernoulliSuffStat {
             self.k -= 1
         }
     }
+
+    fn merge(&mut self, other: Self) {
+        self.n += other.n;
+        self.k += other.k;
+    }
 }
 
 #[cfg(test)]

src/data/stat/beta.rs

@@ -131,6 +131,12 @@ macro_rules! impl_suffstat {
                     self.sum_ln_1mx = 0.0;
                 }
             }
+
+            fn merge(&mut self, other: Self) {
+                self.n += other.n;
+                self.sum_ln_x += other.sum_ln_x;
+                self.sum_ln_1mx += other.sum_ln_1mx;
+            }
         }
     };
 }

src/data/stat/categorical.rs

@@ -118,6 +118,16 @@ impl<X: CategoricalDatum> SuffStat<X> for CategoricalSuffStat {
         self.n -= 1;
         self.counts[ix] -= 1.0;
     }
+
+    fn merge(&mut self, other: Self) {
+        self.n += other.n;
+        self.counts
+            .iter_mut()
+            .zip(other.counts.iter().copied())
+            .for_each(|(ct, ct_o)| {
+                *ct += ct_o;
+            });
+    }
 }
 
 #[cfg(test)]