limma-rust 0.1.0

//! Mean-variance modelling of count data at the observation level. Port of
//! limma's `voom` (`voom.R`) and the span heuristic `chooseLowessSpan`.
//!
//! `voom` transforms counts to log2-counts-per-million, fits a gene-wise linear
//! model, fits a LOWESS trend of the residual standard deviation against the
//! average log-count, and converts that trend into per-observation precision
//! weights for entry to the usual `lmFit` -> `eBayes` pipeline.
//!
//! The current port covers the default path: an optional design and library
//! sizes, `normalize.method = "none"`, no offsets and no blocking/correlation.
//! Per-array prior weights are supported via [`voom_weighted`], which backs
//! [`voom_with_quality_weights`] (voom + gene-by-gene array quality weights).

use anyhow::{bail, Result};
use ndarray::{Array1, Array2, Axis};

use crate::arrayweights::array_weights_gene_by_gene;
use crate::fit::{lmfit, lmfit_weighted, MArrayLM};
use crate::lowess::LowessInterpolator;

/// Choose a LOWESS span as a function of the number of points. Port of
/// `chooseLowessSpan`: larger spans for small datasets, smaller for large ones.
pub fn choose_lowess_span(n: usize, small_n: f64, min_span: f64, power: f64) -> f64 {
    (min_span + (1.0 - min_span) * (small_n / n as f64).powf(power)).min(1.0)
}

/// Result of `voom`: the log-CPM matrix `E`, the per-observation precision
/// `weights`, the design used, the per-sample library sizes, and the span.
#[derive(Clone, Debug)]
pub struct VoomOutput {
    pub e: Array2<f64>,       // n_genes x n_samples, log2-CPM
    pub weights: Array2<f64>, // n_genes x n_samples
    pub design: Array2<f64>,
    pub lib_size: Array1<f64>,
    pub span: f64,
}

/// Result of `vooma`: the per-observation precision `weights`, the design
/// used, and the span.
#[derive(Clone, Debug)]
pub struct VoomaOutput {
    pub weights: Array2<f64>, // n_genes x n_samples
    pub design: Array2<f64>,
    pub span: f64,
}

/// `voom(counts, design, lib.size, span, adaptive.span)` with
/// `normalize.method = "none"`.
///
/// * `counts` — `n_genes x n_samples` non-negative counts (no NaN).
/// * `design` — optional `n_samples x p` design; defaults to a grand-mean
///   intercept.
/// * `lib_size` — optional per-sample library sizes; defaults to column sums.
/// * `span` — LOWESS span used when `adaptive_span` is false.
/// * `adaptive_span` — when true, the span is chosen by [`choose_lowess_span`]
///   (`small.n = 50, min.span = 0.3, power = 1/3`), matching voom's default.
pub fn voom(
    counts: &Array2<f64>,
    design: Option<&Array2<f64>>,
    lib_size: Option<&Array1<f64>>,
    span: f64,
    adaptive_span: bool,
) -> Result<VoomOutput> {
    voom_weighted(counts, design, lib_size, span, adaptive_span, None)
}

