rsomics_edger_exact_test/

lib.rs

//! edgeR exactTest: classic two-group negative-binomial exact test for DE.
//! Method: Robinson & Smyth (2008), Biostatistics 9:321-332 — an overdispersed
//! generalization of Fisher's exact test, conditioning each gene on its total
//! count. Per gene: logFC from per-group one-group NB fits, logCPM from
//! aveLogCPM, and PValue from the doubletail NB convolution after quantile-
//! adjusting both groups to a common library size.

mod special;

use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::Path;

use rsomics_common::{Result, RsomicsError};

const LN2: f64 = std::f64::consts::LN_2;
const PRIOR_COUNT_FC: f64 = 0.125; // edgeR exactTest's logFC prior.count
const AVELOGCPM_PRIOR: f64 = 2.0; // edgeR aveLogCPM default
const AVELOGCPM_DISP: f64 = 0.05; // exactTest's logCPM uses aveLogCPM's own default, not the test dispersion
const BIG_COUNT: f64 = 900.0;

pub struct Matrix {
    pub header: String,
    pub genes: Vec<String>,
    pub counts: Vec<f64>,
    pub n_samples: usize,
}

impl Matrix {
    pub fn load(path: &Path) -> Result<Self> {
        let file = File::open(path)
            .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
        let mut lines = BufReader::new(file).lines();
        let header = lines
            .next()
            .ok_or_else(|| RsomicsError::InvalidInput("empty count matrix".into()))?
            .map_err(RsomicsError::Io)?;
        let n_samples = header.split('\t').count() - 1;
        if n_samples == 0 {
            return Err(RsomicsError::InvalidInput(
                "count matrix has no sample columns".into(),
            ));
        }

        let mut genes = Vec::new();
        let mut counts = Vec::new();
        for line in lines {
            let line = line.map_err(RsomicsError::Io)?;
            if line.is_empty() {
                continue;
            }
            let mut fields = line.split('\t');
            let gene = fields
                .next()
                .ok_or_else(|| RsomicsError::InvalidInput("row without a gene id".into()))?;
            genes.push(gene.to_string());
            let before = counts.len();
            for f in fields {
                counts.push(f.parse::<f64>().map_err(|_| {
                    RsomicsError::InvalidInput(format!("non-numeric count '{f}' for gene {gene}"))
                })?);
            }
            if counts.len() - before != n_samples {
                return Err(RsomicsError::InvalidInput(format!(
                    "gene {gene}: {} values, header has {n_samples} samples",
                    counts.len() - before
                )));
            }
        }
        Ok(Self {
            header,
            genes,
            counts,
            n_samples,
        })
    }

    pub fn n_genes(&self) -> usize {
        self.genes.len()
    }
    fn row(&self, g: usize) -> &[f64] {
        &self.counts[g * self.n_samples..(g + 1) * self.n_samples]
    }
}

pub struct ExactTestOpts {
    pub dispersion: f64,
    pub fdr: bool,
}

/// Split a `--group` spec like `a,a,b,b` into the two sample-index lists. Levels
/// are taken in first-appearance order; logFC is reported level2 over level1.
fn parse_groups(spec: &str, n_samples: usize) -> Result<(Vec<usize>, Vec<usize>)> {
    let labels: Vec<&str> = spec.split(',').map(str::trim).collect();
    if labels.len() != n_samples {
        return Err(RsomicsError::InvalidInput(format!(
            "--group has {} labels but matrix has {n_samples} samples",
            labels.len()
        )));
    }
    let mut levels: Vec<&str> = Vec::new();
    for l in &labels {
        if !levels.contains(l) {
            levels.push(l);
        }
    }
    if levels.len() != 2 {
        return Err(RsomicsError::InvalidInput(format!(
            "exactTest needs exactly two groups, found {}: {levels:?}",
            levels.len()
        )));
    }
    let mut g1 = Vec::new();
    let mut g2 = Vec::new();
    for (i, &lab) in labels.iter().enumerate() {
        if lab == levels[0] { &mut g1 } else { &mut g2 }.push(i);
    }
    Ok((g1, g2))
}

fn load_norm_factors(path: &Path, n_samples: usize) -> Result<Vec<f64>> {
    let file = File::open(path)
        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
    let mut factors = Vec::with_capacity(n_samples);
    for line in BufReader::new(file).lines() {
        let line = line.map_err(RsomicsError::Io)?;
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        let val = line.rsplit('\t').next().unwrap_or(line);
        factors.push(
            val.parse::<f64>().map_err(|_| {
                RsomicsError::InvalidInput(format!("non-numeric norm factor '{val}'"))
            })?,
        );
    }
    if factors.len() != n_samples {
        return Err(RsomicsError::InvalidInput(format!(
            "{} norm factors for {n_samples} samples",
            factors.len()
        )));
    }
    Ok(factors)
}

