samplics 0.1.0

A rust library for fitting statistical models
Documentation
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 stratum: Option<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 stratified: 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 stratified: 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),
            // stratum: None,
            mos: Vec::with_capacity(samp_size),
            samp_size,
            pop_size: None,
            probs: None,
            fpc: 0.0,
            with_replacement,
            // stratified: false,
            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)
            }
        };
    }
}