limma-rust 0.1.0

//! Predictive (empirical-Bayes-shrunk) fold changes. Pure-Rust port of limma's
//! `predFCm` (Phipson & Smyth).
//!
//! Given an [`MArrayLM`] fit, [`pred_fcm`] estimates the proportion of truly
//! differentially expressed genes (`1 - propTrueNull`), re-runs `eBayes` at that
//! proportion, fits a gamma intercept to the squared coefficients to get the
//! prior fold-change variance `v0`, and shrinks each coefficient toward zero by
//! `v0 / (v0 + v)`. When `all_de = false` the shrunk value is further scaled by
//! the posterior probability of differential expression.
//!
//! The `var_indep_of_fc = false` branch reproduces limma's `pmin(v0, 1e-8)`
//! clamp verbatim (which forces `v0` down to `1e-8`), so its output is the same
//! near-zero shrinkage as the reference implementation.

use anyhow::Result;
use ndarray::Array1;

use crate::ebayes::ebayes;
use crate::fit::MArrayLM;
use crate::fitgamma::fit_gamma_intercept;
use crate::proptruenull::{prop_true_null, PropTrueNullMethod};

/// Whether the entries of `v` are not all identical (used to detect the
/// trend / robust eBayes priors, which vary by gene). `Inf == Inf` counts as
/// equal, so an all-`Inf` `df.prior` is treated as constant (non-robust).
fn is_varying(v: &Array1<f64>) -> bool {
    let first = v[0];
    v.iter().any(|&x| x != first)
}

/// Logistic transform of the log-odds, with limma's overflow guard
/// (`lods > 700 -> 1`).
fn prob_de(lods: f64) -> f64 {
    if lods > 700.0 {
        1.0
    } else {
        let e = lods.exp();
        e / (1.0 + e)
    }
}

