Skip to main content

lawkit_core/generate/
poisson.rs

1use 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        // Inject fraud if specified (add non-Poisson patterns)
38        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    // Fraud: add artificially high values or clustering
50    for _ in 0..fraud_count {
51        let index = rng.gen_range(0..numbers.len());
52
53        if rng.gen_bool(0.5) {
54            // Add artificially high values
55            numbers[index] = rng.gen_range(50..100);
56        } else {
57            // Force clustering around specific values
58            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        // Check that mean is approximately lambda
76        let mean = result.iter().sum::<u32>() as f64 / result.len() as f64;
77        assert!((mean - 2.5).abs() < 0.5);
78
79        // Poisson distribution should have variance ≈ mean
80        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}