const MAXIT: usize = 50;
const TOL: f64 = 1e-10;

/// One-group NB fit (edgeR mglmOneGroup / glm_one_group): Fisher scoring for the
/// single coefficient beta where mu[j] = exp(beta + offset[j]). Returns beta on
/// the natural-log scale; -inf when the total count is zero.
fn mglm_one_group(row: &[f64], offset: &[f64], dispersion: f64) -> f64 {
    let total: f64 = row.iter().sum();
    if total == 0.0 {
        return f64::NEG_INFINITY;
    }
    let mean_off = offset.iter().sum::<f64>() / offset.len() as f64;
    let mut beta = (total / row.len() as f64).ln() - mean_off;
    for _ in 0..MAXIT {
        let mut dl = 0.0;
        let mut info = 0.0;
        for (&y, &off) in row.iter().zip(offset) {
            let mu = (beta + off).exp();
            let denom = 1.0 + mu * dispersion;
            dl += (y - mu) / denom;
            info += mu / denom;
        }
        let step = dl / info;
        beta += step;
        if step.abs() < TOL {
            break;
        }
    }
    beta
}

/// aveLogCPM for one gene across all libraries (edgeR aveLogCPM, prior.count=2,
/// dispersion=0.05). Counts are augmented by a per-library-scaled prior; the fit
/// runs against offsets that include twice that prior.
fn ave_log_cpm_gene(row: &[f64], lib: &[f64]) -> f64 {
    let mean_lib = lib.iter().sum::<f64>() / lib.len() as f64;
    let prior: Vec<f64> = lib
        .iter()
        .map(|&l| AVELOGCPM_PRIOR * l / mean_lib)
        .collect();
    let off_aug: Vec<f64> = lib
        .iter()
        .zip(&prior)
        .map(|(&l, &p)| (l + 2.0 * p).ln())
        .collect();
    let aug: Vec<f64> = row.iter().zip(&prior).map(|(&c, &p)| c + p).collect();
    let beta = mglm_one_group(&aug, &off_aug, AVELOGCPM_DISP);
    (beta + 1e6f64.ln()) / LN2
}

/// q2qnbinom: remap count x from an NB with mean `input_mean` to the NB with
/// mean `output_mean` (same dispersion) by averaging a normal-quantile and a
/// gamma-quantile transformation in the matched tail (Robinson & Smyth).
fn q2qnbinom(x: f64, input_mean: f64, output_mean: f64, dispersion: f64) -> f64 {
    let eps = 1e-14;
    let (mut im, mut om) = (input_mean, output_mean);
    if im < eps || om < eps {
        im += 0.25;
        om += 0.25;
    }
    let ri = 1.0 + dispersion * im;
    let vi = im * ri;
    let ro = 1.0 + dispersion * om;
    let vo = om * ro;
    let upper = x >= im;
    let p1 = special::pnorm(x, im, vi.sqrt(), !upper, true);
    let p2 = special::pgamma(x, im / ri, ri, !upper, true);
    let q1 = special::qnorm(p1, om, vo.sqrt(), !upper, true);
    let q2 = special::qgamma(p2, om / ro, ro, !upper, true);
    (q1 + q2) / 2.0
}

/// Doubletail NB exact test on two unrounded group sums (n1 / n2 libraries),
/// conditional on the total. Mirrors edgeR exactTestDoubleTail: Poisson limit
/// via the exact binomial test, beta approximation for large counts (using the
/// unrounded sums), otherwise the convolution of two NB pmfs summed over the
/// equal-or-more-extreme integer tail and doubled.
fn exact_doubletail(raw1: f64, raw2: f64, n1: f64, n2: f64, dispersion: f64) -> f64 {
    let s1 = raw1.round();
    let s2 = raw2.round();
    let s = s1 + s2;
    if s == 0.0 {
        return 1.0;
    }
    if dispersion <= 0.0 {
        return binom_test(s1, s2, n1 / (n1 + n2));
    }
    if s1 > BIG_COUNT && s2 > BIG_COUNT {
        return beta_approx(raw1, raw2, n1, n2, dispersion);
    }
    let mu = s / (n1 + n2);
    let mu1 = n1 * mu;
    let mu2 = n2 * mu;
    let size1 = n1 / dispersion;
    let size2 = n2 / dispersion;
    let p_bot = special::dnbinom_mu(s, (n1 + n2) / dispersion, s);
    let p = if s1 < mu1 {
        let mut top = 0.0;
        let mut x = 0.0;
        while x <= s1 {
            top += special::dnbinom_mu(x, size1, mu1) * special::dnbinom_mu(s - x, size2, mu2);
            x += 1.0;
        }
        2.0 * top / p_bot
    } else if s1 > mu1 {
        let mut top = 0.0;
        let mut x = s1;
        while x <= s {
            top += special::dnbinom_mu(x, size1, mu1) * special::dnbinom_mu(s - x, size2, mu2);
            x += 1.0;
        }
        2.0 * top / p_bot
    } else {
        1.0
    };
    p.min(1.0)
}

