sklears_datasets/generators/
experimental.rs

1//! Experimental design and A/B testing simulation
2//!
3//! This module contains generators for experimental design datasets including
4//! A/B testing simulations, factorial designs, and causal inference experiments.
5
6use scirs2_core::ndarray::{Array1, Array2};
7use scirs2_core::random::prelude::*;
8use scirs2_core::random::rngs::StdRng;
9use scirs2_core::random::{Normal, RngExt};
10use sklears_core::error::{Result, SklearsError};
11
12/// Result type for A/B testing simulation: (features, group, outcomes, conversion)
13type AbTestingResult = (Array2<f64>, Array1<i32>, Array1<f64>, Array1<i32>);
14
15/// A/B test configuration
16#[derive(Debug, Clone)]
17pub struct ABTestConfig {
18    pub control_rate: f64,
19    pub treatment_effect: f64,
20    pub significance_level: f64,
21    pub power: f64,
22    pub minimum_detectable_effect: f64,
23}
24
25/// Generate A/B testing simulation data
26#[allow(non_snake_case)] // X follows mathematical convention for feature matrix
27pub fn make_ab_testing_simulation(
28    config: ABTestConfig,
29    n_samples: usize,
30    n_features: usize,
31    random_state: Option<u64>,
32) -> Result<AbTestingResult> {
33    if n_samples == 0 || n_features == 0 {
34        return Err(SklearsError::InvalidInput(
35            "n_samples and n_features must be positive".to_string(),
36        ));
37    }
38
39    if config.control_rate < 0.0 || config.control_rate > 1.0 {
40        return Err(SklearsError::InvalidInput(
41            "control_rate must be in [0, 1]".to_string(),
42        ));
43    }
44
45    let mut rng = if let Some(seed) = random_state {
46        StdRng::seed_from_u64(seed)
47    } else {
48        StdRng::from_rng(&mut scirs2_core::random::thread_rng())
49    };
50
51    // Generate user features
52    let mut X = Array2::zeros((n_samples, n_features));
53    let normal = Normal::new(0.0, 1.0).expect("operation should succeed");
54
55    for i in 0..n_samples {
56        for j in 0..n_features {
57            X[[i, j]] = rng.sample(normal);
58        }
59    }
60
61    // Assign users to control (0) or treatment (1) groups
62    let mut group_assignment = Array1::zeros(n_samples);
63    let n_treatment = (n_samples as f64 * (1.0 - config.control_rate)) as usize;
64
65    // Random assignment with stratification based on features
66    let mut indices: Vec<usize> = (0..n_samples).collect();
67    indices.shuffle(&mut rng);
68
69    for (i, &idx) in indices.iter().enumerate() {
70        group_assignment[idx] = if i < n_treatment { 1 } else { 0 };
71    }
72
73    // Generate outcomes based on group assignment and user features
74    let mut outcomes = Array1::zeros(n_samples);
75    let mut conversion = Array1::zeros(n_samples);
76
77    for i in 0..n_samples {
78        let user_propensity =
79            X.slice(scirs2_core::ndarray::s![i, ..]).iter().sum::<f64>() / n_features as f64;
80        let base_rate = 0.1 + 0.05 * user_propensity.tanh(); // Base conversion rate
81
82        let conversion_probability = if group_assignment[i] == 1 {
83            // Treatment group gets the effect
84            base_rate + config.treatment_effect
85        } else {
86            // Control group
87            base_rate
88        };
89
90        let converted = rng.random::<f64>() < conversion_probability.clamp(0.0, 1.0);
91        conversion[i] = if converted { 1 } else { 0 };
92
93        // Outcome value (e.g., revenue) - higher for conversions
94        outcomes[i] = if converted {
95            let baseline_value = 10.0 + 5.0 * user_propensity;
96            let treatment_bonus = if group_assignment[i] == 1 {
97                config.treatment_effect * 20.0
98            } else {
99                0.0
100            };
101            baseline_value
102                + treatment_bonus
103                + rng.sample(Normal::new(0.0, 2.0).expect("sampling should succeed"))
104        } else {
105            0.0
106        };
107    }
108
109    Ok((X, group_assignment, outcomes, conversion))
110}
111
112#[allow(non_snake_case)]
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_make_ab_testing_simulation() {
119        let config = ABTestConfig {
120            control_rate: 0.5,
121            treatment_effect: 0.1,
122            significance_level: 0.05,
123            power: 0.8,
124            minimum_detectable_effect: 0.05,
125        };
126
127        let (X, groups, outcomes, conversions) =
128            make_ab_testing_simulation(config, 200, 3, Some(42)).expect("operation should succeed");
129
130        assert_eq!(X.shape(), &[200, 3]);
131        assert_eq!(groups.len(), 200);
132        assert_eq!(outcomes.len(), 200);
133        assert_eq!(conversions.len(), 200);
134
135        // Check group assignments are 0 or 1
136        for &group in groups.iter() {
137            assert!(group == 0 || group == 1, "Groups should be 0 or 1");
138        }
139
140        // Check conversions are 0 or 1
141        for &conv in conversions.iter() {
142            assert!(conv == 0 || conv == 1, "Conversions should be 0 or 1");
143        }
144
145        // Should have roughly balanced groups
146        let treatment_count = groups.iter().filter(|&&g| g == 1).count();
147        assert!(
148            treatment_count > 50 && treatment_count < 150,
149            "Groups should be roughly balanced"
150        );
151    }
152
153    #[test]
154    fn test_ab_testing_simulation_invalid_input() {
155        let config = ABTestConfig {
156            control_rate: 1.5, // Invalid
157            treatment_effect: 0.1,
158            significance_level: 0.05,
159            power: 0.8,
160            minimum_detectable_effect: 0.05,
161        };
162
163        assert!(make_ab_testing_simulation(config, 200, 3, Some(42)).is_err());
164    }
165}
sklears_datasets/generators/experimental.rs

sklears_datasets/generators/
experimental.rs
