use std::collections::HashSet;
use rand::Rng;
use rand::SeedableRng;
use rand::rngs::SmallRng;
use rayon::prelude::*;
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
pub struct BootstrapCI {
pub lower: f64,
pub upper: f64,
pub confidence: f64,
pub prob_positive: f64,
pub std_err: f64,
}
pub fn bootstrap_ci<T, F>(
units: &[T],
n_samples: usize,
seed: u64,
confidence: f64,
statistic: F,
) -> Vec<BootstrapCI>
where
T: Copy + Eq + std::hash::Hash + Sync,
F: Fn(&HashSet<T>) -> Vec<f64> + Sync,
{
assert!(
confidence > 0.0 && confidence < 1.0,
"bootstrap_ci: confidence must be in (0, 1), got {confidence} — \
pass a fraction such as 0.95, not a percentage"
);
let n_units = units.len();
if n_samples == 0 || n_units == 0 {
return Vec::new();
}
let all_samples: Vec<Vec<f64>> = (0..n_samples)
.into_par_iter()
.map(|i| {
let mut rng = SmallRng::seed_from_u64(seed.wrapping_add(i as u64));
let sample: HashSet<T> = (0..n_units)
.map(|_| units[rng.random_range(0..n_units)])
.collect();
statistic(&sample)
})
.collect();
let num_stats = all_samples.first().map_or(0, Vec::len);
let alpha = 1.0 - confidence;
let nb = n_samples;
(0..num_stats)
.map(|m| {
let mut samples: Vec<f64> = all_samples.iter().map(|s| s[m]).collect();
samples.sort_by(f64::total_cmp);
let lo_idx = ((alpha / 2.0) * nb as f64).floor() as usize;
let hi_idx = ((1.0 - alpha / 2.0) * nb as f64).ceil() as usize;
let lower = samples[lo_idx.min(nb - 1)];
let upper = samples[hi_idx.min(nb - 1)];
let pos_count = samples.iter().filter(|&&x| x > 0.0).count();
let prob_positive = pos_count as f64 / nb as f64;
let mean: f64 = samples.iter().sum::<f64>() / nb as f64;
let variance = if nb > 1 {
samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (nb - 1) as f64
} else {
0.0
};
BootstrapCI {
lower,
upper,
confidence,
prob_positive,
std_err: variance.sqrt(),
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn units(n: usize) -> Vec<usize> {
(0..n).collect()
}
#[test]
fn interval_brackets_the_population_mean() {
let values: Vec<f64> = (0..200).map(|i| i as f64).collect();
let truth = values.iter().sum::<f64>() / values.len() as f64;
let cis = bootstrap_ci(&units(values.len()), 500, 42, 0.95, |sample| {
let sum: f64 = sample.iter().map(|&i| values[i]).sum();
vec![sum / sample.len() as f64]
});
assert_eq!(cis.len(), 1);
assert!(
cis[0].lower <= truth && truth <= cis[0].upper,
"truth {truth} outside [{}, {}]",
cis[0].lower,
cis[0].upper
);
}
#[test]
fn statistic_receives_the_units_themselves() {
let image_ids: Vec<u64> = vec![101, 202, 303, 404];
let cis = bootstrap_ci(&image_ids, 20, 1, 0.95, |sample| {
assert!(sample.iter().all(|id| image_ids.contains(id)));
vec![sample.len() as f64]
});
assert_eq!(cis.len(), 1);
}
#[test]
fn same_seed_reproduces_the_same_interval() {
let stat = |sample: &HashSet<usize>| vec![sample.len() as f64];
let a = bootstrap_ci(&units(100), 50, 7, 0.9, stat);
let b = bootstrap_ci(&units(100), 50, 7, 0.9, stat);
assert_eq!(a[0].lower, b[0].lower);
assert_eq!(a[0].upper, b[0].upper);
assert_eq!(a[0].std_err, b[0].std_err);
}
#[test]
fn different_seed_gives_a_different_draw() {
let stat = |sample: &HashSet<usize>| vec![sample.len() as f64];
let a = bootstrap_ci(&units(1000), 50, 1, 0.9, stat);
let b = bootstrap_ci(&units(1000), 50, 2, 0.9, stat);
assert_ne!(a[0].std_err, b[0].std_err);
}
#[test]
fn sample_holds_roughly_63_percent_of_units() {
let cis = bootstrap_ci(&units(10_000), 20, 99, 0.95, |sample| {
vec![sample.len() as f64 / 10_000.0]
});
assert!(
cis[0].lower > 0.60 && cis[0].upper < 0.66,
"expected ~0.632, got [{}, {}]",
cis[0].lower,
cis[0].upper
);
}
#[test]
fn always_positive_statistic_reports_probability_one() {
let cis = bootstrap_ci(&units(50), 100, 3, 0.95, |_| vec![1.0]);
assert_eq!(cis[0].prob_positive, 1.0);
assert_eq!(cis[0].std_err, 0.0);
assert_eq!(cis[0].lower, 1.0);
}
#[test]
fn vector_statistic_yields_one_interval_per_entry() {
let cis = bootstrap_ci(&units(50), 30, 5, 0.95, |s| {
vec![s.len() as f64, -(s.len() as f64), 0.0]
});
assert_eq!(cis.len(), 3);
assert_eq!(cis[1].prob_positive, 0.0);
assert_eq!(cis[2].prob_positive, 0.0);
}
#[test]
fn degenerate_inputs_return_empty_not_a_panic() {
assert!(bootstrap_ci(&units(0), 10, 1, 0.95, |_| vec![1.0]).is_empty());
assert!(bootstrap_ci(&units(10), 0, 1, 0.95, |_| vec![1.0]).is_empty());
}
#[test]
#[should_panic(expected = "confidence must be in (0, 1)")]
fn percentage_confidence_panics_instead_of_min_max() {
bootstrap_ci(&units(10), 10, 1, 95.0, |_| vec![1.0]);
}
#[test]
#[should_panic(expected = "confidence must be in (0, 1)")]
fn confidence_bounds_are_exclusive() {
bootstrap_ci(&units(10), 10, 1, 1.0, |_| vec![1.0]);
}
#[test]
fn nan_statistic_does_not_panic_the_percentile_sort() {
let cis = bootstrap_ci(&units(50), 40, 9, 0.95, |sample| {
let n = sample.len() as f64;
vec![if sample.len() % 3 == 0 { f64::NAN } else { n }]
});
assert_eq!(cis.len(), 1);
assert!(cis[0].lower.is_finite());
}
}