universalmotif 0.2.0

//! Motif enrichment with [`enrich_motifs`].
//!
//! Target and background sequences are scanned for hits, and each motif's
//! hits are compared between the two sets with a one-sided Fisher's exact
//! test, either on the number of sequences with a hit or on the number of
//! hits per scannable position.

use crate::error::{Result, arg_err};
use crate::motif::Motif;
use crate::scan::{ScanOptions, scan_table};
use crate::shuffle::{ShuffleMethod, ShuffleOptions, shuffle_sequences};
use std::f64::consts::LN_10;

use crate::stats::{fisher_greater, ln_fisher_greater, ln_p_adjust_bh, p_adjust_bh};

/// What [`enrich_motifs`] counts.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EnrichTest {
    /// Sequences containing at least one hit.
    Seqs,
    /// Hits, out of all the positions a motif could occupy.
    Sites,
}

impl EnrichTest {
    /// The name [`str::parse`] reads, such as `"seqs"`.
    pub fn as_str(&self) -> &'static str {
        match self {
            EnrichTest::Seqs => "seqs",
            EnrichTest::Sites => "sites",
        }
    }
}

impl std::str::FromStr for EnrichTest {
    type Err = crate::error::Error;

    /// Parse `"seqs"` or `"sites"`.
    fn from_str(s: &str) -> Result<EnrichTest> {
        match s {
            "seqs" => Ok(EnrichTest::Seqs),
            "sites" => Ok(EnrichTest::Sites),
            _ => arg_err(format!("unknown test '{s}'; expected 'seqs' or 'sites'")),
        }
    }
}

impl std::fmt::Display for EnrichTest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

/// Options for [`enrich_motifs`].
#[derive(Debug, Clone, PartialEq)]
pub struct EnrichOptions {
    /// P-value cutoff for hits.
    pub pvalue: f64,
    /// What to count; sequences with hits by default.
    pub test: EnrichTest,
    /// Scan both strands.
    pub rc: bool,
    /// k-mers kept when shuffling the targets into a background.
    pub shuffle_k: usize,
    /// Seed for that shuffle.
    pub seed: u64,
    /// Added to every cell of the 2x2 table.
    pub pseudocount: u64,
    /// Threads to use; 0 means all available.
    pub nthreads: usize,
}

impl Default for EnrichOptions {
    fn default() -> Self {
        EnrichOptions {
            pvalue: 1e-4,
            test: EnrichTest::Seqs,
            rc: true,
            shuffle_k: 2,
            seed: 1,
            pseudocount: 1,
            nthreads: 1,
        }
    }
}

/// Enrichment of one motif.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Enrichment {
    /// Index of the motif.
    pub motif: usize,
    /// The number of target sequences.
    pub target_seq_n: usize,
    /// Target sequences with at least one hit.
    pub target_seq_hits: usize,
    /// Hits in the target sequences.
    pub target_site_hits: usize,
    /// The number of background sequences.
    pub bkg_seq_n: usize,
    /// Background sequences with at least one hit.
    pub bkg_seq_hits: usize,
    /// Hits in the background sequences.
    pub bkg_site_hits: usize,
    /// Ratio of the hit rates, with a floor of half a hit on the background
    /// rate to avoid dividing by zero.
    pub enrichment: f64,
    /// The base 2 logarithm of the enrichment.
    pub log2_enrichment: f64,
    /// One-sided Fisher's exact P-value.
    pub pvalue: f64,
    /// Base 10 log of `pvalue`, which stays finite when `pvalue` is too small
    /// for a double.
    pub log10_pvalue: f64,
    /// Benjamini-Hochberg q-value across all motifs.
    pub qvalue: f64,
    /// Base 10 log of `qvalue`, which stays finite when `qvalue` is too small
    /// for a double.
    pub log10_qvalue: f64,
}

struct Counts {
    seqs: Vec<usize>,
    sites: Vec<usize>,
}

/// Hit counts in the targets and in the background, from one scan of both
/// so that each motif's scoring tables are built once. The first `n_tgt`
/// sequences are the targets.
fn count_hits(
    motifs: &[Motif],
    seqs: &[&str],
    n_tgt: usize,
    opts: &ScanOptions,
) -> Result<(Counts, Counts)> {
    let hits = scan_table(motifs, seqs, opts)?;
    let empty = || Counts {
        seqs: vec![0; motifs.len()],
        sites: vec![0; motifs.len()],
    };
    let (mut tgt, mut neg) = (empty(), empty());
    // Hits come ordered by motif and sequence, so a new sequence for a motif
    // shows up as a change from the previous row.
    let mut last: Option<(u32, u32)> = None;
    for (&m, &s) in hits.motif.iter().zip(&hits.sequence) {
        let c = if (s as usize) < n_tgt {
            &mut tgt
        } else {
            &mut neg
        };
        c.sites[m as usize] += 1;
        if last != Some((m, s)) {
            c.seqs[m as usize] += 1;
            last = Some((m, s));
        }
    }
    Ok((tgt, neg))
}

