v_utils 2.15.41

My utils crate
Documentation
use rand::{
	SeedableRng,
	distr::{Distribution, weighted::WeightedIndex},
	rngs::StdRng,
};

pub struct ReimanZeta {
	pub alpha: f64,
	pub weights: Vec<f64>,
	pub normalization_constant: f64,
}

impl ReimanZeta {
	pub fn new(alpha: f64, max_k: usize) -> ReimanZeta {
		let normalization_constant = (1..=max_k).map(|k| 1.0 / (k as f64).powf(alpha)).sum();

		let mut weights = Vec::with_capacity(max_k);
		for k in 1..=max_k {
			weights.push((k as f64).powf(-alpha) / normalization_constant);
		}

		ReimanZeta {
			alpha,
			weights,
			normalization_constant,
		}
	}

	pub fn sample(&self, seed: Option<u64>) -> usize {
		let mut rng = match seed {
			Some(s) => StdRng::seed_from_u64(s),
			None => StdRng::from_rng(&mut rand::rng()),
		};

		let dist = WeightedIndex::new(&self.weights).unwrap();
		dist.sample(&mut rng) + 1
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::internal_utils::simplified_snapshots::SnapshotP;

	#[test]
	fn test_reiman_zeta() {
		let alpha = 1.0;
		let max_k = 1000;
		let zeta = ReimanZeta::new(alpha, max_k);

		let mut samples = (0..1000).map(|i| zeta.sample(Some(i)) as u32).collect::<Vec<u32>>();
		samples.sort_by(|a, b| b.cmp(a));
		let plot = SnapshotP::build(&samples).draw();
		insta::assert_snapshot!(plot, @r###"
  █                                                                                         991
  ██▄                                                                                          
  ████▄                                                                                        
  █████▇▂                                                                                      
  ███████▅                                                                                     
  ████████▇                                                                                    
  ██████████▄▂                                                                                 
  ████████████▇▄▁                                                                              
  ███████████████▅▂▁                                                                           
  ██████████████████▇▄▃▁                                                                       
  ██████████████████████▇▆▅▄▃▂▁                                                                
  ██████████████████████████████▇▇▇▆▆▅▅▄▄▄▄▃▃▃▃▃▃▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁1
  "###);
	}
}