/// `voom` with optional per-array prior weights (the `weights=` argument used by
/// [`voom_with_quality_weights`]). When `array_weights` is `Some(aw)` the
/// internal `lmFit` is weighted by the broadcast matrix `W[g,j] = aw[j]`, which
/// changes `fit.sigma` (hence the trend) and the fitted values; the returned
/// `weights` are still `1 / trend(fitted_logcount)^4` and do not themselves
/// include `aw`.
pub(crate) fn voom_weighted(
    counts: &Array2<f64>,
    design: Option<&Array2<f64>>,
    lib_size: Option<&Array1<f64>>,
    span: f64,
    adaptive_span: bool,
    array_weights: Option<&Array1<f64>>,
) -> Result<VoomOutput> {
    let ngenes = counts.nrows();
    let nsamples = counts.ncols();
    if ngenes < 2 {
        bail!("Need at least two genes to fit a mean-variance trend");
    }
    let mut min_count = f64::INFINITY;
    for &c in counts.iter() {
        if c.is_nan() {
            bail!("NA counts not allowed");
        }
        min_count = min_count.min(c);
    }
    if min_count < 0.0 {
        bail!("Negative counts not allowed");
    }

    let design_owned;
    let design = match design {
        Some(d) => {
            if d.nrows() != nsamples {
                bail!(
                    "design rows ({}) does not match number of samples ({})",
                    d.nrows(),
                    nsamples
                );
            }
            d
        }
        None => {
            design_owned = Array2::<f64>::ones((nsamples, 1));
            &design_owned
        }
    };
    let p = design.ncols();

    let lib_size: Array1<f64> = match lib_size {
        Some(l) => {
            if l.len() != nsamples {
                bail!(
                    "lib_size length ({}) does not match number of samples ({})",
                    l.len(),
                    nsamples
                );
            }
            l.clone()
        }
        None => counts.sum_axis(Axis(0)),
    };

    let span = if adaptive_span {
        choose_lowess_span(ngenes, 50.0, 0.3, 1.0 / 3.0)
    } else {
        span
    };

    // Log2-counts-per-million (normalize.method = "none").
    let mut y = Array2::<f64>::zeros((ngenes, nsamples));
    for g in 0..ngenes {
        for s in 0..nsamples {
            y[[g, s]] = ((counts[[g, s]] + 0.5) / (lib_size[s] + 1.0) * 1e6).log2();
        }
    }

    let gene_names: Vec<String> = (0..ngenes).map(|i| i.to_string()).collect();
    let coef_names: Vec<String> = (0..p).map(|j| j.to_string()).collect();
    let fit = match array_weights {
        Some(aw) => {
            let mut wmat = Array2::<f64>::zeros((ngenes, nsamples));
            for g in 0..ngenes {
                for s in 0..nsamples {
                    wmat[[g, s]] = aw[s];
                }
            }
            lmfit_weighted(&y, design, &wmat, gene_names, coef_names)?
        }
        None => lmfit(&y, design, gene_names, coef_names)?,
    };

    // With too little replication for a trend, all weights are 1.
    let n_with_reps = fit.df_residual.iter().filter(|&&d| d > 0.0).count();
    if n_with_reps < 2 {
        return Ok(VoomOutput {
            e: y,
            weights: Array2::ones((ngenes, nsamples)),
            design: design.clone(),
            lib_size,
            span,
        });
    }

    // Trend of sqrt-standard-deviation against average log-count.
    let mean_log_lib = lib_size.iter().map(|&l| (l + 1.0).log2()).sum::<f64>() / nsamples as f64;
    let log2_1e6 = 1e6_f64.log2();
    let sx_all: Vec<f64> = (0..ngenes)
        .map(|g| fit.amean[g] + mean_log_lib - log2_1e6)
        .collect();
    let sy_all: Vec<f64> = (0..ngenes).map(|g| fit.sigma[g].sqrt()).collect();

    // Genes with all-zero counts are excluded from the trend fit (they still
    // receive weights below).
    let (sx, sy): (Vec<f64>, Vec<f64>) = (0..ngenes)
        .filter(|&g| counts.row(g).sum() != 0.0)
        .map(|g| (sx_all[g], sy_all[g]))
        .unzip();
    let trend = LowessInterpolator::fit(&sx, &sy, span, 3);

    // Apply the trend to each fitted observation.
    let fitted = fit.coefficients.dot(&design.t()); // n_genes x n_samples
    let mut weights = Array2::<f64>::zeros((ngenes, nsamples));
    for g in 0..ngenes {
        for s in 0..nsamples {
            let fitted_cpm = 2f64.powf(fitted[[g, s]]);
            let fitted_count = 1e-6 * fitted_cpm * (lib_size[s] + 1.0);
            let pred = trend.eval(fitted_count.log2());
            weights[[g, s]] = 1.0 / pred.powi(4);
        }
    }

    Ok(VoomOutput {
        e: y,
        weights,
        design: design.clone(),
        lib_size,
        span,
    })
}

/// Result of [`voom_with_quality_weights`]: a [`VoomOutput`]-style payload whose
/// `weights` already fold in the per-sample `sample_weights`.
#[derive(Clone, Debug)]
pub struct VoomQualityWeights {
    pub e: Array2<f64>,       // n_genes x n_samples, log2-CPM
    pub weights: Array2<f64>, // n_genes x n_samples (voom weights scaled by sample_weights)
    pub design: Array2<f64>,
    pub lib_size: Array1<f64>,
    pub span: f64,
    pub sample_weights: Array1<f64>, // n_samples array quality weights (aw)
}

