use rand::{seq::SliceRandom, Rng, SeedableRng};
use rand_chacha;
use std::collections::HashMap;
pub trait EqualProbability<T>
where
T: Copy,
{
fn new(samp_size: usize, with_replacement: bool, pps_method: Option<PPS>) -> Self;
fn select(&mut self, pop: &Population<T>, seed: Option<u64>);
fn calculate_probs(&mut self);
}
pub trait UnequalProbability<T>
where
T: Copy,
{
fn new(samp_size: HashMap<T, usize>, with_replacement: bool, pps_method: Option<PPS>) -> Self;
fn select(&mut self, pop: &Population<T>, seed: Option<u64>);
fn calculate_probs(&mut self);
}
#[derive(Debug)]
pub struct Population<T>
where
T: Copy,
{
pub y: Vec<T>,
pub stratum: Option<Vec<T>>,
pub mos: Option<Vec<f64>>,
}
impl<T> Population<T>
where
T: Copy,
{
pub fn new(y: Vec<T>, stratum: Option<Vec<T>>, mos: Option<Vec<f64>>) -> Self {
Self { y, stratum, mos }
}
}
pub struct UnstratifiedSample<T> {
pub index: Vec<usize>,
pub y: Vec<T>,
pub mos: Vec<f64>,
pub samp_size: usize,
pub pop_size: Option<usize>,
pub probs: Option<Vec<f64>>,
pub fpc: f64,
pub with_replacement: bool,
pub pps_method: Option<PPS>,
pub seed: Option<u64>,
}
pub struct StratifiedSample<T> {
pub index: Option<Vec<usize>>,
pub y: Option<Vec<T>>,
pub stratum: Option<Vec<T>>,
pub mos: Option<Vec<f64>>,
pub samp_size: HashMap<T, usize>,
pub pop_size: HashMap<T, usize>,
pub probs: Vec<f64>,
pub fpc: HashMap<T, f64>,
pub with_replacement: bool,
pub pps_method: Option<PPS>,
pub seed: Option<u64>,
}
pub enum PPS {
Systematic,
Brewer,
Murphy,
RaoSampford,
HanuravVijayan,
}
impl<T> EqualProbability<T> for UnstratifiedSample<T>
where
T: Copy,
{
fn new(samp_size: usize, with_replacement: bool, pps_method: Option<PPS>) -> Self {
Self {
index: Vec::with_capacity(samp_size),
y: Vec::with_capacity(samp_size),
mos: Vec::with_capacity(samp_size),
samp_size,
pop_size: None,
probs: None,
fpc: 0.0,
with_replacement,
pps_method,
seed: None,
}
}
fn select(&mut self, pop: &Population<T>, seed: Option<u64>)
where
T: Copy,
{
self.pop_size = Some(pop.y.len());
self.fpc = self.samp_size as f64 / self.pop_size.unwrap() as f64;
self.seed = seed;
let mut rng = match self.seed {
Some(x) => rand_chacha::ChaCha12Rng::seed_from_u64(x),
None => rand_chacha::ChaCha12Rng::from_entropy(),
};
if self.with_replacement {
for i in 0..self.samp_size {
self.index.push(rng.gen_range(0, self.pop_size.unwrap()));
self.y.push(pop.y[self.index[i]]);
}
} else {
let mut pop_indexes: Vec<usize> = (0..self.pop_size.unwrap()).collect();
pop_indexes.shuffle(&mut rng);
for i in 0..self.samp_size {
self.index.push(pop_indexes[i]);
self.y.push(pop.y[self.index[i]]);
}
}
}
fn calculate_probs(&mut self) {
self.probs = match self.with_replacement {
true => Some(vec![
self.samp_size as f64 / self.pop_size.unwrap() as f64;
self.samp_size
]),
false => {
let mut probs_without_replacement: Vec<f64> = Vec::with_capacity(self.samp_size);
for i in 0..self.samp_size {
probs_without_replacement
.push((self.samp_size - i) as f64 / self.pop_size.unwrap() as f64)
}
Some(probs_without_replacement)
}
};
}
}