Add KS test for IndexedRandom::choose_weighted

distr_test/tests/weighted.rs

@@ -8,7 +8,8 @@
 
 mod ks;
 use ks::test_discrete;
-use rand::distr::WeightedIndex;
+use rand::distr::{Distribution, WeightedIndex};
+use rand::seq::IndexedRandom;
 use rand_distr::{WeightedAliasIndex, WeightedTreeIndex};
 
 fn make_cdf(num: usize, f: impl Fn(i64) -> f64) -> impl Fn(i64) -> f64 {
@@ -75,3 +76,24 @@ fn weighted_tree_index() {
     test_weights(100, |i| (i as f64).powi(3));
     test_weights(100, |i| 1.0 / ((i + 1) as f64));
 }
+
+#[test]
+fn choose_weighted_indexed() {
+    struct Adapter<F: Fn(i64) -> f64>(Vec<i64>, F);
+    impl<F: Fn(i64) -> f64> Distribution<i64> for Adapter<F> {
+        fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> i64 {
+            *IndexedRandom::choose_weighted(&self.0[..], rng, |i| (self.1)(*i)).unwrap()
+        }
+    }
+
+    fn test_weights(num: usize, weight: impl Fn(i64) -> f64) {
+        let distr = Adapter((0..num).map(|i| i as i64).collect(), &weight);
+        test_discrete(0, distr, make_cdf(num, |i| weight(i)));
+    }
+
+    test_weights(100, |_| 1.0);
+    test_weights(100, |i| ((i + 1) as f64).ln());
+    test_weights(100, |i| i as f64);
+    test_weights(100, |i| (i as f64).powi(3));
+    test_weights(100, |i| 1.0 / ((i + 1) as f64));
+}