/// `voomWithQualityWeights(counts, design, span, adaptive.span, var.design,
/// method="genebygene")` — combine voom observation weights with array quality
/// weights estimated gene-by-gene.
///
/// Mirrors limma's orchestration: voom → [`array_weights_gene_by_gene`] (using
/// the voom weights as observation weights) → voom weighted by those array
/// weights → re-estimate array weights → scale each sample column of the voom
/// weights by the final array weights. The `plot`, `var.group`, REML method and
/// `DGEList`/`EList` inputs are out of scope; this is the matrix +
/// `method="genebygene"` path.
///
/// * `var_design` — optional centred, full-rank variance design `Z2`; defaults
///   to `contr.sum(narrays)`.
/// * `prior_n` — prior support for the array weights (limma default `10`).
pub fn voom_with_quality_weights(
    counts: &Array2<f64>,
    design: Option<&Array2<f64>>,
    lib_size: Option<&Array1<f64>>,
    span: f64,
    adaptive_span: bool,
    var_design: Option<&Array2<f64>>,
    prior_n: f64,
) -> Result<VoomQualityWeights> {
    // 1. voom without array weights.
    let v1 = voom_weighted(counts, design, lib_size, span, adaptive_span, None)?;

    // 2. Array weights on top of the voom weights (voom weights act as the
    //    observation weights for the gene-by-gene estimator). The resolved voom
    //    design matches what arrayWeights would build for a full-rank design.
    let aw1 = array_weights_gene_by_gene(&v1.e, &v1.design, Some(&v1.weights), var_design, prior_n);

    // 3. voom again, now weighted by the array weights.
    let v2 = voom_weighted(counts, design, lib_size, span, adaptive_span, Some(&aw1))?;

    // 4. Re-estimate array weights on the reweighted voom fit.
    let aw2 = array_weights_gene_by_gene(&v2.e, &v2.design, Some(&v2.weights), var_design, prior_n);

    // 5. Scale each sample column of the voom weights by the final array weights.
    let mut weights = v2.weights;
    let (ng, ns) = weights.dim();
    for g in 0..ng {
        for s in 0..ns {
            weights[[g, s]] *= aw2[s];
        }
    }

    Ok(VoomQualityWeights {
        e: v2.e,
        weights,
        design: v2.design,
        lib_size: v2.lib_size,
        span: v2.span,
        sample_weights: aw2,
    })
}

/// `vooma(y, design, span, legacy.span)` — the continuous-data analog of
/// `voom`, with no predictor and no blocking.
///
/// * `y` — `n_genes x n_samples` log-expression matrix (must be finite; R's
///   vooma disallows NAs).
/// * `design` — optional `n_samples x p` design; defaults to a grand-mean
///   intercept.
/// * `span` — LOWESS span; when `None` it is chosen adaptively
///   (`small.n = 50, power = 1/3`, or `small.n = 10, power = 1/2` when
///   `legacy_span`).
pub fn vooma(
    y: &Array2<f64>,
    design: Option<&Array2<f64>>,
    span: Option<f64>,
    legacy_span: bool,
) -> Result<VoomaOutput> {
    let ngenes = y.nrows();
    let nsamples = y.ncols();
    if ngenes < 1 {
        bail!("y has no rows");
    }

    let design_owned;
    let design = match design {
        Some(d) => {
            if d.nrows() != nsamples {
                bail!(
                    "design rows ({}) does not match number of samples ({})",
                    d.nrows(),
                    nsamples
                );
            }
            d
        }
        None => {
            design_owned = Array2::<f64>::ones((nsamples, 1));
            &design_owned
        }
    };
    let p = design.ncols();

    let gene_names: Vec<String> = (0..ngenes).map(|i| i.to_string()).collect();
    let coef_names: Vec<String> = (0..p).map(|j| j.to_string()).collect();
    let fit = lmfit(y, design, gene_names, coef_names)?;

    if fit.amean.iter().any(|a| a.is_nan()) {
        bail!("y contains entirely NA rows");
    }

    let sx: Vec<f64> = fit.amean.to_vec();
    let sy: Vec<f64> = (0..ngenes).map(|g| fit.sigma[g].sqrt()).collect();

    let span = match span {
        Some(s) => s,
        None => {
            if legacy_span {
                choose_lowess_span(ngenes, 10.0, 0.3, 0.5)
            } else {
                choose_lowess_span(ngenes, 50.0, 0.3, 1.0 / 3.0)
            }
        }
    };
    let trend = LowessInterpolator::fit(&sx, &sy, span, 3);

    let mu = fit.coefficients.dot(&design.t()); // n_genes x n_samples
    let mut weights = Array2::<f64>::zeros((ngenes, nsamples));
    for g in 0..ngenes {
        for s in 0..nsamples {
            weights[[g, s]] = 1.0 / trend.eval(mu[[g, s]]).powi(4);
        }
    }

    Ok(VoomaOutput {
        weights,
        design: design.clone(),
        span,
    })
}

/// Result of [`vooma_by_group`].
#[derive(Clone, Debug)]
pub struct VoomaByGroupOutput {
    /// Observation weights (`n_genes x n_samples`), assembled from the per-group
    /// vooma fits; a singleton group's column borrows the global fit's weights.
    pub weights: Array2<f64>,
    /// The design used: the supplied design, or the `~0+group` indicator.
    pub design: Array2<f64>,
    /// Reported LOWESS span: the last non-singleton group's span, falling back
    /// to the global fit's span.
    pub span: f64,
}

