Skip to main content

edge_rs/limma/
array_weights.rs

1//! limma's `arrayWeights` and `duplicateCorrelation`.
2//!
3//! [`array_weights`] gives one precision weight per sample, so a bad array is
4//! down-weighted everywhere rather than dropped. [`duplicate_correlation`]
5//! gives the consensus intra-block correlation for a repeated-measures design.
6//!
7//! ### Array weights
8//!
9//! Sample `j` has variance `sigma^2 / w_j` with `log(1 / w_j) = z_j' gamma`,
10//! where `z_j` is a row of the variance design (`contr.sum(n_samples)` by
11//! default, so the log weights sum to zero). Both methods solve the same REML
12//! score equations and differ in how they sweep the genes:
13//!
14//! * [`ArrayWeightMethod::GeneByGene`]: one Fisher scoring step per gene,
15//!   weights updated as it goes. One pass, no convergence test, copes with
16//!   missing values and gene-level weights.
17//! * [`ArrayWeightMethod::Reml`]: score and information accumulated over all
18//!   genes, iterated to convergence. Needs a complete matrix.
19//!
20//! `prior_n` is a ridge on the information: the number of notional genes that
21//! saw every weight equal to one.
22//!
23//! ### Duplicate correlation
24//!
25//! A per-gene two-component mixed model through `statmod::mixedModel2Fit`, then
26//! a trimmed mean of the Fisher z-transformed correlations. The per-gene fit
27//! projects out the fixed effects, takes an SVD of the projected block design,
28//! and fits a gamma GLM of the squared rotated residuals on the squared
29//! singular values. The Levenberg damping of that GLM is reproduced step for
30//! step, since the answer depends on where it stops. edgePython uses a one-way
31//! ANOVA moment estimator instead. See `UPSTREAM_DEVIATIONS.md` A14.
32//!
33//! ### Implementation notes
34//!
35//! * No scalar optimiser: `arrayWeights` is Fisher scoring with a closed-form
36//!   step and `glmgam.fit` is its own damped scoring, so
37//!   [`crate::numeric::optimise`] would stop somewhere else.
38//! * The dense linear algebra is local: a LINPACK `dqrdc2` QR and a
39//!   partial-pivot elimination, so the rank rule and pivot order are R's. faer
40//!   is used only for the SVD in the mixed model.
41//! * `f64` only, not generic. The matrices scale with sample count, so there is
42//!   no memory to save.
43//!
44//! ### References
45//!
46//! Ritchie, Diyagama, Neilson, van Laar, Dobrovic, Holloway and Smyth, BMC
47//! Bioinformatics, 2006 (array weights)
48//!
49//! Smyth, Michaud and Scott, Bioinformatics, 2005 (duplicate correlation)
50
51use rayon::prelude::*;
52
53use crate::numeric::stats::trimmed_mean;
54use crate::prelude::*;
55
56////////////
57// Consts //
58////////////
59
60/// Relative tolerance below which a column counts as linearly dependent.
61const LM_RANK_TOL: f64 = 1e-7;
62
63/// Residual variances below this are treated as structurally zero.
64///
65/// limma skips such genes: they are constant or perfectly fitted and carry no
66/// information on the sample variances.
67const MIN_RESIDUAL_VAR: f64 = 1e-15;
68
69/// Residual degrees of freedom a gene needs before it can contribute.
70///
71/// One residual df gives no leverage on how variance splits across samples.
72const MIN_RESIDUAL_DF: usize = 2;
73
74/// Largest intra-block correlation `duplicateCorrelation` will report.
75///
76/// Clipping keeps the Fisher z-transform finite.
77const RHO_MAX: f64 = 0.99;
78
79/// Slack added to the lower bound on the correlation.
80///
81/// A block of size `m` cannot have a correlation below `1 / (1 - m)` without
82/// losing positive definiteness. limma clips to that bound plus this.
83const RHO_MIN_SLACK: f64 = 0.01;
84
85/// Below this, the block factor counts as already spanned by the design.
86///
87/// limma gives up if no block indicator survives projection onto the residual
88/// space, as there is no within-block replication left.
89const BLOCK_ABSORBED_TOL: f64 = 1e-8;
90
91/// Squared singular values above this count as non-zero.
92///
93/// `mixedModel2Fit` runs the gamma GLM only with at least two of them.
94const SINGULAR_VALUE_TOL: f64 = 1e-15;
95
96/// Convergence tolerance of `statmod::glmgam.fit` on the score-step product.
97const GLMGAM_TOL: f64 = 1e-6;
98
99/// Iteration budget `mixedModel2Fit` gives `glmgam.fit` (`maxit = 20`).
100///
101/// `glmgam.fit` tests `iter > maxit` at the bottom of the loop, so this permits
102/// 21 scoring steps.
103const GLMGAM_MAX_ITER: usize = 20;
104
105/// Ratio by which the gamma variance is floored inside `glmgam.fit`.
106///
107/// `v <- pmax(mu^2, max(mu^2) / 1e3)`.
108const GLMGAM_VARIANCE_FLOOR: f64 = 1e3;
109
110/// Damping ratio at which `glmgam.fit` declares the step hopeless and stops.
111const GLMGAM_MAX_DAMPING: f64 = 1e15;
112
113/// Factor by which Levenberg damping is relaxed after a step that worked.
114const GLMGAM_DAMPING_RELAX: f64 = 10.0;
115
116/// Relative size below which a deviance counts as converged regardless.
117const GLMGAM_DEVIANCE_TOL: f64 = 1e-15;
118
119/// Below this, an observation and its fitted value both count as zero.
120///
121/// `statmod`'s gamma deviance drops such terms rather than evaluating
122/// `log(0 / 0)`.
123const GLMGAM_ZERO_TOL: f64 = 1e-15;
124
125////////////////////////
126// Small dense helper //
127////////////////////////
128
129/// Euclidean norm of a slice.
130///
131/// ### Params
132///
133/// * `x` - Values
134///
135/// ### Returns
136///
137/// `sqrt(sum(x^2))`.
138fn euclidean_norm(x: &[f64]) -> f64 {
139    x.iter().map(|v| v * v).sum::<f64>().sqrt()
140}
141
142/// Solves a small dense system by Gaussian elimination with partial pivoting.
143///
144/// This is LAPACK's `dgesv` and therefore R's `solve`. The systems are
145/// `ngam` by `ngam`, tens of rows at most.
146///
147/// ### Params
148///
149/// * `a` - Column-major `n * n` coefficient matrix, copied before use
150/// * `b` - Right hand side of length `n`
151/// * `n` - Order of the system
152///
153/// ### Returns
154///
155/// The solution, or [`EdgeErrors::SolveFailed`] if a pivot vanishes or the
156/// result is not finite.
157fn solve_linear(a: &[f64], b: &[f64], n: usize) -> Result<Vec<f64>, EdgeErrors> {
158    let mut m = a.to_vec();
159    let mut x = b.to_vec();
160
161    for k in 0..n {
162        let (pivot_row, pivot) = (k..n).fold((k, 0.0), |(best, mag), i| {
163            let v = m[k * n + i].abs();
164            if v > mag { (i, v) } else { (best, mag) }
165        });
166        if pivot == 0.0 || !pivot.is_finite() {
167            return Err(EdgeErrors::SolveFailed(format!(
168                "the array weight information matrix is singular at pivot {k}"
169            )));
170        }
171        if pivot_row != k {
172            for c in 0..n {
173                m.swap(c * n + k, c * n + pivot_row);
174            }
175            x.swap(k, pivot_row);
176        }
177        for i in (k + 1)..n {
178            let factor = m[k * n + i] / m[k * n + k];
179            if factor == 0.0 {
180                continue;
181            }
182            for c in (k + 1)..n {
183                m[c * n + i] -= factor * m[c * n + k];
184            }
185            x[i] -= factor * x[k];
186        }
187    }
188
189    for k in (0..n).rev() {
190        let mut acc = x[k];
191        for c in (k + 1)..n {
192            acc -= m[c * n + k] * x[c];
193        }
194        x[k] = acc / m[k * n + k];
195    }
196
197    if x.iter().any(|v| !v.is_finite()) {
198        return Err(EdgeErrors::SolveFailed(
199            "the array weight update was not finite".to_string(),
200        ));
201    }
202    Ok(x)
203}
204
205/// Transposes a row-major matrix into column-major order.
206///
207/// ### Params
208///
209/// * `x` - Row-major `n_rows * n_cols` values
210/// * `n_rows` - Number of rows
211/// * `n_cols` - Number of columns
212///
213/// ### Returns
214///
215/// The same matrix, column-major.
216fn to_column_major(x: &[f64], n_rows: usize, n_cols: usize) -> Vec<f64> {
217    let mut out = vec![0.0; n_rows * n_cols];
218    for i in 0..n_rows {
219        for j in 0..n_cols {
220            out[j * n_rows + i] = x[i * n_cols + j];
221        }
222    }
223    out
224}
225
226/// R's `contr.sum(n)`.
227///
228/// ### Params
229///
230/// * `n` - Number of levels
231///
232/// ### Returns
233///
234/// Column-major `n * (n - 1)`: an identity block over the first `n - 1` rows
235/// and a row of `-1` underneath. Constraining the log weights against this
236/// basis is what makes them sum to zero.
237fn contr_sum(n: usize) -> Vec<f64> {
238    let mut out = vec![0.0; n * (n - 1)];
239    for k in 0..(n - 1) {
240        out[k * n + k] = 1.0;
241        out[k * n + (n - 1)] = -1.0;
242    }
243    out
244}
245
246////////////////////
247// Householder QR //
248////////////////////
249
250/// A LINPACK `dqrdc2` Householder QR: the factorisation R's `lm.fit` uses.
251///
252/// This is not faer's QR, deliberately. `dqrdc2` walks the columns left to right
253/// and moves a column aside only when the part of it orthogonal to its
254/// predecessors has lost a factor [`LM_RANK_TOL`] of its length. Survivors keep
255/// their original order, which is what makes `QR$pivot[1:rank]` in `arrayWeights`
256/// an order-preserving subset.
257///
258/// Rejected columns stay in place instead of being cycled to the end as LINPACK
259/// does; nothing reads them, and it keeps [`LinpackQr::pivot`] ascending.
260struct LinpackQr {
261    /// Column-major `n * p` working store. For each accepted step `l` sitting in
262    /// original column `j`, row `l` holds `R[l, l]` and rows `l + 1..n` hold the
263    /// tail of that step's Householder vector.
264    qr: Vec<f64>,
265    /// Leading element of each accepted step's Householder vector, LINPACK's
266    /// `qraux`. Kept out of `qr` because `R`'s diagonal needs that slot.
267    qraux: Vec<f64>,
268    /// Original index of each accepted column, ascending. Its length is the rank.
269    pivot: Vec<usize>,
270    /// Number of rows.
271    n: usize,
272}
273
274impl LinpackQr {
275    /// Factorises a column-major matrix.
276    ///
277    /// ### Params
278    ///
279    /// * `a` - Column-major `n * p` values, consumed as the working store
280    /// * `n` - Number of rows
281    /// * `p` - Number of columns
282    ///
283    /// ### Returns
284    ///
285    /// The factorisation. A zero column is always rejected, since its original
286    /// norm is replaced by one before the comparison, exactly as `dqrdc2` does.
287    fn new(mut a: Vec<f64>, n: usize, p: usize) -> Self {
288        let original: Vec<f64> = (0..p)
289            .map(|j| {
290                let norm = euclidean_norm(&a[j * n..(j + 1) * n]);
291                if norm == 0.0 { 1.0 } else { norm }
292            })
293            .collect();
294
295        let mut qraux = Vec::with_capacity(p.min(n));
296        let mut pivot = Vec::with_capacity(p.min(n));
297
298        for j in 0..p {
299            let l = pivot.len();
300            if l >= n {
301                break;
302            }
303            let remaining = euclidean_norm(&a[j * n + l..(j + 1) * n]);
304            if remaining < original[j] * LM_RANK_TOL {
305                continue;
306            }
307
308            // LINPACK's reflector: normalise by the signed norm, then bump the
309            // leading element by one so that `v[0]` doubles as the scale.
310            let mut nrmxl = remaining;
311            if a[j * n + l] != 0.0 {
312                nrmxl = nrmxl.copysign(a[j * n + l]);
313            }
314            for i in l..n {
315                a[j * n + i] /= nrmxl;
316            }
317            a[j * n + l] += 1.0;
318            let v0 = a[j * n + l];
319
320            for c in (j + 1)..p {
321                let mut t = 0.0;
322                for i in l..n {
323                    t += a[j * n + i] * a[c * n + i];
324                }
325                let t = -t / v0;
326                for i in l..n {
327                    a[c * n + i] += t * a[j * n + i];
328                }
329            }
330
331            qraux.push(v0);
332            a[j * n + l] = -nrmxl;
333            pivot.push(j);
334        }
335
336        LinpackQr {
337            qr: a,
338            qraux,
339            pivot,
340            n,
341        }
342    }
343
344    /// Numerical rank, in R's sense.
345    ///
346    /// ### Returns
347    ///
348    /// The number of columns that survived the [`LM_RANK_TOL`] test.
349    fn rank(&self) -> usize {
350        self.pivot.len()
351    }
352
353    /// Applies `Q'` in place, giving R's `effects`.
354    ///
355    /// ### Params
356    ///
357    /// * `y` - Vector of length `n`, overwritten with `Q' y`
358    fn qty(&self, y: &mut [f64]) {
359        for l in 0..self.rank() {
360            self.reflect(l, y);
361        }
362    }
363
364    /// Applies `Q` in place.
365    ///
366    /// Same reflectors as [`LinpackQr::qty`], applied in the opposite order:
367    /// `Q = H_0 ... H_{r-1}` and `Q' = H_{r-1} ... H_0`, each reflector being its
368    /// own inverse.
369    ///
370    /// ### Params
371    ///
372    /// * `y` - Vector of length `n`, overwritten with `Q y`
373    fn qy(&self, y: &mut [f64]) {
374        for l in (0..self.rank()).rev() {
375            self.reflect(l, y);
376        }
377    }
378
379    /// Applies one Householder reflector in place.
380    ///
381    /// ### Params
382    ///
383    /// * `l` - Index of the accepted step whose reflector to apply
384    /// * `y` - Vector of length `n`, overwritten
385    fn reflect(&self, l: usize, y: &mut [f64]) {
386        let v = &self.qr[self.pivot[l] * self.n..(self.pivot[l] + 1) * self.n];
387        let mut t = self.qraux[l] * y[l];
388        for (vi, yi) in v[l + 1..].iter().zip(y[l + 1..].iter()) {
389            t += vi * yi;
390        }
391        let t = -t / self.qraux[l];
392        y[l] += t * self.qraux[l];
393        for (vi, yi) in v[l + 1..].iter().zip(y[l + 1..].iter_mut()) {
394            *yi += t * vi;
395        }
396    }
397
398    /// The leading columns of `Q`, R's `qr.qy(qr, diag(1, n, k))`.
399    ///
400    /// ### Params
401    ///
402    /// * `k` - Number of columns wanted
403    ///
404    /// ### Returns
405    ///
406    /// Column-major `n * k`.
407    fn leading_q(&self, k: usize) -> Vec<f64> {
408        let mut out = vec![0.0; self.n * k];
409        for c in 0..k {
410            out[c * self.n + c] = 1.0;
411            let (_, column) = out.split_at_mut(c * self.n);
412            self.qy(&mut column[..self.n]);
413        }
414        out
415    }
416
417    /// Leverages of the rows, R's `hat(qr)`.
418    ///
419    /// ### Returns
420    ///
421    /// The hat matrix diagonal, one value per row.
422    fn hat(&self) -> Vec<f64> {
423        let rank = self.rank();
424        let q = self.leading_q(rank);
425        (0..self.n)
426            .map(|i| (0..rank).map(|c| q[c * self.n + i].powi(2)).sum())
427            .collect()
428    }
429
430    /// Back-substitutes `R b = effects` over the accepted columns.
431    ///
432    /// ### Params
433    ///
434    /// * `effects` - `Q' y`, of which only the leading `rank` entries are read
435    ///
436    /// ### Returns
437    ///
438    /// The coefficients of the accepted columns, in [`LinpackQr::pivot`] order.
439    fn solve_r(&self, effects: &[f64]) -> Vec<f64> {
440        let rank = self.rank();
441        let mut b = vec![0.0; rank];
442        for l in (0..rank).rev() {
443            let mut acc = effects[l];
444            for (&column, coefficient) in self.pivot[l + 1..rank].iter().zip(b[l + 1..rank].iter())
445            {
446                acc -= self.qr[column * self.n + l] * coefficient;
447            }
448            b[l] = acc / self.qr[self.pivot[l] * self.n + l];
449        }
450        b
451    }
452}
453
454///////////////////////
455// ArrayWeightMethod //
456///////////////////////
457
458/// How [`array_weights`] sweeps the genes.
459#[derive(Clone, Copy, Debug, PartialEq, Eq)]
460pub enum ArrayWeightMethod {
461    /// One Fisher scoring step per gene, weights updated as the sweep proceeds.
462    ///
463    /// limma's `.arrayWeightsGeneByGene`. Accepts missing values and gene-level
464    /// weights; `method = "auto"` selects it when either is present.
465    /// `method = "auto"` selects whenever either is present.
466    GeneByGene,
467    /// Score and information pooled over all genes, iterated to convergence.
468    ///
469    /// limma's `.arrayWeightsREML`, or `.arrayWeightsPrWtsREML` with gene-level
470    /// weights. `method = "auto"` selects it for a complete, unweighted matrix.
471    Reml,
472}
473
474///////////////////////
475// ArrayWeightParams //
476///////////////////////
477
478/// Knobs for [`array_weights`].
479#[derive(Clone, Debug)]
480pub struct ArrayWeightParams {
481    /// Which sweep to run. See [`ArrayWeightMethod`].
482    pub method: ArrayWeightMethod,
483    /// Prior number of genes pulling the weights towards one.
484    ///
485    /// Enters as a ridge `prior_n * Z2' Z2` on the information and, for the REML
486    /// sweep, a matching `prior_n * (w - 1)` on the score. limma's default is 10.
487    pub prior_n: f64,
488    /// Iteration budget for the REML sweep. Ignored by the gene-by-gene sweep.
489    pub max_iter: usize,
490    /// Convergence tolerance for the REML sweep, on the score-step product
491    /// scaled by the number of variance coefficients and effective gene count.
492    /// limma's default is `1e-5`.
493    pub tol: f64,
494    /// Variance design, row-major `n_samples * n_var_coef`.
495    ///
496    /// `None` uses `contr.sum(n_samples)`: one free weight per sample, logs
497    /// summing to zero. A coarser basis, such as group indicators, pools samples.
498    pub var_design: Option<Vec<f64>>,
499    /// Number of columns in `var_design`. Ignored when `var_design` is `None`.
500    pub n_var_coef: usize,
501}
502
503impl Default for ArrayWeightParams {
504    /// limma's defaults, with the method fixed at [`ArrayWeightMethod::Reml`].
505    ///
506    /// limma's `method = "auto"` resolves to `"reml"` for a complete, unweighted
507    /// matrix and `"genebygene"` otherwise. There is no `Auto` variant, so the
508    /// default matches a bare `arrayWeights(y, design)`; pass
509    /// [`ArrayWeightMethod::GeneByGene`] for weights or missing values.
510    fn default() -> Self {
511        ArrayWeightParams {
512            method: ArrayWeightMethod::Reml,
513            prior_n: 10.0,
514            max_iter: 50,
515            tol: 1e-5,
516            var_design: None,
517            n_var_coef: 0,
518        }
519    }
520}
521
522impl ArrayWeightParams {
523    /// Builds a parameter set.
524    ///
525    /// ### Params
526    ///
527    /// * `method` - Which sweep to run
528    /// * `prior_n` - Prior gene count shrinking the weights towards one, limma's 10
529    /// * `max_iter` - REML iteration budget, limma's 50
530    /// * `tol` - REML convergence tolerance, limma's 1e-5
531    /// * `var_design` - Optional variance design, row-major `n_samples * n_var_coef`
532    /// * `n_var_coef` - Number of columns in `var_design`
533    ///
534    /// ### Returns
535    ///
536    /// The parameters. Validation happens in [`array_weights`], which knows the
537    /// sample count these have to agree with.
538    pub fn new(
539        method: ArrayWeightMethod,
540        prior_n: f64,
541        max_iter: usize,
542        tol: f64,
543        var_design: Option<Vec<f64>>,
544        n_var_coef: usize,
545    ) -> Self {
546        ArrayWeightParams {
547            method,
548            prior_n,
549            max_iter,
550            tol,
551            var_design,
552            n_var_coef,
553        }
554    }
555}
556
557/////////////////////
558// Shared plumbing //
559/////////////////////
560
561/// Checks the shape of an expression matrix and its design.
562///
563/// ### Params
564///
565/// * `y` - Row-major `n_genes * n_samples` expression values
566/// * `n_genes` - Number of genes
567/// * `n_samples` - Number of samples
568/// * `design` - Row-major `n_samples * n_coef` design
569/// * `n_coef` - Number of design columns
570///
571/// ### Returns
572///
573/// `Ok(())`, or [`EdgeErrors::EmptyCounts`], [`EdgeErrors::MustBePositive`] or
574/// [`EdgeErrors::LengthMismatch`].
575fn validate_shapes(
576    y: &[f64],
577    n_genes: usize,
578    n_samples: usize,
579    design: &[f64],
580    n_coef: usize,
581) -> Result<(), EdgeErrors> {
582    if n_genes == 0 || n_samples == 0 {
583        return Err(EdgeErrors::EmptyCounts { n_genes, n_samples });
584    }
585    if n_coef == 0 {
586        return Err(EdgeErrors::MustBePositive("n_coef".to_string()));
587    }
588    if y.len() != n_genes * n_samples {
589        return Err(EdgeErrors::LengthMismatch {
590            name: "y",
591            expected: n_genes * n_samples,
592            got: y.len(),
593        });
594    }
595    if design.len() != n_samples * n_coef {
596        return Err(EdgeErrors::LengthMismatch {
597            name: "design",
598            expected: n_samples * n_coef,
599            got: design.len(),
600        });
601    }
602    Ok(())
603}
604
605/// Drops the linearly dependent columns of a design.
606///
607/// limma keeps `design[, QR$pivot[1:QR$rank]]`, so a rank deficient design loses
608/// columns silently.
609///
610/// ### Params
611///
612/// * `design` - Row-major `n_samples * n_coef`
613/// * `n_samples` - Number of samples
614/// * `n_coef` - Number of columns
615///
616/// ### Returns
617///
618/// The retained columns as a column-major `n_samples * rank` matrix, and the
619/// rank.
620fn reduce_design(design: &[f64], n_samples: usize, n_coef: usize) -> (Vec<f64>, usize) {
621    let column_major = to_column_major(design, n_samples, n_coef);
622    let qr = LinpackQr::new(column_major.clone(), n_samples, n_coef);
623    let rank = qr.rank();
624    if rank == n_coef {
625        return (column_major, n_coef);
626    }
627    let mut out = Vec::with_capacity(n_samples * rank);
628    for &j in &qr.pivot {
629        out.extend_from_slice(&column_major[j * n_samples..(j + 1) * n_samples]);
630    }
631    (out, rank)
632}
633
634/// Builds the variance design `Z2`.
635///
636/// With no user basis this is `contr.sum(n_samples)`. Otherwise limma centres
637/// every column and drops the linearly dependent ones, so a user-supplied
638/// intercept becomes a zero column and fails the [`LM_RANK_TOL`] test.
639///
640/// ### Params
641///
642/// * `n_samples` - Number of samples
643/// * `var_design` - Optional row-major `n_samples * n_var_coef` basis
644/// * `n_var_coef` - Number of columns in `var_design`
645///
646/// ### Returns
647///
648/// Column-major `n_samples * ngam` and `ngam`, or [`EdgeErrors`] on a shape
649/// mismatch.
650fn prepare_var_design(
651    n_samples: usize,
652    var_design: Option<&[f64]>,
653    n_var_coef: usize,
654) -> Result<(Vec<f64>, usize), EdgeErrors> {
655    let Some(vd) = var_design else {
656        return Ok((contr_sum(n_samples), n_samples - 1));
657    };
658    if n_var_coef == 0 {
659        return Err(EdgeErrors::MustBePositive("n_var_coef".to_string()));
660    }
661    if vd.len() != n_samples * n_var_coef {
662        return Err(EdgeErrors::LengthMismatch {
663            name: "var_design",
664            expected: n_samples * n_var_coef,
665            got: vd.len(),
666        });
667    }
668
669    let mut z2 = to_column_major(vd, n_samples, n_var_coef);
670    for j in 0..n_var_coef {
671        let column = &mut z2[j * n_samples..(j + 1) * n_samples];
672        let mean = column.iter().sum::<f64>() / n_samples as f64;
673        for v in column.iter_mut() {
674            *v -= mean;
675        }
676    }
677
678    let qr = LinpackQr::new(z2.clone(), n_samples, n_var_coef);
679    let ngam = qr.rank();
680    if ngam == n_var_coef {
681        return Ok((z2, ngam));
682    }
683    let mut out = Vec::with_capacity(n_samples * ngam);
684    for &j in &qr.pivot {
685        out.extend_from_slice(&z2[j * n_samples..(j + 1) * n_samples]);
686    }
687    Ok((out, ngam))
688}
689
690/// Scales each design column by the per-observation square-root weight.
691///
692/// ### Params
693///
694/// * `design` - Column-major `n * p` design
695/// * `root_w` - Square-root weights, length `n`
696/// * `n` - Number of observations
697/// * `p` - Number of columns
698///
699/// ### Returns
700///
701/// The weighted design, column-major `n * p`.
702fn weight_columns(design: &[f64], root_w: &[f64], n: usize, p: usize) -> Vec<f64> {
703    let mut xw = vec![0.0; n * p];
704    for j in 0..p {
705        for i in 0..n {
706            xw[j * n + i] = root_w[i] * design[j * n + i];
707        }
708    }
709    xw
710}
711
712/// Reconstructs unweighted residuals `y - X b` from a pivoted QR's solve.
713///
714/// ### Params
715///
716/// * `qr` - The pivoted QR the coefficients were solved from
717/// * `y` - Response of length `n`
718/// * `design` - Column-major `n * p` design, unweighted
719/// * `coef` - Coefficients in pivot order, as returned by [`LinpackQr::solve_r`]
720/// * `n` - Number of observations
721///
722/// ### Returns
723///
724/// The residuals, length `n`.
725fn qr_residuals(qr: &LinpackQr, y: &[f64], design: &[f64], coef: &[f64], n: usize) -> Vec<f64> {
726    let mut residuals = y.to_vec();
727    for (m, &j) in qr.pivot.iter().enumerate() {
728        for i in 0..n {
729            residuals[i] -= design[j * n + i] * coef[m];
730        }
731    }
732    residuals
733}
734
735/// One weighted least squares fit, R's `lm.wfit`.
736struct WeightedFit {
737    /// Unweighted residuals `y - X b`, one per observation.
738    residuals: Vec<f64>,
739    /// Leverages of the weighted design, R's `hat(fit$qr)`.
740    hat: Vec<f64>,
741    /// Numerical rank of the weighted design.
742    rank: usize,
743    /// Weighted residual sum of squares over `nobs - rank`, R's
744    /// `mean(fit$effects[-(1:rank)]^2)`.
745    s2: f64,
746}
747
748/// Fits `y ~ X` by weighted least squares.
749///
750/// ### Params
751///
752/// * `x` - Column-major `n * p` design
753/// * `y` - Response of length `n`
754/// * `w` - Weights of length `n`, all non-negative
755/// * `n` - Number of observations
756/// * `p` - Number of columns
757///
758/// ### Returns
759///
760/// The residuals, leverages, rank and residual variance. `s2` is `NaN` when the
761/// fit has no residual degrees of freedom, matching R's `mean(numeric(0))`.
762fn weighted_fit(x: &[f64], y: &[f64], w: &[f64], n: usize, p: usize) -> WeightedFit {
763    let root_w: Vec<f64> = w.iter().map(|v| v.sqrt()).collect();
764    let xw = weight_columns(x, &root_w, n, p);
765    let qr = LinpackQr::new(xw, n, p);
766    let rank = qr.rank();
767
768    let mut effects: Vec<f64> = (0..n).map(|i| root_w[i] * y[i]).collect();
769    qr.qty(&mut effects);
770    let coef = qr.solve_r(&effects);
771
772    let residuals = qr_residuals(&qr, y, x, &coef, n);
773
774    let s2 = if n > rank {
775        effects[rank..].iter().map(|v| v * v).sum::<f64>() / (n - rank) as f64
776    } else {
777        f64::NAN
778    };
779
780    WeightedFit {
781        residuals,
782        hat: qr.hat(),
783        rank,
784        s2,
785    }
786}
787
788/// Accumulates `info[-1, -1] - info[-1, 1] info[1, -1] / info[1, 1]` into `info2`.
789///
790/// `Z = [1 | Z2]`, so the leading row and column of `info` are the intercept's;
791/// sweeping them out leaves the information for `gamma`. The pieces are kept
792/// separate to avoid allocating a `(1 + ngam)^2` matrix per gene.
793///
794/// ### Params
795///
796/// * `info2` - Column-major `ngam * ngam` accumulator, updated in place
797/// * `corner` - `info[1, 1]`
798/// * `edge` - `info[1, -1]`, length `ngam`
799/// * `block` - `info[-1, -1]`, column-major `ngam * ngam`
800/// * `ngam` - Number of variance coefficients
801///
802/// ### Returns
803///
804/// `true` when the contribution was added, `false` when `corner` was
805/// non-positive or NaN. limma has no such guard; it never triggers for admitted
806/// genes and only avoids `NaN` weights.
807/// alternative was `NaN` weights.
808fn sweep_intercept(
809    info2: &mut [f64],
810    corner: f64,
811    edge: &[f64],
812    block: &[f64],
813    ngam: usize,
814) -> bool {
815    if corner.is_nan() || corner <= 0.0 {
816        return false;
817    }
818    for b in 0..ngam {
819        for a in 0..ngam {
820            info2[b * ngam + a] += block[b * ngam + a] - edge[a] * edge[b] / corner;
821        }
822    }
823    true
824}
825
826///////////////////
827// Array weights //
828///////////////////
829
830/// Per-sample quality weights, limma's `arrayWeights`.
831///
832/// Sample `j` has residual variance `sigma_g^2 / w_j`, shared across genes. A
833/// weight below one marks a noisier array; the logs of the weights sum to zero
834/// over the variance design. The estimator is REML on the variance components,
835/// solved per gene or pooled (see [`ArrayWeightMethod`]). Both return all-ones
836/// weights when there is nothing to estimate from: fewer than two genes, or
837/// fewer than two residual degrees of freedom.
838///
839/// Zero weights become missing expression values and unit weights, as in limma,
840/// so the gene-by-gene sweep skips those observations.
841///
842/// ### Params
843///
844/// * `y` - Row-major `n_genes * n_samples` log-expression values. Non-finite
845///   entries are treated as missing.
846/// * `n_genes` - Number of genes
847/// * `n_samples` - Number of samples
848/// * `design` - Row-major `n_samples * n_coef` design matrix
849/// * `n_coef` - Number of design columns. A rank deficient design is reduced to
850///   its independent columns rather than rejected.
851/// * `weights` - Optional gene by sample observation weights, non-negative
852/// * `params` - See [`ArrayWeightParams`]; `None` takes the defaults
853///
854/// ### Returns
855///
856/// One weight per sample, or [`EdgeErrors`] on a shape mismatch, an invalid
857/// weight, an out-of-domain parameter, or a singular information matrix.
858///
859/// ### edgePython
860///
861/// `edgepython/voom_lmfit.py:834` differs from limma in five places (no
862/// `.arrayWeightsPrWtsREML`, fully pivoted QR, `pinv` for `solve`, clamping,
863/// looser gene admission). This follows limma. See `UPSTREAM_DEVIATIONS.md` A14.
864///
865/// ### References
866///
867/// Ritchie, Diyagama, Neilson, van Laar, Dobrovic, Holloway and Smyth, BMC
868/// Bioinformatics 7:261, 2006
869pub fn array_weights(
870    y: &[f64],
871    n_genes: usize,
872    n_samples: usize,
873    design: &[f64],
874    n_coef: usize,
875    weights: Option<&Recycled<f64>>,
876    params: Option<ArrayWeightParams>,
877) -> Result<Vec<f64>, EdgeErrors> {
878    validate_shapes(y, n_genes, n_samples, design, n_coef)?;
879    let params = params.unwrap_or_default();
880
881    if !params.prior_n.is_finite() || params.prior_n < 0.0 {
882        return Err(EdgeErrors::InvalidArgument(format!(
883            "prior_n must be finite and non-negative; got {}.",
884            params.prior_n
885        )));
886    }
887    if !params.tol.is_finite() || params.tol <= 0.0 {
888        return Err(EdgeErrors::InvalidArgument(format!(
889            "tol must be finite and positive; got {}.",
890            params.tol
891        )));
892    }
893    if params.max_iter == 0 {
894        return Err(EdgeErrors::MustBePositive("max_iter".to_string()));
895    }
896
897    let ones = vec![1.0; n_samples];
898    if n_genes < 2 || n_samples < 2 {
899        return Ok(ones);
900    }
901
902    let (design_cm, rank) = reduce_design(design, n_samples, n_coef);
903    if n_samples < rank + MIN_RESIDUAL_DF {
904        return Ok(ones);
905    }
906
907    // Zero weights become missing observations (as in limma), so the expression
908    // matrix is materialised before the sweep.
909    let mut expression = y.to_vec();
910    let mut weight_matrix: Option<Vec<f64>> = None;
911    if let Some(w) = weights {
912        w.validate(n_genes, n_samples)?;
913        let mut dense = w.expand(n_genes, n_samples);
914        if dense.iter().any(|v| !v.is_finite()) {
915            return Err(EdgeErrors::InvalidArgument(
916                "arrayWeights does not accept missing or infinite weights.".to_string(),
917            ));
918        }
919        if dense.iter().any(|v| *v < 0.0) {
920            return Err(EdgeErrors::InvalidArgument(
921                "arrayWeights does not accept negative weights.".to_string(),
922            ));
923        }
924        for (e, wv) in expression.iter_mut().zip(dense.iter_mut()) {
925            if *wv == 0.0 {
926                *e = f64::NAN;
927                *wv = 1.0;
928            }
929        }
930        weight_matrix = Some(dense);
931    }
932
933    let (z2, ngam) =
934        prepare_var_design(n_samples, params.var_design.as_deref(), params.n_var_coef)?;
935    if ngam == 0 {
936        return Ok(ones);
937    }
938
939    match params.method {
940        ArrayWeightMethod::GeneByGene => array_weights_gene_by_gene(
941            &expression,
942            n_genes,
943            n_samples,
944            &design_cm,
945            rank,
946            weight_matrix.as_deref(),
947            &z2,
948            ngam,
949            params.prior_n,
950        ),
951        ArrayWeightMethod::Reml => {
952            // limma drops rows that are wholly or partly missing before the REML
953            // sweep, since it has no per-observation branch.
954            let complete: Vec<usize> = (0..n_genes)
955                .filter(|g| {
956                    expression[g * n_samples..(g + 1) * n_samples]
957                        .iter()
958                        .all(|v| v.is_finite())
959                })
960                .collect();
961            if complete.len() < 2 {
962                return Ok(ones);
963            }
964            let kept: Vec<f64> = complete
965                .iter()
966                .flat_map(|&g| {
967                    expression[g * n_samples..(g + 1) * n_samples]
968                        .iter()
969                        .copied()
970                })
971                .collect();
972
973            match weight_matrix {
974                None => array_weights_reml(
975                    &kept,
976                    complete.len(),
977                    n_samples,
978                    &design_cm,
979                    rank,
980                    &z2,
981                    ngam,
982                    params.prior_n,
983                    params.max_iter,
984                    params.tol,
985                ),
986                Some(w) => {
987                    let kept_w: Vec<f64> = complete
988                        .iter()
989                        .flat_map(|&g| w[g * n_samples..(g + 1) * n_samples].iter().copied())
990                        .collect();
991                    array_weights_prior_reml(
992                        &kept,
993                        &kept_w,
994                        complete.len(),
995                        n_samples,
996                        &design_cm,
997                        rank,
998                        &z2,
999                        ngam,
1000                        params.prior_n,
1001                        params.max_iter,
1002                        params.tol,
1003                    )
1004                }
1005            }
1006        }
1007    }
1008}
1009
1010/// The prior ridge `prior_n * Z2' Z2`, the starting information for all three
1011/// estimators.
1012///
1013/// ### Params
1014///
1015/// * `z2` - Column-major `n_samples * ngam` variance design
1016/// * `n_samples` - Number of samples
1017/// * `ngam` - Number of variance coefficients
1018/// * `prior_n` - Prior gene count
1019///
1020/// ### Returns
1021///
1022/// The `ngam * ngam` ridge matrix, row-major.
1023fn prior_ridge(z2: &[f64], n_samples: usize, ngam: usize, prior_n: f64) -> Vec<f64> {
1024    let mut m = vec![0.0; ngam * ngam];
1025    for b in 0..ngam {
1026        for a in 0..ngam {
1027            let dot: f64 = (0..n_samples)
1028                .map(|i| z2[a * n_samples + i] * z2[b * n_samples + i])
1029                .sum();
1030            m[b * ngam + a] = prior_n * dot;
1031        }
1032    }
1033    m
1034}
1035
1036/// limma's `.arrayWeightsGeneByGene`.
1037///
1038/// One Fisher scoring step per gene with weights refreshed after each step. Gene
1039/// `i + 1` is fitted under the weights gene `i` produced, so this is the one
1040/// estimator in the module that runs on a single thread.
1041///
1042/// ### Params
1043///
1044/// * `expression` - Row-major `n_genes * n_samples`, non-finite meaning missing
1045/// * `n_genes` - Number of genes
1046/// * `n_samples` - Number of samples
1047/// * `design` - Column-major `n_samples * p`, already reduced to full rank
1048/// * `p` - Number of design columns
1049/// * `weights` - Optional row-major `n_genes * n_samples` observation weights
1050/// * `z2` - Column-major `n_samples * ngam` variance design
1051/// * `ngam` - Number of variance coefficients
1052/// * `prior_n` - Prior gene count
1053///
1054/// ### Returns
1055///
1056/// One weight per sample.
1057#[allow(clippy::too_many_arguments)]
1058fn array_weights_gene_by_gene(
1059    expression: &[f64],
1060    n_genes: usize,
1061    n_samples: usize,
1062    design: &[f64],
1063    p: usize,
1064    weights: Option<&[f64]>,
1065    z2: &[f64],
1066    ngam: usize,
1067    prior_n: f64,
1068) -> Result<Vec<f64>, EdgeErrors> {
1069    let mut gam = vec![0.0; ngam];
1070    let mut aw = vec![1.0; n_samples];
1071
1072    // info2 starts at the prior ridge `prior_n * Z2' Z2`.
1073    let mut info2 = prior_ridge(z2, n_samples, ngam, prior_n);
1074
1075    let mut edge = vec![0.0; ngam];
1076    let mut block = vec![0.0; ngam * ngam];
1077
1078    for gene in 0..n_genes {
1079        let row = &expression[gene * n_samples..(gene + 1) * n_samples];
1080        let w: Vec<f64> = match weights {
1081            None => aw.clone(),
1082            Some(m) => (0..n_samples)
1083                .map(|j| aw[j] * m[gene * n_samples + j])
1084                .collect(),
1085        };
1086
1087        let mut d = vec![0.0; n_samples];
1088        let mut h1 = vec![0.0; n_samples];
1089        let s2;
1090
1091        if row.iter().all(|v| v.is_finite()) {
1092            // No residual degrees of freedom check on this branch: limma has
1093            // none either, and the caller has already guaranteed at least two.
1094            let fit = weighted_fit(design, row, &w, n_samples, p);
1095            for j in 0..n_samples {
1096                d[j] = w[j] * fit.residuals[j] * fit.residuals[j];
1097                h1[j] = 1.0 - fit.hat[j];
1098            }
1099            s2 = fit.s2;
1100        } else {
1101            let observed: Vec<usize> = (0..n_samples).filter(|&j| row[j].is_finite()).collect();
1102            let n_obs = observed.len();
1103            if n_obs <= MIN_RESIDUAL_DF {
1104                continue;
1105            }
1106            let mut sub_x = vec![0.0; n_obs * p];
1107            for c in 0..p {
1108                for (k, &j) in observed.iter().enumerate() {
1109                    sub_x[c * n_obs + k] = design[c * n_samples + j];
1110                }
1111            }
1112            let sub_y: Vec<f64> = observed.iter().map(|&j| row[j]).collect();
1113            let sub_w: Vec<f64> = observed.iter().map(|&j| w[j]).collect();
1114            let fit = weighted_fit(&sub_x, &sub_y, &sub_w, n_obs, p);
1115            if n_obs < fit.rank + MIN_RESIDUAL_DF {
1116                continue;
1117            }
1118            for (k, &j) in observed.iter().enumerate() {
1119                d[j] = sub_w[k] * fit.residuals[k] * fit.residuals[k];
1120                h1[j] = 1.0 - fit.hat[k];
1121            }
1122            s2 = fit.s2;
1123        }
1124
1125        if !s2.is_finite() || s2 < MIN_RESIDUAL_VAR {
1126            continue;
1127        }
1128
1129        // info = Z' diag(h1) Z with Z = [1 | Z2], kept in its three blocks.
1130        let corner: f64 = h1.iter().sum();
1131        for a in 0..ngam {
1132            edge[a] = (0..n_samples).map(|i| h1[i] * z2[a * n_samples + i]).sum();
1133        }
1134        for b in 0..ngam {
1135            for a in 0..=b {
1136                let v: f64 = (0..n_samples)
1137                    .map(|i| h1[i] * z2[a * n_samples + i] * z2[b * n_samples + i])
1138                    .sum();
1139                block[b * ngam + a] = v;
1140                block[a * ngam + b] = v;
1141            }
1142        }
1143        if !sweep_intercept(&mut info2, corner, &edge, &block, ngam) {
1144            continue;
1145        }
1146
1147        let score: Vec<f64> = (0..ngam)
1148            .map(|a| {
1149                (0..n_samples)
1150                    .map(|i| z2[a * n_samples + i] * (d[i] / s2 - h1[i]))
1151                    .sum()
1152            })
1153            .collect();
1154        let step = solve_linear(&info2, &score, ngam)?;
1155        for a in 0..ngam {
1156            gam[a] += step[a];
1157        }
1158        for (j, w) in aw.iter_mut().enumerate() {
1159            let eta: f64 = (0..ngam).map(|a| z2[a * n_samples + j] * gam[a]).sum();
1160            *w = (-eta).exp();
1161        }
1162    }
1163
1164    Ok(aw)
1165}
1166
1167/// Builds limma's `Q2` and the leverages from a fitted QR.
1168///
1169/// `Q2` holds every product `Q[, i] * Q[, j]` for `i <= j` over the leading `p`
1170/// columns of `Q`, off-diagonal blocks scaled by `sqrt(2)`, so `Q2' Q2` gives the
1171/// second moments of the residual projection (the second-derivative term of the
1172/// REML information).
1173///
1174/// ### Params
1175///
1176/// * `qr` - Factorisation of the weighted design
1177/// * `n_samples` - Number of samples
1178/// * `p` - Number of design columns
1179///
1180/// ### Returns
1181///
1182/// Column-major `n_samples * (p * (p + 1) / 2)` and the leverages, which are the
1183/// row sums of the first `p` columns of `Q2`.
1184fn residual_moments(qr: &LinpackQr, n_samples: usize, p: usize) -> (Vec<f64>, Vec<f64>) {
1185    let p2 = p * (p + 1) / 2;
1186    let q = qr.leading_q(p);
1187    let mut q2 = vec![0.0; n_samples * p2];
1188    let mut j0 = 0;
1189    for k in 0..p {
1190        for c in 0..(p - k) {
1191            for i in 0..n_samples {
1192                q2[(j0 + c) * n_samples + i] = q[c * n_samples + i] * q[(c + k) * n_samples + i];
1193            }
1194        }
1195        j0 += p - k;
1196    }
1197    let root_two = 2.0_f64.sqrt();
1198    for v in q2[p * n_samples..].iter_mut() {
1199        *v *= root_two;
1200    }
1201    let hat: Vec<f64> = (0..n_samples)
1202        .map(|i| (0..p).map(|c| q2[c * n_samples + i]).sum())
1203        .collect();
1204    (q2, hat)
1205}
1206
1207/// The REML information for `gamma`, `crossprod(Z, (1 - 2h) Z) + crossprod(crossprod(Q2, Z))`.
1208///
1209/// ### Params
1210///
1211/// * `q2` - Column-major `n_samples * p2` second-moment matrix
1212/// * `p2` - Number of columns in `q2`
1213/// * `hat` - Leverages, one per sample
1214/// * `z2` - Column-major `n_samples * ngam` variance design
1215/// * `ngam` - Number of variance coefficients
1216/// * `n_samples` - Number of samples
1217///
1218/// ### Returns
1219///
1220/// `info[1, 1]`, `info[1, -1]` and `info[-1, -1]` of the `(1 + ngam)` square
1221/// information, ready for [`sweep_intercept`].
1222fn reml_information(
1223    q2: &[f64],
1224    p2: usize,
1225    hat: &[f64],
1226    z2: &[f64],
1227    ngam: usize,
1228    n_samples: usize,
1229) -> (f64, Vec<f64>, Vec<f64>) {
1230    let u: Vec<f64> = hat.iter().map(|h| 1.0 - 2.0 * h).collect();
1231
1232    // g = Q2' Z, of shape p2 by (1 + ngam), stored column-major.
1233    let m = 1 + ngam;
1234    let mut g = vec![0.0; p2 * m];
1235    for r in 0..p2 {
1236        g[r] = (0..n_samples).map(|i| q2[r * n_samples + i]).sum();
1237        for c in 0..ngam {
1238            g[(1 + c) * p2 + r] = (0..n_samples)
1239                .map(|i| q2[r * n_samples + i] * z2[c * n_samples + i])
1240                .sum();
1241        }
1242    }
1243
1244    let corner = u.iter().sum::<f64>() + (0..p2).map(|r| g[r] * g[r]).sum::<f64>();
1245    let edge: Vec<f64> = (0..ngam)
1246        .map(|a| {
1247            let direct: f64 = (0..n_samples).map(|i| u[i] * z2[a * n_samples + i]).sum();
1248            let moment: f64 = (0..p2).map(|r| g[r] * g[(1 + a) * p2 + r]).sum();
1249            direct + moment
1250        })
1251        .collect();
1252    let mut block = vec![0.0; ngam * ngam];
1253    for b in 0..ngam {
1254        for a in 0..=b {
1255            let direct: f64 = (0..n_samples)
1256                .map(|i| u[i] * z2[a * n_samples + i] * z2[b * n_samples + i])
1257                .sum();
1258            let moment: f64 = (0..p2)
1259                .map(|r| g[(1 + a) * p2 + r] * g[(1 + b) * p2 + r])
1260                .sum();
1261            block[b * ngam + a] = direct + moment;
1262            block[a * ngam + b] = direct + moment;
1263        }
1264    }
1265    (corner, edge, block)
1266}
1267
1268/// Turns the current `gamma` into weights, `exp(-Z2 gamma)`.
1269///
1270/// ### Params
1271///
1272/// * `z2` - Column-major `n_samples * ngam` variance design
1273/// * `gam` - Current variance coefficients
1274/// * `n_samples` - Number of samples
1275/// * `ngam` - Number of variance coefficients
1276///
1277/// ### Returns
1278///
1279/// One weight per sample.
1280fn weights_from_gamma(z2: &[f64], gam: &[f64], n_samples: usize, ngam: usize) -> Vec<f64> {
1281    (0..n_samples)
1282        .map(|j| {
1283            let eta: f64 = (0..ngam).map(|a| z2[a * n_samples + j] * gam[a]).sum();
1284            (-eta).exp()
1285        })
1286        .collect()
1287}
1288
1289/// limma's `.arrayWeightsREML`, the unweighted pooled sweep.
1290///
1291/// All genes share one weighted design, so the QR is factorised once per
1292/// iteration. Genes are chunked and reduced in a fixed order so the result does
1293/// not depend on rayon's split.
1294///
1295/// ### Params
1296///
1297/// * `y` - Row-major `n_genes * n_samples`, complete
1298/// * `n_genes` - Number of genes
1299/// * `n_samples` - Number of samples
1300/// * `design` - Column-major `n_samples * p`, full rank
1301/// * `p` - Number of design columns
1302/// * `z2` - Column-major `n_samples * ngam` variance design
1303/// * `ngam` - Number of variance coefficients
1304/// * `prior_n` - Prior gene count
1305/// * `max_iter` - Iteration budget
1306/// * `tol` - Convergence tolerance
1307///
1308/// ### Returns
1309///
1310/// One weight per sample.
1311#[allow(clippy::too_many_arguments)]
1312fn array_weights_reml(
1313    y: &[f64],
1314    n_genes: usize,
1315    n_samples: usize,
1316    design: &[f64],
1317    p: usize,
1318    z2: &[f64],
1319    ngam: usize,
1320    prior_n: f64,
1321    max_iter: usize,
1322    tol: f64,
1323) -> Result<Vec<f64>, EdgeErrors> {
1324    // limma drops genes with no residual variation once, from the unweighted
1325    // fit, before the iteration starts.
1326    let unit = vec![1.0; n_samples];
1327    let (_, first_s2, _) = reml_gene_pass(y, n_genes, n_samples, design, p, &unit);
1328    let usable: Vec<usize> = (0..n_genes)
1329        .filter(|&g| first_s2[g] >= MIN_RESIDUAL_VAR)
1330        .collect();
1331    let (kept, n_kept): (Vec<f64>, usize) = if usable.len() == n_genes {
1332        (y.to_vec(), n_genes)
1333    } else {
1334        if usable.len() < 2 {
1335            return Ok(vec![1.0; n_samples]);
1336        }
1337        (
1338            usable
1339                .iter()
1340                .flat_map(|&g| y[g * n_samples..(g + 1) * n_samples].iter().copied())
1341                .collect(),
1342            usable.len(),
1343        )
1344    };
1345
1346    let p2 = p * (p + 1) / 2;
1347    let ridge: Vec<f64> = prior_ridge(z2, n_samples, ngam, prior_n);
1348
1349    let mut gam = vec![0.0; ngam];
1350    let mut w: Vec<f64> = vec![1.0; n_samples];
1351    let mut last_convergence = f64::INFINITY;
1352    let genes = n_kept as f64;
1353
1354    for _ in 0..max_iter {
1355        let (mean_scaled, _, qr) = reml_gene_pass(&kept, n_kept, n_samples, design, p, &w);
1356        let (q2, hat) = residual_moments(&qr, n_samples, p);
1357        let (corner, edge, block) = reml_information(&q2, p2, &hat, z2, ngam, n_samples);
1358        if corner.is_nan() || corner <= 0.0 {
1359            break;
1360        }
1361
1362        let mut info2 = vec![0.0; ngam * ngam];
1363        sweep_intercept(&mut info2, corner, &edge, &block, ngam);
1364        for (slot, prior) in info2.iter_mut().zip(ridge.iter()) {
1365            *slot = genes * *slot + prior;
1366        }
1367
1368        let z: Vec<f64> = (0..n_samples)
1369            .map(|j| genes * (mean_scaled[j] - (1.0 - hat[j])) + prior_n * (w[j] - 1.0))
1370            .collect();
1371        let score: Vec<f64> = (0..ngam)
1372            .map(|a| (0..n_samples).map(|i| z2[a * n_samples + i] * z[i]).sum())
1373            .collect();
1374        let step = solve_linear(&info2, &score, ngam)?;
1375
1376        let convergence = score
1377            .iter()
1378            .zip(step.iter())
1379            .map(|(a, b)| a * b)
1380            .sum::<f64>()
1381            / ngam as f64
1382            / (genes + prior_n);
1383        // limma stops rather than cycling once the criterion stops improving.
1384        if !convergence.is_finite() || convergence >= last_convergence {
1385            break;
1386        }
1387        last_convergence = convergence;
1388
1389        for a in 0..ngam {
1390            gam[a] += step[a];
1391        }
1392        w = weights_from_gamma(z2, &gam, n_samples, ngam);
1393        if convergence < tol {
1394            break;
1395        }
1396    }
1397
1398    Ok(w)
1399}
1400
1401/// One pooled pass over the genes at fixed sample weights.
1402///
1403/// Computes the mean over genes of `w_j r_gj^2 / s2_g` at each sample and the
1404/// per-gene residual variances. The QR of the weighted design is shared, so
1405/// there is no factorisation per gene. Genes are split into fixed chunks whose
1406/// partial sums are added in index order, so results are bit-reproducible.
1407///
1408/// ### Params
1409///
1410/// * `y` - Row-major `n_genes * n_samples`
1411/// * `n_genes` - Number of genes
1412/// * `n_samples` - Number of samples
1413/// * `design` - Column-major `n_samples * p`
1414/// * `p` - Number of design columns
1415/// * `w` - Current sample weights
1416///
1417/// ### Returns
1418///
1419/// The per-sample mean of `w_j r^2 / s2`, the per-gene `s2`, and the shared
1420/// factorisation of the weighted design (reused for the REML information).
1421fn reml_gene_pass(
1422    y: &[f64],
1423    n_genes: usize,
1424    n_samples: usize,
1425    design: &[f64],
1426    p: usize,
1427    w: &[f64],
1428) -> (Vec<f64>, Vec<f64>, LinpackQr) {
1429    let root_w: Vec<f64> = w.iter().map(|v| v.sqrt()).collect();
1430    let xw = weight_columns(design, &root_w, n_samples, p);
1431    let qr = LinpackQr::new(xw, n_samples, p);
1432    let rank = qr.rank();
1433    let df = (n_samples - rank) as f64;
1434
1435    let chunk = n_genes.div_ceil(rayon::current_num_threads().max(1)).max(1);
1436    let partials: Vec<(Vec<f64>, Vec<f64>)> = y
1437        .par_chunks(chunk * n_samples)
1438        .map(|slab| {
1439            let mut acc = vec![0.0; n_samples];
1440            let mut s2 = Vec::with_capacity(slab.len() / n_samples);
1441            for row in slab.chunks_exact(n_samples) {
1442                let mut effects: Vec<f64> = (0..n_samples).map(|i| root_w[i] * row[i]).collect();
1443                qr.qty(&mut effects);
1444                let coef = qr.solve_r(&effects);
1445                let variance = effects[rank..].iter().map(|v| v * v).sum::<f64>() / df;
1446                s2.push(variance);
1447                let residual = qr_residuals(&qr, row, design, &coef, n_samples);
1448                for i in 0..n_samples {
1449                    acc[i] += w[i] * residual[i] * residual[i] / variance;
1450                }
1451            }
1452            (acc, s2)
1453        })
1454        .collect();
1455
1456    let mut total = vec![0.0; n_samples];
1457    let mut s2 = Vec::with_capacity(n_genes);
1458    for (acc, part) in partials {
1459        for (t, a) in total.iter_mut().zip(acc.iter()) {
1460            *t += a;
1461        }
1462        s2.extend(part);
1463    }
1464    for t in total.iter_mut() {
1465        *t /= n_genes as f64;
1466    }
1467    (total, s2, qr)
1468}
1469
1470/// limma's `.arrayWeightsPrWtsREML`, the pooled sweep with gene-level weights.
1471///
1472/// Each gene has its own weighted design, so the QR is refactorised per gene per
1473/// iteration. Score and information are accumulated by chunk and summed in index
1474/// order for reproducibility.
1475///
1476/// Unlike [`array_weights_reml`], convergence is tested only for finiteness and
1477/// falling below `tol`, not for improvement. That is limma's behaviour.
1478///
1479/// ### Params
1480///
1481/// * `y` - Row-major `n_genes * n_samples`, complete
1482/// * `weights` - Row-major `n_genes * n_samples` observation weights
1483/// * `n_genes` - Number of genes
1484/// * `n_samples` - Number of samples
1485/// * `design` - Column-major `n_samples * p`, full rank
1486/// * `p` - Number of design columns
1487/// * `z2` - Column-major `n_samples * ngam` variance design
1488/// * `ngam` - Number of variance coefficients
1489/// * `prior_n` - Prior gene count
1490/// * `max_iter` - Iteration budget
1491/// * `tol` - Convergence tolerance
1492///
1493/// ### Returns
1494///
1495/// One weight per sample.
1496#[allow(clippy::too_many_arguments)]
1497fn array_weights_prior_reml(
1498    y: &[f64],
1499    weights: &[f64],
1500    n_genes: usize,
1501    n_samples: usize,
1502    design: &[f64],
1503    p: usize,
1504    z2: &[f64],
1505    ngam: usize,
1506    prior_n: f64,
1507    max_iter: usize,
1508    tol: f64,
1509) -> Result<Vec<f64>, EdgeErrors> {
1510    let p2 = p * (p + 1) / 2;
1511    let ridge: Vec<f64> = prior_ridge(z2, n_samples, ngam, prior_n);
1512
1513    let mut gam = vec![0.0; ngam];
1514    let mut w = vec![1.0; n_samples];
1515    let chunk = n_genes.div_ceil(rayon::current_num_threads().max(1)).max(1);
1516
1517    for _ in 0..max_iter {
1518        let partials: Vec<(Vec<f64>, Vec<f64>)> = y
1519            .par_chunks(chunk * n_samples)
1520            .zip(weights.par_chunks(chunk * n_samples))
1521            .map(|(slab, wslab)| {
1522                let mut info2 = vec![0.0; ngam * ngam];
1523                let mut z = vec![0.0; n_samples];
1524                for (row, wrow) in slab
1525                    .chunks_exact(n_samples)
1526                    .zip(wslab.chunks_exact(n_samples))
1527                {
1528                    let full_w: Vec<f64> = (0..n_samples).map(|j| w[j] * wrow[j]).collect();
1529                    let root_w: Vec<f64> = full_w.iter().map(|v| v.sqrt()).collect();
1530                    let xw = weight_columns(design, &root_w, n_samples, p);
1531                    let qr = LinpackQr::new(xw, n_samples, p);
1532                    let rank = qr.rank();
1533                    let mut effects: Vec<f64> =
1534                        (0..n_samples).map(|i| root_w[i] * row[i]).collect();
1535                    qr.qty(&mut effects);
1536                    let coef = qr.solve_r(&effects);
1537                    let s2 = effects[rank..].iter().map(|v| v * v).sum::<f64>()
1538                        / (n_samples - rank) as f64;
1539
1540                    let (q2, hat) = residual_moments(&qr, n_samples, p);
1541                    let (corner, edge, block) =
1542                        reml_information(&q2, p2, &hat, z2, ngam, n_samples);
1543                    let added = sweep_intercept(&mut info2, corner, &edge, &block, ngam);
1544
1545                    // Skip the score too when the information was skipped, so
1546                    // info2 and z stay in step (as in array_weights_gene_by_gene).
1547                    if added && s2 > MIN_RESIDUAL_VAR {
1548                        let residual = qr_residuals(&qr, row, design, &coef, n_samples);
1549                        for i in 0..n_samples {
1550                            z[i] += full_w[i] * residual[i] * residual[i] / s2 - (1.0 - hat[i]);
1551                        }
1552                    }
1553                }
1554                (info2, z)
1555            })
1556            .collect();
1557
1558        let mut info2 = ridge.clone();
1559        let mut z: Vec<f64> = w.iter().map(|v| prior_n * (v - 1.0)).collect();
1560        for (part_info, part_z) in partials {
1561            for (slot, v) in info2.iter_mut().zip(part_info.iter()) {
1562                *slot += v;
1563            }
1564            for (slot, v) in z.iter_mut().zip(part_z.iter()) {
1565                *slot += v;
1566            }
1567        }
1568        let scale = n_genes as f64 + prior_n;
1569        for slot in info2.iter_mut() {
1570            *slot /= scale;
1571        }
1572        for slot in z.iter_mut() {
1573            *slot /= scale;
1574        }
1575
1576        let score: Vec<f64> = (0..ngam)
1577            .map(|a| (0..n_samples).map(|i| z2[a * n_samples + i] * z[i]).sum())
1578            .collect();
1579        let step = solve_linear(&info2, &score, ngam)?;
1580        for a in 0..ngam {
1581            gam[a] += step[a];
1582        }
1583        w = weights_from_gamma(z2, &gam, n_samples, ngam);
1584
1585        let convergence = score
1586            .iter()
1587            .zip(step.iter())
1588            .map(|(a, b)| a * b)
1589            .sum::<f64>()
1590            / scale
1591            / ngam as f64;
1592        if !convergence.is_finite() || convergence < tol {
1593            break;
1594        }
1595    }
1596
1597    Ok(w)
1598}
1599
1600///////////////////////////
1601// Duplicate correlation //
1602///////////////////////////
1603
1604/// Consensus intra-block correlation, limma's `duplicateCorrelation`.
1605///
1606/// Fits `y_g = X beta_g + Z u_g + e_g` for each gene, with `u_g` a random block
1607/// effect, and reports `tanh` of the trimmed mean of the Fisher z-transformed
1608/// per-gene correlations `sigma_block^2 / (sigma_block^2 + sigma_e^2)`.
1609///
1610/// Fits `y_g = X beta_g + Z u_g + e_g` per gene with `u_g` a random block
1611/// effect, and reports `tanh` of the trimmed mean of the Fisher z-transformed
1612/// per-gene correlations `sigma_block^2 / (sigma_block^2 + sigma_e^2)`.
1613///
1614/// Returns exactly zero, as limma does, when every block has a single member or
1615/// the block factor already lies in the column space of the design.
1616///
1617/// A gene contributes only with more than `n_coef + 2` observations, more than
1618/// one block and fewer blocks than `n_obs - 1`; others are dropped from the
1619/// trimmed mean, not counted as zero.
1620///
1621/// edgePython replaces the per-gene mixed model with a one-way ANOVA moment
1622/// estimator and drops several of limma's admission rules. This follows limma.
1623/// See `UPSTREAM_DEVIATIONS.md` A14.
1624///
1625/// ### Params
1626///
1627/// * `y` - Row-major `n_genes * n_samples` log-expression values
1628/// * `n_genes` - Number of genes
1629/// * `n_samples` - Number of samples
1630/// * `design` - Row-major `n_samples * n_coef` design matrix
1631/// * `n_coef` - Number of design columns
1632/// * `block` - Block label per sample; only equality of labels matters
1633/// * `weights` - Optional gene by sample observation weights. Non-positive or
1634///   non-finite weights mark the observation missing, as limma does.
1635/// * `trim` - Fraction trimmed from each tail of the z-transformed
1636///   correlations, in `[0, 0.5)`. limma's default is 0.15.
1637///
1638/// ### Returns
1639///
1640/// The consensus correlation, or [`EdgeErrors`] on a shape mismatch, an
1641/// out-of-range `trim`, or when no gene yielded a correlation at all.
1642///
1643/// ### References
1644///
1645/// Smyth, Michaud and Scott, Bioinformatics 21:2067-2075, 2005
1646#[allow(clippy::too_many_arguments)]
1647pub fn duplicate_correlation(
1648    y: &[f64],
1649    n_genes: usize,
1650    n_samples: usize,
1651    design: &[f64],
1652    n_coef: usize,
1653    block: &[usize],
1654    weights: Option<&Recycled<f64>>,
1655    trim: f64,
1656) -> Result<f64, EdgeErrors> {
1657    validate_shapes(y, n_genes, n_samples, design, n_coef)?;
1658    if block.len() != n_samples {
1659        return Err(EdgeErrors::LengthMismatch {
1660            name: "block",
1661            expected: n_samples,
1662            got: block.len(),
1663        });
1664    }
1665    if !(0.0..0.5).contains(&trim) {
1666        return Err(EdgeErrors::InvalidArgument(format!(
1667            "trim must lie in [0, 0.5); got {trim}."
1668        )));
1669    }
1670
1671    let mut levels: Vec<usize> = block.to_vec();
1672    levels.sort_unstable();
1673    levels.dedup();
1674    let max_block = levels
1675        .iter()
1676        .map(|lv| block.iter().filter(|b| *b == lv).count())
1677        .max()
1678        .unwrap_or(0);
1679    if max_block <= 1 {
1680        return Ok(0.0);
1681    }
1682
1683    let design_cm = to_column_major(design, n_samples, n_coef);
1684    let qr = LinpackQr::new(design_cm.clone(), n_samples, n_coef);
1685    let rank = qr.rank();
1686
1687    // Treatment-coded block indicators: if none of them survives the projection
1688    // onto the design's residual space, the blocks are already in the model.
1689    let mut absorbed = true;
1690    for lv in levels.iter().skip(1) {
1691        let mut column: Vec<f64> = block
1692            .iter()
1693            .map(|b| if b == lv { 1.0 } else { 0.0 })
1694            .collect();
1695        qr.qty(&mut column);
1696        if column[rank..].iter().any(|v| v.abs() >= BLOCK_ABSORBED_TOL) {
1697            absorbed = false;
1698            break;
1699        }
1700    }
1701    if absorbed {
1702        return Ok(0.0);
1703    }
1704
1705    let weight_matrix = match weights {
1706        None => None,
1707        Some(w) => {
1708            w.validate(n_genes, n_samples)?;
1709            Some(w.expand(n_genes, n_samples))
1710        }
1711    };
1712
1713    let rho_min = 1.0 / (1.0 - max_block as f64) + RHO_MIN_SLACK;
1714    let correlations: Vec<Option<f64>> = (0..n_genes)
1715        .into_par_iter()
1716        .map(|gene| {
1717            let row = &y[gene * n_samples..(gene + 1) * n_samples];
1718            let wrow = weight_matrix
1719                .as_ref()
1720                .map(|m| &m[gene * n_samples..(gene + 1) * n_samples]);
1721            let observed: Vec<usize> = (0..n_samples)
1722                .filter(|&j| {
1723                    row[j].is_finite() && wrow.is_none_or(|w| w[j].is_finite() && w[j] > 0.0)
1724                })
1725                .collect();
1726            let n_obs = observed.len();
1727            if n_obs <= n_coef + 2 {
1728                return None;
1729            }
1730            let mut present: Vec<usize> = observed.iter().map(|&j| block[j]).collect();
1731            present.sort_unstable();
1732            present.dedup();
1733            let n_blocks = present.len();
1734            if n_blocks <= 1 || n_blocks >= n_obs - 1 {
1735                return None;
1736            }
1737
1738            let mut x = vec![0.0; n_obs * n_coef];
1739            for c in 0..n_coef {
1740                for (k, &j) in observed.iter().enumerate() {
1741                    x[c * n_obs + k] = design_cm[c * n_samples + j];
1742                }
1743            }
1744            let response: Vec<f64> = observed.iter().map(|&j| row[j]).collect();
1745            let z: Vec<f64> = present
1746                .iter()
1747                .flat_map(|lv| {
1748                    observed
1749                        .iter()
1750                        .map(move |&j| if block[j] == *lv { 1.0 } else { 0.0 })
1751                })
1752                .collect();
1753            let w = wrow.map(|w| observed.iter().map(|&j| w[j]).collect::<Vec<f64>>());
1754
1755            let varcomp =
1756                mixed_model_varcomp(&response, &x, &z, w.as_deref(), n_obs, n_coef, n_blocks)?;
1757            let total = varcomp[0] + varcomp[1];
1758            let rho = varcomp[1] / total;
1759            rho.is_finite().then_some(rho)
1760        })
1761        .collect();
1762
1763    let z: Vec<f64> = correlations
1764        .into_iter()
1765        .flatten()
1766        .map(|rho| rho.clamp(rho_min, RHO_MAX).atanh())
1767        .collect();
1768    if z.is_empty() {
1769        return Err(EdgeErrors::InvalidArgument(
1770            "no gene yielded an intra-block correlation; check the block structure and design."
1771                .to_string(),
1772        ));
1773    }
1774    Ok(trimmed_mean(&z, trim)?.tanh())
1775}
1776
1777/// `statmod::mixedModel2Fit(..., only.varcomp = TRUE)`.
1778///
1779/// Projects the response and block design onto the residual space of the fixed
1780/// effects, rotates by the SVD of the projected block design, and fits a gamma
1781/// GLM of the squared rotated residuals on the squared singular values. Its two
1782/// coefficients are the residual and block variance components.
1783///
1784/// `w` scales the response and fixed effects but not the block design, as in
1785/// `statmod`. Left singular vectors beyond the rank are an arbitrary orthonormal
1786/// completion (faer's differs from LAPACK's), which does not matter: those rows
1787/// share a squared singular value of zero and enter the fit only through their
1788/// sum.
1789///
1790/// ### Params
1791///
1792/// * `y` - Response of length `n_obs`
1793/// * `x` - Column-major `n_obs * n_coef` fixed effects
1794/// * `z` - Column-major `n_obs * n_blocks` block indicators
1795/// * `w` - Optional observation weights of length `n_obs`
1796/// * `n_obs` - Number of observations
1797/// * `n_coef` - Number of fixed effects
1798/// * `n_blocks` - Number of blocks
1799///
1800/// ### Returns
1801///
1802/// The residual and block variance components, or `None` when the projection
1803/// leaves no residual space or the fit is degenerate.
1804fn mixed_model_varcomp(
1805    y: &[f64],
1806    x: &[f64],
1807    z: &[f64],
1808    w: Option<&[f64]>,
1809    n_obs: usize,
1810    n_coef: usize,
1811    n_blocks: usize,
1812) -> Option<[f64; 2]> {
1813    let mut xw = x.to_vec();
1814    let mut yw = y.to_vec();
1815    if let Some(w) = w {
1816        for i in 0..n_obs {
1817            let root = w[i].sqrt();
1818            yw[i] *= root;
1819            for c in 0..n_coef {
1820                xw[c * n_obs + i] *= root;
1821            }
1822        }
1823    }
1824
1825    let qr = LinpackQr::new(xw, n_obs, n_coef);
1826    let rank = qr.rank();
1827    let mq = n_obs - rank;
1828    if mq == 0 {
1829        return None;
1830    }
1831
1832    // Q' Z restricted to the residual rows, column-major mq by n_blocks.
1833    let mut qtz = vec![0.0; mq * n_blocks];
1834    for c in 0..n_blocks {
1835        let mut column = z[c * n_obs..(c + 1) * n_obs].to_vec();
1836        qr.qty(&mut column);
1837        qtz[c * mq..(c + 1) * mq].copy_from_slice(&column[rank..]);
1838    }
1839    qr.qty(&mut yw);
1840    let effects = &yw[rank..];
1841
1842    let svd = faer::MatRef::from_column_major_slice(&qtz, mq, n_blocks)
1843        .svd()
1844        .ok()?;
1845    let singular = svd.S().column_vector();
1846    let u = svd.U();
1847
1848    let mut d = vec![0.0; mq];
1849    for i in 0..singular.nrows().min(mq) {
1850        d[i] = singular[i] * singular[i];
1851    }
1852    let dy: Vec<f64> = (0..mq)
1853        .map(|c| {
1854            let projected: f64 = (0..mq).map(|i| u[(i, c)] * effects[i]).sum();
1855            projected * projected
1856        })
1857        .collect();
1858
1859    // Start from the ordinary least squares fit of dy on (1, d).
1860    let mut dx = Vec::with_capacity(mq * 2);
1861    dx.extend(std::iter::repeat_n(1.0, mq));
1862    dx.extend_from_slice(&d);
1863    let dqr = LinpackQr::new(dx.clone(), mq, 2);
1864    if dqr.rank() < 2 {
1865        return None;
1866    }
1867    let mut dy_effects = dy.clone();
1868    dqr.qty(&mut dy_effects);
1869    let start = dqr.solve_r(&dy_effects);
1870    let fitted: Vec<f64> = (0..mq).map(|i| start[0] + start[1] * d[i]).collect();
1871
1872    let non_zero = d.iter().filter(|v| v.abs() > SINGULAR_VALUE_TOL).count();
1873    let mean_d = d.iter().sum::<f64>() / mq as f64;
1874    let var_d = d.iter().map(|v| (v - mean_d).powi(2)).sum::<f64>() / (mq as f64 - 1.0);
1875    if !(mq > 2 && non_zero > 1 && var_d > SINGULAR_VALUE_TOL) {
1876        return Some([start[0], start[1]]);
1877    }
1878
1879    let coef_start = if fitted.iter().all(|v| *v >= 0.0) {
1880        [start[0], start[1]]
1881    } else {
1882        [dy.iter().sum::<f64>() / mq as f64, 0.0]
1883    };
1884    glmgam_fit(&dx, &dy, coef_start, mq)
1885}
1886
1887/// `statmod::glmgam.fit` with a supplied starting value.
1888///
1889/// A gamma GLM with identity link, fitted by Levenberg-damped Fisher scoring on
1890/// the deviance. The damping schedule is reproduced verbatim because the
1891/// objective is nearly flat at the optimum, so where the iteration stops decides
1892/// the answer.
1893///
1894/// ### Params
1895///
1896/// * `x` - Column-major `n * 2` design, an intercept and the squared singular values
1897/// * `y` - Squared rotated residuals, non-negative
1898/// * `start` - Starting coefficients, from the linear fit
1899/// * `n` - Number of observations
1900///
1901/// ### Returns
1902///
1903/// The fitted coefficients, or `None` if the damping ran away or the result was
1904/// not finite.
1905fn glmgam_fit(x: &[f64], y: &[f64], start: [f64; 2], n: usize) -> Option<[f64; 2]> {
1906    let max_y = y.iter().fold(0.0_f64, |acc, v| acc.max(*v));
1907    if max_y == 0.0 {
1908        return Some([0.0, 0.0]);
1909    }
1910
1911    let fitted = |beta: [f64; 2]| -> Vec<f64> {
1912        (0..n)
1913            .map(|i| x[i] * beta[0] + x[n + i] * beta[1])
1914            .collect()
1915    };
1916
1917    let mut beta = start;
1918    let mut mu = fitted(beta);
1919    if mu.iter().any(|v| *v < 0.0) {
1920        return None;
1921    }
1922    let mut deviance = gamma_deviance(y, &mu);
1923
1924    let mut lambda = 0.0;
1925    let mut max_information;
1926    let mut iter = 0;
1927    loop {
1928        iter += 1;
1929        let v: Vec<f64> = mu.iter().map(|m| m * m).collect();
1930        let floor = v.iter().fold(0.0_f64, |acc, s| acc.max(*s)) / GLMGAM_VARIANCE_FLOOR;
1931        let v: Vec<f64> = v.into_iter().map(|s| s.max(floor)).collect();
1932
1933        // Fisher information X' V^-1 X, symmetric two by two.
1934        let mut xvx = [0.0_f64; 4];
1935        for i in 0..n {
1936            let (a, b) = (x[i], x[n + i]);
1937            xvx[0] += a * a / v[i];
1938            xvx[1] += a * b / v[i];
1939            xvx[3] += b * b / v[i];
1940        }
1941        xvx[2] = xvx[1];
1942        max_information = xvx[0].max(xvx[3]);
1943        if iter == 1 {
1944            lambda = ((xvx[0] + xvx[3]) / 2.0).abs() / 2.0;
1945        }
1946
1947        let mut score = [0.0_f64; 2];
1948        for i in 0..n {
1949            let r = (y[i] - mu[i]) / v[i];
1950            score[0] += x[i] * r;
1951            score[1] += x[n + i] * r;
1952        }
1953
1954        let old_beta = beta;
1955        let old_deviance = deviance;
1956        let mut level = 0;
1957        let mut step;
1958        let mut overdamped = false;
1959        loop {
1960            level += 1;
1961            let a = xvx[0] + lambda;
1962            let b = xvx[1];
1963            let c = xvx[3] + lambda;
1964            let det = a * c - b * b;
1965            if det.is_nan() || det <= 0.0 {
1966                return None;
1967            }
1968            step = [
1969                (c * score[0] - b * score[1]) / det,
1970                (a * score[1] - b * score[0]) / det,
1971            ];
1972            beta = [old_beta[0] + step[0], old_beta[1] + step[1]];
1973            mu = fitted(beta);
1974            deviance = gamma_deviance(y, &mu);
1975            let peak = mu.iter().fold(f64::NEG_INFINITY, |acc, v| acc.max(*v));
1976            if deviance <= old_deviance || deviance / peak < GLMGAM_DEVIANCE_TOL {
1977                break;
1978            }
1979            if lambda / max_information > GLMGAM_MAX_DAMPING {
1980                beta = old_beta;
1981                overdamped = true;
1982                break;
1983            }
1984            lambda *= 2.0;
1985        }
1986        if overdamped || lambda / max_information > GLMGAM_MAX_DAMPING {
1987            break;
1988        }
1989        if level == 1 {
1990            lambda /= GLMGAM_DAMPING_RELAX;
1991        }
1992        let peak = mu.iter().fold(f64::NEG_INFINITY, |acc, v| acc.max(*v));
1993        if score[0] * step[0] + score[1] * step[1] < GLMGAM_TOL
1994            || deviance / peak < GLMGAM_DEVIANCE_TOL
1995        {
1996            break;
1997        }
1998        if iter > GLMGAM_MAX_ITER {
1999            break;
2000        }
2001    }
2002
2003    beta.iter().all(|v| v.is_finite()).then_some(beta)
2004}
2005
2006/// Gamma deviance with `statmod`'s treatment of joint zeros.
2007///
2008/// ### Params
2009///
2010/// * `y` - Observed values, non-negative
2011/// * `mu` - Fitted values
2012///
2013/// ### Returns
2014///
2015/// `2 * sum((y - mu) / mu - log(y / mu))` over the terms where `y` and `mu` are
2016/// not both negligible, `Inf` if any fitted value is negative.
2017fn gamma_deviance(y: &[f64], mu: &[f64]) -> f64 {
2018    if mu.iter().any(|v| *v < 0.0) {
2019        return f64::INFINITY;
2020    }
2021    let kept: Vec<usize> = (0..y.len())
2022        .filter(|&i| !(y[i] < GLMGAM_ZERO_TOL && mu[i] < GLMGAM_ZERO_TOL))
2023        .collect();
2024    if kept.is_empty() {
2025        return 0.0;
2026    }
2027    2.0 * kept
2028        .iter()
2029        .map(|&i| (y[i] - mu[i]) / mu[i] - (y[i] / mu[i]).ln())
2030        .sum::<f64>()
2031}
2032
2033///////////
2034// Tests //
2035///////////
2036
2037#[cfg(test)]
2038mod tests {
2039    use super::*;
2040    use approx::assert_relative_eq;
2041
2042    /// Reference values come from limma 3.66 on R 4.5 with `statmod` 1.5. Inputs
2043    /// are embedded alongside the outputs so nothing depends on an RNG stream:
2044    ///
2045    /// ```r
2046    /// set.seed(1); y <- matrix(rnorm(200), 20, 10)
2047    /// cat(sprintf("%.17g", as.vector(t(y))), sep = ",\n")
2048    /// ```
2049    ///
2050    /// Row-major, 20 genes by 10 samples.
2051    ///
2052    /// ### Agreement achieved
2053    ///
2054    /// [`array_weights`] matches limma to better than `1e-16` relative on every
2055    /// fixture (both methods, weighted and not); assertions sit at `1e-14`.
2056    ///
2057    /// [`duplicate_correlation`] reaches `5e-15` at worst, bound by
2058    /// `glmgam.fit`'s own `tol = 1e-6` and SVD rounding rather than the
2059    /// arithmetic; assertions sit at `1e-13`.
2060    const Y: [f64; 200] = [
2061        -0.6264538107423324,
2062        0.9189773716082182,
2063        -0.1645235962535869,
2064        2.401617760504777,
2065        -0.568668732818502,
2066        -0.620366677224124,
2067        -0.5059574621142573,
2068        -1.914359425680012,
2069        0.42510037737244794,
2070        -1.231323421558044,
2071        0.18364332422208227,
2072        0.782136300731067,
2073        -0.2533616801365075,
2074        -0.039240002733169244,
2075        -0.13517861512383206,
2076        0.04211587314423524,
2077        1.3430388251704113,
2078        1.1765833120185623,
2079        -0.2386471009130328,
2080        0.9838955700533794,
2081        -0.8356286124100472,
2082        0.0745649833651906,
2083        0.6969633754047375,
2084        0.6897393624507766,
2085        1.1780869965732044,
2086        -0.9109216485524455,
2087        -0.21457940854686872,
2088        -1.6649724362120033,
2089        1.0584830487090204,
2090        0.21992480366065134,
2091        1.5952808021377916,
2092        -1.9893516958633728,
2093        0.5566631986736573,
2094        0.028002158780666062,
2095        -1.523566800429762,
2096        0.15802877240407498,
2097        -0.17955653004338712,
2098        -0.46353040147238567,
2099        0.886422651374936,
2100        -1.467250029092243,
2101        0.3295077718153605,
2102        0.6198257478947102,
2103        -0.6887556945495199,
2104        -0.7432732088824053,
2105        0.5939461876284216,
2106        -0.6545846439188178,
2107        -0.10019074121356196,
2108        -1.1159201050428456,
2109        -0.619243048231147,
2110        0.5210227426481385,
2111        -0.8204683841180155,
2112        -0.056128739529000764,
2113        -0.7074951569621196,
2114        0.18879229951434284,
2115        0.3329503712135183,
2116        1.7672872693726465,
2117        0.7126663070514053,
2118        -0.7508190011934479,
2119        2.2061024645404674,
2120        -0.15875460471601593,
2121        0.4874290524284853,
2122        -0.1557955067053293,
2123        0.3645819621368303,
2124        -1.8049586288910378,
2125        1.0630998372763627,
2126        0.7167074760172057,
2127        -0.07356440412632632,
2128        2.087166545628347,
2129        -0.25502703014101524,
2130        1.4645873119697974,
2131        0.7383247051292172,
2132        -1.4707523838992744,
2133        0.7685329245154158,
2134        1.4655548615628862,
2135        -0.30418392363430063,
2136        0.9101742294952271,
2137        -0.03763417146704791,
2138        0.017395619693251662,
2139        -1.4244946502128084,
2140        -0.7660819996046648,
2141        0.5757813516534921,
2142        -0.4781500551086204,
2143        -0.11234621215022805,
2144        0.1532533382118977,
2145        0.37001880991628816,
2146        0.3841853578263446,
2147        -0.6816604787556568,
2148        -1.2863005304343256,
2149        -0.14439960195421936,
2150        -0.43021175392854644,
2151        -0.305388387156356,
2152        0.4179415601997024,
2153        0.881107726454215,
2154        2.1726116703621523,
2155        0.26709879077223103,
2156        1.6821760805194184,
2157        -0.32427027224631916,
2158        -1.6406055344185784,
2159        0.20753833923234472,
2160        -0.9261094973774372,
2161        1.5117811684508478,
2162        1.3586795515290442,
2163        0.398105880367068,
2164        0.4755095288996625,
2165        -0.5425200309916505,
2166        -0.6357364539489768,
2167        0.06016044043451516,
2168        0.4501871012726556,
2169        2.3079783990593614,
2170        -0.17710396143654025,
2171        0.38984323641143104,
2172        -0.10278772734299552,
2173        -0.6120263932507712,
2174        -0.7099464309218142,
2175        1.2078678059831724,
2176        -0.461644730360566,
2177        -0.5888944862596638,
2178        -0.018559832714637976,
2179        0.10580236789371146,
2180        0.4020117794863379,
2181        -0.6212405805418039,
2182        0.38767161155936924,
2183        0.34111969142442483,
2184        0.6107263534890547,
2185        1.1604026156949516,
2186        1.4322822385416627,
2187        0.5314961926325724,
2188        -0.3180683745438444,
2189        0.4569988054234134,
2190        -0.7317481731196062,
2191        -2.2146998871775,
2192        -0.05380504058290512,
2193        -1.1293630960807926,
2194        -0.9340976316442515,
2195        0.7002136495149983,
2196        -0.6506963533103668,
2197        -1.5183940817867871,
2198        -0.9293621474537023,
2199        -0.077152935356531,
2200        0.8303731679816739,
2201        1.1249309181431082,
2202        -1.3770595568286068,
2203        1.4330237017010372,
2204        -1.2536334002391023,
2205        1.5868334545408456,
2206        -0.20738074360196543,
2207        0.3065578607897656,
2208        -1.4874603101414847,
2209        -0.3340008423665445,
2210        -1.208082786304465,
2211        -0.04493360901523085,
2212        -0.41499456329967976,
2213        1.98039989850586,
2214        0.29144623551746296,
2215        0.558486425565304,
2216        -0.3928079294419839,
2217        -1.536449823537586,
2218        -1.0751922966156808,
2219        -0.034726028311276184,
2220        -1.0479844128077418,
2221        -0.016190263098946084,
2222        -0.3942899537103493,
2223        -0.36722147646650916,
2224        -0.4432918732184329,
2225        -1.2765922084580363,
2226        -0.3199928685485067,
2227        -0.30097612683661074,
2228        1.000028803713914,
2229        0.787639605630162,
2230        1.441157706844281,
2231        0.9438362106852992,
2232        -0.05931339671118567,
2233        -1.0441346263165303,
2234        0.0011053516316241311,
2235        -0.5732654142368862,
2236        -0.27911330297655895,
2237        -0.5282799044450062,
2238        -0.6212666947968234,
2239        2.075245008652285,
2240        -1.0158474653046492,
2241        0.8212211950980886,
2242        1.1000253719838828,
2243        0.5697196274424129,
2244        0.07434132415166406,
2245        -1.2246126148983558,
2246        0.494188331267827,
2247        -0.6520947806809989,
2248        -1.3844268473844912,
2249        1.0273924387637678,
2250        0.4119747123175149,
2251        0.5939013212175087,
2252        0.7631757484575442,
2253        -0.13505460388082438,
2254        -0.589520946188072,
2255        -0.4734006364393116,
2256        -0.17733048226960638,
2257        -0.05689677784739257,
2258        1.8692906224235806,
2259        1.2079083983867038,
2260        -0.38107605110891957,
2261    ];
2262
2263    /// Observation weights, row-major 20 by 10:
2264    ///
2265    /// ```r
2266    /// set.seed(42); W <- matrix(runif(200, 0.2, 3), 20, 10)
2267    /// cat(sprintf("%.17g", as.vector(t(W))), sep = ",\n")
2268    /// ```
2269    const W: [f64; 200] = [
2270        2.761456921789795,
2271        2.731287884432822,
2272        1.26276587350294,
2273        2.091700368653983,
2274        1.8284912070259451,
2275        1.9534869646653532,
2276        1.1988215981982648,
2277        0.6156504987739027,
2278        2.8175042705610394,
2279        2.770557968318462,
2280        2.823811157234013,
2281        0.5883884696289897,
2282        1.4201604379341006,
2283        2.9518881541676816,
2284        0.6421345829032361,
2285        0.8080415550619363,
2286        1.349778352305293,
2287        0.42474050642922523,
2288        1.7413834362290799,
2289        2.615363375749439,
2290        1.001190697401762,
2291        2.9688968409784136,
2292        0.3048068920150399,
2293        2.3267239491455256,
2294        1.2052792564034462,
2295        0.8063884707167744,
2296        1.8057325170375407,
2297        1.4993947437033057,
2298        1.8849454591982067,
2299        1.0875306658446788,
2300        2.525253352988511,
2301        2.8506710511632263,
2302        2.9259117585606873,
2303        1.7861675874330103,
2304        2.007769259531051,
2305        1.2890460802242159,
2306        1.85109924999997,
2307        2.382230852078646,
2308        0.751584566757083,
2309        0.9259296127595007,
2310        1.9968874529004097,
2311        0.430825162678957,
2312        1.4089034968987106,
2313        2.579131211992353,
2314        2.372305415384471,
2315        2.838875937554985,
2316        2.2150404185988006,
2317        2.2538782867603,
2318        1.6986625095829366,
2319        2.2783460654318333,
2320        1.6534686575643718,
2321        1.6397929961793123,
2322        2.881214470602572,
2323        0.7305270191282034,
2324        1.778211156371981,
2325        2.895302438549697,
2326        1.3059245266020298,
2327        2.488245244324207,
2328        0.7027560699731111,
2329        2.292611127067357,
2330        2.2624472809955476,
2331        1.292569707892835,
2332        2.685713735409081,
2333        0.959602521173656,
2334        0.8543695161119104,
2335        2.2715947818011046,
2336        2.773771001584828,
2337        0.6764549475163221,
2338        1.4652821843512356,
2339        2.7701312952674924,
2340        0.5770664722658694,
2341        2.736066766548902,
2342        1.991940554510802,
2343        2.518843758665025,
2344        0.451945445779711,
2345        2.253088536020368,
2346        2.8951968220062554,
2347        2.84521691147238,
2348        1.08774938499555,
2349        2.4209353857673706,
2350        2.0395784131251276,
2351        1.4515149586834013,
2352        2.918706509005278,
2353        2.1409734970889986,
2354        0.4397137817926705,
2355        1.7001316119916736,
2356        0.8538658714853227,
2357        1.0221467558294535,
2358        0.5252890771254897,
2359        0.5733229311183095,
2360        2.174181395303458,
2361        2.5408119279891253,
2362        1.932746980525553,
2363        0.8735252710059285,
2364        1.054611434508115,
2365        0.20636430503800512,
2366        2.228593279980123,
2367        0.6174017468467354,
2368        0.7210860389284789,
2369        1.0056993060745298,
2370        1.4816769734956323,
2371        2.2652677297592163,
2372        1.1335961915552615,
2373        0.32036862885579465,
2374        2.0687942410819233,
2375        1.905024867132306,
2376        2.7301766703836616,
2377        2.214260055590421,
2378        2.243244270607829,
2379        0.7450932026840746,
2380        2.213514304626733,
2381        2.4709543955512343,
2382        1.1708950949832797,
2383        0.5933414635248482,
2384        0.20066891042515636,
2385        2.5430443664081395,
2386        1.8897274373099207,
2387        1.1074406669475139,
2388        1.3532417993992567,
2389        2.3955062714405355,
2390        2.8170822920277714,
2391        1.286703191883862,
2392        1.3157591519877314,
2393        0.805879162158817,
2394        0.7839958794414997,
2395        2.3042631755582987,
2396        1.9682204369455576,
2397        2.3806665958836675,
2398        1.359339108876884,
2399        0.5608420528471469,
2400        0.9152007081545889,
2401        2.1184752423316238,
2402        2.3971397719345986,
2403        1.5423159797675907,
2404        2.812495556566864,
2405        1.467648403160274,
2406        2.8246803790330883,
2407        1.3044348050840198,
2408        1.5448683612048626,
2409        0.5614499938674271,
2410        1.4944199031218885,
2411        0.21105534862726927,
2412        0.30902217514812946,
2413        0.7527489583939314,
2414        2.7918052961118516,
2415        1.7002119827084243,
2416        2.581351701915264,
2417        2.1000600301660595,
2418        1.3969845036976039,
2419        0.4023087115958333,
2420        2.8320406637154516,
2421        2.532165024708956,
2422        2.2966270812787113,
2423        2.214196345489472,
2424        2.255464042816311,
2425        1.7046547468751667,
2426        1.8234985173679887,
2427        2.3723101196810603,
2428        0.5821730083785951,
2429        0.3487625537440181,
2430        2.9390339994803067,
2431        0.22053561126813293,
2432        2.096375124529004,
2433        0.22207726845517756,
2434        1.1326015535742044,
2435        0.20386636201292277,
2436        2.499930986948311,
2437        0.7260333232581615,
2438        2.5091023378074166,
2439        1.68924842197448,
2440        0.5289646126329899,
2441        0.7814451239071786,
2442        0.6795401250943541,
2443        1.25137190092355,
2444        1.6421773235313593,
2445        1.1958646707236766,
2446        0.5184121054597199,
2447        0.2814402929507196,
2448        1.8584518791176379,
2449        0.514463076647371,
2450        1.5299918283708394,
2451        2.73848394183442,
2452        0.9310462987050414,
2453        1.6403415832668542,
2454        2.283129009697586,
2455        1.9139726524241267,
2456        2.3406217242591083,
2457        0.5799986317753791,
2458        2.424311537388712,
2459        2.2809256152249873,
2460        1.7689316894859075,
2461        1.9129802016541362,
2462        1.640356217045337,
2463        0.20439755180850627,
2464        1.933645872119814,
2465        2.5210379655472934,
2466        1.9461176801472901,
2467        2.104459698777646,
2468        2.3532907919958235,
2469        2.2476833363994957,
2470    ];
2471
2472    /// A three-coefficient design, row-major 10 by 3:
2473    ///
2474    /// ```r
2475    /// set.seed(7); X3 <- cbind(1, rep(0:1, each = 5), round(rnorm(10), 6))
2476    /// ```
2477    const X3: [f64; 30] = [
2478        1.0, 0.0, 2.287247, 1.0, 0.0, -1.196772, 1.0, 0.0, -0.694293, 1.0, 0.0, -0.412293, 1.0,
2479        0.0, -0.970673, 1.0, 1.0, -0.94728, 1.0, 1.0, 0.748139, 1.0, 1.0, -0.116955, 1.0, 1.0,
2480        0.152658, 1.0, 1.0, 2.189978,
2481    ];
2482
2483    const N_GENES: usize = 20;
2484    const N_SAMPLES: usize = 10;
2485
2486    /// `cbind(1, rep(0:1, each = 5))`, row-major.
2487    fn two_group_design() -> Vec<f64> {
2488        let mut d = Vec::with_capacity(20);
2489        for i in 0..10 {
2490            d.push(1.0);
2491            d.push(if i < 5 { 0.0 } else { 1.0 });
2492        }
2493        d
2494    }
2495
2496    /// `cbind(1, rep(0:1, each = 5), rep(1:0, each = 5))`, row-major. Rank two.
2497    fn rank_deficient_design() -> Vec<f64> {
2498        let mut d = Vec::with_capacity(30);
2499        for i in 0..10 {
2500            d.push(1.0);
2501            d.push(if i < 5 { 0.0 } else { 1.0 });
2502            d.push(if i < 5 { 1.0 } else { 0.0 });
2503        }
2504        d
2505    }
2506
2507    fn assert_weights(got: &[f64], want: &[f64], tolerance: f64) {
2508        assert_eq!(got.len(), want.len());
2509        for (g, w) in got.iter().zip(want.iter()) {
2510            assert_relative_eq!(g, w, max_relative = tolerance);
2511        }
2512    }
2513
2514    //////////////////////////////
2515    // arrayWeights, no weights //
2516    //////////////////////////////
2517
2518    /// `Rscript -e 'suppressMessages(library(limma)); set.seed(1); y <- matrix(rnorm(200), 20, 10); X <- cbind(1, rep(0:1, each=5)); cat(arrayWeights(y, X), "\n")'`
2519    ///
2520    /// `method = "auto"` resolves to `"reml"` here, so this is also the default
2521    /// path of [`ArrayWeightParams`].
2522    #[test]
2523    fn test_array_weights_reml_matches_limma() {
2524        let design = two_group_design();
2525        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, None).unwrap();
2526        let want = [
2527            0.9742805411895777,
2528            1.026017651831166,
2529            1.3455650214699084,
2530            1.0480584447646801,
2531            0.8463342237363761,
2532            0.9747809021137841,
2533            1.4807872631702028,
2534            0.8161127881463421,
2535            0.683967423080732,
2536            1.0402644635138032,
2537        ];
2538        assert_weights(&got, &want, 1e-14);
2539    }
2540
2541    /// `arrayWeights(y, X, method = "genebygene")` on the same fixture.
2542    #[test]
2543    fn test_array_weights_gene_by_gene_matches_limma() {
2544        let design = two_group_design();
2545        let params = ArrayWeightParams {
2546            method: ArrayWeightMethod::GeneByGene,
2547            ..Default::default()
2548        };
2549        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2550        let want = [
2551            0.9855865468275133,
2552            1.0354449545106663,
2553            1.3167194637260566,
2554            0.9892096743664166,
2555            0.883308082435258,
2556            0.9875617156255206,
2557            1.4399727557165392,
2558            0.8087646791684824,
2559            0.7206377138806224,
2560            1.027606206031394,
2561        ];
2562        assert_weights(&got, &want, 1e-14);
2563    }
2564
2565    /// `arrayWeights(y, X, method = "genebygene", prior.n = 5)`. Halving the
2566    /// prior widens the spread of the weights.
2567    #[test]
2568    fn test_array_weights_prior_n_loosens_the_shrinkage() {
2569        let design = two_group_design();
2570        let params = ArrayWeightParams::new(ArrayWeightMethod::GeneByGene, 5.0, 50, 1e-5, None, 0);
2571        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2572        let want = [
2573            0.9753130283019934,
2574            1.0345246343437244,
2575            1.4463951214443538,
2576            0.9581734676776122,
2577            0.86147110226916,
2578            0.982408539384646,
2579            1.6499951468701075,
2580            0.7624393728516808,
2581            0.6570573074282369,
2582            1.0222567198392511,
2583        ];
2584        assert_weights(&got, &want, 1e-14);
2585    }
2586
2587    /// `arrayWeights(y, method = "reml")` and `method = "genebygene"`, that is
2588    /// an intercept-only design.
2589    #[test]
2590    fn test_array_weights_intercept_only_design() {
2591        let design = vec![1.0; N_SAMPLES];
2592        let reml = array_weights(&Y, N_GENES, N_SAMPLES, &design, 1, None, None).unwrap();
2593        let want_reml = [
2594            1.0039401405346158,
2595            1.2457559559217304,
2596            1.2269924234759986,
2597            0.9319055327134608,
2598            0.8102086600806064,
2599            1.2349234755723884,
2600            1.4595139406006172,
2601            0.6290074491230409,
2602            0.8344850279179324,
2603            0.9122753178175574,
2604        ];
2605        assert_weights(&reml, &want_reml, 1e-14);
2606
2607        let params = ArrayWeightParams {
2608            method: ArrayWeightMethod::GeneByGene,
2609            ..Default::default()
2610        };
2611        let gbg = array_weights(&Y, N_GENES, N_SAMPLES, &design, 1, None, Some(params)).unwrap();
2612        let want_gbg = [
2613            1.0332474602716566,
2614            1.2339092930392186,
2615            1.206630141064309,
2616            0.8772804188552625,
2617            0.8376944311823541,
2618            1.218771000735954,
2619            1.4282811697295135,
2620            0.6395457423258586,
2621            0.8434756288804333,
2622            0.9419637959843594,
2623        ];
2624        assert_weights(&gbg, &want_gbg, 1e-14);
2625    }
2626
2627    /// `arrayWeights(y, X3, method = "reml")` and `"genebygene"` on an intercept,
2628    /// group indicator and continuous covariate. Three coefficients give `Q2`
2629    /// six columns, which a two-column design never exercises.
2630    #[test]
2631    fn test_array_weights_three_coefficient_design() {
2632        let reml = array_weights(&Y, N_GENES, N_SAMPLES, &X3, 3, None, None).unwrap();
2633        let want_reml = [
2634            1.1123858654314625,
2635            0.8725880455354021,
2636            1.3674933744743083,
2637            1.1067965225608778,
2638            0.8492159040121154,
2639            1.1449733538385882,
2640            1.414815203220406,
2641            0.6507714947593773,
2642            0.7456578248308494,
2643            1.0196730581383637,
2644        ];
2645        assert_weights(&reml, &want_reml, 1e-14);
2646
2647        let params = ArrayWeightParams {
2648            method: ArrayWeightMethod::GeneByGene,
2649            ..Default::default()
2650        };
2651        let gbg = array_weights(&Y, N_GENES, N_SAMPLES, &X3, 3, None, Some(params)).unwrap();
2652        let want_gbg = [
2653            1.0789676466626188,
2654            0.9176081655063308,
2655            1.3301594166920214,
2656            1.039568426333845,
2657            0.8938386978754715,
2658            1.1276155513517505,
2659            1.388971000002176,
2660            0.673140420315832,
2661            0.7647038723639742,
2662            1.0135966976839954,
2663        ];
2664        assert_weights(&gbg, &want_gbg, 1e-14);
2665    }
2666
2667    /// A design whose third column is `1 - ` its second is reduced to the first
2668    /// two, so the answer has to equal the two-column fit exactly. limma does
2669    /// the same through `QR$pivot[1:QR$rank]`.
2670    #[test]
2671    fn test_array_weights_reduces_a_rank_deficient_design() {
2672        let got = array_weights(
2673            &Y,
2674            N_GENES,
2675            N_SAMPLES,
2676            &rank_deficient_design(),
2677            3,
2678            None,
2679            None,
2680        )
2681        .unwrap();
2682        let want = [
2683            0.9742805411895777,
2684            1.026017651831166,
2685            1.3455650214699084,
2686            1.0480584447646801,
2687            0.8463342237363761,
2688            0.9747809021137841,
2689            1.4807872631702028,
2690            0.8161127881463421,
2691            0.683967423080732,
2692            1.0402644635138032,
2693        ];
2694        assert_weights(&got, &want, 1e-14);
2695    }
2696
2697    ////////////////////////////////
2698    // arrayWeights, with weights //
2699    ////////////////////////////////
2700
2701    /// `arrayWeights(y, X, weights = W, method = "genebygene")`. This is also
2702    /// what `method = "auto"` picks once weights are present, and the branch
2703    /// `voomLmFit` always takes.
2704    #[test]
2705    fn test_array_weights_gene_by_gene_with_observation_weights() {
2706        let design = two_group_design();
2707        let weights = Recycled::full(W.to_vec(), N_GENES, N_SAMPLES).unwrap();
2708        let params = ArrayWeightParams {
2709            method: ArrayWeightMethod::GeneByGene,
2710            ..Default::default()
2711        };
2712        let got = array_weights(
2713            &Y,
2714            N_GENES,
2715            N_SAMPLES,
2716            &design,
2717            2,
2718            Some(&weights),
2719            Some(params),
2720        )
2721        .unwrap();
2722        let want = [
2723            0.9165352875614898,
2724            1.0605424640617098,
2725            1.3409356509781636,
2726            0.9427478239354696,
2727            0.996662010628232,
2728            1.0484128636731012,
2729            1.1908686059553828,
2730            0.8437343085959492,
2731            0.7664520368017649,
2732            1.0113104325677869,
2733        ];
2734        assert_weights(&got, &want, 1e-14);
2735    }
2736
2737    /// `arrayWeights(y, X, weights = W, method = "reml")`, which limma routes to
2738    /// `.arrayWeightsPrWtsREML` rather than `.arrayWeightsREML`.
2739    #[test]
2740    fn test_array_weights_reml_with_observation_weights() {
2741        let design = two_group_design();
2742        let weights = Recycled::full(W.to_vec(), N_GENES, N_SAMPLES).unwrap();
2743        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap();
2744        let want = [
2745            0.9209037806194705,
2746            1.0645265066532537,
2747            1.37274621213818,
2748            0.9947452835222503,
2749            0.9441035850402303,
2750            1.040449837341745,
2751            1.1976947015007404,
2752            0.8374281704069466,
2753            0.7330195053976797,
2754            1.0343728257277731,
2755        ];
2756        assert_weights(&got, &want, 1e-14);
2757    }
2758
2759    /// Both weighted branches on the three-coefficient design.
2760    #[test]
2761    fn test_array_weights_weighted_three_coefficient_design() {
2762        let weights = Recycled::full(W.to_vec(), N_GENES, N_SAMPLES).unwrap();
2763        let params = ArrayWeightParams {
2764            method: ArrayWeightMethod::GeneByGene,
2765            ..Default::default()
2766        };
2767        let gbg =
2768            array_weights(&Y, N_GENES, N_SAMPLES, &X3, 3, Some(&weights), Some(params)).unwrap();
2769        let want_gbg = [
2770            1.1997764476593915,
2771            0.9528873823812497,
2772            1.3052954876962655,
2773            0.9659096989828025,
2774            0.9271573637545184,
2775            1.1130783426736774,
2776            1.078354693852502,
2777            0.7064691660204454,
2778            0.8130233115307034,
2779            1.0853647189535636,
2780        ];
2781        assert_weights(&gbg, &want_gbg, 1e-14);
2782
2783        let reml = array_weights(&Y, N_GENES, N_SAMPLES, &X3, 3, Some(&weights), None).unwrap();
2784        let want_reml = [
2785            1.2370489318901068,
2786            0.9124220797886592,
2787            1.3534944344493567,
2788            1.0292176614054709,
2789            0.8823440119643678,
2790            1.1423166038779065,
2791            1.077558416900025,
2792            0.6755922802424658,
2793            0.7926814918065541,
2794            1.093465231695838,
2795        ];
2796        assert_weights(&reml, &want_reml, 1e-14);
2797    }
2798
2799    /// A zero weight becomes a missing observation and a weight of one, which is
2800    /// what forces the gene-by-gene sweep down its subsetting branch:
2801    ///
2802    /// ```r
2803    /// W0 <- W; W0[1, 1] <- 0; W0[3, 7] <- 0
2804    /// arrayWeights(y, X, weights = W0, method = "genebygene")
2805    /// ```
2806    #[test]
2807    fn test_array_weights_zero_weights_become_missing() {
2808        let design = two_group_design();
2809        let mut w = W.to_vec();
2810        w[0] = 0.0;
2811        w[2 * N_SAMPLES + 6] = 0.0;
2812        let weights = Recycled::full(w, N_GENES, N_SAMPLES).unwrap();
2813        let params = ArrayWeightParams {
2814            method: ArrayWeightMethod::GeneByGene,
2815            ..Default::default()
2816        };
2817        let got = array_weights(
2818            &Y,
2819            N_GENES,
2820            N_SAMPLES,
2821            &design,
2822            2,
2823            Some(&weights),
2824            Some(params),
2825        )
2826        .unwrap();
2827        let want = [
2828            0.9468801432504194,
2829            1.0719407282350122,
2830            1.3187893121618048,
2831            0.987724517154794,
2832            0.9497467768369868,
2833            1.047094852664723,
2834            1.1467989187355798,
2835            0.8571863121305534,
2836            0.7667221536949717,
2837            1.0090869707631613,
2838        ];
2839        assert_weights(&got, &want, 1e-14);
2840    }
2841
2842    /// Missing expression values, reached without any weights at all:
2843    ///
2844    /// ```r
2845    /// yNA <- y; yNA[2, 3] <- NA; yNA[5, 9] <- NA
2846    /// arrayWeights(yNA, X, method = "genebygene")
2847    /// ```
2848    #[test]
2849    fn test_array_weights_handles_missing_values() {
2850        let design = two_group_design();
2851        let mut y = Y.to_vec();
2852        y[N_SAMPLES + 2] = f64::NAN;
2853        y[4 * N_SAMPLES + 8] = f64::NAN;
2854        let params = ArrayWeightParams {
2855            method: ArrayWeightMethod::GeneByGene,
2856            ..Default::default()
2857        };
2858        let got = array_weights(&y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2859        let want = [
2860            0.9844172815741963,
2861            1.0605350684234676,
2862            1.3055898280414249,
2863            0.9877067956763591,
2864            0.8794361790533194,
2865            0.9842652252777853,
2866            1.4500921121935295,
2867            0.8113169367939767,
2868            0.6976973359698356,
2869            1.0454211225192833,
2870        ];
2871        assert_weights(&got, &want, 1e-14);
2872    }
2873
2874    /// A non-finite value makes the REML sweep drop that gene, as limma does
2875    /// before `.arrayWeightsREML`. Asserts the result is finite, normalised and
2876    /// differs from the complete-matrix fit.
2877    #[test]
2878    fn test_array_weights_reml_drops_incomplete_genes() {
2879        let design = two_group_design();
2880        let mut y = Y.to_vec();
2881        y[3] = f64::INFINITY;
2882        let got = array_weights(&y, N_GENES, N_SAMPLES, &design, 2, None, None).unwrap();
2883        assert!(got.iter().all(|v| v.is_finite() && *v > 0.0));
2884        let full = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, None).unwrap();
2885        assert!(
2886            got.iter()
2887                .zip(full.iter())
2888                .any(|(a, b)| (a - b).abs() > 1e-6)
2889        );
2890    }
2891
2892    //////////////////////////////
2893    // arrayWeights, var.design //
2894    //////////////////////////////
2895
2896    /// A two-level variance design pools the samples into two weights:
2897    ///
2898    /// ```r
2899    /// VD <- cbind(rep(c(1, -1), each = 5))
2900    /// arrayWeights(y, X, var.design = VD, method = "reml")
2901    /// arrayWeights(y, X, var.design = VD, method = "genebygene")
2902    /// ```
2903    #[test]
2904    fn test_array_weights_with_a_var_design() {
2905        let design = two_group_design();
2906        let var_design: Vec<f64> = (0..10).map(|i| if i < 5 { 1.0 } else { -1.0 }).collect();
2907
2908        let params = ArrayWeightParams {
2909            var_design: Some(var_design.clone()),
2910            n_var_coef: 1,
2911            ..Default::default()
2912        };
2913        let reml = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2914        let want_reml = [
2915            1.0380239930063324,
2916            1.0380239930063324,
2917            1.0380239930063324,
2918            1.0380239930063324,
2919            1.0380239930063324,
2920            0.9633688688676578,
2921            0.9633688688676578,
2922            0.9633688688676578,
2923            0.9633688688676578,
2924            0.9633688688676578,
2925        ];
2926        assert_weights(&reml, &want_reml, 1e-14);
2927
2928        let params = ArrayWeightParams {
2929            method: ArrayWeightMethod::GeneByGene,
2930            var_design: Some(var_design),
2931            n_var_coef: 1,
2932            ..Default::default()
2933        };
2934        let gbg = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2935        let want_gbg = [
2936            1.035442258871988,
2937            1.035442258871988,
2938            1.035442258871988,
2939            1.035442258871988,
2940            1.035442258871988,
2941            0.9657708978281427,
2942            0.9657708978281427,
2943            0.9657708978281427,
2944            0.9657708978281427,
2945            0.9657708978281427,
2946        ];
2947        assert_weights(&gbg, &want_gbg, 1e-14);
2948    }
2949
2950    /// Centring kills an intercept column in the variance design and the rank
2951    /// reduction then drops it, so `cbind(1, g)` must give the same answer as
2952    /// `cbind(g)`:
2953    ///
2954    /// ```r
2955    /// arrayWeights(y, X, var.design = cbind(1, rep(c(1, -1), each = 5)), method = "reml")
2956    /// ```
2957    #[test]
2958    fn test_var_design_intercept_column_is_dropped() {
2959        let design = two_group_design();
2960        let mut var_design = Vec::with_capacity(20);
2961        for i in 0..10 {
2962            var_design.push(1.0);
2963            var_design.push(if i < 5 { 1.0 } else { -1.0 });
2964        }
2965        let params = ArrayWeightParams {
2966            var_design: Some(var_design),
2967            n_var_coef: 2,
2968            ..Default::default()
2969        };
2970        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
2971        let want = [
2972            1.0380239930063324,
2973            1.0380239930063324,
2974            1.0380239930063324,
2975            1.0380239930063324,
2976            1.0380239930063324,
2977            0.9633688688676578,
2978            0.9633688688676578,
2979            0.9633688688676578,
2980            0.9633688688676578,
2981            0.9633688688676578,
2982        ];
2983        assert_weights(&got, &want, 1e-14);
2984    }
2985
2986    ///////////////////////////////
2987    // arrayWeights, early exits //
2988    ///////////////////////////////
2989
2990    #[test]
2991    fn test_array_weights_returns_ones_for_a_single_gene() {
2992        let design = two_group_design();
2993        let got = array_weights(&Y[..N_SAMPLES], 1, N_SAMPLES, &design, 2, None, None).unwrap();
2994        assert_eq!(got, vec![1.0; N_SAMPLES]);
2995    }
2996
2997    /// Three samples against a two-column design leaves one residual degree of
2998    /// freedom, which limma refuses to work with.
2999    #[test]
3000    fn test_array_weights_returns_ones_without_residual_df() {
3001        let y: Vec<f64> = (0..12).map(|i| i as f64 * 0.37).collect();
3002        let design = vec![1.0, 0.0, 1.0, 1.0, 1.0, 2.0];
3003        let got = array_weights(&y, 4, 3, &design, 2, None, None).unwrap();
3004        assert_eq!(got, vec![1.0; 3]);
3005    }
3006
3007    //////////////////////////
3008    // arrayWeights, errors //
3009    //////////////////////////
3010
3011    #[test]
3012    fn test_array_weights_rejects_empty_input() {
3013        let err = array_weights(&[], 0, 4, &[1.0; 4], 1, None, None).unwrap_err();
3014        assert!(matches!(err, EdgeErrors::EmptyCounts { .. }));
3015        let err = array_weights(&[], 4, 0, &[1.0; 4], 1, None, None).unwrap_err();
3016        assert!(matches!(err, EdgeErrors::EmptyCounts { .. }));
3017    }
3018
3019    #[test]
3020    fn test_array_weights_rejects_a_shape_mismatch() {
3021        let err = array_weights(&[1.0; 7], 2, 4, &[1.0; 4], 1, None, None).unwrap_err();
3022        assert!(matches!(err, EdgeErrors::LengthMismatch { name: "y", .. }));
3023        let err = array_weights(&[1.0; 8], 2, 4, &[1.0; 3], 1, None, None).unwrap_err();
3024        assert!(matches!(
3025            err,
3026            EdgeErrors::LengthMismatch { name: "design", .. }
3027        ));
3028    }
3029
3030    #[test]
3031    fn test_array_weights_rejects_a_zero_coefficient_design() {
3032        let err = array_weights(&[1.0; 8], 2, 4, &[], 0, None, None).unwrap_err();
3033        assert!(matches!(err, EdgeErrors::MustBePositive(_)));
3034    }
3035
3036    #[test]
3037    fn test_array_weights_rejects_bad_weights() {
3038        let design = two_group_design();
3039        let mut w = W.to_vec();
3040        w[3] = -1.0;
3041        let weights = Recycled::full(w, N_GENES, N_SAMPLES).unwrap();
3042        let err =
3043            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap_err();
3044        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3045
3046        let mut w = W.to_vec();
3047        w[3] = f64::INFINITY;
3048        let weights = Recycled::full(w, N_GENES, N_SAMPLES).unwrap();
3049        let err =
3050            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap_err();
3051        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3052
3053        let weights = Recycled::by_sample(vec![1.0; 3]);
3054        let err =
3055            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap_err();
3056        assert!(matches!(err, EdgeErrors::LengthMismatch { .. }));
3057    }
3058
3059    #[test]
3060    fn test_array_weights_rejects_bad_parameters() {
3061        let design = two_group_design();
3062        let bad_prior = ArrayWeightParams {
3063            prior_n: -1.0,
3064            ..Default::default()
3065        };
3066        let err =
3067            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(bad_prior)).unwrap_err();
3068        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3069
3070        let bad_tol = ArrayWeightParams {
3071            tol: 0.0,
3072            ..Default::default()
3073        };
3074        let err =
3075            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(bad_tol)).unwrap_err();
3076        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3077
3078        let bad_iter = ArrayWeightParams {
3079            max_iter: 0,
3080            ..Default::default()
3081        };
3082        let err =
3083            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(bad_iter)).unwrap_err();
3084        assert!(matches!(err, EdgeErrors::MustBePositive(_)));
3085    }
3086
3087    #[test]
3088    fn test_array_weights_rejects_a_bad_var_design() {
3089        let design = two_group_design();
3090        let short = ArrayWeightParams {
3091            var_design: Some(vec![1.0; 7]),
3092            n_var_coef: 1,
3093            ..Default::default()
3094        };
3095        let err = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(short)).unwrap_err();
3096        assert!(matches!(
3097            err,
3098            EdgeErrors::LengthMismatch {
3099                name: "var_design",
3100                ..
3101            }
3102        ));
3103
3104        let zero_cols = ArrayWeightParams {
3105            var_design: Some(vec![1.0; 10]),
3106            n_var_coef: 0,
3107            ..Default::default()
3108        };
3109        let err =
3110            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(zero_cols)).unwrap_err();
3111        assert!(matches!(err, EdgeErrors::MustBePositive(_)));
3112    }
3113
3114    /// A variance design that centres away to nothing leaves no coefficients to
3115    /// estimate, so the weights stay at one.
3116    #[test]
3117    fn test_array_weights_constant_var_design_gives_unit_weights() {
3118        let design = two_group_design();
3119        let params = ArrayWeightParams {
3120            var_design: Some(vec![1.0; 10]),
3121            n_var_coef: 1,
3122            ..Default::default()
3123        };
3124        let got = array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, None, Some(params)).unwrap();
3125        assert_eq!(got, vec![1.0; N_SAMPLES]);
3126    }
3127
3128    /// The parallel reductions are chunked in index order, so repeated calls
3129    /// have to agree bit for bit.
3130    #[test]
3131    fn test_array_weights_is_reproducible() {
3132        let design = two_group_design();
3133        let weights = Recycled::full(W.to_vec(), N_GENES, N_SAMPLES).unwrap();
3134        let first =
3135            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap();
3136        let second =
3137            array_weights(&Y, N_GENES, N_SAMPLES, &design, 2, Some(&weights), None).unwrap();
3138        assert_eq!(first, second);
3139    }
3140
3141    //////////////////////////
3142    // duplicateCorrelation //
3143    //////////////////////////
3144
3145    /// `Rscript -e 'suppressMessages(library(limma)); set.seed(1); y <- matrix(rnorm(200), 20, 10); X <- cbind(1, rep(0:1, each=5)); b <- rep(1:5, 2); cat(duplicateCorrelation(y, X, block=b)$consensus.correlation, "\n")'`
3146    #[test]
3147    fn test_duplicate_correlation_balanced_blocks() {
3148        let design = two_group_design();
3149        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3150        let got =
3151            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3152        assert_relative_eq!(got, 0.07936899121798541, max_relative = 1e-14);
3153    }
3154
3155    /// The same fixture with the trimming turned off, `trim = 0`.
3156    #[test]
3157    fn test_duplicate_correlation_untrimmed() {
3158        let design = two_group_design();
3159        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3160        let got =
3161            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.0).unwrap();
3162        assert_relative_eq!(got, 0.05971830681411544, max_relative = 1e-13);
3163    }
3164
3165    /// `block = c(1,1,1,2,2,3,3,3,4,4)`, so the blocks are of size 3, 2, 3 and
3166    /// 2. The lower clip becomes `1 / (1 - 3) + 0.01`, not `-0.99`.
3167    #[test]
3168    fn test_duplicate_correlation_unbalanced_blocks() {
3169        let design = two_group_design();
3170        let block = vec![0, 0, 0, 1, 1, 2, 2, 2, 3, 3];
3171        let trimmed =
3172            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3173        assert_relative_eq!(trimmed, -0.15040825061556953, max_relative = 1e-14);
3174
3175        let untrimmed =
3176            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.0).unwrap();
3177        assert_relative_eq!(untrimmed, -0.09752275604037172, max_relative = 1e-14);
3178    }
3179
3180    /// `duplicateCorrelation(y, X, block = b, weights = W)`, balanced and
3181    /// unbalanced. `mixedModel2Fit` scales the response and fixed effects by
3182    /// `sqrt(w)` but not the block design.
3183    #[test]
3184    fn test_duplicate_correlation_with_weights() {
3185        let design = two_group_design();
3186        let weights = Recycled::full(W.to_vec(), N_GENES, N_SAMPLES).unwrap();
3187
3188        let balanced: Vec<usize> = (0..10).map(|i| i % 5).collect();
3189        let got = duplicate_correlation(
3190            &Y,
3191            N_GENES,
3192            N_SAMPLES,
3193            &design,
3194            2,
3195            &balanced,
3196            Some(&weights),
3197            0.15,
3198        )
3199        .unwrap();
3200        assert_relative_eq!(got, 0.04173698661489739, max_relative = 1e-14);
3201
3202        let unbalanced = vec![0, 0, 0, 1, 1, 2, 2, 2, 3, 3];
3203        let got = duplicate_correlation(
3204            &Y,
3205            N_GENES,
3206            N_SAMPLES,
3207            &design,
3208            2,
3209            &unbalanced,
3210            Some(&weights),
3211            0.15,
3212        )
3213        .unwrap();
3214        assert_relative_eq!(got, -0.1610960088669458, max_relative = 1e-14);
3215    }
3216
3217    /// `duplicateCorrelation(y, X3, block = b)`, three fixed effects.
3218    #[test]
3219    fn test_duplicate_correlation_three_coefficient_design() {
3220        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3221        let got =
3222            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &X3, 3, &block, None, 0.15).unwrap();
3223        assert_relative_eq!(got, 0.15590493686350534, max_relative = 1e-14);
3224    }
3225
3226    /// `duplicateCorrelation(y, matrix(1, 10, 1), block = b)`.
3227    #[test]
3228    fn test_duplicate_correlation_intercept_only() {
3229        let design = vec![1.0; N_SAMPLES];
3230        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3231        let got =
3232            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 1, &block, None, 0.15).unwrap();
3233        assert_relative_eq!(got, 0.031669365881093696, max_relative = 1e-14);
3234    }
3235
3236    /// Per-gene correlations on the Fisher scale against
3237    /// `duplicateCorrelation(y, X, block = b)$atanh.correlations`. Catches a wrong
3238    /// gene the trimmed mean would hide.
3239    #[test]
3240    fn test_duplicate_correlation_per_gene_values() {
3241        let design = two_group_design();
3242        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3243        let want = [
3244            0.908548407046912,
3245            0.21099075616381313,
3246            0.3440500133251655,
3247            0.4185794031987016,
3248            0.9957853699803254,
3249            0.0565836402590664,
3250            0.9173034946434048,
3251            -0.19437110982329026,
3252            0.787175958405252,
3253            -0.2506510886181334,
3254            -0.20227495727455905,
3255            0.2062792827285404,
3256            -1.2213369921806791,
3257            0.3954884812282991,
3258            -0.7169799798471564,
3259            -0.05498686558164328,
3260            -0.8010392798825152,
3261            0.23348428595384552,
3262            -0.4464403717163369,
3263            -0.3903994527358152,
3264        ];
3265        // One gene at a time reproduces the per-gene value, since a single gene
3266        // makes the trimmed mean the identity.
3267        for (gene, expected) in want.iter().enumerate() {
3268            let row = &Y[gene * N_SAMPLES..(gene + 1) * N_SAMPLES];
3269            let got =
3270                duplicate_correlation(row, 1, N_SAMPLES, &design, 2, &block, None, 0.0).unwrap();
3271            assert_relative_eq!(got.atanh(), expected, max_relative = 1e-13);
3272        }
3273    }
3274
3275    /// Blocks of size one carry no within-block replication, so limma warns and
3276    /// returns exactly zero.
3277    #[test]
3278    fn test_duplicate_correlation_singleton_blocks_are_zero() {
3279        let design = two_group_design();
3280        let block: Vec<usize> = (0..10).collect();
3281        let got =
3282            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3283        assert_eq!(got, 0.0);
3284    }
3285
3286    /// A block factor that is already a column of the design has been absorbed
3287    /// into the fixed effects, so again limma returns exactly zero.
3288    #[test]
3289    fn test_duplicate_correlation_absorbed_block_is_zero() {
3290        let design = two_group_design();
3291        let block: Vec<usize> = (0..10).map(|i| usize::from(i >= 5)).collect();
3292        let got =
3293            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3294        assert_eq!(got, 0.0);
3295    }
3296
3297    #[test]
3298    fn test_duplicate_correlation_rejects_a_bad_block() {
3299        let design = two_group_design();
3300        let err = duplicate_correlation(
3301            &Y,
3302            N_GENES,
3303            N_SAMPLES,
3304            &design,
3305            2,
3306            &[0, 0, 1, 1],
3307            None,
3308            0.15,
3309        )
3310        .unwrap_err();
3311        assert!(matches!(
3312            err,
3313            EdgeErrors::LengthMismatch { name: "block", .. }
3314        ));
3315    }
3316
3317    #[test]
3318    fn test_duplicate_correlation_rejects_a_bad_trim() {
3319        let design = two_group_design();
3320        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3321        for trim in [-0.1, 0.5, f64::NAN] {
3322            let err = duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, trim)
3323                .unwrap_err();
3324            assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3325        }
3326    }
3327
3328    #[test]
3329    fn test_duplicate_correlation_rejects_a_shape_mismatch() {
3330        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3331        let err = duplicate_correlation(
3332            &Y,
3333            N_GENES,
3334            N_SAMPLES,
3335            &two_group_design()[..18],
3336            2,
3337            &block,
3338            None,
3339            0.15,
3340        )
3341        .unwrap_err();
3342        assert!(matches!(
3343            err,
3344            EdgeErrors::LengthMismatch { name: "design", .. }
3345        ));
3346        let err = duplicate_correlation(&[], 0, 10, &two_group_design(), 2, &block, None, 0.15)
3347            .unwrap_err();
3348        assert!(matches!(err, EdgeErrors::EmptyCounts { .. }));
3349    }
3350
3351    /// Every gene fails the admission rules, so there is nothing to average.
3352    #[test]
3353    fn test_duplicate_correlation_errors_when_no_gene_qualifies() {
3354        // Four samples against a two-column design: `n_obs > n_coef + 2` fails.
3355        let y = vec![0.3, -0.7, 1.1, 0.2, -0.4, 0.9, 0.1, -0.2];
3356        let design = vec![1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0];
3357        let block = vec![0, 1, 0, 1];
3358        let err = duplicate_correlation(&y, 2, 4, &design, 2, &block, None, 0.15).unwrap_err();
3359        assert!(matches!(err, EdgeErrors::InvalidArgument(_)));
3360    }
3361
3362    #[test]
3363    fn test_duplicate_correlation_is_reproducible() {
3364        let design = two_group_design();
3365        let block: Vec<usize> = (0..10).map(|i| i % 5).collect();
3366        let first =
3367            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3368        let second =
3369            duplicate_correlation(&Y, N_GENES, N_SAMPLES, &design, 2, &block, None, 0.15).unwrap();
3370        assert_eq!(first, second);
3371    }
3372
3373    ////////////////////////
3374    // Internal machinery //
3375    ////////////////////////
3376
3377    /// The QR has to reproduce R's, so its leverages must match
3378    /// `hat(qr(X))` and its rank must match `qr(X)$rank`.
3379    #[test]
3380    fn test_linpack_qr_leverages_and_rank() {
3381        // Column-major two-group design: each group of five gives a leverage of
3382        // 1/5, and the leverages sum to the rank.
3383        let mut a = vec![1.0; 10];
3384        a.extend((0..10).map(|i| if i < 5 { 0.0 } else { 1.0 }));
3385        let qr = LinpackQr::new(a, 10, 2);
3386        assert_eq!(qr.rank(), 2);
3387        let hat = qr.hat();
3388        for h in &hat {
3389            assert_relative_eq!(h, &0.2, max_relative = 1e-14);
3390        }
3391        assert_relative_eq!(hat.iter().sum::<f64>(), 2.0, max_relative = 1e-14);
3392    }
3393
3394    /// A duplicated column is dropped, and the survivors keep their original
3395    /// order rather than being sorted by magnitude.
3396    #[test]
3397    fn test_linpack_qr_drops_dependent_columns_in_place() {
3398        let mut a = vec![1.0; 6];
3399        a.extend([0.0, 0.0, 0.0, 1.0, 1.0, 1.0]);
3400        a.extend([1.0, 1.0, 1.0, 0.0, 0.0, 0.0]);
3401        let qr = LinpackQr::new(a, 6, 3);
3402        assert_eq!(qr.rank(), 2);
3403        assert_eq!(qr.pivot, vec![0, 1]);
3404    }
3405
3406    /// `Q' Q = I` on the leading columns, which is what every downstream
3407    /// projection assumes.
3408    #[test]
3409    fn test_linpack_qr_leading_q_is_orthonormal() {
3410        let a = vec![
3411            1.0, 1.0, 1.0, 1.0, 1.0, 0.3, -1.2, 0.7, 2.1, -0.4, 1.7, 0.2, -0.9, 0.5, 1.1,
3412        ];
3413        let qr = LinpackQr::new(a, 5, 3);
3414        let q = qr.leading_q(3);
3415        for a in 0..3 {
3416            for b in 0..3 {
3417                let dot: f64 = (0..5).map(|i| q[a * 5 + i] * q[b * 5 + i]).sum();
3418                let want = if a == b { 1.0 } else { 0.0 };
3419                assert_relative_eq!(dot, want, epsilon = 1e-14);
3420            }
3421        }
3422    }
3423
3424    #[test]
3425    fn test_solve_linear_round_trips() {
3426        // Column-major [[4, 1], [1, 3]].
3427        let a = [4.0, 1.0, 1.0, 3.0];
3428        let x = solve_linear(&a, &[1.0, 2.0], 2).unwrap();
3429        assert_relative_eq!(x[0], 1.0 / 11.0, max_relative = 1e-14);
3430        assert_relative_eq!(x[1], 7.0 / 11.0, max_relative = 1e-14);
3431    }
3432
3433    #[test]
3434    fn test_solve_linear_rejects_a_singular_system() {
3435        let a = [1.0, 2.0, 2.0, 4.0];
3436        assert!(matches!(
3437            solve_linear(&a, &[1.0, 1.0], 2),
3438            Err(EdgeErrors::SolveFailed(_))
3439        ));
3440    }
3441
3442    /// `contr.sum(4)` in R, column-major.
3443    #[test]
3444    fn test_contr_sum_matches_r() {
3445        let z = contr_sum(4);
3446        assert_eq!(
3447            z,
3448            vec![
3449                1.0, 0.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0, 0.0, 0.0, 1.0, -1.0
3450            ]
3451        );
3452    }
3453}