Skip to main content

use_probability/
bernoulli.rs

1use crate::Probability;
2
3/// A Bernoulli distribution with a single success probability.
4#[derive(Debug, Clone, Copy, PartialEq)]
5pub struct Bernoulli {
6    success: Probability,
7}
8
9impl Bernoulli {
10    /// Creates a Bernoulli model from a validated success probability.
11    #[must_use]
12    pub const fn new(success: Probability) -> Self {
13        Self { success }
14    }
15
16    /// Returns the success probability.
17    #[must_use]
18    pub const fn success_probability(&self) -> Probability {
19        self.success
20    }
21
22    /// Returns the failure probability.
23    #[must_use]
24    pub const fn failure_probability(&self) -> Probability {
25        self.success.complement()
26    }
27
28    /// Returns the mean of the Bernoulli model.
29    #[must_use]
30    pub const fn mean(&self) -> f64 {
31        self.success.value()
32    }
33
34    /// Returns the variance of the Bernoulli model.
35    #[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    /// Returns the probability mass for `success`.
44    #[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}