/// Drop all-zero columns of `m`, preserving order. Reduces a `~0+group`
/// per-group subdesign — whose only non-zero column is that group's indicator —
/// to an intercept, reproducing limma's rank-deficient `lm.fit` fitted values
/// (the group mean) without a pivoted QR.
fn drop_zero_columns(m: &Array2<f64>) -> Array2<f64> {
    let keep: Vec<usize> = (0..m.ncols())
        .filter(|&j| m.column(j).iter().any(|&v| v != 0.0))
        .collect();
    m.select(Axis(1), &keep)
}

/// `voomaByGroup(y, group, design, span, legacy.span)` — fit a separate vooma
/// mean-variance trend within each level of `group` and assemble the per-group
/// observation weights into one `n_genes x n_samples` matrix for entry to
/// `lmFit`.
///
/// * `y` — `n_genes x n_samples` log-expression (finite; no NAs, as in vooma).
/// * `group` — length-`n_samples` label per sample; levels are taken in sorted
///   order (matching R's factor levels).
/// * `design` — optional `n_samples x p` design. When `None` it defaults to the
///   `~0+group` indicator (one column per level, in sorted order); each
///   per-group subdesign then reduces to an intercept (the group mean), matching
///   limma's rank-deficient `lm.fit`.
/// * `span`, `legacy_span` — forwarded to each per-group [`vooma`].
///
/// Singleton groups (a single sample) borrow their weights from a global vooma
/// fit over all samples, matching limma. The plot-only `voom.xy` trend points
/// are not returned, and blocking/`correlation` is not supported.
pub fn vooma_by_group(
    y: &Array2<f64>,
    group: &[usize],
    design: Option<&Array2<f64>>,
    span: Option<f64>,
    legacy_span: bool,
) -> Result<VoomaByGroupOutput> {
    let ngenes = y.nrows();
    let narrays = y.ncols();
    if group.len() != narrays {
        bail!(
            "length of group ({}) must equal number of samples ({})",
            group.len(),
            narrays
        );
    }

    // Sorted unique levels; map each sample to its level index.
    let mut levels: Vec<usize> = group.to_vec();
    levels.sort_unstable();
    levels.dedup();
    let ngroups = levels.len();
    let intgroup: Vec<usize> = group
        .iter()
        .map(|g| levels.iter().position(|l| l == g).unwrap())
        .collect();

    // Design: supplied, or the ~0+group indicator (sorted level order).
    let design: Array2<f64> = match design {
        Some(d) => {
            if d.nrows() != narrays {
                bail!(
                    "design rows ({}) does not match number of samples ({})",
                    d.nrows(),
                    narrays
                );
            }
            d.clone()
        }
        None => {
            let mut ind = Array2::<f64>::zeros((narrays, ngroups));
            for (s, &gi) in intgroup.iter().enumerate() {
                ind[[s, gi]] = 1.0;
            }
            ind
        }
    };

    // Sample columns belonging to each group level.
    let mut cols_of: Vec<Vec<usize>> = vec![Vec::new(); ngroups];
    for (s, &gi) in intgroup.iter().enumerate() {
        cols_of[gi].push(s);
    }
    let has_singleton = cols_of.iter().any(|c| c.len() == 1);

    // Global fit, used to supply weights for any singleton group.
    let voomall = if has_singleton {
        Some(vooma(y, Some(&design), span, legacy_span)?)
    } else {
        None
    };

    let mut weights = Array2::<f64>::zeros((ngenes, narrays));
    let mut last_span: Option<f64> = None;
    for cols in &cols_of {
        if cols.len() == 1 {
            let va = voomall.as_ref().unwrap();
            let s = cols[0];
            for g in 0..ngenes {
                weights[[g, s]] = va.weights[[g, s]];
            }
        } else {
            let yi = y.select(Axis(1), cols);
            let di = drop_zero_columns(&design.select(Axis(0), cols));
            let voomi = vooma(&yi, Some(&di), span, legacy_span)?;
            for (k, &s) in cols.iter().enumerate() {
                for g in 0..ngenes {
                    weights[[g, s]] = voomi.weights[[g, k]];
                }
            }
            last_span = Some(voomi.span);
        }
    }

    let out_span = last_span
        .or_else(|| voomall.as_ref().map(|v| v.span))
        .unwrap_or(0.0);

    Ok(VoomaByGroupOutput {
        weights,
        design,
        span: out_span,
    })
}

/// Result of [`vooma_lm_fit`].
#[derive(Clone, Debug)]
pub struct VoomaLmFit {
    /// Weighted least-squares fit carrying the vooma weights, ready for `eBayes`.
    pub fit: MArrayLM,
    /// vooma observation weights (`n_genes x n_samples`) used in the final fit.
    pub weights: Array2<f64>,
    /// LOWESS span used for the mean-variance trend.
    pub span: f64,
}

