universalmotif 0.2.0

//! Motif P-values and score thresholds.
//!
//! A motif's P-value for a score is the probability that a random site, drawn
//! from the background, scores at least that much. It is computed exactly from
//! the motif's score distribution, as described in [`crate::cdf`].

use rayon::prelude::*;

use crate::cdf::{ScoreDistribution, grid_value};
use crate::convert::convert_type;
use crate::error::{Error, Result, arg_err};
use crate::matrix::Matrix;
use crate::motif::{Motif, MotifType};
use crate::ops::with_pseudocount;

/// Options for [`motif_pvalue`] and [`pvalue_to_score`].
#[derive(Debug, Clone, PartialEq)]
pub struct PvalueOptions {
    /// Background frequencies: none to use each motif's own, one set for all
    /// motifs, or one set per motif.
    pub bkg: Vec<Vec<f64>>,
    /// Threads to use; 0 means all available.
    pub nthreads: usize,
}

impl Default for PvalueOptions {
    fn default() -> Self {
        PvalueOptions {
            bkg: Vec::new(),
            nthreads: 1,
        }
    }
}

/// Run `f` on a thread pool of `nthreads` threads, or all available for 0.
/// Parallel iterators inside `f` use this pool, so they never spread to more
/// threads than asked for.
pub(crate) fn with_threads<T: Send>(nthreads: usize, f: impl FnOnce() -> T + Send) -> Result<T> {
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(nthreads)
        .build()
        .map_err(|e| Error::Argument(format!("could not start threads: {e}")))?;
    Ok(pool.install(f))
}

/// The PWM and background a motif is scored with.
#[derive(Debug, Clone, PartialEq)]
pub struct ScoringMatrix {
    /// The motif as a PWM of log2 odds, with no infinite values.
    pub pwm: Matrix,
    /// The background frequencies the PWM is relative to.
    pub bkg: Vec<f64>,
    /// Whether the pseudocount had to be reapplied because the PWM had
    /// infinite values.
    pub pseudocount_added: bool,
}

/// Prepare the PWM a motif is scored with, as `motif_pvalue()` does.
///
/// Against its own background, which is the default, the motif's PWM is used
/// as scanning uses it, because converting back and forth can move entries
/// across integer score boundaries. Otherwise the motif is converted to a PPM
/// and then to a PWM against `bkg`. Either way, if the motif's PWM would have
/// infinite values the pseudocount is first reapplied with
/// [`normalize`](crate::ops::normalize) (through the probabilities of a CWM).
pub fn scoring_matrix(motif: &Motif, bkg: Option<&[f64]>) -> Result<ScoringMatrix> {
    let bkg = match bkg {
        Some(b) => {
            if b.len() != motif.alphabet.len() {
                return arg_err(format!(
                    "the background must have {} values, one per letter",
                    motif.alphabet.len()
                ));
            }
            b.to_vec()
        }
        None => motif.bkg.clone(),
    };
    // Against its own background the motif's PWM is used as scanning uses
    // it, since converting back and forth can move entries across integer
    // score boundaries.
    let pwm = convert_type(motif, MotifType::Pwm, None, false)?;
    let infinite = pwm.matrix.as_slice().iter().any(|v| v.is_infinite());
    if bkg == motif.bkg {
        let pwm = if infinite {
            convert_type(&with_pseudocount(motif)?, MotifType::Pwm, None, false)?
        } else {
            pwm
        };
        return Ok(ScoringMatrix {
            pwm: pwm.into_parts().matrix,
            bkg,
            pseudocount_added: infinite,
        });
    }
    let base = if infinite {
        with_pseudocount(motif)?
    } else {
        motif.clone()
    };
    let ppm = convert_type(&base, MotifType::Ppm, None, false)?;
    let pseudocount_added = infinite;
    let mut parts = ppm.into_parts();
    parts.bkg = bkg.clone();
    let ppm = Motif::from_parts(parts)?;
    let pwm = convert_type(&ppm, MotifType::Pwm, None, false)?;
    Ok(ScoringMatrix {
        pwm: pwm.into_parts().matrix,
        bkg,
        pseudocount_added,
    })
}

/// The score distribution of each motif, built in parallel.
pub fn score_distributions(
    motifs: &[Motif],
    opts: &PvalueOptions,
) -> Result<Vec<ScoreDistribution>> {
    if opts.bkg.len() > 1 && opts.bkg.len() != motifs.len() {
        return arg_err("give one background for all motifs or one per motif");
    }
    let mats = motifs
        .iter()
        .enumerate()
        .map(|(i, m)| {
            let bkg = match opts.bkg.len() {
                0 => None,
                1 => Some(opts.bkg[0].as_slice()),
                _ => Some(opts.bkg[i].as_slice()),
            };
            scoring_matrix(m, bkg)
        })
        .collect::<Result<Vec<_>>>()?;
    with_threads(opts.nthreads, || {
        mats.par_iter()
            .map(|m| ScoreDistribution::new(&m.pwm, &m.bkg))
            .collect()
    })
}

fn check_queries(n: usize, queries: &[(usize, f64)]) -> Result<()> {
    for &(i, v) in queries {
        if i >= n {
            return arg_err(format!("query refers to motif {i}, but there are {n}"));
        }
        if !v.is_finite() {
            return arg_err("scores and P-values must be finite numbers");
        }
    }
    Ok(())
}