/// Exact two-sided binomial test (Poisson-limit case), edgeR binomTest with
/// p != 0.5: sum the probabilities of all outcomes no more likely than observed.
fn binom_test(y1: f64, y2: f64, p: f64) -> f64 {
    let y1 = y1.round();
    let size = (y1 + y2.round()) as i64;
    if size == 0 {
        return 1.0;
    }
    let mut d = vec![0.0f64; (size + 1) as usize];
    for (k, dk) in d.iter_mut().enumerate() {
        *dk = special::dbinom(k as f64, size as f64, p);
    }
    let observed = d[y1 as usize];
    let tol = observed * (1.0 + 1e-7);
    let pv: f64 = d.iter().filter(|&&v| v <= tol).sum();
    pv.min(1.0)
}

/// Beta approximation for large counts (edgeR exactTestBetaApprox).
fn beta_approx(s1: f64, s2: f64, n1: f64, n2: f64, dispersion: f64) -> f64 {
    let y = s1 + s2;
    if y <= 0.0 {
        return 1.0;
    }
    let mu = y / (n1 + n2);
    let alpha1 = n1 * mu / (1.0 + dispersion * mu);
    let alpha2 = n2 / n1 * alpha1;
    let med = special::qbeta(0.5, alpha1, alpha2);
    if (s1 + 0.5) / y < med {
        (2.0 * special::pbeta((s1 + 0.5) / y, alpha1, alpha2, true, false)).min(1.0)
    } else if (s1 - 0.5) / y > med {
        (2.0 * special::pbeta((s1 - 0.5) / y, alpha1, alpha2, false, false)).min(1.0)
    } else {
        1.0
    }
}

fn bh_fdr(pvals: &[f64]) -> Vec<f64> {
    let n = pvals.len();
    let mut order: Vec<usize> = (0..n).collect();
    order.sort_by(|&a, &b| pvals[b].partial_cmp(&pvals[a]).unwrap());
    let mut adj = vec![0.0f64; n];
    let mut cummin = f64::INFINITY;
    for (rank, &i) in order.iter().enumerate() {
        let m = n - rank;
        let v = (pvals[i] * n as f64 / m as f64).min(1.0);
        cummin = cummin.min(v);
        adj[i] = cummin;
    }
    adj
}