/// `voomaLmFit(y, design, prior.weights, span, legacy.span)` — vooma followed
/// by a weighted `lmFit`, returning a fit ready for `eBayes`.
///
/// Implements the common, non-iterating path. An optional `prior_weights`
/// matrix is supported: as in limma the first fit uses the prior weights, and
/// the final weights are `vooma_weights * prior_weights`. The block-correlation,
/// sample-weight (`sample.weights`/`var.design`/`var.group`) and precision
/// `predictor` options are not implemented — they require
/// duplicateCorrelation/gls.series or arrayWeights iteration.
pub fn vooma_lm_fit(
    y: &Array2<f64>,
    design: Option<&Array2<f64>>,
    prior_weights: Option<&Array2<f64>>,
    span: Option<f64>,
    legacy_span: bool,
    gene_names: Vec<String>,
    coef_names: Vec<String>,
) -> Result<VoomaLmFit> {
    let ngenes = y.nrows();
    let nsamples = y.ncols();
    if nsamples < 2 {
        bail!("Too few samples");
    }
    if ngenes < 2 {
        bail!("Need multiple rows");
    }

    let design_owned;
    let design = match design {
        Some(d) => {
            if d.nrows() != nsamples {
                bail!(
                    "design rows ({}) does not match number of samples ({})",
                    d.nrows(),
                    nsamples
                );
            }
            d
        }
        None => {
            design_owned = Array2::<f64>::ones((nsamples, 1));
            &design_owned
        }
    };

    if let Some(pw) = prior_weights {
        if pw.nrows() != ngenes || pw.ncols() != nsamples {
            bail!(
                "prior_weights dimensions ({}x{}) must match y ({}x{})",
                pw.nrows(),
                pw.ncols(),
                ngenes,
                nsamples
            );
        }
    }

    // First fit: weighted when prior weights are supplied, otherwise OLS.
    let fit0 = match prior_weights {
        Some(pw) => lmfit_weighted(y, design, pw, gene_names.clone(), coef_names.clone())?,
        None => lmfit(y, design, gene_names.clone(), coef_names.clone())?,
    };
    if fit0.amean.iter().any(|a| a.is_nan()) {
        bail!("y contains entirely NA rows");
    }

    // Mean-variance trend on (average expression, sqrt residual SD).
    let sx: Vec<f64> = fit0.amean.to_vec();
    let sy: Vec<f64> = (0..ngenes).map(|g| fit0.sigma[g].sqrt()).collect();
    let span = match span {
        Some(s) => s,
        None => {
            if legacy_span {
                choose_lowess_span(ngenes, 10.0, 0.3, 0.5)
            } else {
                choose_lowess_span(ngenes, 50.0, 0.3, 1.0 / 3.0)
            }
        }
    };
    let trend = LowessInterpolator::fit(&sx, &sy, span, 3);

    // vooma weights from the fitted values, then combine with prior weights.
    let mu = fit0.coefficients.dot(&design.t());
    let mut weights = Array2::<f64>::zeros((ngenes, nsamples));
    for g in 0..ngenes {
        for s in 0..nsamples {
            weights[[g, s]] = 1.0 / trend.eval(mu[[g, s]]).powi(4);
        }
    }
    if let Some(pw) = prior_weights {
        weights = &weights * pw;
    }

    let fit = lmfit_weighted(y, design, &weights, gene_names, coef_names)?;
    Ok(VoomaLmFit { fit, weights, span })
}

#[cfg(test)]
#[allow(clippy::excessive_precision)]
mod tests {
    use super::*;
    use ndarray::array;

    fn counts_fixture() -> Array2<f64> {
        array![
            [10., 12., 8., 11.],
            [100., 110., 95., 105.],
            [5., 7., 6., 4.],
            [200., 190., 210., 205.],
            [50., 45., 55., 60.],
            [1., 2., 0., 3.],
            [1000., 1100., 900., 950.],
            [30., 25., 35., 40.],
        ]
    }

    fn design_fixture() -> Array2<f64> {
        // model.matrix(~group), group = A A B B.
        array![[1., 0.], [1., 0.], [1., 1.], [1., 1.]]
    }

    #[test]
    fn choose_span_matches_r() {
        // limma:::chooseLowessSpan(8, 50, 0.3, 1/3) clamps to 1.
        assert_eq!(choose_lowess_span(8, 50.0, 0.3, 1.0 / 3.0), 1.0);
        // A large dataset gives a span below 1.
        let s = choose_lowess_span(1000, 50.0, 0.3, 1.0 / 3.0);
        let expect = 0.3 + 0.7 * (50.0_f64 / 1000.0).powf(1.0 / 3.0);
        assert!((s - expect).abs() < 1e-12 && s < 1.0);
    }

