limma/

beadcountweights.rs

//! `beadCountWeights` (beadCountWeights.R): bead-count quality weights for
//! Illumina BeadChips. Each probe's variance is split into a *technical*
//! component predicted from the bead-level coefficient of variation and bead
//! counts, and a *biological* component recovered from the residual variance
//! after fitting the design. The returned weights are `1 / (tech + bio)`,
//! row-normalised to mean 1.
//!
//! The biological variance is found per probe by Newton's method, exactly as in
//! limma's internal `.ilmn.biological.variance`.

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

use crate::linalg::qr_econ;

/// How the bead-level dispersion is supplied. limma accepts either the
/// per-probe standard deviation directly, or the standard error of the bead
/// mean (in which case `stdev = stderr * sqrt(nbeads)`).
pub enum BeadDispersion<'a> {
    /// Bead-level standard deviation (`BEAD_STDEV`).
    Stdev(&'a Array2<f64>),
    /// Standard error of the bead mean (`BEAD_STDERR`); converted internally.
    Stderr(&'a Array2<f64>),
}

/// Result of [`bead_count_weights`].
pub struct BeadCountWeights {
    /// `n_probes x n_arrays` observation weights, row-normalised to mean 1
    /// (and optionally rescaled by `rowMeans(1/tech)` when `scale` is set).
    pub weights: Array2<f64>,
    /// Per-probe biological variance component (length `n_probes`).
    pub var_biological: Array1<f64>,
    /// `n_probes x n_arrays` technical variance component.
    pub var_technical: Array2<f64>,
    /// Per-array trimmed coefficient of variation (length `n_arrays`).
    pub cv_array: Array1<f64>,
    /// Single trimmed coefficient of variation pooled over all observations.
    pub cv_constant: f64,
}

/// `beadCountWeights(y, x, design, nbeads, bead.stdev|bead.stderr, array.cv, scale)`.
///
/// * `e_norm` — `n_probes x n_arrays` normalised (log-scale) expression (`y`).
/// * `e_raw` — `n_probes x n_arrays` non-normalised, unlogged expression (`x`).
/// * `design` — `n_arrays x p` design matrix.
/// * `nbeads` — `n_probes x n_arrays` bead counts.
/// * `dispersion` — bead-level [`BeadDispersion`].
/// * `array_cv` — `TRUE` uses a per-array CV; `FALSE` uses one pooled CV.
/// * `scale` — when `true`, multiply each probe's weights by `mean(1/tech)`.
#[allow(clippy::too_many_arguments)]
pub fn bead_count_weights(
    e_norm: &Array2<f64>,
    e_raw: &Array2<f64>,
    design: &Array2<f64>,
    nbeads: &Array2<f64>,
    dispersion: BeadDispersion,
    array_cv: bool,
    scale: bool,
) -> Result<BeadCountWeights> {
    let p = e_norm.nrows();
    let a = e_norm.ncols();
    if e_raw.dim() != (p, a) {
        bail!("dimensions don't match");
    }
    if design.nrows() != a {
        bail!("row dimension of design doesn't match column dimension of data");
    }
    if nbeads.dim() != (p, a) {
        bail!("dimensions don't match");
    }

    let bead_stdev = match dispersion {
        BeadDispersion::Stdev(s) => {
            if s.dim() != (p, a) {
                bail!("dimensions don't match");
            }
            s.to_owned()
        }
        BeadDispersion::Stderr(se) => {
            if se.dim() != (p, a) {
                bail!("dimensions don't match");
            }
            Array2::from_shape_fn((p, a), |(i, j)| se[[i, j]] * nbeads[[i, j]].sqrt())
        }
    };

    // Coefficient of variation of bead-level observations: cv = stdev / raw.
    let sqrt_cv =
        Array2::from_shape_fn((p, a), |(i, j)| (bead_stdev[[i, j]] / e_raw[[i, j]]).sqrt());

    // Per-array (trimmed mean down each column) and pooled CVs, both squared.
    let cv_array: Array1<f64> = (0..a)
        .map(|j| {
            let col: Vec<f64> = (0..p).map(|i| sqrt_cv[[i, j]]).collect();
            let m = trimmed_mean_narm(&col, 0.125);
            m * m
        })
        .collect();
    let pooled: Vec<f64> = sqrt_cv.iter().copied().collect();
    let cm = trimmed_mean_narm(&pooled, 0.125);
    let cv_constant = cm * cm;

    // Spread the chosen CV across the matrix.
    let cv = if array_cv {
        Array2::from_shape_fn((p, a), |(_, j)| cv_array[j])
    } else {
        Array2::from_elem((p, a), cv_constant)
    };

    // Predicted technical variance of normalised probe values:
    //   tv = log(cv^2 / nbeads + 1) / log(2)^2.
    let ln2_sq = std::f64::consts::LN_2 * std::f64::consts::LN_2;
    let var_technical = Array2::from_shape_fn((p, a), |(i, j)| {
        (cv[[i, j]] * cv[[i, j]] / nbeads[[i, j]] + 1.0).ln() / ln2_sq
    });

    // Leverages and squared residuals from the design fit (qr.Q / qr.resid).
    let (q, _r) = qr_econ(design); // a x k, orthonormal columns
    let k = q.ncols();
    let h: Vec<f64> = (0..a)
        .map(|i| (0..k).map(|c| q[[i, c]] * q[[i, c]]).sum())
        .collect();
    let mut r2 = Array2::<f64>::zeros((p, a));
    for g in 0..p {
        let y: Vec<f64> = (0..a).map(|j| e_norm[[g, j]]).collect();
        let qty: Vec<f64> = (0..k)
            .map(|c| (0..a).map(|j| q[[j, c]] * y[j]).sum())
            .collect();
        for j in 0..a {
            let fitted: f64 = (0..k).map(|c| q[[j, c]] * qty[c]).sum();
            let res = y[j] - fitted;
            r2[[g, j]] = res * res;
        }
    }

    let var_biological = ilmn_biological_variance(&var_technical, &r2, &h)?;

    // Weights = 1 / (technical + biological), row-normalised to mean 1.
    let mut weights = Array2::from_shape_fn((p, a), |(g, j)| {
        1.0 / (var_technical[[g, j]] + var_biological[g])
    });
    for g in 0..p {
        let rm: f64 = (0..a).map(|j| weights[[g, j]]).sum::<f64>() / a as f64;
        for j in 0..a {
            weights[[g, j]] /= rm;
        }
    }
    if scale {
        for g in 0..p {
            let rs: f64 = (0..a).map(|j| 1.0 / var_technical[[g, j]]).sum::<f64>() / a as f64;
            for j in 0..a {
                weights[[g, j]] *= rs;
            }
        }
    }

    Ok(BeadCountWeights {
        weights,
        var_biological,
        var_technical,
        cv_array,
        cv_constant,
    })
}

/// `.ilmn.biological.variance`: per-probe biological variance via Newton's
/// method. `h` holds the per-array leverages (constant down each column of the
/// implicit `n_probes x n_arrays` leverage matrix).
fn ilmn_biological_variance(tv: &Array2<f64>, r2: &Array2<f64>, h: &[f64]) -> Result<Array1<f64>> {
    let p = tv.nrows();
    let a = tv.ncols();
    if tv.iter().any(|&v| v < 0.0) || r2.iter().any(|&v| v < 0.0) {
        bail!("negative variances not allowed");
    }

    let mut bv = vec![0.0f64; p];
    let mut fvec = vec![0.0f64; p];
    for g in 0..p {
        let mut acc = 0.0;
        for j in 0..a {
            let t = tv[[g, j]];
            acc += r2[[g, j]] / (2.0 * t * t) - (1.0 - h[j]) / (2.0 * t);
        }
        fvec[g] = acc / a as f64;
    }
    let mut active: Vec<bool> = (0..p).map(|g| fvec[g] > 0.0).collect();

    let mut iter = 0u32;
    while active.iter().any(|&x| x) {
        iter += 1;
        if iter > 200 {
            break; // limma warns and returns the current estimate
        }
        for g in 0..p {
            if !active[g] {
                continue;
            }
            let mut fd = 0.0;
            for j in 0..a {
                let denom = bv[g] + tv[[g, j]];
                fd += r2[[g, j]] / denom.powi(3) - (1.0 - h[j]) / (2.0 * denom * denom);
            }
            let fdash = -(fd / a as f64);
            let step = -fvec[g] / fdash;
            bv[g] += step;
            let mut nf = 0.0;
            for j in 0..a {
                let denom = bv[g] + tv[[g, j]];
                nf += r2[[g, j]] / (2.0 * denom * denom) - (1.0 - h[j]) / (2.0 * denom);
            }
            fvec[g] = nf / a as f64;
            active[g] = step > 1e-5;
        }
    }

    Ok(Array1::from(bv))
}

/// R `mean(x, trim, na.rm = TRUE)`: drop `NaN`, then trim `floor(trim*n)` order
/// statistics from each tail and average the rest.
fn trimmed_mean_narm(xs: &[f64], trim: f64) -> f64 {
    let mut v: Vec<f64> = xs.iter().copied().filter(|x| !x.is_nan()).collect();
    let n = v.len();
    v.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let lo = (trim * n as f64).floor() as usize;
    let hi = n - lo;
    let s = &v[lo..hi];
    s.iter().sum::<f64>() / s.len() as f64
}

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

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

    /// 8x4 fixtures matching `scratch/beadcountweights_ref.R`.
    fn fixture() -> (Array2<f64>, Array2<f64>, Array2<f64>, Array2<f64>) {
        let (p, a) = (8usize, 4usize);
        let mut enorm = Array2::zeros((p, a));
        let mut eraw = Array2::zeros((p, a));
        let mut nbeads = Array2::zeros((p, a));
        let mut bstdev = Array2::zeros((p, a));
        for g0 in 0..p {
            for j0 in 0..a {
                let (g, j) = (g0 as i64, j0 as i64);
                enorm[[g0, j0]] =
                    5.0 + (g % 4) as f64 * 0.3 + ((g * 3 + j * 2) % 5 - 2) as f64 * 0.15;
                eraw[[g0, j0]] =
                    50.0 + (g % 5) as f64 * 10.0 + j as f64 * 5.0 + ((g * 2 + j) % 7) as f64;
                nbeads[[g0, j0]] = 20.0 + ((g + j) % 10) as f64;
                bstdev[[g0, j0]] = 5.0 + (g % 3) as f64 * 1.5 + (j % 2) as f64 * 0.5;
            }
        }
        (enorm, eraw, nbeads, bstdev)
    }

    fn design4() -> Array2<f64> {
        Array2::from_shape_vec((4, 2), vec![1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0]).unwrap()
    }

    fn assert_mat(out: &Array2<f64>, exp: &[[f64; 4]; 8], label: &str) {
        for g in 0..8 {
            for j in 0..4 {
                assert!(
                    rclose(out[[g, j]], exp[g][j]),
                    "{label}[{g},{j}]: {} vs {}",
                    out[[g, j]],
                    exp[g][j]
                );
            }
        }
    }

    #[test]
    fn bead_count_weights_array_cv() {
        let (e, x, nb, sd) = fixture();
        let out = bead_count_weights(
            &e,
            &x,
            &design4(),
            &nb,
            BeadDispersion::Stdev(&sd),
            true,
            false,
        )
        .unwrap();
        let w = [
            [
                0.99782734558068886,
                0.99858220483100923,
                1.0017640081022272,
                1.0018264414860747,
            ],
            [
                0.99795142109299873,
                0.99864815134429352,
                1.001677550700262,
                1.001722876862446,
            ],
            [
                0.99686487564372694,
                0.99790796071791255,
                1.0025886936831236,
                1.0026384699552373,
            ],
            [
                0.99868108910280529,
                0.99911248623705162,
                1.0010967180710966,
                1.0011097065890469,
            ],
            [
                0.99718642392248746,
                0.99808778549730581,
                1.0023570293212056,
                1.0023687612590013,
            ],
            [
                0.99835495001703034,
                0.99887346538654376,
                1.0013868359620184,
                1.0013847486344079,
            ],
            [
                0.99842913135349864,
                0.99891590226037741,
                1.0013325480313711,
                1.0013224183547531,
            ],
            [
                0.99880592784455446,
                0.99954887746182364,
                1.0033290822752707,
                0.99831611241835094,
            ],
        ];
        assert_mat(&out.weights, &w, "bcw A");

        let bio = [
            0.072508375912570672,
            0.072343975538660149,
            0.04434429071263505,
            0.10062096340379205,
            0.044396127280804984,
            0.072621855179473602,
            0.072490525871691081,
            0.044404095553014722,
        ];
        for (g, &b) in bio.iter().enumerate() {
            assert!(rclose(out.var_biological[g], b), "bio[{g}]");
        }
        let cva = [
            0.091824134833759868,
            0.091067645980364265,
            0.078902066838947571,
            0.080362689765811415,
        ];
        for (j, &c) in cva.iter().enumerate() {
            assert!(rclose(out.cv_array[j], c), "cv.array[{j}]");
        }
        assert!(rclose(out.cv_constant, 0.085513043675327097), "cv.constant");

        // Technical variance, first and last rows.
        assert!(rclose(out.var_technical[[0, 0]], 0.00087728608895408207));
        assert!(rclose(out.var_technical[[7, 3]], 0.00067198240487293759));
    }

    #[test]
    fn bead_count_weights_constant_cv() {
        let (e, x, nb, sd) = fixture();
        let out = bead_count_weights(
            &e,
            &x,
            &design4(),
            &nb,
            BeadDispersion::Stdev(&sd),
            false,
            false,
        )
        .unwrap();
        let w = [
            [
                0.99930095702677646,
                0.99979572350253321,
                1.0002459440283038,
                1.0006573754423866,
            ],
            [
                0.99936238568389368,
                0.99981251521384185,
                1.000223863553922,
                1.0006012355483425,
            ],
            [
                0.99905212502205476,
                0.99971950893033579,
                1.0003320697666715,
                1.0008962962809378,
            ],
            [
                0.99961286553655182,
                0.99988496416876738,
                1.0001354289590132,
                1.0003667413356676,
            ],
            [
                0.99919714223094125,
                0.9997600888508188,
                1.0002803012872656,
                1.0007624676309743,
            ],
            [
                0.99954310018584613,
                0.9998629549019481,
                1.0001593024646029,
                1.0004346424476027,
            ],
            [
                0.99957600343922848,
                0.99987231147702482,
                1.0001476147611972,
                1.0004040703225496,
            ],
            [
                1.0007641431549719,
                1.0012122896222646,
                1.0016298942856547,
                0.9963936729371089,
            ],
        ];
        assert_mat(&out.weights, &w, "bcw B");
    }

    #[test]
    fn bead_count_weights_scaled() {
        let (e, x, nb, sd) = fixture();
        let out = bead_count_weights(
            &e,
            &x,
            &design4(),
            &nb,
            BeadDispersion::Stdev(&sd),
            true,
            true,
        )
        .unwrap();
        let w = [
            [
                1438.3935054717908,
                1439.4816542863277,
                1444.0683046535266,
                1444.1583039647787,
            ],
            [
                1505.0568251182424,
                1506.1075963259523,
                1510.6763740042441,
                1510.7447324917509,
            ],
            [
                1569.8302304364456,
                1571.4728467251539,
                1578.843912040946,
                1578.9222980876727,
            ],
            [
                1639.2234167489569,
                1639.9315069412025,
                1643.1884017819241,
                1643.2097209828235,
            ],
            [
                1703.2035875341246,
                1704.7431213976697,
                1712.0350291317538,
                1712.0550673894586,
            ],
            [
                1771.7107857931012,
                1772.6309587964511,
                1777.0912619753831,
                1777.0875577408913,
            ],
            [
                1838.3587171634942,
                1839.2549846218658,
                1843.7046362595136,
                1843.685984981493,
            ],
            [
                1719.8292154910105,
                1721.108489493894,
                1727.617568482578,
                1718.9858095232457,
            ],
        ];
        assert_mat(&out.weights, &w, "bcw C");
    }

    #[test]
    fn bead_count_weights_stderr_path() {
        let (e, x, nb, sd) = fixture();
        // stderr = stdev / sqrt(nbeads) so the resolved stdev equals `sd`.
        let stderr = Array2::from_shape_fn((8, 4), |(i, j)| sd[[i, j]] / nb[[i, j]].sqrt());
        let out = bead_count_weights(
            &e,
            &x,
            &design4(),
            &nb,
            BeadDispersion::Stderr(&stderr),
            true,
            false,
        )
        .unwrap();
        // Must equal the array.cv case A weights exactly (same resolved stdev).
        let ref_out = bead_count_weights(
            &e,
            &x,
            &design4(),
            &nb,
            BeadDispersion::Stdev(&sd),
            true,
            false,
        )
        .unwrap();
        for g in 0..8 {
            for j in 0..4 {
                assert!(
                    rclose(out.weights[[g, j]], ref_out.weights[[g, j]]),
                    "stderr[{g},{j}]"
                );
            }
        }
    }
}