chromy_core/
allocate.rs

1//! From sizes to expected hits, and the allocation that equalises CVs.
2//!
3//! A size measure is not yet a design. [`expected_hits`] turns sizes into
4//! the expected hits the selection consumes: proportional to size, with
5//! units so large that proportionality would push them past one resolved
6//! into certainties by iteration — or, under minimum replacement, left
7//! above one and allowed several hits. [`eqcv`] closes the loop for
8//! multivariate designs: it searches for the per-variable sample sizes
9//! whose maximal measure gives every auxiliary variable the same
10//! anticipated coefficient of variation, so no estimate is quietly
11//! sacrificed to another.
12
13use crate::error::Error;
14use crate::size::mpps;
15
16/// Expected hits proportional to size.
17///
18/// `n` is the sample size. With `min_replacement` false (the usual case),
19/// units whose proportional share reaches one are made certainties and the
20/// remainder is re-shared until no unit crosses one; the result is a vector
21/// of inclusion probabilities in `[0, 1]`, and `n` may not exceed the
22/// number of units with positive size. With `min_replacement` true the
23/// expected hits are simply `n` times the size shares, values above one
24/// meaning several hits, which is what Chromy's procedure is for.
25pub fn expected_hits(size: &[f64], n: u64, min_replacement: bool) -> Result<Vec<f64>, Error> {
26    for &s in size {
27        if !s.is_finite() || s < 0.0 {
28            return Err(Error::Input(format!(
29                "sizes must be finite and non-negative; found {s}"
30            )));
31        }
32    }
33    let total: f64 = size.iter().sum();
34    if total <= 0.0 {
35        return Err(Error::Input("the sizes sum to zero".into()));
36    }
37    if n < 1 {
38        return Err(Error::Input("the sample size must be at least one".into()));
39    }
40    let positive = size.iter().filter(|&&s| s > 0.0).count() as u64;
41    if !min_replacement && n > positive {
42        return Err(Error::Input(format!(
43            "a sample of {n} without replacement needs at least {n} units of \
44             positive size; the frame has {positive}"
45        )));
46    }
47
48    let n_f = n as f64;
49    if min_replacement {
50        return Ok(size.iter().map(|&s| n_f * s / total).collect());
51    }
52    if n == positive {
53        return Ok(size.iter().map(|&s| f64::from(s > 0.0)).collect());
54    }
55
56    let n_units = size.len();
57    let mut certain = vec![false; n_units];
58    let mut lambda = vec![0.0_f64; n_units];
59    loop {
60        let n_certain = certain.iter().filter(|&&c| c).count();
61        let remaining = n_f - n_certain as f64;
62        if remaining <= 0.0 && certain.iter().any(|&c| !c) {
63            return Err(Error::Internal(
64                "certainty resolution consumed the whole sample".into(),
65            ));
66        }
67        let rest_total: f64 = size
68            .iter()
69            .zip(&certain)
70            .filter(|&(_, &c)| !c)
71            .map(|(&s, _)| s)
72            .sum();
73        for i in 0..n_units {
74            lambda[i] = if certain[i] {
75                1.0
76            } else if remaining > 0.0 {
77                remaining * size[i] / rest_total
78            } else {
79                0.0
80            };
81        }
82        let mut grew = false;
83        for i in 0..n_units {
84            if !certain[i] && lambda[i] >= 1.0 {
85                certain[i] = true;
86                grew = true;
87            }
88        }
89        if !grew {
90            break;
91        }
92    }
93    Ok(lambda.iter().map(|&l| l.min(1.0)).collect())
94}
95
96/// How [`eqcv`] is to search.
97#[derive(Debug, Clone, Copy)]
98pub struct EqcvOptions {
99    /// The most iterations the search may take.
100    pub max_iterations: usize,
101    /// Convergence: the largest anticipated CV may exceed the smallest by
102    /// at most this relative amount.
103    pub tolerance: f64,
104    /// Damping of each multiplicative step, in `(0, 1]`; one is undamped.
105    pub damping: f64,
106}
107
108impl Default for EqcvOptions {
109    fn default() -> Self {
110        EqcvOptions {
111            max_iterations: 300,
112            tolerance: 1e-3,
113            damping: 0.5,
114        }
115    }
116}
117
118/// The outcome of an [`eqcv`] search.
119#[derive(Debug, Clone)]
120pub struct EqcvAllocation {
121    /// The per-variable sample-size factors, one per column of `x`,
122    /// normalised to geometric mean one over the informative columns.
123    pub n_j: Vec<f64>,
124    /// The anticipated CV of each informative column at the returned
125    /// factors, in column order.
126    pub cv: Vec<f64>,
127    /// Which columns were informative (summed to more than zero).
128    pub used: Vec<bool>,
129    /// Iterations actually taken.
130    pub iterations: usize,
131    /// Whether the CVs met the tolerance. A false is a report, not an
132    /// error: the factors returned are still the best the search reached.
133    pub converged: bool,
134}
135
136/// Search for the sample-size factors that equalise anticipated CVs.
137///
138/// `x` is the auxiliary matrix, column-major with `n_rows` units, and `n`
139/// the sample size (at most the number of units). At each step the
140/// anticipated CV of the total of every informative column is computed
141/// under the design that [`mpps`] and [`expected_hits`] would produce, and
142/// the factors move against the spread of the CVs, damped, until the
143/// spread falls within tolerance.
144pub fn eqcv(
145    x: &[f64],
146    n_rows: usize,
147    n: u64,
148    options: EqcvOptions,
149) -> Result<EqcvAllocation, Error> {
150    if n_rows == 0 || x.len() % n_rows != 0 {
151        return Err(Error::Input(format!(
152            "a matrix of {} values cannot have {} rows",
153            x.len(),
154            n_rows
155        )));
156    }
157    if n as usize > n_rows {
158        return Err(Error::Input(format!(
159            "equalising CVs needs n at most the number of units; n = {n} \
160             exceeds {n_rows}. With n beyond the frame use minimum \
161             replacement and equal or fixed factors"
162        )));
163    }
164    let n_cols = x.len() / n_rows;
165
166    let mut col_sums = vec![0.0_f64; n_cols];
167    for c in 0..n_cols {
168        for r in 0..n_rows {
169            col_sums[c] += x[c * n_rows + r];
170        }
171    }
172    let used_cols: Vec<usize> = (0..n_cols).filter(|&c| col_sums[c] > 0.0).collect();
173    if used_cols.is_empty() {
174        return Err(Error::Input(
175            "every auxiliary column sums to zero; there is nothing to equalise".into(),
176        ));
177    }
178
179    // The shares each informative column assigns to each unit, fixed for
180    // the whole search.
181    let mut shares = vec![0.0_f64; n_rows * used_cols.len()];
182    for (uc, &c) in used_cols.iter().enumerate() {
183        for r in 0..n_rows {
184            shares[uc * n_rows + r] = x[c * n_rows + r] / col_sums[c];
185        }
186    }
187
188    let mut n_j = vec![1.0_f64; n_cols];
189
190    let anticipated_cv = |n_j: &[f64]| -> Result<Vec<f64>, Error> {
191        let size = mpps(x, n_rows, n_j)?;
192        let pi = expected_hits(&size, n, false)?;
193        let a: Vec<f64> = pi.iter().map(|&p| p * (1.0 - p)).collect();
194        let a_total: f64 = a.iter().sum();
195        if a_total == 0.0 {
196            return Ok(vec![0.0; used_cols.len()]);
197        }
198        let mut cv = vec![0.0_f64; used_cols.len()];
199        for uc in 0..used_cols.len() {
200            let mut bias = 0.0_f64;
201            for r in 0..n_rows {
202                bias += (1.0 - pi[r]) * shares[uc * n_rows + r];
203            }
204            let centre = bias / a_total;
205            let mut ss = 0.0_f64;
206            for r in 0..n_rows {
207                let d = shares[uc * n_rows + r] / pi[r] - centre;
208                ss += a[r] * d * d;
209            }
210            cv[uc] = ss.sqrt();
211        }
212        Ok(cv)
213    };
214
215    let geometric_mean = |values: &[f64]| -> f64 {
216        let mut s = 0.0_f64;
217        for &v in values {
218            s += v.ln();
219        }
220        (s / values.len() as f64).exp()
221    };
222
223    let spread_ok = |cv: &[f64], tol: f64| -> bool {
224        let pos: Vec<f64> = cv.iter().copied().filter(|&v| v > 0.0).collect();
225        if pos.is_empty() {
226            return true;
227        }
228        let max = pos.iter().copied().fold(f64::NEG_INFINITY, f64::max);
229        let min = pos.iter().copied().fold(f64::INFINITY, f64::min);
230        max / min - 1.0 < tol
231    };
232
233    let mut cv = anticipated_cv(&n_j)?;
234    let mut iterations = 0_usize;
235    while iterations < options.max_iterations {
236        if spread_ok(&cv, options.tolerance) {
237            break;
238        }
239        let positive: Vec<f64> = cv.iter().copied().filter(|&v| v > 0.0).collect();
240        let g = geometric_mean(&positive);
241        for (uc, &c) in used_cols.iter().enumerate() {
242            if cv[uc] > 0.0 {
243                n_j[c] *= (cv[uc] / g).powf(options.damping);
244            }
245        }
246        let norm = geometric_mean(&used_cols.iter().map(|&c| n_j[c]).collect::<Vec<_>>());
247        for f in &mut n_j {
248            *f /= norm;
249        }
250        cv = anticipated_cv(&n_j)?;
251        iterations += 1;
252    }
253
254    let converged = spread_ok(&cv, options.tolerance);
255    let used: Vec<bool> = (0..n_cols).map(|c| col_sums[c] > 0.0).collect();
256    Ok(EqcvAllocation {
257        n_j,
258        cv,
259        used,
260        iterations,
261        converged,
262    })
263}
chromy_core/allocate.rs

chromy_core/
allocate.rs
