use_probability/
bernoulli.rs1use crate::Probability;
2
3#[derive(Debug, Clone, Copy, PartialEq)]
5pub struct Bernoulli {
6 success: Probability,
7}
8
9impl Bernoulli {
10 #[must_use]
12 pub const fn new(success: Probability) -> Self {
13 Self { success }
14 }
15
16 #[must_use]
18 pub const fn success_probability(&self) -> Probability {
19 self.success
20 }
21
22 #[must_use]
24 pub const fn failure_probability(&self) -> Probability {
25 self.success.complement()
26 }
27
28 #[must_use]
30 pub const fn mean(&self) -> f64 {
31 self.success.value()
32 }
33
34 #[must_use]
36 pub fn variance(&self) -> f64 {
37 let success = self.success.value();
38 let failure = self.failure_probability().value();
39
40 success * failure
41 }
42
43 #[must_use]
45 pub const fn pmf(&self, success: bool) -> Probability {
46 if success {
47 self.success
48 } else {
49 self.failure_probability()
50 }
51 }
52}
53
54#[cfg(test)]
55mod tests {
56 use super::Bernoulli;
57 use crate::{Probability, ProbabilityError};
58
59 fn assert_close(left: f64, right: f64, tolerance: f64) {
60 assert!(
61 (left - right).abs() <= tolerance,
62 "expected {left} to be within {tolerance} of {right}"
63 );
64 }
65
66 #[test]
67 fn exposes_success_and_failure_probabilities() -> Result<(), ProbabilityError> {
68 let model = Bernoulli::new(Probability::from_fraction(1, 4)?);
69
70 assert_eq!(model.success_probability(), Probability::try_new(0.25)?);
71 assert_eq!(model.failure_probability(), Probability::try_new(0.75)?);
72
73 Ok(())
74 }
75
76 #[test]
77 fn computes_mean_variance_and_mass_values() -> Result<(), ProbabilityError> {
78 let model = Bernoulli::new(Probability::from_fraction(1, 4)?);
79
80 assert_close(model.mean(), 0.25, 1.0e-12);
81 assert_close(model.variance(), 0.1875, 1.0e-12);
82 assert_eq!(model.pmf(true), Probability::try_new(0.25)?);
83 assert_eq!(model.pmf(false), Probability::try_new(0.75)?);
84
85 Ok(())
86 }
87}