lawkit_core/generate/
poisson.rs1use super::{DataGenerator, GenerateConfig};
2use crate::error::Result;
3use rand::prelude::*;
4use rand_distr::{Distribution, Poisson};
5
6#[derive(Debug, Clone)]
7pub struct PoissonGenerator {
8 pub lambda: f64,
9 pub time_series: bool,
10}
11
12impl PoissonGenerator {
13 pub fn new(lambda: f64, time_series: bool) -> Self {
14 Self {
15 lambda,
16 time_series,
17 }
18 }
19}
20
21impl DataGenerator for PoissonGenerator {
22 type Output = Vec<u32>;
23
24 fn generate(&self, config: &GenerateConfig) -> Result<Self::Output> {
25 let mut rng = config.create_rng();
26 let mut numbers = Vec::with_capacity(config.samples);
27
28 let poisson = Poisson::new(self.lambda).map_err(|e| {
29 crate::error::BenfError::ParseError(format!("Invalid lambda parameter: {e}"))
30 })?;
31
32 for _ in 0..config.samples {
33 let value = poisson.sample(&mut rng) as u32;
34 numbers.push(value);
35 }
36
37 if config.fraud_rate > 0.0 {
39 inject_poisson_fraud(&mut numbers, config.fraud_rate, &mut rng);
40 }
41
42 Ok(numbers)
43 }
44}
45
46fn inject_poisson_fraud(numbers: &mut [u32], fraud_rate: f64, rng: &mut impl Rng) {
47 let fraud_count = (numbers.len() as f64 * fraud_rate) as usize;
48
49 for _ in 0..fraud_count {
51 let index = rng.gen_range(0..numbers.len());
52
53 if rng.gen_bool(0.5) {
54 numbers[index] = rng.gen_range(50..100);
56 } else {
57 numbers[index] = if rng.gen_bool(0.3) { 0 } else { 1 };
59 }
60 }
61}
62
63#[cfg(test)]
64mod tests {
65 use super::*;
66
67 #[test]
68 fn test_poisson_generator() {
69 let generator = PoissonGenerator::new(2.5, false);
70 let config = GenerateConfig::new(1000).with_seed(42);
71
72 let result = generator.generate(&config).unwrap();
73 assert_eq!(result.len(), 1000);
74
75 let mean = result.iter().sum::<u32>() as f64 / result.len() as f64;
77 assert!((mean - 2.5).abs() < 0.5);
78
79 let variance = result
81 .iter()
82 .map(|&x| (x as f64 - mean).powi(2))
83 .sum::<f64>()
84 / result.len() as f64;
85
86 assert!((variance - mean).abs() < 1.0);
87 }
88}