/// Answer `(motif index, value)` queries one motif at a time, so only the
/// distributions in flight are held in memory.
fn per_motif(
    motifs: &[Motif],
    queries: &[(usize, f64)],
    opts: &PvalueOptions,
    answer: impl Fn(&ScoreDistribution, f64) -> f64 + Sync,
) -> Result<Vec<f64>> {
    if opts.bkg.len() > 1 && opts.bkg.len() != motifs.len() {
        return arg_err("give one background for all motifs or one per motif");
    }
    let mut by_motif: Vec<Vec<usize>> = vec![Vec::new(); motifs.len()];
    for (q, &(i, _)) in queries.iter().enumerate() {
        by_motif[i].push(q);
    }
    let answers: Vec<Vec<(usize, f64)>> = with_threads(opts.nthreads, || {
        by_motif
            .par_iter()
            .enumerate()
            .filter(|(_, qs)| !qs.is_empty())
            .map(|(i, qs)| {
                let bkg = match opts.bkg.len() {
                    0 => None,
                    1 => Some(opts.bkg[0].as_slice()),
                    _ => Some(opts.bkg[i].as_slice()),
                };
                let m = scoring_matrix(&motifs[i], bkg)?;
                let dist = ScoreDistribution::new(&m.pwm, &m.bkg);
                Ok(qs
                    .iter()
                    .map(|&q| (q, answer(&dist, queries[q].1)))
                    .collect())
            })
            .collect::<Result<Vec<_>>>()
    })??;
    let mut out = vec![0.0; queries.len()];
    for (q, v) in answers.into_iter().flatten() {
        out[q] = v;
    }
    Ok(out)
}

/// P-values for `(motif index, score)` pairs.
///
/// # Example
///
/// ```
/// use universalmotif::create::{CreateOptions, from_strings};
/// use universalmotif::pvalue::{PvalueOptions, motif_pvalue, pvalue_to_score};
///
/// let opts = CreateOptions { nsites: Some(20.0), pseudocount: 1.0, ..Default::default() };
/// let m = from_strings(&["TGACTCA"], None, &opts)?;
/// // The score a random site reaches with probability 1e-4, and back again.
/// let score = pvalue_to_score(std::slice::from_ref(&m), &[(0, 1e-4)], &PvalueOptions::default())?[0];
/// let p = motif_pvalue(&[m], &[(0, score)], &PvalueOptions::default())?[0];
/// assert!(p <= 1e-4);
/// # Ok::<(), universalmotif::Error>(())
/// ```
pub fn motif_pvalue(
    motifs: &[Motif],
    queries: &[(usize, f64)],
    opts: &PvalueOptions,
) -> Result<Vec<f64>> {
    check_queries(motifs.len(), queries)?;
    per_motif(motifs, queries, opts, |d, s| d.pvalue(s))
}

/// Score thresholds for `(motif index, P-value)` pairs: the lowest score a
/// random site reaches with at most that probability, or infinity if no
/// score is that unlikely.
pub fn pvalue_to_score(
    motifs: &[Motif],
    queries: &[(usize, f64)],
    opts: &PvalueOptions,
) -> Result<Vec<f64>> {
    check_queries(motifs.len(), queries)?;
    per_motif(motifs, queries, opts, |d, p| d.score(p))
}

/// The smallest P-value each motif can give, that of its highest possible
/// score. A P-value cutoff below it leaves the motif without hits.
///
/// No score distribution is needed: the highest score takes, in every
/// column, one of the letters with that column's best score on the integer
/// grid, so its probability is the product over columns of their background
/// probabilities.
pub fn min_pvalues(motifs: &[Motif], opts: &PvalueOptions) -> Result<Vec<f64>> {
    if opts.bkg.len() > 1 && opts.bkg.len() != motifs.len() {
        return arg_err("give one background for all motifs or one per motif");
    }
    motifs
        .iter()
        .enumerate()
        .map(|(i, motif)| {
            let bkg = match opts.bkg.len() {
                0 => None,
                1 => Some(opts.bkg[0].as_slice()),
                _ => Some(opts.bkg[i].as_slice()),
            };
            let m = scoring_matrix(motif, bkg)?;
            let (mut top, mut total) = (1.0, 1.0);
            for col in m.pwm.columns() {
                let grid: Vec<i32> = col.iter().map(|&v| grid_value(v)).collect();
                let best = *grid.iter().max().unwrap();
                top *= grid
                    .iter()
                    .zip(&m.bkg)
                    .filter(|(g, _)| **g == best)
                    .map(|(_, b)| b)
                    .sum::<f64>();
                total *= m.bkg.iter().sum::<f64>();
            }
            Ok(if total > 0.0 {
                (top / total).min(1.0)
            } else {
                0.0
            })
        })
        .collect()
}

/// The natural logarithm of a P-value written as text, such as `"1e-400"`.
/// Values too small for a double are handled through their mantissa and
/// exponent.
pub fn log_pvalue_text(pvalue: &str) -> Result<f64> {
    let s = pvalue.trim();
    let bad = || Error::Argument(format!("'{pvalue}' is not a P-value"));
    let value: f64 = s.parse().map_err(|_| bad())?;
    if value.is_nan() || value < 0.0 {
        return Err(bad());
    }
    if value >= f64::MIN_POSITIVE {
        return Ok(value.ln());
    }
    // Zero, or too small for a double.
    match s.split_once(['e', 'E']) {
        Some((m, e)) => {
            let m: f64 = m.parse().map_err(|_| bad())?;
            let e = e.parse::<i64>().map_err(|_| bad())? as f64;
            Ok(m.ln() + e * std::f64::consts::LN_10)
        }
        None => Ok(value.ln()),
    }
}