/// Test motifs for enrichment in `sequences` relative to `background`, which by
/// default is the targets shuffled with the Euler method keeping
/// `shuffle_k`-lets. Returns one result per motif, in input order and
/// unfiltered.
pub fn enrich_motifs<S: AsRef<str> + Sync>(
    motifs: &[Motif],
    sequences: &[S],
    background: Option<&[S]>,
    opts: &EnrichOptions,
) -> Result<Vec<Enrichment>> {
    let Some(first) = motifs.first() else {
        return arg_err("no motifs given");
    };
    if motifs.iter().any(|m| m.alphabet != first.alphabet) {
        return arg_err("all motifs must share the same alphabet");
    }
    if !(opts.pvalue > 0.0 && opts.pvalue < 1.0) {
        return arg_err("pvalue must be in (0, 1)");
    }
    if sequences.is_empty() {
        return arg_err("no sequences given");
    }
    let shuffled: Vec<String>;
    let bkg: Vec<&str> = match background {
        Some(b) => b.iter().map(|s| s.as_ref()).collect(),
        None => {
            shuffled = shuffle_sequences(
                sequences,
                &ShuffleOptions {
                    k: opts.shuffle_k,
                    method: ShuffleMethod::Euler,
                    seed: opts.seed,
                    window: None,
                    nthreads: opts.nthreads,
                },
            )?;
            shuffled.iter().map(|s| s.as_str()).collect()
        }
    };
    if bkg.is_empty() {
        return arg_err("no background sequences");
    }
    let scan = ScanOptions {
        pvalue: opts.pvalue,
        rc: opts.rc,
        no_overlaps: None,
        nthreads: opts.nthreads,
        ..Default::default()
    };
    let both: Vec<&str> = sequences
        .iter()
        .map(|s| s.as_ref())
        .chain(bkg.iter().copied())
        .collect();
    let (tgt, neg) = count_hits(motifs, &both, sequences.len(), &scan)?;
    let (n_tgt, n_bkg) = (sequences.len(), bkg.len());
    let tgt_lens: Vec<usize> = sequences.iter().map(|s| s.as_ref().len()).collect();
    let bkg_lens: Vec<usize> = bkg.iter().map(|s| s.len()).collect();
    // Only DNA and RNA have a reverse strand to scan.
    let strands = if opts.rc && first.alphabet.is_nucleotide() {
        2
    } else {
        1
    };
    let positions = |lens: &[usize], w: usize| -> usize {
        lens.iter()
            .map(|&l| (l + 1).saturating_sub(w))
            .sum::<usize>()
            * strands
    };

    let pc = opts.pseudocount;
    let mut out: Vec<Enrichment> = Vec::with_capacity(motifs.len());
    for (i, m) in motifs.iter().enumerate() {
        let (a, b, c, d, eff) = match opts.test {
            EnrichTest::Seqs => {
                let (a, b) = (tgt.seqs[i], neg.seqs[i]);
                let pos_rate = a as f64 / n_tgt as f64;
                let neg_rate = b as f64 / n_bkg as f64;
                let neg_den = if neg_rate > 1e-9 {
                    neg_rate
                } else {
                    0.5 / n_bkg as f64
                };
                (a, b, n_tgt - a, n_bkg - b, pos_rate / neg_den)
            }
            EnrichTest::Sites => {
                let w = m.width();
                let (ppos, npos) = (positions(&tgt_lens, w), positions(&bkg_lens, w));
                let (a, b) = (tgt.sites[i], neg.sites[i]);
                let pos_rate = if ppos > 0 {
                    a as f64 / ppos as f64
                } else {
                    0.0
                };
                let neg_den = if npos > 0 && b > 0 {
                    b as f64 / npos as f64
                } else {
                    0.5 / npos.max(1) as f64
                };
                let eff = if pos_rate > 0.0 {
                    pos_rate / neg_den
                } else {
                    0.0
                };
                (a, b, ppos.saturating_sub(a), npos.saturating_sub(b), eff)
            }
        };
        let cells = (a as u64 + pc, b as u64 + pc, c as u64 + pc, d as u64 + pc);
        let pvalue = fisher_greater(cells.0, cells.1, cells.2, cells.3);
        let log10_pvalue = ln_fisher_greater(cells.0, cells.1, cells.2, cells.3) / LN_10;
        out.push(Enrichment {
            motif: i,
            target_seq_n: n_tgt,
            target_seq_hits: tgt.seqs[i],
            target_site_hits: tgt.sites[i],
            bkg_seq_n: n_bkg,
            bkg_seq_hits: neg.seqs[i],
            bkg_site_hits: neg.sites[i],
            enrichment: eff,
            log2_enrichment: if eff > 0.0 {
                eff.log2()
            } else {
                f64::NEG_INFINITY
            },
            pvalue,
            log10_pvalue,
            qvalue: 0.0,
            log10_qvalue: 0.0,
        });
    }
    let q = p_adjust_bh(&out.iter().map(|e| e.pvalue).collect::<Vec<_>>());
    let ln_q = ln_p_adjust_bh(
        &out.iter()
            .map(|e| e.log10_pvalue * LN_10)
            .collect::<Vec<_>>(),
        out.len() as f64,
    );
    for ((e, q), ln_q) in out.iter_mut().zip(q).zip(ln_q) {
        e.qvalue = q;
        // Where the q-value is representable, its log is taken directly so
        // that the two columns agree exactly.
        e.log10_qvalue = if q > 1e-300 { q.log10() } else { ln_q / LN_10 };
    }
    Ok(out)
}