1use std::collections::VecDeque;
36use thiserror::Error;
37
38#[derive(Debug, Error, Clone, PartialEq)]
44pub enum BayesError {
45 #[error(
47 "prior/observation mismatch: cannot update {prior_type} prior with {obs_type} observation"
48 )]
49 PriorObservationMismatch {
50 prior_type: String,
52 obs_type: String,
54 },
55
56 #[error("invalid parameters: {0}")]
58 InvalidParameters(String),
59
60 #[error("unsupported operation: {0}")]
62 UnsupportedOperation(String),
63
64 #[error("numerical error: {0}")]
66 NumericalError(String),
67}
68
69#[derive(Debug, Clone, PartialEq)]
75pub enum Prior {
76 Beta {
78 alpha: f64,
80 beta: f64,
82 },
83 Gaussian {
86 mean: f64,
88 variance: f64,
90 },
91 Dirichlet {
93 alphas: Vec<f64>,
95 },
96 Gamma {
98 shape: f64,
100 rate: f64,
102 },
103}
104
105impl Prior {
106 fn type_name(&self) -> &'static str {
108 match self {
109 Prior::Beta { .. } => "Beta",
110 Prior::Gaussian { .. } => "Gaussian",
111 Prior::Dirichlet { .. } => "Dirichlet",
112 Prior::Gamma { .. } => "Gamma",
113 }
114 }
115
116 fn validate(&self) -> Result<(), BayesError> {
118 match self {
119 Prior::Beta { alpha, beta } => {
120 if *alpha <= 0.0 || alpha.is_nan() {
121 return Err(BayesError::InvalidParameters(format!(
122 "Beta alpha must be > 0, got {alpha}"
123 )));
124 }
125 if *beta <= 0.0 || beta.is_nan() {
126 return Err(BayesError::InvalidParameters(format!(
127 "Beta beta must be > 0, got {beta}"
128 )));
129 }
130 }
131 Prior::Gaussian { variance, .. } => {
132 if *variance <= 0.0 || variance.is_nan() {
133 return Err(BayesError::InvalidParameters(format!(
134 "Gaussian variance must be > 0, got {variance}"
135 )));
136 }
137 }
138 Prior::Dirichlet { alphas } => {
139 if alphas.is_empty() {
140 return Err(BayesError::InvalidParameters(
141 "Dirichlet alphas must be non-empty".to_string(),
142 ));
143 }
144 for (i, &a) in alphas.iter().enumerate() {
145 if a <= 0.0 || a.is_nan() {
146 return Err(BayesError::InvalidParameters(format!(
147 "Dirichlet alpha[{i}] must be > 0, got {a}"
148 )));
149 }
150 }
151 }
152 Prior::Gamma { shape, rate } => {
153 if *shape <= 0.0 || shape.is_nan() {
154 return Err(BayesError::InvalidParameters(format!(
155 "Gamma shape must be > 0, got {shape}"
156 )));
157 }
158 if *rate <= 0.0 || rate.is_nan() {
159 return Err(BayesError::InvalidParameters(format!(
160 "Gamma rate must be > 0, got {rate}"
161 )));
162 }
163 }
164 }
165 Ok(())
166 }
167}
168
169#[derive(Debug, Clone, PartialEq)]
171pub enum Observation {
172 Bernoulli {
174 successes: u64,
176 trials: u64,
178 },
179 Gaussian {
181 sample_mean: f64,
183 sample_variance: f64,
185 n: u64,
187 },
188 Categorical {
190 counts: Vec<u64>,
192 },
193 Poisson {
195 total_events: u64,
197 total_time: f64,
199 },
200}
201
202impl Observation {
203 fn type_name(&self) -> &'static str {
205 match self {
206 Observation::Bernoulli { .. } => "Bernoulli",
207 Observation::Gaussian { .. } => "Gaussian",
208 Observation::Categorical { .. } => "Categorical",
209 Observation::Poisson { .. } => "Poisson",
210 }
211 }
212}
213
214#[derive(Debug, Clone)]
216pub struct Posterior {
217 pub prior: Prior,
219 pub likelihood_type: String,
221 pub updated: Prior,
223 pub log_marginal: f64,
225}
226
227#[derive(Debug, Clone, PartialEq)]
229pub struct CredibleInterval {
230 pub lower: f64,
232 pub upper: f64,
234 pub probability: f64,
236}
237
238fn ln_gamma(x: f64) -> f64 {
245 if x < 0.5 {
247 return ln_gamma(x + 1.0) - x.ln();
250 }
251 if x < 7.0 {
252 return ln_gamma(x + 1.0) - x.ln();
254 }
255 let half_ln_two_pi = 0.918_938_533_204_672_8_f64; let inv_x = 1.0 / x;
258 let inv_x2 = inv_x * inv_x;
259 (x - 0.5) * x.ln() - x
260 + half_ln_two_pi
261 + inv_x * (1.0 / 12.0 - inv_x2 * (1.0 / 360.0 - inv_x2 / 1260.0))
262}
263
264fn log_beta(a: f64, b: f64) -> f64 {
266 ln_gamma(a) + ln_gamma(b) - ln_gamma(a + b)
267}
268
269fn log_dirichlet_norm(alphas: &[f64]) -> f64 {
272 let sum: f64 = alphas.iter().sum();
273 let sum_lg: f64 = alphas.iter().map(|&a| ln_gamma(a)).sum();
274 sum_lg - ln_gamma(sum)
275}
276
277fn digamma(x: f64) -> f64 {
282 if x < 6.0 {
283 return digamma(x + 1.0) - 1.0 / x;
285 }
286 let inv_x = 1.0 / x;
288 let inv_x2 = inv_x * inv_x;
289 x.ln() - 0.5 * inv_x - inv_x2 * (1.0 / 12.0 - inv_x2 * (1.0 / 120.0 - inv_x2 / 252.0))
290}
291
292fn z_score(probability: f64) -> f64 {
297 let p = (1.0 + probability) / 2.0;
299 if (p - 0.5).abs() < 1e-10 {
301 return 0.0;
302 }
303 let t = (-2.0 * (1.0 - p).ln()).sqrt();
304 let c0 = 2.515_517;
305 let c1 = 0.802_853;
306 let c2 = 0.010_328;
307 let d1 = 1.432_788;
308 let d2 = 0.189_269;
309 let d3 = 0.001_308;
310 let numer = c0 + c1 * t + c2 * t * t;
311 let denom = 1.0 + d1 * t + d2 * t * t + d3 * t * t * t;
312 t - numer / denom
313}
314
315fn check_finite(val: f64, context: &str) -> Result<f64, BayesError> {
317 if val.is_finite() {
318 Ok(val)
319 } else {
320 Err(BayesError::NumericalError(format!(
321 "{context}: computed non-finite value {val}"
322 )))
323 }
324}
325
326#[derive(Debug)]
335pub struct BayesianUpdateEngine {
336 history: VecDeque<Posterior>,
338 max_history: usize,
340}
341
342impl BayesianUpdateEngine {
343 pub fn new(max_history: usize) -> Self {
345 Self {
346 history: VecDeque::with_capacity(max_history.min(1024)),
347 max_history,
348 }
349 }
350
351 pub fn update(
358 &mut self,
359 prior: Prior,
360 observation: &Observation,
361 ) -> Result<Posterior, BayesError> {
362 prior.validate()?;
363
364 let posterior = match (&prior, observation) {
365 (Prior::Beta { alpha, beta }, Observation::Bernoulli { successes, trials }) => {
367 if successes > trials {
368 return Err(BayesError::InvalidParameters(format!(
369 "successes ({successes}) cannot exceed trials ({trials})"
370 )));
371 }
372 let s = *successes as f64;
373 let f = (*trials - *successes) as f64;
374 let alpha_post = alpha + s;
375 let beta_post = beta + f;
376 let log_marginal = check_finite(
377 log_beta(alpha_post, beta_post) - log_beta(*alpha, *beta),
378 "Beta-Bernoulli log_marginal",
379 )?;
380 Posterior {
381 prior: prior.clone(),
382 likelihood_type: "Bernoulli".to_string(),
383 updated: Prior::Beta {
384 alpha: alpha_post,
385 beta: beta_post,
386 },
387 log_marginal,
388 }
389 }
390
391 (
393 Prior::Gaussian {
394 mean: prior_mean,
395 variance: prior_var,
396 },
397 Observation::Gaussian {
398 sample_mean,
399 sample_variance,
400 n,
401 },
402 ) => {
403 if *sample_variance <= 0.0 || sample_variance.is_nan() {
404 return Err(BayesError::InvalidParameters(format!(
405 "sample_variance must be > 0, got {sample_variance}"
406 )));
407 }
408 if *n == 0 {
409 return Err(BayesError::InvalidParameters(
410 "n must be > 0 for Gaussian observation".to_string(),
411 ));
412 }
413 let n_f = *n as f64;
414 let post_prec = 1.0 / prior_var + n_f / sample_variance;
416 let post_var = 1.0 / post_prec;
417 let post_mean =
418 post_var * (prior_mean / prior_var + n_f * sample_mean / sample_variance);
419
420 let effective_var = prior_var + sample_variance / n_f;
422 let log_marginal = check_finite(
423 -0.5 * (std::f64::consts::TAU * effective_var).ln(),
424 "Gaussian-Gaussian log_marginal",
425 )?;
426
427 Posterior {
428 prior: prior.clone(),
429 likelihood_type: "Gaussian".to_string(),
430 updated: Prior::Gaussian {
431 mean: post_mean,
432 variance: post_var,
433 },
434 log_marginal,
435 }
436 }
437
438 (Prior::Dirichlet { alphas }, Observation::Categorical { counts }) => {
440 if alphas.len() != counts.len() {
441 return Err(BayesError::InvalidParameters(format!(
442 "Dirichlet dim {} != Categorical counts dim {}",
443 alphas.len(),
444 counts.len()
445 )));
446 }
447 let alphas_post: Vec<f64> = alphas
448 .iter()
449 .zip(counts.iter())
450 .map(|(&a, &c)| a + c as f64)
451 .collect();
452
453 let log_marginal = check_finite(
454 log_dirichlet_norm(&alphas_post) - log_dirichlet_norm(alphas),
455 "Dirichlet-Categorical log_marginal",
456 )?;
457
458 Posterior {
459 prior: prior.clone(),
460 likelihood_type: "Categorical".to_string(),
461 updated: Prior::Dirichlet {
462 alphas: alphas_post,
463 },
464 log_marginal,
465 }
466 }
467
468 (
470 Prior::Gamma { shape, rate },
471 Observation::Poisson {
472 total_events,
473 total_time,
474 },
475 ) => {
476 if *total_time <= 0.0 || total_time.is_nan() {
477 return Err(BayesError::InvalidParameters(format!(
478 "total_time must be > 0, got {total_time}"
479 )));
480 }
481 let k = *total_events as f64;
482 let shape_post = shape + k;
483 let rate_post = rate + total_time;
484
485 let log_marginal = check_finite(
488 ln_gamma(shape_post) - ln_gamma(*shape) + shape * rate.ln()
489 - shape_post * rate_post.ln(),
490 "Gamma-Poisson log_marginal",
491 )?;
492
493 Posterior {
494 prior: prior.clone(),
495 likelihood_type: "Poisson".to_string(),
496 updated: Prior::Gamma {
497 shape: shape_post,
498 rate: rate_post,
499 },
500 log_marginal,
501 }
502 }
503
504 _ => {
506 return Err(BayesError::PriorObservationMismatch {
507 prior_type: prior.type_name().to_string(),
508 obs_type: observation.type_name().to_string(),
509 });
510 }
511 };
512
513 if self.history.len() >= self.max_history && self.max_history > 0 {
515 self.history.pop_front();
516 }
517 if self.max_history > 0 {
518 self.history.push_back(posterior.clone());
519 }
520
521 Ok(posterior)
522 }
523
524 pub fn sequential_update(
531 &mut self,
532 prior: Prior,
533 observations: &[Observation],
534 ) -> Result<Posterior, BayesError> {
535 if observations.is_empty() {
536 return Err(BayesError::InvalidParameters(
537 "observations slice must not be empty".to_string(),
538 ));
539 }
540
541 let mut current_prior = prior;
542 let mut last_posterior: Option<Posterior> = None;
543
544 for obs in observations {
545 let posterior = self.update(current_prior, obs)?;
546 current_prior = posterior.updated.clone();
547 last_posterior = Some(posterior);
548 }
549
550 last_posterior.ok_or_else(|| {
552 BayesError::NumericalError("unexpected empty observation sequence".to_string())
553 })
554 }
555
556 pub fn credible_interval(
567 posterior: &Prior,
568 probability: f64,
569 ) -> Result<CredibleInterval, BayesError> {
570 if !(0.0 < probability && probability < 1.0) {
571 return Err(BayesError::InvalidParameters(format!(
572 "probability must be in (0, 1), got {probability}"
573 )));
574 }
575 posterior.validate()?;
576
577 let z = z_score(probability);
578
579 match posterior {
580 Prior::Beta { alpha, beta } => {
581 let n = alpha + beta;
582 let center = alpha / n;
583 let half_width = z * (center * (1.0 - center) / n).sqrt();
584 let lower = (center - half_width).max(0.0);
585 let upper = (center + half_width).min(1.0);
586 Ok(CredibleInterval {
587 lower,
588 upper,
589 probability,
590 })
591 }
592
593 Prior::Gaussian { mean, variance } => {
594 let half_width = z * variance.sqrt();
595 Ok(CredibleInterval {
596 lower: mean - half_width,
597 upper: mean + half_width,
598 probability,
599 })
600 }
601
602 Prior::Gamma { shape, rate } => {
603 let mean = shape / rate;
605 let std_dev = (shape / (rate * rate)).sqrt();
606 let half_width = z * std_dev;
607 let lower = (mean - half_width).max(0.0);
608 let upper = mean + half_width;
609 Ok(CredibleInterval {
610 lower,
611 upper,
612 probability,
613 })
614 }
615
616 Prior::Dirichlet { alphas } => {
617 let sum: f64 = alphas.iter().sum();
619 let max_alpha = alphas.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
620 let center = max_alpha / sum;
621 let half_width = z * (center * (1.0 - center) / (sum + 1.0)).sqrt();
622 let lower = (center - half_width).max(0.0);
623 let upper = (center + half_width).min(1.0);
624 Ok(CredibleInterval {
625 lower,
626 upper,
627 probability,
628 })
629 }
630 }
631 }
632
633 pub fn map_estimate(posterior: &Prior) -> f64 {
644 match posterior {
645 Prior::Beta { alpha, beta } => {
646 if *alpha > 1.0 && *beta > 1.0 {
647 (alpha - 1.0) / (alpha + beta - 2.0)
648 } else {
649 alpha / (alpha + beta)
650 }
651 }
652 Prior::Gaussian { mean, .. } => *mean,
653 Prior::Gamma { shape, rate } => {
654 if *shape > 1.0 {
655 (shape - 1.0) / rate
656 } else {
657 0.0
658 }
659 }
660 Prior::Dirichlet { alphas } => {
661 let sum: f64 = alphas.iter().sum();
662 let max_alpha = alphas.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
663 max_alpha / sum
664 }
665 }
666 }
667
668 pub fn kl_divergence(p: &Prior, q: &Prior) -> Result<f64, BayesError> {
676 match (p, q) {
677 (
678 Prior::Beta {
679 alpha: ap,
680 beta: bp,
681 },
682 Prior::Beta {
683 alpha: aq,
684 beta: bq,
685 },
686 ) => {
687 let psi_ap = digamma(*ap);
692 let psi_bp = digamma(*bp);
693 let psi_ap_bp = digamma(ap + bp);
694 let kl = log_beta(*aq, *bq) - log_beta(*ap, *bp)
695 + (ap - aq) * psi_ap
696 + (bp - bq) * psi_bp
697 - ((ap + bp) - (aq + bq)) * psi_ap_bp;
698 check_finite(kl, "KL(Beta‖Beta)")
699 }
700
701 (
702 Prior::Gaussian {
703 mean: mp,
704 variance: vp,
705 },
706 Prior::Gaussian {
707 mean: mq,
708 variance: vq,
709 },
710 ) => {
711 let kl = 0.5 * ((vq / vp).ln() + vp / vq + (mp - mq) * (mp - mq) / vq - 1.0);
714 check_finite(kl, "KL(Gaussian‖Gaussian)")
715 }
716
717 _ => Err(BayesError::UnsupportedOperation(format!(
718 "KL divergence not implemented for {} vs {}",
719 p.type_name(),
720 q.type_name()
721 ))),
722 }
723 }
724
725 pub fn history(&self) -> &VecDeque<Posterior> {
729 &self.history
730 }
731
732 pub fn clear_history(&mut self) {
734 self.history.clear();
735 }
736}
737
738#[cfg(test)]
743mod tests {
744 use super::{
745 digamma, ln_gamma, log_beta, log_dirichlet_norm, z_score, BayesError, BayesianUpdateEngine,
746 Observation, Prior,
747 };
748
749 #[test]
752 fn ln_gamma_integer_values() {
753 assert!((ln_gamma(1.0)).abs() < 1e-8);
756 assert!((ln_gamma(2.0)).abs() < 1e-8);
758 assert!((ln_gamma(3.0) - 2.0_f64.ln()).abs() < 1e-8);
760 assert!((ln_gamma(5.0) - 24.0_f64.ln()).abs() < 1e-7);
762 }
763
764 #[test]
765 fn ln_gamma_half() {
766 let expected = 0.5 * std::f64::consts::PI.ln();
768 assert!((ln_gamma(0.5) - expected).abs() < 1e-6);
769 }
770
771 #[test]
772 fn log_beta_symmetry() {
773 let diff = (log_beta(2.0, 5.0) - log_beta(5.0, 2.0)).abs();
775 assert!(diff < 1e-12);
776 }
777
778 #[test]
779 fn log_beta_known_value() {
780 assert!(
782 log_beta(1.0, 1.0).abs() < 1e-6,
783 "log_beta(1,1) = {}",
784 log_beta(1.0, 1.0)
785 );
786 }
787
788 #[test]
789 fn log_dirichlet_norm_two_dim_equals_log_beta() {
790 let a = 3.0_f64;
792 let b = 7.0_f64;
793 let dir = log_dirichlet_norm(&[a, b]);
794 let lb = log_beta(a, b);
795 assert!((dir - lb).abs() < 1e-10);
796 }
797
798 #[test]
799 fn digamma_known_value() {
800 let expected = -0.577_215_664_9_f64;
802 assert!((digamma(1.0) - expected).abs() < 1e-4);
803 }
804
805 #[test]
806 fn digamma_recurrence_property() {
807 let x = 4.5_f64;
809 let diff = digamma(x + 1.0) - digamma(x);
810 assert!((diff - 1.0 / x).abs() < 1e-8);
811 }
812
813 #[test]
814 fn z_score_95_percent() {
815 let z = z_score(0.95);
817 assert!((z - 1.96).abs() < 0.01);
818 }
819
820 #[test]
821 fn z_score_99_percent() {
822 let z = z_score(0.99);
824 assert!((z - 2.576).abs() < 0.01);
825 }
826
827 #[test]
830 fn beta_bernoulli_uniform_prior() {
831 let mut engine = BayesianUpdateEngine::new(10);
832 let prior = Prior::Beta {
833 alpha: 1.0,
834 beta: 1.0,
835 };
836 let obs = Observation::Bernoulli {
837 successes: 7,
838 trials: 10,
839 };
840 let post = engine
841 .update(prior, &obs)
842 .expect("test: TD update should succeed");
843 match &post.updated {
844 Prior::Beta { alpha, beta } => {
845 assert!((alpha - 8.0).abs() < 1e-10);
846 assert!((beta - 4.0).abs() < 1e-10);
847 }
848 _ => panic!("wrong variant"),
849 }
850 }
851
852 #[test]
853 fn beta_bernoulli_all_successes() {
854 let mut engine = BayesianUpdateEngine::new(10);
855 let prior = Prior::Beta {
856 alpha: 2.0,
857 beta: 3.0,
858 };
859 let obs = Observation::Bernoulli {
860 successes: 5,
861 trials: 5,
862 };
863 let post = engine
864 .update(prior, &obs)
865 .expect("test: TD update should succeed");
866 match post.updated {
867 Prior::Beta { alpha, beta } => {
868 assert!((alpha - 7.0).abs() < 1e-10);
869 assert!((beta - 3.0).abs() < 1e-10);
870 }
871 _ => panic!("wrong variant"),
872 }
873 }
874
875 #[test]
876 fn beta_bernoulli_zero_successes() {
877 let mut engine = BayesianUpdateEngine::new(10);
878 let prior = Prior::Beta {
879 alpha: 1.0,
880 beta: 1.0,
881 };
882 let obs = Observation::Bernoulli {
883 successes: 0,
884 trials: 5,
885 };
886 let post = engine
887 .update(prior, &obs)
888 .expect("test: TD update should succeed");
889 match post.updated {
890 Prior::Beta { alpha, beta } => {
891 assert!((alpha - 1.0).abs() < 1e-10);
892 assert!((beta - 6.0).abs() < 1e-10);
893 }
894 _ => panic!("wrong variant"),
895 }
896 }
897
898 #[test]
899 fn beta_bernoulli_log_marginal_finite() {
900 let mut engine = BayesianUpdateEngine::new(10);
901 let prior = Prior::Beta {
902 alpha: 2.0,
903 beta: 2.0,
904 };
905 let obs = Observation::Bernoulli {
906 successes: 3,
907 trials: 6,
908 };
909 let post = engine
910 .update(prior, &obs)
911 .expect("test: TD update should succeed");
912 assert!(post.log_marginal.is_finite());
913 }
914
915 #[test]
916 fn beta_bernoulli_successes_exceed_trials_error() {
917 let mut engine = BayesianUpdateEngine::new(10);
918 let prior = Prior::Beta {
919 alpha: 1.0,
920 beta: 1.0,
921 };
922 let obs = Observation::Bernoulli {
923 successes: 11,
924 trials: 10,
925 };
926 let result = engine.update(prior, &obs);
927 assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
928 }
929
930 #[test]
933 fn gaussian_gaussian_update_basic() {
934 let mut engine = BayesianUpdateEngine::new(10);
935 let prior = Prior::Gaussian {
936 mean: 0.0,
937 variance: 1.0,
938 };
939 let obs = Observation::Gaussian {
940 sample_mean: 2.0,
941 sample_variance: 1.0,
942 n: 1,
943 };
944 let post = engine
945 .update(prior, &obs)
946 .expect("test: TD update should succeed");
947 match post.updated {
948 Prior::Gaussian { mean, variance } => {
949 assert!((variance - 0.5).abs() < 1e-10);
951 assert!((mean - 1.0).abs() < 1e-10);
952 }
953 _ => panic!("wrong variant"),
954 }
955 }
956
957 #[test]
958 fn gaussian_gaussian_large_n_pulls_to_sample() {
959 let mut engine = BayesianUpdateEngine::new(10);
960 let prior = Prior::Gaussian {
961 mean: 0.0,
962 variance: 100.0,
963 };
964 let obs = Observation::Gaussian {
965 sample_mean: 5.0,
966 sample_variance: 1.0,
967 n: 1000,
968 };
969 let post = engine
970 .update(prior, &obs)
971 .expect("test: TD update should succeed");
972 match post.updated {
973 Prior::Gaussian { mean, .. } => {
974 assert!((mean - 5.0).abs() < 0.1);
976 }
977 _ => panic!("wrong variant"),
978 }
979 }
980
981 #[test]
982 fn gaussian_gaussian_log_marginal_finite() {
983 let mut engine = BayesianUpdateEngine::new(10);
984 let prior = Prior::Gaussian {
985 mean: 1.0,
986 variance: 2.0,
987 };
988 let obs = Observation::Gaussian {
989 sample_mean: 1.5,
990 sample_variance: 0.5,
991 n: 10,
992 };
993 let post = engine
994 .update(prior, &obs)
995 .expect("test: TD update should succeed");
996 assert!(post.log_marginal.is_finite());
997 }
998
999 #[test]
1000 fn gaussian_gaussian_zero_n_error() {
1001 let mut engine = BayesianUpdateEngine::new(10);
1002 let prior = Prior::Gaussian {
1003 mean: 0.0,
1004 variance: 1.0,
1005 };
1006 let obs = Observation::Gaussian {
1007 sample_mean: 1.0,
1008 sample_variance: 1.0,
1009 n: 0,
1010 };
1011 let result = engine.update(prior, &obs);
1012 assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1013 }
1014
1015 #[test]
1018 fn dirichlet_categorical_update_basic() {
1019 let mut engine = BayesianUpdateEngine::new(10);
1020 let prior = Prior::Dirichlet {
1021 alphas: vec![1.0, 1.0, 1.0],
1022 };
1023 let obs = Observation::Categorical {
1024 counts: vec![3, 2, 5],
1025 };
1026 let post = engine
1027 .update(prior, &obs)
1028 .expect("test: TD update should succeed");
1029 match post.updated {
1030 Prior::Dirichlet { alphas } => {
1031 assert!((alphas[0] - 4.0).abs() < 1e-10);
1032 assert!((alphas[1] - 3.0).abs() < 1e-10);
1033 assert!((alphas[2] - 6.0).abs() < 1e-10);
1034 }
1035 _ => panic!("wrong variant"),
1036 }
1037 }
1038
1039 #[test]
1040 fn dirichlet_categorical_dim_mismatch_error() {
1041 let mut engine = BayesianUpdateEngine::new(10);
1042 let prior = Prior::Dirichlet {
1043 alphas: vec![1.0, 1.0],
1044 };
1045 let obs = Observation::Categorical {
1046 counts: vec![1, 2, 3],
1047 };
1048 let result = engine.update(prior, &obs);
1049 assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1050 }
1051
1052 #[test]
1053 fn dirichlet_categorical_log_marginal_finite() {
1054 let mut engine = BayesianUpdateEngine::new(10);
1055 let prior = Prior::Dirichlet {
1056 alphas: vec![2.0, 3.0, 5.0],
1057 };
1058 let obs = Observation::Categorical {
1059 counts: vec![10, 15, 25],
1060 };
1061 let post = engine
1062 .update(prior, &obs)
1063 .expect("test: TD update should succeed");
1064 assert!(post.log_marginal.is_finite());
1065 }
1066
1067 #[test]
1068 fn dirichlet_categorical_zero_counts_no_change() {
1069 let mut engine = BayesianUpdateEngine::new(10);
1070 let prior = Prior::Dirichlet {
1071 alphas: vec![2.0, 3.0],
1072 };
1073 let obs = Observation::Categorical { counts: vec![0, 0] };
1074 let post = engine
1075 .update(prior, &obs)
1076 .expect("test: TD update should succeed");
1077 match post.updated {
1078 Prior::Dirichlet { alphas } => {
1079 assert!((alphas[0] - 2.0).abs() < 1e-10);
1080 assert!((alphas[1] - 3.0).abs() < 1e-10);
1081 }
1082 _ => panic!("wrong variant"),
1083 }
1084 }
1085
1086 #[test]
1089 fn gamma_poisson_update_basic() {
1090 let mut engine = BayesianUpdateEngine::new(10);
1091 let prior = Prior::Gamma {
1092 shape: 1.0,
1093 rate: 1.0,
1094 };
1095 let obs = Observation::Poisson {
1096 total_events: 5,
1097 total_time: 2.0,
1098 };
1099 let post = engine
1100 .update(prior, &obs)
1101 .expect("test: TD update should succeed");
1102 match post.updated {
1103 Prior::Gamma { shape, rate } => {
1104 assert!((shape - 6.0).abs() < 1e-10);
1105 assert!((rate - 3.0).abs() < 1e-10);
1106 }
1107 _ => panic!("wrong variant"),
1108 }
1109 }
1110
1111 #[test]
1112 fn gamma_poisson_log_marginal_finite() {
1113 let mut engine = BayesianUpdateEngine::new(10);
1114 let prior = Prior::Gamma {
1115 shape: 2.0,
1116 rate: 0.5,
1117 };
1118 let obs = Observation::Poisson {
1119 total_events: 10,
1120 total_time: 5.0,
1121 };
1122 let post = engine
1123 .update(prior, &obs)
1124 .expect("test: TD update should succeed");
1125 assert!(post.log_marginal.is_finite());
1126 }
1127
1128 #[test]
1129 fn gamma_poisson_zero_time_error() {
1130 let mut engine = BayesianUpdateEngine::new(10);
1131 let prior = Prior::Gamma {
1132 shape: 1.0,
1133 rate: 1.0,
1134 };
1135 let obs = Observation::Poisson {
1136 total_events: 5,
1137 total_time: 0.0,
1138 };
1139 let result = engine.update(prior, &obs);
1140 assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1141 }
1142
1143 #[test]
1146 fn mismatch_beta_gaussian_obs() {
1147 let mut engine = BayesianUpdateEngine::new(10);
1148 let prior = Prior::Beta {
1149 alpha: 1.0,
1150 beta: 1.0,
1151 };
1152 let obs = Observation::Gaussian {
1153 sample_mean: 0.5,
1154 sample_variance: 1.0,
1155 n: 10,
1156 };
1157 let result = engine.update(prior, &obs);
1158 assert!(matches!(
1159 result,
1160 Err(BayesError::PriorObservationMismatch { .. })
1161 ));
1162 }
1163
1164 #[test]
1165 fn mismatch_gaussian_bernoulli_obs() {
1166 let mut engine = BayesianUpdateEngine::new(10);
1167 let prior = Prior::Gaussian {
1168 mean: 0.0,
1169 variance: 1.0,
1170 };
1171 let obs = Observation::Bernoulli {
1172 successes: 3,
1173 trials: 5,
1174 };
1175 let result = engine.update(prior, &obs);
1176 assert!(matches!(
1177 result,
1178 Err(BayesError::PriorObservationMismatch { .. })
1179 ));
1180 }
1181
1182 #[test]
1185 fn sequential_update_equivalent_to_batch() {
1186 let mut engine = BayesianUpdateEngine::new(64);
1188 let prior = Prior::Beta {
1189 alpha: 1.0,
1190 beta: 1.0,
1191 };
1192 let obs = vec![
1193 Observation::Bernoulli {
1194 successes: 3,
1195 trials: 5,
1196 },
1197 Observation::Bernoulli {
1198 successes: 2,
1199 trials: 4,
1200 },
1201 ];
1202 let seq_post = engine
1203 .sequential_update(prior.clone(), &obs)
1204 .expect("test: should succeed");
1205
1206 let mut engine2 = BayesianUpdateEngine::new(64);
1208 let batch_obs = Observation::Bernoulli {
1209 successes: 5,
1210 trials: 9,
1211 };
1212 let batch_post = engine2
1213 .update(prior, &batch_obs)
1214 .expect("test: TD update should succeed");
1215
1216 match (&seq_post.updated, &batch_post.updated) {
1217 (
1218 Prior::Beta {
1219 alpha: a1,
1220 beta: b1,
1221 },
1222 Prior::Beta {
1223 alpha: a2,
1224 beta: b2,
1225 },
1226 ) => {
1227 assert!((a1 - a2).abs() < 1e-10);
1228 assert!((b1 - b2).abs() < 1e-10);
1229 }
1230 _ => panic!("wrong variant"),
1231 }
1232 }
1233
1234 #[test]
1235 fn sequential_update_empty_error() {
1236 let mut engine = BayesianUpdateEngine::new(10);
1237 let prior = Prior::Beta {
1238 alpha: 1.0,
1239 beta: 1.0,
1240 };
1241 let result = engine.sequential_update(prior, &[]);
1242 assert!(matches!(result, Err(BayesError::InvalidParameters(_))));
1243 }
1244
1245 #[test]
1248 fn credible_interval_beta_bounds() {
1249 let post = Prior::Beta {
1250 alpha: 8.0,
1251 beta: 4.0,
1252 };
1253 let ci =
1254 BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1255 assert!(ci.lower >= 0.0);
1256 assert!(ci.upper <= 1.0);
1257 assert!(ci.lower < ci.upper);
1258 assert!((ci.probability - 0.95).abs() < 1e-10);
1259 }
1260
1261 #[test]
1262 fn credible_interval_gaussian_symmetric() {
1263 let post = Prior::Gaussian {
1264 mean: 5.0,
1265 variance: 1.0,
1266 };
1267 let ci =
1268 BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1269 let center = (ci.lower + ci.upper) / 2.0;
1270 assert!((center - 5.0).abs() < 1e-8);
1271 let hw = (ci.upper - ci.lower) / 2.0;
1273 assert!((hw - 1.96).abs() < 0.01);
1274 }
1275
1276 #[test]
1277 fn credible_interval_invalid_probability() {
1278 let post = Prior::Gaussian {
1279 mean: 0.0,
1280 variance: 1.0,
1281 };
1282 assert!(BayesianUpdateEngine::credible_interval(&post, 0.0).is_err());
1283 assert!(BayesianUpdateEngine::credible_interval(&post, 1.0).is_err());
1284 assert!(BayesianUpdateEngine::credible_interval(&post, -0.1).is_err());
1285 }
1286
1287 #[test]
1288 fn credible_interval_gamma() {
1289 let post = Prior::Gamma {
1290 shape: 9.0,
1291 rate: 3.0,
1292 };
1293 let ci =
1294 BayesianUpdateEngine::credible_interval(&post, 0.95).expect("test: should succeed");
1295 assert!(ci.lower >= 0.0);
1297 assert!(ci.upper > ci.lower);
1298 }
1299
1300 #[test]
1301 fn credible_interval_dirichlet() {
1302 let post = Prior::Dirichlet {
1303 alphas: vec![10.0, 2.0, 3.0],
1304 };
1305 let ci =
1306 BayesianUpdateEngine::credible_interval(&post, 0.90).expect("test: should succeed");
1307 assert!(ci.lower >= 0.0);
1308 assert!(ci.upper <= 1.0);
1309 }
1310
1311 #[test]
1314 fn map_beta_mode() {
1315 let p = Prior::Beta {
1317 alpha: 3.0,
1318 beta: 3.0,
1319 };
1320 let map = BayesianUpdateEngine::map_estimate(&p);
1321 assert!((map - 0.5).abs() < 1e-10);
1322 }
1323
1324 #[test]
1325 fn map_beta_uniform_fallback() {
1326 let p = Prior::Beta {
1328 alpha: 1.0,
1329 beta: 1.0,
1330 };
1331 let map = BayesianUpdateEngine::map_estimate(&p);
1332 assert!((map - 0.5).abs() < 1e-10);
1333 }
1334
1335 #[test]
1336 fn map_gaussian_is_mean() {
1337 let p = Prior::Gaussian {
1338 mean: 3.7,
1339 variance: 2.0,
1340 };
1341 assert!((BayesianUpdateEngine::map_estimate(&p) - 3.7).abs() < 1e-10);
1342 }
1343
1344 #[test]
1345 fn map_gamma_mode() {
1346 let p = Prior::Gamma {
1348 shape: 5.0,
1349 rate: 2.0,
1350 };
1351 assert!((BayesianUpdateEngine::map_estimate(&p) - 2.0).abs() < 1e-10);
1352 }
1353
1354 #[test]
1355 fn map_gamma_shape_one_gives_zero() {
1356 let p = Prior::Gamma {
1357 shape: 1.0,
1358 rate: 2.0,
1359 };
1360 assert!((BayesianUpdateEngine::map_estimate(&p) - 0.0).abs() < 1e-10);
1361 }
1362
1363 #[test]
1364 fn map_dirichlet_argmax_proportion() {
1365 let p = Prior::Dirichlet {
1366 alphas: vec![1.0, 5.0, 2.0],
1367 };
1368 let expected = 5.0 / 8.0;
1370 assert!((BayesianUpdateEngine::map_estimate(&p) - expected).abs() < 1e-10);
1371 }
1372
1373 #[test]
1376 fn kl_beta_self_is_zero() {
1377 let p = Prior::Beta {
1378 alpha: 3.0,
1379 beta: 5.0,
1380 };
1381 let kl = BayesianUpdateEngine::kl_divergence(&p, &p).expect("test: should succeed");
1382 assert!(kl.abs() < 1e-8);
1383 }
1384
1385 #[test]
1386 fn kl_gaussian_self_is_zero() {
1387 let p = Prior::Gaussian {
1388 mean: 2.0,
1389 variance: 3.0,
1390 };
1391 let kl = BayesianUpdateEngine::kl_divergence(&p, &p).expect("test: should succeed");
1392 assert!(kl.abs() < 1e-10);
1393 }
1394
1395 #[test]
1396 fn kl_gaussian_asymmetry() {
1397 let p = Prior::Gaussian {
1398 mean: 0.0,
1399 variance: 1.0,
1400 };
1401 let q = Prior::Gaussian {
1402 mean: 1.0,
1403 variance: 2.0,
1404 };
1405 let kl_pq = BayesianUpdateEngine::kl_divergence(&p, &q).expect("test: should succeed");
1406 let kl_qp = BayesianUpdateEngine::kl_divergence(&q, &p).expect("test: should succeed");
1407 assert!((kl_pq - kl_qp).abs() > 1e-6);
1409 assert!(kl_pq >= 0.0);
1411 assert!(kl_qp >= 0.0);
1412 }
1413
1414 #[test]
1415 fn kl_known_gaussian_value() {
1416 let p = Prior::Gaussian {
1418 mean: 0.0,
1419 variance: 1.0,
1420 };
1421 let q = Prior::Gaussian {
1422 mean: 1.0,
1423 variance: 1.0,
1424 };
1425 let kl = BayesianUpdateEngine::kl_divergence(&p, &q).expect("test: should succeed");
1426 assert!((kl - 0.5).abs() < 1e-10);
1427 }
1428
1429 #[test]
1430 fn kl_unsupported_pair_error() {
1431 let p = Prior::Beta {
1432 alpha: 1.0,
1433 beta: 1.0,
1434 };
1435 let q = Prior::Gamma {
1436 shape: 1.0,
1437 rate: 1.0,
1438 };
1439 let result = BayesianUpdateEngine::kl_divergence(&p, &q);
1440 assert!(matches!(result, Err(BayesError::UnsupportedOperation(_))));
1441 }
1442
1443 #[test]
1446 fn history_bounded_by_max() {
1447 let mut engine = BayesianUpdateEngine::new(3);
1448 for i in 0..5_u64 {
1449 let prior = Prior::Beta {
1450 alpha: 1.0,
1451 beta: 1.0,
1452 };
1453 let obs = Observation::Bernoulli {
1454 successes: i % 3,
1455 trials: 5,
1456 };
1457 engine
1458 .update(prior, &obs)
1459 .expect("test: TD update should succeed");
1460 }
1461 assert_eq!(engine.history().len(), 3);
1462 }
1463
1464 #[test]
1465 fn history_clear() {
1466 let mut engine = BayesianUpdateEngine::new(10);
1467 let prior = Prior::Beta {
1468 alpha: 1.0,
1469 beta: 1.0,
1470 };
1471 let obs = Observation::Bernoulli {
1472 successes: 3,
1473 trials: 5,
1474 };
1475 engine
1476 .update(prior, &obs)
1477 .expect("test: TD update should succeed");
1478 assert!(!engine.history().is_empty());
1479 engine.clear_history();
1480 assert!(engine.history().is_empty());
1481 }
1482
1483 #[test]
1484 fn history_zero_capacity_no_store() {
1485 let mut engine = BayesianUpdateEngine::new(0);
1486 let prior = Prior::Beta {
1487 alpha: 1.0,
1488 beta: 1.0,
1489 };
1490 let obs = Observation::Bernoulli {
1491 successes: 3,
1492 trials: 5,
1493 };
1494 engine
1495 .update(prior, &obs)
1496 .expect("test: TD update should succeed");
1497 assert!(engine.history().is_empty());
1498 }
1499
1500 #[test]
1503 fn invalid_beta_alpha_zero() {
1504 let mut engine = BayesianUpdateEngine::new(10);
1505 let prior = Prior::Beta {
1506 alpha: 0.0,
1507 beta: 1.0,
1508 };
1509 let obs = Observation::Bernoulli {
1510 successes: 1,
1511 trials: 2,
1512 };
1513 assert!(matches!(
1514 engine.update(prior, &obs),
1515 Err(BayesError::InvalidParameters(_))
1516 ));
1517 }
1518
1519 #[test]
1520 fn invalid_gamma_rate_negative() {
1521 let mut engine = BayesianUpdateEngine::new(10);
1522 let prior = Prior::Gamma {
1523 shape: 1.0,
1524 rate: -1.0,
1525 };
1526 let obs = Observation::Poisson {
1527 total_events: 5,
1528 total_time: 1.0,
1529 };
1530 assert!(matches!(
1531 engine.update(prior, &obs),
1532 Err(BayesError::InvalidParameters(_))
1533 ));
1534 }
1535
1536 #[test]
1539 fn likelihood_type_labels() {
1540 let mut engine = BayesianUpdateEngine::new(10);
1541
1542 let p1 = engine
1543 .update(
1544 Prior::Beta {
1545 alpha: 1.0,
1546 beta: 1.0,
1547 },
1548 &Observation::Bernoulli {
1549 successes: 1,
1550 trials: 2,
1551 },
1552 )
1553 .expect("test: should succeed");
1554 assert_eq!(p1.likelihood_type, "Bernoulli");
1555
1556 let p2 = engine
1557 .update(
1558 Prior::Gaussian {
1559 mean: 0.0,
1560 variance: 1.0,
1561 },
1562 &Observation::Gaussian {
1563 sample_mean: 1.0,
1564 sample_variance: 1.0,
1565 n: 5,
1566 },
1567 )
1568 .expect("test: should succeed");
1569 assert_eq!(p2.likelihood_type, "Gaussian");
1570
1571 let p3 = engine
1572 .update(
1573 Prior::Dirichlet {
1574 alphas: vec![1.0, 1.0],
1575 },
1576 &Observation::Categorical { counts: vec![3, 2] },
1577 )
1578 .expect("test: should succeed");
1579 assert_eq!(p3.likelihood_type, "Categorical");
1580
1581 let p4 = engine
1582 .update(
1583 Prior::Gamma {
1584 shape: 1.0,
1585 rate: 1.0,
1586 },
1587 &Observation::Poisson {
1588 total_events: 3,
1589 total_time: 1.0,
1590 },
1591 )
1592 .expect("test: should succeed");
1593 assert_eq!(p4.likelihood_type, "Poisson");
1594 }
1595}