pub fn exact_test(
    counts_path: &Path,
    group_spec: &str,
    norm_factors_path: Option<&Path>,
    opts: &ExactTestOpts,
    output: &mut dyn Write,
) -> Result<u64> {
    let m = Matrix::load(counts_path)?;
    let (g1, g2) = parse_groups(group_spec, m.n_samples)?;

    let norm_factors = match norm_factors_path {
        Some(p) => load_norm_factors(p, m.n_samples)?,
        None => vec![1.0; m.n_samples],
    };

    let mut lib = vec![0.0f64; m.n_samples];
    for row in m.counts.chunks_exact(m.n_samples) {
        for (s, &c) in lib.iter_mut().zip(row) {
            *s += c;
        }
    }
    let eff_lib: Vec<f64> = lib
        .iter()
        .zip(&norm_factors)
        .map(|(&l, &f)| l * f)
        .collect();
    let offset: Vec<f64> = eff_lib.iter().map(|&l| l.ln()).collect();
    let lib_average = (offset.iter().sum::<f64>() / offset.len() as f64).exp();

    let mean_eff = eff_lib.iter().sum::<f64>() / eff_lib.len() as f64;
    let prior_fc: Vec<f64> = eff_lib
        .iter()
        .map(|&l| PRIOR_COUNT_FC * l / mean_eff)
        .collect();
    let offset_aug: Vec<f64> = eff_lib
        .iter()
        .zip(&prior_fc)
        .map(|(&l, &p)| (l + 2.0 * p).ln())
        .collect();

    let disp = opts.dispersion;
    let n1 = g1.len() as f64;
    let n2 = g2.len() as f64;

    let gene_col = m.header.split('\t').next().unwrap_or("gene");
    let mut logfc = Vec::with_capacity(m.n_genes());
    let mut logcpm = Vec::with_capacity(m.n_genes());
    let mut pvals = Vec::with_capacity(m.n_genes());

    let mut buf1 = vec![0.0f64; g1.len()];
    let mut buf2 = vec![0.0f64; g2.len()];
    let mut off1 = vec![0.0f64; g1.len()];
    let mut off2 = vec![0.0f64; g2.len()];
    for (j, &s) in g1.iter().enumerate() {
        off1[j] = offset_aug[s];
    }
    for (j, &s) in g2.iter().enumerate() {
        off2[j] = offset_aug[s];
    }
    let off1_g: Vec<f64> = g1.iter().map(|&s| offset[s]).collect();
    let off2_g: Vec<f64> = g2.iter().map(|&s| offset[s]).collect();
    let mut all_off = vec![0.0f64; g1.len() + g2.len()];
    all_off[..g1.len()].copy_from_slice(&off1_g);
    all_off[g1.len()..].copy_from_slice(&off2_g);
    let mut all_row = vec![0.0f64; g1.len() + g2.len()];

    for g in 0..m.n_genes() {
        let row = m.row(g);

        for (j, &s) in g1.iter().enumerate() {
            buf1[j] = row[s] + prior_fc[s];
        }
        for (j, &s) in g2.iter().enumerate() {
            buf2[j] = row[s] + prior_fc[s];
        }
        let ab1 = mglm_one_group(&buf1, &off1, disp);
        let ab2 = mglm_one_group(&buf2, &off2, disp);
        logfc.push((ab2 - ab1) / LN2);

        for (j, &s) in g1.iter().enumerate() {
            all_row[j] = row[s];
        }
        for (j, &s) in g2.iter().enumerate() {
            all_row[g1.len() + j] = row[s];
        }
        logcpm.push(ave_log_cpm_gene(&all_row, &eff_lib));

        let abundance = mglm_one_group(&all_row, &all_off, disp);
        let e = abundance.exp();
        let output_mean = e * lib_average;
        let mut s1 = 0.0;
        for &si in &g1 {
            let im = e * eff_lib[si];
            s1 += q2qnbinom(row[si], im, output_mean, disp);
        }
        let mut s2 = 0.0;
        for &si in &g2 {
            let im = e * eff_lib[si];
            s2 += q2qnbinom(row[si], im, output_mean, disp);
        }
        pvals.push(exact_doubletail(s1, s2, n1, n2, disp));
    }

    let fdr = if opts.fdr { Some(bh_fdr(&pvals)) } else { None };

    let mut out = BufWriter::new(output);
    if let Some(_f) = &fdr {
        writeln!(out, "{gene_col}\tlogFC\tlogCPM\tPValue\tFDR").map_err(RsomicsError::Io)?;
    } else {
        writeln!(out, "{gene_col}\tlogFC\tlogCPM\tPValue").map_err(RsomicsError::Io)?;
    }
    for g in 0..m.n_genes() {
        write!(
            out,
            "{}\t{}\t{}\t{}",
            m.genes[g],
            fmt(logfc[g]),
            fmt(logcpm[g]),
            fmt_p(pvals[g])
        )
        .map_err(RsomicsError::Io)?;
        if let Some(f) = &fdr {
            write!(out, "\t{}", fmt_p(f[g])).map_err(RsomicsError::Io)?;
        }
        writeln!(out).map_err(RsomicsError::Io)?;
    }
    out.flush().map_err(RsomicsError::Io)?;
    Ok(m.n_genes() as u64)
}

fn fmt(v: f64) -> String {
    format!("{v:.6}")
}

