chromy-core 0.2.0

Chromy's sequential sampling with probability minimum replacement: serpentine sorting, multivariate size measures, expected hits, equal-CV allocation, and the multiplicity estimator
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

//! What the method promises, checked by simulation.
//!
//! The golden tests pin the crate to the reference implementation; these
//! pin it to the method itself: hits within one of the expected hits, the
//! sample size always exact, expected hits recovered in expectation, and
//! the serpentine order never jumping.

use chromy_core::*;

/// SplitMix64: a tiny deterministic generator for the tests, so the crate
/// itself stays free of dependencies and the tests free of flakiness.
struct Rng(u64);

impl Rng {
    fn next_u64(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
        z ^ (z >> 31)
    }

    fn uniform(&mut self) -> f64 {
        (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
    }

    fn uniforms(&mut self, n: usize) -> Vec<f64> {
        (0..n).map(|_| self.uniform()).collect()
    }
}

/// Expected hits that sum exactly to an integer, built from random sizes.
fn random_lambda(rng: &mut Rng, n_units: usize, n: u64) -> Vec<f64> {
    let size: Vec<f64> = (0..n_units).map(|_| rng.uniform() + 0.01).collect();
    expected_hits(&size, n, true).unwrap()
}

#[test]
fn hits_stay_within_one_of_the_expected_hits() {
    let mut rng = Rng(1);
    for trial in 0..200 {
        let n_units = 5 + (trial % 60);
        let n = 1 + (trial % 17) as u64;
        let lambda = random_lambda(&mut rng, n_units, n);
        let uniforms = rng.uniforms(n_units);
        let priority: Vec<usize> = (0..n_units).collect();
        let s = pmr(&lambda, &uniforms, &priority).unwrap();

        assert_eq!(s.hits.iter().sum::<u64>(), n, "the sample size is exact");
        for (i, &h) in s.hits.iter().enumerate() {
            let lo = s.effective_lambda[i].floor() as u64;
            let hi = s.effective_lambda[i].ceil() as u64;
            assert!(
                h == lo || h == hi,
                "unit {i}: {h} hits against expected {}",
                s.effective_lambda[i]
            );
        }
    }
}

#[test]
fn expected_hits_are_recovered_in_expectation() {
    let mut rng = Rng(2);
    let lambda = vec![0.3, 1.7, 0.5, 0.2, 0.8, 1.5];
    let n: u64 = lambda.iter().sum::<f64>().round() as u64;
    assert_eq!(n, 5);
    let priority: Vec<usize> = (0..lambda.len()).collect();

    let replications = 200_000;
    let mut mean = vec![0.0_f64; lambda.len()];
    for _ in 0..replications {
        let s = pmr(&lambda, &rng.uniforms(lambda.len()), &priority).unwrap();
        for (m, &h) in mean.iter_mut().zip(&s.hits) {
            *m += h as f64;
        }
    }
    for (i, m) in mean.iter().enumerate() {
        let observed = m / replications as f64;
        // Three standard errors of a binomial-like proportion.
        let se = (lambda[i].fract() * (1.0 - lambda[i].fract()) / replications as f64).sqrt();
        assert!(
            (observed - lambda[i]).abs() < 3.5 * se.max(1e-4),
            "unit {i}: observed {observed}, expected {}",
            lambda[i]
        );
    }
}

#[test]
fn certainty_resolution_sums_to_the_sample_size() {
    let mut rng = Rng(3);
    for trial in 0..100 {
        let n_units = 10 + (trial % 50);
        let mut size: Vec<f64> = (0..n_units).map(|_| rng.uniform() * 10.0).collect();
        size[0] = 500.0; // force a certainty
        let n = 2 + (trial % (n_units - 2)) as u64;
        let lambda = expected_hits(&size, n, false).unwrap();
        let total: f64 = lambda.iter().sum();
        assert!(
            (total - n as f64).abs() < 1e-9,
            "trial {trial}: probabilities sum to {total}, not {n}"
        );
        assert!(lambda.iter().all(|&l| (0.0..=1.0).contains(&l)));
    }
}

#[test]
fn the_serpentine_order_never_jumps_within_a_stratum() {
    // Two strata of sizes; within each stratum the serpentine must run
    // monotonically, upward in one, downward in the next.
    let stratum: Vec<String> = ["a", "a", "a", "a", "b", "b", "b", "b"]
        .iter()
        .map(|s| s.to_string())
        .collect();
    let size = vec![4.0, 2.0, 3.0, 1.0, 30.0, 10.0, 40.0, 20.0];
    let order = serpentine(&[Key::Categorical(&stratum), Key::Numeric(&size)]).unwrap();

    let listed: Vec<f64> = order.iter().map(|&u| size[u]).collect();
    assert_eq!(&listed[..4], &[1.0, 2.0, 3.0, 4.0], "first stratum ascends");
    assert_eq!(
        &listed[4..],
        &[40.0, 30.0, 20.0, 10.0],
        "second stratum descends: no seam at the boundary"
    );
}

#[test]
fn a_full_selection_is_reproducible_and_exact() {
    let mut rng = Rng(4);
    let n_units = 40;
    let size: Vec<f64> = (0..n_units).map(|_| rng.uniform() * 5.0 + 0.1).collect();
    let lambda = expected_hits(&size, 12, true).unwrap();
    let order: Vec<usize> = (0..n_units).collect();

    let u_head = rng.uniform();
    let uniforms = rng.uniforms(required_uniforms(&lambda, &order, Head::All));
    let a = select(&lambda, &order, Head::All, u_head, &uniforms).unwrap();
    let b = select(&lambda, &order, Head::All, u_head, &uniforms).unwrap();

    assert_eq!(a.hits, b.hits, "the same uniforms give the same sample");
    assert_eq!(a.hits.iter().sum::<u64>(), 12);
    assert_eq!(a.list.len(), n_units);

    // Residual mode fixes the integer part and randomises the rest.
    let lambda_mixed: Vec<f64> = lambda.iter().map(|&l| l + 1.0).collect(); // all have a fixed hit
    let needed = required_uniforms(&lambda_mixed, &order, Head::Residual);
    let c = select(
        &lambda_mixed,
        &order,
        Head::Residual,
        rng.uniform(),
        &rng.uniforms(needed),
    )
    .unwrap();
    for (p, &u) in c.list.iter().enumerate() {
        assert!(
            c.hits[p] >= lambda_mixed[u].floor() as u64,
            "a fixed hit was lost at position {p}"
        );
    }
}

#[test]
fn refusals_are_loud() {
    assert!(
        pmr(&[0.5, 0.6], &[0.1, 0.2], &[0, 1]).is_err(),
        "drift beyond tolerance"
    );
    assert!(
        pmr(&[0.5, 0.5], &[0.1], &[0, 1]).is_err(),
        "missing uniforms"
    );
    assert!(
        pmr(&[-0.5, 1.5], &[0.1, 0.2], &[0, 1]).is_err(),
        "negative expected hits"
    );
    assert!(
        expected_hits(&[1.0, 1.0], 3, false).is_err(),
        "n beyond the frame"
    );
    assert!(
        expected_hits(&[0.0, 0.0], 1, false).is_err(),
        "sizes sum to zero"
    );
    assert!(
        mpps(&[1.0, 2.0, 3.0], 2, &[1.0]).is_err(),
        "a ragged matrix"
    );
    assert!(
        head_position(&[0.0, 0.0], 0.5).is_err(),
        "weights sum to zero"
    );
}