    #[test]
    fn voom_matches_r() {
        let counts = counts_fixture();
        let design = design_fixture();
        let v = voom(&counts, Some(&design), None, 0.5, false).unwrap();

        // Library sizes are column sums.
        let lib = array![1396.0, 1491.0, 1309.0, 1378.0];
        for (a, b) in v.lib_size.iter().zip(lib.iter()) {
            assert!((a - b).abs() < 1e-9);
        }

        // E = log2((counts+0.5)/(lib+1)*1e6) is a deterministic closed form.
        for g in 0..counts.nrows() {
            for s in 0..counts.ncols() {
                let expect = ((counts[[g, s]] + 0.5) / (lib[s] + 1.0) * 1e6).log2();
                assert!((v.e[[g, s]] - expect).abs() < 1e-9);
            }
        }

        // Weights from R limma 3.68.3 (scratch/voom_ref.R).
        let want = array![
            [
                28.935889399719,
                28.935889399719,
                10.326303259626,
                17.013134202459
            ],
            [
                78.797663894764,
                80.528056879469,
                77.731587915647,
                79.054873411272
            ],
            [
                0.641934244773,
                0.877223560596,
                0.507411168766,
                0.507411168766
            ],
            [
                97.087653321213,
                99.514156468690,
                100.012320428675,
                107.314600669557
            ],
            [
                50.936936306788,
                57.095789543676,
                65.164440442125,
                66.225455485966
            ],
            [
                0.507411168766,
                0.507411168766,
                0.507411168766,
                0.507411168766
            ],
            [
                2193.696131719829,
                2193.696131719829,
                1703.605796348448,
                1968.434108844519
            ],
            [
                28.935889399719,
                28.935889399719,
                35.153539352782,
                38.113189099620
            ],
        ];
        for (a, b) in v.weights.iter().zip(want.iter()) {
            let rel = (a - b).abs() / b.abs();
            assert!(rel < 1e-9, "weight {a} vs {b} (rel {rel:e})");
        }
    }

    #[test]
    fn vooma_matches_r() {
        // Continuous log-expression fixture (scratch/vooma_ref.R).
        let y = array![
            [5.10, 4.80, 6.20, 5.50],
            [2.30, 3.10, 2.80, 3.50],
            [7.70, 7.20, 8.10, 6.90],
            [1.10, 0.90, 1.40, 1.20],
            [9.30, 9.10, 8.80, 9.50],
            [4.40, 4.90, 5.20, 4.10],
            [6.60, 6.20, 6.90, 6.40],
            [3.30, 3.80, 3.10, 2.90],
            [8.10, 7.60, 8.40, 8.00],
            [0.50, 1.20, 0.80, 0.30],
        ];
        let design = design_fixture();
        let v = vooma(&y, Some(&design), None, false).unwrap();
        assert_eq!(v.span, 1.0); // ngenes=10 -> adaptive span clamps to 1

        let want = array![
            [
                5.769009938122,
                5.769009938122,
                5.710681439260,
                5.710681439260
            ],
            [
                7.764204399417,
                7.764204399417,
                7.299187981925,
                7.299187981925
            ],
            [
                6.131523915696,
                6.131523915696,
                6.146966664045,
                6.146966664045
            ],
            [
                10.027067353155,
                10.027067353155,
                9.549770567334,
                9.549770567334
            ],
            [
                6.732012491201,
                6.732012491201,
                6.722541890163,
                6.722541890163
            ],
            [
                5.873692933054,
                5.873692933054,
                5.873692933054,
                5.873692933054
            ],
            [
                5.828720439884,
                5.828720439884,
                5.892290676457,
                5.892290676457
            ],
            [
                6.892732120849,
                6.892732120849,
                7.443835424811,
                7.443835424811
            ],
            [
                6.257090085442,
                6.257090085442,
                6.374681498763,
                6.374681498763
            ],
            [
                10.292360128274,
                10.292360128274,
                10.566485009140,
                10.566485009140
            ],
        ];
        for (a, b) in v.weights.iter().zip(want.iter()) {
            let rel = (a - b).abs() / b.abs();
            assert!(rel < 1e-9, "weight {a} vs {b} (rel {rel:e})");
        }
    }

    fn rclose(a: f64, b: f64) -> bool {
        (a - b).abs() <= 1e-7 * (1.0 + b.abs())
    }