/// Predictive fold changes for coefficient `coef` (0-based) of `fit`. Port of
/// `predFCm`. `method` selects the `propTrueNull` estimator (limma's default is
/// [`PropTrueNullMethod::Lfdr`]); the histogram estimator uses limma's default
/// of 20 bins. Returns one shrunken fold change per gene.
pub fn pred_fcm(
    fit: &MArrayLM,
    coef: usize,
    var_indep_of_fc: bool,
    all_de: bool,
    method: PropTrueNullMethod,
) -> Result<Vec<f64>> {
    // Ensure eBayes has been run (default proportion/trend/robust) so that
    // p.value, s2.prior and df.prior are available.
    let mut base = fit.clone();
    if base.p_value.is_none() {
        ebayes(&mut base, 0.01, (0.1, 4.0), false, false)?;
    }
    let ng = base.coefficients.nrows();

    // p = 1 - propTrueNull(p.value[, coef]); clamp away from exactly zero.
    let pcol: Vec<f64> = base.p_value.as_ref().unwrap().column(coef).to_vec();
    let mut p = 1.0 - prop_true_null(&pcol, method, 20);
    if p == 0.0 {
        p = 1e-8;
    }

    // trend/robust are inferred from whether the priors vary by gene. (limma
    // tests `length(s2.prior)==1`; our MArrayLM always stores a full-length
    // vector, so a constant vector is the non-trend / non-robust case.)
    let trend = base.s2_prior.as_ref().is_some_and(is_varying);
    let robust = base.df_prior.as_ref().is_some_and(is_varying);

    // Re-run eBayes at the estimated proportion (on the original lmFit fit).
    let mut f = fit.clone();
    ebayes(&mut f, p, (0.1, 4.0), trend, robust)?;

    let v = f.cov_coefficients[[coef, coef]];
    let beta: Array1<f64> = f.coefficients.column(coef).to_owned();
    let s2post = f.s2_post.as_ref().expect("eBayes fills s2_post");

    let a = p / (1.0 - p);

    let pfc: Vec<f64> = if var_indep_of_fc {
        let y2: Vec<f64> = beta.iter().map(|&b| b * b).collect();
        let offset: Vec<f64> = s2post.iter().map(|&s| v * s).collect();
        let mut v0 = fit_gamma_intercept(&y2, &offset, 1000);
        if v0 < 0.0 {
            v0 = 1e-8;
        }
        let base = beta
            .iter()
            .zip(s2post.iter())
            .map(|(&b, &s)| b * v0 / (v0 + v * s));
        if all_de {
            base.collect()
        } else {
            base.zip(beta.iter().zip(s2post.iter()))
                .map(|(pf, (&bb, &s))| {
                    let vs = v * s;
                    let bfac = (vs / (vs + v0)).sqrt();
                    let cfac = (bb * bb * v0 / (2.0 * v * v * s * s + 2.0 * v * v0 * s)).exp();
                    pf * prob_de((a * bfac * cfac).ln())
                })
                .collect()
        }
    } else {
        let b2: Vec<f64> = beta
            .iter()
            .zip(s2post.iter())
            .map(|(&b, &s)| b * b / s)
            .collect();
        let offset = vec![v; ng];
        // limma clamps with pmin (a minimum), forcing v0 down to <= 1e-8.
        let v0 = fit_gamma_intercept(&b2, &offset, 1000).min(1e-8);
        let base = beta.iter().map(|&b| b * v0 / (v0 + v));
        if all_de {
            base.collect()
        } else {
            let bfac = (v / (v + v0)).sqrt();
            base.zip(beta.iter().zip(s2post.iter()))
                .map(|(pf, (&bb, &s))| {
                    let cfac = (bb * bb * v0 / (2.0 * v * v * s + 2.0 * v * v0 * s)).exp();
                    pf * prob_de((a * bfac * cfac).ln())
                })
                .collect()
        }
    };

    Ok(pfc)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fit::lmfit;
    use ndarray::{array, Array2};

    fn build_fit() -> MArrayLM {
        let y: Array2<f64> = array![
            [
                2.28724716134052,
                0.839750359624071,
                1.21855053450735,
                2.57637147386057,
                2.8425853501473,
                3.01905836527295
            ],
            [
                -1.19677168222235,
                0.7053418309055,
                -0.699317078685514,
                -1.84084472173727,
                -1.99575176374787,
                -1.41246029509811
            ],
            [
                -0.694292510435459,
                1.30596472081169,
                -0.285432751528726,
                2.3436741847097,
                1.82921984200165,
                1.72066693894479
            ],
            [
                -0.412292951136803,
                -1.38799621659285,
                -1.31155267260939,
                -0.795192647386953,
                -1.8934234291213,
                -2.67436101486938
            ],
            [
                -0.970673341119483,
                1.27291686425524,
                -0.391012431449258,
                3.31896914255711,
                2.20729543721107,
                3.30872211763837
            ],
            [
                -0.947279945228108,
                0.184192771235767,
                -0.401526613094972,
                -1.49075021102894,
                -2.91170186529517,
                -4.20387854268349
            ],
            [
                0.748139340290551,
                0.752279895740033,
                1.35051758092295,
                0.769154194657112,
                -0.34606859178086,
                0.991289625052712
            ],
            [
                -0.116955225887152,
                0.591745052462727,
                0.591190027089221,
                1.15347367477894,
                -0.304607588245324,
                1.02322044470037
            ],
            [
                0.152657626282234,
                -0.983052595771021,
                0.100525455628569,
                1.26068350268094,
                -1.78589348744452,
                0.840145438883483
            ],
            [
                2.18997810732938,
                -0.276063955112006,
                0.931071995520097,
                0.700623506572324,
                0.58727467185797,
                0.12007860795572
            ],
            [
                0.356986230329022,
                -0.870851022568591,
                -0.262742348566532,
                0.432627160845955,
                1.63579443444659,
                -0.426255055445707
            ],
            [
                2.71675178313072,
                0.718710553084245,
                -0.00766810471266396,
                -0.922601718256921,
                -0.645423473634397,
                0.458926243504027
            ],
            [
                2.28145192598956,
                0.110652877769336,
                0.367153006545634,
                -0.615584206630919,
                0.61899216878734,
                0.645047947693915
            ],
            [
                0.324020540138516,
                -0.0784667679717042,
                1.70716254513761,
                -0.866659688251375,
                0.236393598401322,
                0.611530549290199
            ],
            [
                1.89606706680993,
                -0.420490459341998,
                0.72374026252838,
                -1.63951708718114,
                0.846500898751643,
                -0.889211293868794
            ],
            [
                0.467680511321698,
                -0.562125876285266,
                0.481036048707917,
                -1.32583924384341,
                -0.573645738849693,
                1.54389234869222
            ],
            [
                -0.893800723085444,
                0.997513444755305,
                -1.56786824422525,
                -0.88903672763797,
                1.11799320399617,
                -1.24176360488504
            ],
            [
                -0.307328299537195,
                -1.10513005881326,
                0.318250283480828,
                -0.557602330302113,
                -1.54000113193302,
                1.10344734039128
            ],
            [
                -0.00482242226757041,
                -0.142287830774585,
                0.16599145067735,
                -0.0624023088383481,
                -0.438123899300085,
                0.982772356675575
            ],
            [
                0.988164149499945,
                0.314994904887913,
                -0.899907629628172,
                2.42269297715943,
                -0.150672970896448,
                0.304327174033201
            ],
        ];
        let design: Array2<f64> = array![
            [1.0, 0.0],
            [1.0, 0.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [1.0, 1.0],
            [1.0, 1.0],
        ];
        lmfit(&y, &design, vec![String::new(); 20], vec![String::new(); 2]).unwrap()
    }

    fn assert_close(got: &[f64], want: &[f64]) {
        assert_eq!(got.len(), want.len());
        for (g, w) in got.iter().zip(want.iter()) {
            assert!((g - w).abs() <= 1e-7 * w.abs() + 1e-13, "got {g}, want {w}");
        }
    }

    #[test]
    fn pred_fcm_matches_r() {
        let fit = build_fit();

        // var.indep.of.fc = TRUE, all.de = TRUE (defaults).
        let vi1_ad1 = [
            0.806899941410833,
            -0.800165248392108,
            1.09769272440758,
            -0.443849844795292,
            1.75947115517052,
            -1.46726184458905,
            -0.28324270522318,
            0.158937728669327,
            0.206001182143709,
            -0.283330991201421,
            0.476902630464977,
            -0.894526143117239,
            -0.416180722805507,
            -0.38870547681697,
            -0.765312904916374,
            -0.14633414417854,
            0.0889910610860675,
            0.0197269550114635,
            0.0913603309470277,
            0.428463022546148,
        ];
        let got = pred_fcm(&fit, 1, true, true, PropTrueNullMethod::Lfdr).unwrap();
        assert_close(&got, &vi1_ad1);

        // all.de = TRUE makes the result independent of the proportion method.
        let got_mean = pred_fcm(&fit, 1, true, true, PropTrueNullMethod::Mean).unwrap();
        assert_close(&got_mean, &vi1_ad1);

        // var.indep.of.fc = TRUE, all.de = FALSE.
        let vi1_ad0 = [
            0.270478766277483,
            -0.265522699426538,
            0.57398255588566,
            -0.093477821284407,
            1.65297623275505,
            -1.18175391914552,
            -0.0523086424303369,
            0.0275529195325538,
            0.0364280464603553,
            -0.0523279383852657,
            0.103863229886988,
            -0.342837057870809,
            -0.085364720344582,
            -0.0777824277550811,
            -0.241150197241235,
            -0.0252550990394734,
            0.0151196325151701,
            0.00332235440767211,
            0.0155298924292991,
            0.0889042770881764,
        ];
        let got = pred_fcm(&fit, 1, true, false, PropTrueNullMethod::Lfdr).unwrap();
        assert_close(&got, &vi1_ad0);

        // var.indep.of.fc = FALSE, all.de = TRUE (pmin clamp -> ~1e-8 scale).
        let vi0_ad1 = [
            2.04623353621094e-08,
            -2.02915489485312e-08,
            2.78366071164941e-08,
            -1.12556760863579e-08,
            4.46187773593186e-08,
            -3.72085836014727e-08,
            -7.18280783738075e-09,
            4.03053332738792e-09,
            5.22402476154024e-09,
            -7.18504669898158e-09,
            1.20938682218567e-08,
            -2.2684465559181e-08,
            -1.05540093439608e-08,
            -9.85725914146247e-09,
            -1.94077214703572e-08,
            -3.71091653306243e-09,
            2.25674193629162e-09,
            5.00259757625001e-10,
            2.31682471975736e-09,
            1.08654786147003e-08,
        ];
        let got = pred_fcm(&fit, 1, false, true, PropTrueNullMethod::Lfdr).unwrap();
        assert_close(&got, &vi0_ad1);

        // var.indep.of.fc = FALSE, all.de = FALSE.
        let vi0_ad0 = [
            4.92171467490767e-09,
            -4.88063617568995e-09,
            6.69541561943082e-09,
            -2.70727776232544e-09,
            1.07319571258011e-08,
            -8.94961577592611e-09,
            -1.72764885189419e-09,
            9.69446269885002e-10,
            1.25651146134861e-09,
            -1.72818735532886e-09,
            2.90888439514644e-09,
            -5.45619380724241e-09,
            -2.53850897651131e-09,
            -2.3709227445188e-09,
            -4.66805306961375e-09,
            -8.92570262603364e-10,
            5.42804108910441e-10,
            1.20325256308392e-10,
            5.57255553814274e-10,
            2.61342529709022e-09,
        ];
        let got = pred_fcm(&fit, 1, false, false, PropTrueNullMethod::Lfdr).unwrap();
        assert_close(&got, &vi0_ad0);
    }
}