/// Scientific notation with R's `formatC(format="e")` exponent shape: a sign and
/// at least two exponent digits (`4.676393e-02`), so goldens diff cleanly.
fn fmt_p(v: f64) -> String {
    let s = format!("{v:.6e}");
    let (mantissa, exp) = s.split_once('e').unwrap();
    let (sign, digits) = match exp.strip_prefix('-') {
        Some(d) => ('-', d),
        None => ('+', exp),
    };
    format!("{mantissa}e{sign}{digits:0>2}")
}

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

    // Probed in R: exactTest(d, dispersion=0.1) on the 5-gene fixture.
    const EXPECT_LOGFC: [f64; 5] = [
        2.07695720785,
        0.08819639348,
        5.71691109651,
        -2.08892448998,
        -0.13684694279,
    ];
    const EXPECT_LOGCPM: [f64; 5] = [
        15.60092586,
        17.29409726,
        13.03042119,
        12.44757336,
        19.58295511,
    ];
    const EXPECT_PVAL: [f64; 5] = [
        0.0003140189807,
        0.8671983532895,
        0.0010940667164,
        0.2522215919834,
        0.7703694926792,
    ];

    fn fixture() -> (Vec<String>, Vec<f64>, usize) {
        let genes = (1..=5).map(|i| format!("g{i}")).collect();
        let counts = vec![
            10.0, 12.0, 40.0, 55.0, 100.0, 90.0, 95.0, 110.0, 0.0, 0.0, 5.0, 8.0, 3.0, 2.0, 0.0,
            1.0, 500.0, 520.0, 480.0, 460.0,
        ];
        (genes, counts, 4)
    }

    #[test]
    fn matches_edger_probe() {
        let (genes, counts, ns) = fixture();
        let m = Matrix {
            header: "gene\ta1\ta2\tb1\tb2".into(),
            genes,
            counts,
            n_samples: ns,
        };
        let (g1, g2) = parse_groups("a,a,b,b", ns).unwrap();
        let disp = 0.1;

        let mut lib = vec![0.0; ns];
        for r in m.counts.chunks_exact(ns) {
            for (s, &c) in lib.iter_mut().zip(r) {
                *s += c;
            }
        }
        let offset: Vec<f64> = lib.iter().map(|&l| l.ln()).collect();
        let lib_average = (offset.iter().sum::<f64>() / ns as f64).exp();
        let mean_eff = lib.iter().sum::<f64>() / ns as f64;
        let prior_fc: Vec<f64> = lib.iter().map(|&l| PRIOR_COUNT_FC * l / mean_eff).collect();
        let offset_aug: Vec<f64> = lib
            .iter()
            .zip(&prior_fc)
            .map(|(&l, &p)| (l + 2.0 * p).ln())
            .collect();
        let off1: Vec<f64> = g1.iter().map(|&s| offset_aug[s]).collect();
        let off2: Vec<f64> = g2.iter().map(|&s| offset_aug[s]).collect();
        let all_off: Vec<f64> = (0..ns).map(|s| offset[s]).collect();

        for g in 0..5 {
            let row = m.row(g);
            let b1: Vec<f64> = g1.iter().map(|&s| row[s] + prior_fc[s]).collect();
            let b2: Vec<f64> = g2.iter().map(|&s| row[s] + prior_fc[s]).collect();
            let lfc = (mglm_one_group(&b2, &off2, disp) - mglm_one_group(&b1, &off1, disp)) / LN2;
            assert!((lfc - EXPECT_LOGFC[g]).abs() < 1e-6, "logFC g{g}: {lfc}");

            let lcpm = ave_log_cpm_gene(row, &lib);
            assert!(
                (lcpm - EXPECT_LOGCPM[g]).abs() < 1e-5,
                "logCPM g{g}: {lcpm}"
            );

            let abundance = mglm_one_group(row, &all_off, disp);
            let e = abundance.exp();
            let om = e * lib_average;
            let s1: f64 = g1
                .iter()
                .map(|&s| q2qnbinom(row[s], e * lib[s], om, disp))
                .sum();
            let s2: f64 = g2
                .iter()
                .map(|&s| q2qnbinom(row[s], e * lib[s], om, disp))
                .sum();
            let p = exact_doubletail(s1, s2, 2.0, 2.0, disp);
            let rel = (p - EXPECT_PVAL[g]).abs() / (1.0 + EXPECT_PVAL[g].abs());
            assert!(rel < 1e-5, "PValue g{g}: {p} vs {}", EXPECT_PVAL[g]);
        }
    }

    #[test]
    fn bh_matches_r_padjust() {
        let p = [
            0.0003140189807,
            0.8671983532895,
            0.0010940667164,
            0.2522215919834,
            0.7703694926792,
        ];
        let adj = bh_fdr(&p);
        // p.adjust(p, "BH") in R:
        let expect = [
            0.001570095,
            0.867198353,
            0.002735167,
            0.420369320,
            0.867198353,
        ];
        for (a, e) in adj.iter().zip(expect) {
            assert!((a - e).abs() < 1e-6, "{a} vs {e}");
        }
    }
}