    /// Rebuild scratch/voomalmfit_ref.R's 50x6 expression and prior-weight
    /// matrices from the same purely rational, 0-indexed formula.
    fn vlf_fixture() -> (Array2<f64>, Array2<f64>, Array2<f64>) {
        let ngenes = 50usize;
        let narrays = 6usize;
        let group = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0];
        let mut y = Array2::<f64>::zeros((ngenes, narrays));
        let mut pw = Array2::<f64>::zeros((ngenes, narrays));
        for g in 0..ngenes {
            let gi = g as i64;
            let base = ((gi % 7) - 3) as f64;
            let extra = (((gi * 5) % 11) - 5) as f64;
            let lvl = 5.0 + (gi % 10) as f64 * 0.2;
            let eff = base * 0.25;
            for k in 0..narrays {
                let ki = k as i64;
                let noise = (((gi * 13 + ki * 17) % 19) - 9) as f64 * 0.04;
                y[[g, k]] = lvl + eff * group[k] + extra * 0.05 + noise;
                pw[[g, k]] = 0.6 + (((gi * 3 + ki * 7) % 9) as f64) * 0.1;
            }
        }
        let mut design = Array2::<f64>::zeros((narrays, 2));
        for k in 0..narrays {
            design[[k, 0]] = 1.0;
            design[[k, 1]] = group[k];
        }
        (y, design, pw)
    }

    /// Rebuild scratch/vwqw_ref.R's fixture B: integer counts (10x6) from a
    /// 0-indexed rational formula, plus the `~group` design.
    fn vwqw_counts_fixture() -> (Array2<f64>, Array2<f64>) {
        let ngenes = 10usize;
        let narrays = 6usize;
        let grp = [0i64, 0, 0, 1, 1, 1];
        let mut counts = Array2::<f64>::zeros((ngenes, narrays));
        for g in 0..ngenes {
            let gi = g as i64;
            for j in 0..narrays {
                let ji = j as i64;
                let c = 20 + gi * 15 + ji * 3 + ((gi * 5 + ji * 7) % 13) * 4 + grp[j] * gi * 2;
                counts[[g, j]] = c as f64;
            }
        }
        let mut design = Array2::<f64>::zeros((narrays, 2));
        for j in 0..narrays {
            design[[j, 0]] = 1.0;
            design[[j, 1]] = grp[j] as f64;
        }
        (counts, design)
    }

    #[test]
    fn voom_with_quality_weights_matches_r() {
        let (counts, design) = vwqw_counts_fixture();
        let v = voom_with_quality_weights(&counts, Some(&design), None, 0.5, false, None, 10.0)
            .unwrap();

        // Final per-sample array quality weights.
        let aw_want = [
            1.1428522224351862,
            0.69948202996388553,
            1.1916899290606626,
            1.1563249275179759,
            0.74938331550111592,
            1.2113961366163049,
        ];
        for (g, x) in v.sample_weights.iter().zip(aw_want.iter()) {
            assert!(rclose(*g, *x), "sample weight: got {g}, want {x}");
        }

        // Final voom weights (voom weights scaled by the array weights).
        assert!(rclose(v.weights[[0, 0]], 4.6250384357077206));
        assert!(rclose(v.weights[[0, 3]], 9.0644002874417744));
        assert!(rclose(v.weights[[0, 5]], 10.34938541502423));
        assert!(rclose(v.weights[[4, 2]], 28.109033884186356));
        assert!(rclose(v.weights[[5, 1]], 16.86267872361735));
        assert!(rclose(v.weights[[9, 0]], 75.782464040995038));
        assert!(rclose(v.weights[[9, 5]], 117.79877846147826));

        // E is the deterministic log2-CPM, identical to plain voom.
        assert!(rclose(v.e[[0, 0]], 14.185832765530236));
        assert!(rclose(v.e[[9, 5]], 17.166249779108728));
    }

    #[test]
    fn vooma_lm_fit_matches_r() {
        let (y, design, pw) = vlf_fixture();
        let gn: Vec<String> = (0..50).map(|i| format!("g{i}")).collect();
        let cn = vec!["Intercept".to_string(), "group".to_string()];

        // Scenario 1: no prior weights (first fit is OLS).
        let r = vooma_lm_fit(&y, Some(&design), None, None, false, gn.clone(), cn.clone()).unwrap();
        assert_eq!(r.span, 1.0); // ngenes=50 == small.n -> span clamps to 1
        assert!(rclose(r.fit.coefficients[[0, 0]], 4.8166666666666673));
        assert!(rclose(r.fit.coefficients[[0, 1]], -0.73666666666666669));
        assert!(rclose(r.fit.coefficients[[2, 1]], 0.016666666666666802));
        assert!(rclose(r.fit.sigma[0], 2.0001523415288656));
        assert!(rclose(r.fit.sigma[4], 0.50739386270805098));
        assert!(rclose(r.fit.stdev_unscaled[[0, 0]], 0.07765105731999547));
        assert!(rclose(r.fit.stdev_unscaled[[0, 1]], 0.10897521719515758));
        assert_eq!(r.fit.df_residual[0], 4.0);
        // Within an intercept+group design, weights take two values per gene.
        assert!(rclose(r.weights[[0, 0]], 55.282032012091854));
        assert!(rclose(r.weights[[0, 3]], 57.019909902580309));
        assert!(rclose(r.weights[[24, 0]], 43.694024586743168));
        assert!(rclose(r.weights[[49, 0]], 33.652368280339203));
        assert!(rclose(r.weights[[49, 3]], 48.718872300466991));

        // Scenario 2: with prior weights (first fit is weighted).
        let r2 = vooma_lm_fit(&y, Some(&design), Some(&pw), None, false, gn, cn).unwrap();
        assert_eq!(r2.span, 1.0);
        assert!(rclose(r2.fit.coefficients[[0, 0]], 4.9046666666666683));
        assert!(rclose(r2.fit.coefficients[[0, 1]], -0.83800000000000119));
        assert!(rclose(r2.fit.sigma[0], 1.8170323609512815));
        assert!(rclose(r2.fit.stdev_unscaled[[0, 0]], 0.07413980745014162));
        assert!(rclose(r2.weights[[0, 0]], 36.385394507083909));
        assert!(rclose(r2.weights[[0, 5]], 91.56829234031899));
        assert!(rclose(r2.weights[[49, 0]], 30.459503135870523));
    }

    #[test]
    fn vooma_by_group_matches_r() {
        // See scratch/voomabygroup_ref.R.
        let ngenes = 12usize;
        let mky = |narrays: usize| -> Array2<f64> {
            Array2::from_shape_fn((ngenes, narrays), |(g0, s0)| {
                5.0 + 0.25 * g0 as f64 + 0.1 * s0 as f64 + 0.15 * ((g0 * (s0 + 1)) % 5) as f64
            })
        };
        let rclose = |a: f64, b: f64| (a - b).abs() <= 1e-7 * (1.0 + b.abs());

        // Per-group weights: column k is the weight shared by every sample in
        // sorted group k. Within-group columns are identical (intercept fit).
        let gw = array![
            [128.57974235629533, 201.43265373150086, 169.52805974815644],
            [51.72862401222929, 18.636960776554076, 37.75223112586292],
            [37.559193159543049, 20.439481407168092, 13.083239709395334],
            [37.614806241386034, 146.28421635333993, 8.9183551719300596],
            [106.95854312349614, 172.92065689564916, 12.144674253781014],
            [69.276054650291854, 206.54103286812787, 10.92966941384552],
            [63.332917660356472, 18.636960776554002, 15.654008874403512],
            [32.388566989657896, 20.439481407168017, 12.99773820977671],
            [37.212711780045019, 146.28421635334118, 9.8481582955355105],
            [100.30030326001432, 172.92065689564888, 15.212910649940969],
            [70.172943221746678, 206.54103286812716, 14.112001483913859],
            [124.95596005142109, 32.391260048340605, 24.542006411260687],
        ];

        // Scenario A: three groups of two, default ~0+group design, span 0.5.
        let ya = mky(6);
        let group_a = [0usize, 0, 1, 1, 2, 2];
        let va = vooma_by_group(&ya, &group_a, None, Some(0.5), false).unwrap();
        assert!((va.span - 0.5).abs() < 1e-12);
        for ((g, s), _) in va.weights.indexed_iter() {
            let e = gw[[g, group_a[s]]];
            assert!(rclose(va.weights[[g, s]], e), "A[{g},{s}]");
        }

        // Scenario B: singleton last group (sample 4) borrows the global fit.
        let single_b = [
            92.423854171536973,
            46.676048745081182,
            36.257599360043535,
            52.633157173375849,
            97.417756778627549,
            67.100924018777832,
            43.554031111104514,
            36.052167162893426,
            52.218302840608807,
            99.089796796342043,
            71.514137307408575,
            49.250767908710742,
        ];
        let yb = mky(5);
        let group_b = [0usize, 0, 1, 1, 2];
        let vb = vooma_by_group(&yb, &group_b, None, Some(0.5), false).unwrap();
        assert!((vb.span - 0.5).abs() < 1e-12);
        for g in 0..ngenes {
            // groups 0 and 1 reuse the per-group weights from scenario A
            assert!(rclose(vb.weights[[g, 0]], gw[[g, 0]]), "B[{g},0]");
            assert!(rclose(vb.weights[[g, 1]], gw[[g, 0]]), "B[{g},1]");
            assert!(rclose(vb.weights[[g, 2]], gw[[g, 1]]), "B[{g},2]");
            assert!(rclose(vb.weights[[g, 3]], gw[[g, 1]]), "B[{g},3]");
            assert!(rclose(vb.weights[[g, 4]], single_b[g]), "B[{g},4]");
        }
    }
}