pop_prob/
lib.rs

1use crate::fac::Fac;
2use crate::pop_prob_error::PopProbError;
3use crate::snsk::Snsk;
4use num_bigint::ToBigUint;
5use num_traits::cast::ToPrimitive;
6
7mod fac;
8mod pop_prob_error;
9mod snsk;
10
11pub struct Calculator {
12    fac: Fac,
13    snsk: Snsk,
14}
15
16impl Calculator {
17    pub fn new() -> Self {
18        Self {
19            fac: Fac::new(),
20            snsk: Snsk::new(),
21        }
22    }
23
24    /// Calculates the most likely population size if `sample` elements are selected randomly
25    /// with replacement from the population and `unique` of these elements are unique.
26    pub fn pop(&mut self, sample: u32, unique: u32) -> Result<u32, PopProbError> {
27        let mut prev_prev = None;
28        let mut prev: Option<(u32, f64)> = None;
29        for pop in (1_u32..).map(|x| unique + 2_u32.pow(x - 1) - 1) {
30            let v = (pop, self.prob_unique(pop, sample, unique)?);
31            if let Some(prev) = prev {
32                if v.1 < prev.1 {
33                    return if let Some(prev_prev) = prev_prev {
34                        self.find_peak(sample, unique, prev_prev, prev, v)
35                    } else {
36                        Ok(prev.0)
37                    };
38                }
39            }
40            prev_prev = prev;
41            prev = Some(v);
42        }
43        Err(PopProbError::Generic)
44    }
45
46    fn find_peak(
47        &mut self,
48        sample: u32,
49        unique: u32,
50        lbound: (u32, f64),
51        pivot: (u32, f64),
52        ubound: (u32, f64),
53    ) -> Result<u32, PopProbError> {
54        if ubound.0 == pivot.0 + 1 {
55            return Ok(pivot.0);
56        }
57        let (lbound, pivot, ubound) = if ubound.0 - pivot.0 > pivot.0 - lbound.0 {
58            let x = (pivot.0 + ubound.0) / 2;
59            let v = (x, self.prob_unique(x, sample, unique)?);
60            if v.1 < pivot.1 {
61                (lbound, pivot, v)
62            } else {
63                (pivot, v, ubound)
64            }
65        } else {
66            let x = (lbound.0 + pivot.0) / 2;
67            let v = (x, self.prob_unique(x, sample, unique)?);
68            if pivot.1 < v.1 {
69                (lbound, v, pivot)
70            } else {
71                (v, pivot, ubound)
72            }
73        };
74        self.find_peak(sample, unique, lbound, pivot, ubound)
75    }
76
77    /// Calculates the probability that a population contains exactly `size` unique elements if
78    /// `sample` elements are selected randomly with replacement from the population and
79    /// `unique` of these elements are unique.
80    pub fn prob_pop(&mut self, sample: u32, unique: u32, size: u32) -> Result<f64, PopProbError> {
81        let numerator = self.prob_unique(size, sample, unique)?;
82        let mut denominator = numerator;
83        let mut latest_prob = numerator;
84        let mut low_size = size - 1;
85        let mut low_prob = self.prob_unique(low_size, sample, unique)?;
86        let mut high_size = size + 1;
87        let mut high_prob = self.prob_unique(high_size, sample, unique)?;
88        while latest_prob * 10000.0 > denominator {
89            if low_prob > high_prob {
90                latest_prob = low_prob;
91                denominator += low_prob;
92                low_size -= 1;
93                low_prob = if low_size < unique {
94                    0.0
95                } else {
96                    self.prob_unique(low_size, sample, unique)?
97                };
98            } else {
99                latest_prob = high_prob;
100                denominator += high_prob;
101                high_size += 1;
102                high_prob = self.prob_unique(high_size, sample, unique)?;
103            };
104        }
105        Ok(numerator / denominator)
106    }
107
108    /// Calculates the probability that `unique` unique elements will be selected after
109    /// `sample` elements are selected randomly with replacement
110    /// from a population containing exactly `size` unique elements
111    pub fn prob_unique(
112        &mut self,
113        size: u32,
114        sample: u32,
115        unique: u32,
116    ) -> Result<f64, PopProbError> {
117        if unique > size {
118            return Err(PopProbError::UniqueGreaterThanSize((unique, size)));
119        }
120        if unique > sample {
121            return Err(PopProbError::UniqueGreaterThanSample((unique, sample)));
122        }
123        if unique == 0 {
124            return Err(PopProbError::UniqueZero);
125        }
126
127        let numerator = {
128            let snsk = self.snsk.calc(sample, unique)?;
129            let fac = self.fac.fac(size).clone();
130            snsk * fac
131        };
132
133        let denominator = {
134            let pow = size.to_biguint().unwrap().pow(sample as u32);
135            let fac = self.fac.fac(size - unique).clone();
136            pow * fac
137        };
138
139        Ok((numerator * u32::max_value() / denominator)
140            .to_f64()
141            .unwrap()
142            / u32::max_value() as f64)
143    }
144}
pop_prob/lib.rs

pop_prob/
lib.rs
