fastx-io 0.3.0

//! Summary statistics over a set of records — the `seqkit stats` equivalent.

use std::fmt;

use crate::qual::{self, PHRED33};
use crate::record::Sequence;
use crate::seq::{self, BaseCounts};

/// Accumulates length, composition and quality statistics.
///
/// ```
/// use fastx::{FastxReader, SeqStats};
///
/// let data = b">a\nACGTACGTAC\n>b\nGGCC\n";
/// let mut stats = SeqStats::default();
/// FastxReader::new(&data[..]).for_each_record(|r| { stats.push(r); Ok(()) })?;
///
/// assert_eq!(stats.count, 2);
/// assert_eq!(stats.total_length, 14);
/// assert_eq!(stats.min_length, Some(4));
/// assert_eq!(stats.max_length, Some(10));
/// assert_eq!(stats.n50(), Some(10));
/// # Ok::<(), fastx::Error>(())
/// ```
#[derive(Debug, Clone, Default)]
pub struct SeqStats {
    /// Number of records seen.
    pub count: u64,
    /// Sum of all sequence lengths.
    pub total_length: u64,
    /// Shortest record, if any.
    pub min_length: Option<u64>,
    /// Longest record, if any.
    pub max_length: Option<u64>,
    /// Base composition across all records.
    pub bases: BaseCounts,
    /// Every observed length, kept so that N50 and the median can be computed.
    lengths: Vec<u64>,
    /// Counts per Phred score, from which every quality figure is derived.
    quality: QualityHistogram,
}

/// Number of distinct Phred scores tracked: 0..=93 covers all printable ASCII.
pub const PHRED_SCORES: usize = 94;

/// Highest Phred score the histogram can hold.
const MAX_SCORE: u8 = (PHRED_SCORES - 1) as u8;

/// Counts of each Phred score, indexed by the score itself.
///
/// Accumulating a histogram keeps the hot loop to one increment per base — no
/// floating point at all — and every quality statistic falls out of it
/// afterwards at a cost that does not depend on the size of the input.
#[derive(Debug, Clone)]
struct QualityHistogram([u64; PHRED_SCORES]);

impl Default for QualityHistogram {
    fn default() -> Self {
        QualityHistogram([0; PHRED_SCORES])
    }
}

impl SeqStats {
    /// An empty accumulator.
    pub fn new() -> SeqStats {
        SeqStats::default()
    }

    /// Fold one record in.
    pub fn push(&mut self, record: &Sequence) {
        let len = record.len() as u64;
        self.count += 1;
        self.total_length += len;
        self.min_length = Some(self.min_length.map_or(len, |m| m.min(len)));
        self.max_length = Some(self.max_length.map_or(len, |m| m.max(len)));
        self.bases.merge(&record.base_counts());
        self.lengths.push(len);
        if let Some(quality) = &record.quality {
            let histogram = &mut self.quality.0;
            for &c in quality {
                // `min` also keeps the index provably in range, which removes
                // the bounds check from the loop.
                histogram[qual::score(c, PHRED33).min(MAX_SCORE) as usize] += 1;
            }
        }
    }

    /// Merge another accumulator, for parallel or per-file aggregation.
    pub fn merge(&mut self, other: &SeqStats) {
        self.count += other.count;
        self.total_length += other.total_length;
        self.min_length = min_option(self.min_length, other.min_length);
        self.max_length = max_option(self.max_length, other.max_length);
        self.bases.merge(&other.bases);
        self.lengths.extend_from_slice(&other.lengths);
        for (slot, count) in self.quality.0.iter_mut().zip(other.quality.0.iter()) {
            *slot += count;
        }
    }

    /// True when no records have been seen.
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Mean sequence length.
    pub fn mean_length(&self) -> Option<f64> {
        if self.count == 0 {
            None
        } else {
            Some(self.total_length as f64 / self.count as f64)
        }
    }

    /// Median sequence length (lower median for an even count).
    pub fn median_length(&self) -> Option<u64> {
        if self.lengths.is_empty() {
            return None;
        }
        let mut lengths = self.lengths.clone();
        lengths.sort_unstable();
        Some(lengths[(lengths.len() - 1) / 2])
    }

    /// N50: contigs of at least this length cover half the total.
    pub fn n50(&self) -> Option<u64> {
        seq::n50(&mut self.lengths.clone())
    }

    /// N90.
    pub fn n90(&self) -> Option<u64> {
        seq::nx(&mut self.lengths.clone(), 0.9)
    }

    /// L50: the number of contigs that make up the N50.
    pub fn l50(&self) -> Option<u64> {
        let n50 = self.n50()?;
        let mut lengths = self.lengths.clone();
        lengths.sort_unstable_by(|a, b| b.cmp(a));
        Some(lengths.iter().take_while(|&&l| l >= n50).count() as u64)
    }

    /// GC fraction over unambiguous bases.
    pub fn gc_content(&self) -> Option<f64> {
        self.bases.gc_content()
    }

    /// Number of quality characters seen across all records.
    pub fn quality_bases(&self) -> u64 {
        self.quality.0.iter().sum()
    }

    /// Counts per Phred score, indexed by score — the same data FastQC plots.
    ///
    /// ```
    /// # use fastx::{SeqStats, Sequence};
    /// let mut stats = SeqStats::new();
    /// stats.push(&Sequence::fastq("r", b"ACGT", b"IIII")?);
    /// assert_eq!(stats.quality_histogram()[40], 4); // 'I' is Q40
    /// # Ok::<(), fastx::Error>(())
    /// ```
    pub fn quality_histogram(&self) -> &[u64; PHRED_SCORES] {
        &self.quality.0
    }

    /// Expected number of wrong bases across all records.
    pub fn expected_errors(&self) -> f64 {
        self.quality
            .0
            .iter()
            .enumerate()
            .map(|(score, &count)| count as f64 * qual::error_probability(score as u8))
            .sum()
    }

    /// Mean quality as a Phred score, derived from the mean error rate.
    pub fn mean_quality(&self) -> Option<f64> {
        let bases = self.quality_bases();
        if bases == 0 {
            return None;
        }
        let mean_p = self.expected_errors() / bases as f64;
        Some(-10.0 * mean_p.log10())
    }

    /// Fraction of bases at Q20 or better.
    pub fn q20_fraction(&self) -> Option<f64> {
        self.fraction_at_least(20)
    }

    /// Fraction of bases at Q30 or better.
    pub fn q30_fraction(&self) -> Option<f64> {
        self.fraction_at_least(30)
    }

    /// Fraction of bases whose Phred score is at least `score`.
    pub fn fraction_at_least(&self, score: u8) -> Option<f64> {
        let bases = self.quality_bases();
        if bases == 0 {
            return None;
        }
        let at_least: u64 = self.quality.0[(score as usize).min(PHRED_SCORES)..]
            .iter()
            .sum();
        Some(at_least as f64 / bases as f64)
    }

    /// All observed lengths, in the order the records were seen.
    pub fn lengths(&self) -> &[u64] {
        &self.lengths
    }

    /// Render the statistics as a JSON object, for pipelines that parse output.
    ///
    /// The object is on one line, so a run over several files is valid
    /// line-delimited JSON and streams straight into `jq`. Pipe through `jq .`
    /// if you want it laid out.
    ///
    /// Absent figures — quality for FASTA, N50 after
    /// [`SeqStats::forget_lengths`] — come out as `null` rather than being
    /// omitted, so the shape of the object never changes.
    ///
    /// ```
    /// # use fastx::{SeqStats, Sequence};
    /// let mut stats = SeqStats::new();
    /// stats.push(&Sequence::fasta("a", b"ACGT"));
    /// let json = stats.to_json();
    /// assert!(json.starts_with('{') && json.ends_with('}'));
    /// assert!(!json.contains('\n'), "must stay on one line");
    /// assert!(json.contains("\"records\":1"));
    /// assert!(json.contains("\"mean_quality\":null"));
    /// ```
    pub fn to_json(&self) -> String {
        fn number(value: Option<f64>, decimals: usize) -> String {
            match value {
                // JSON has no NaN or Infinity, so anything not finite is null.
                Some(v) if v.is_finite() => format!("{v:.decimals$}"),
                _ => "null".to_string(),
            }
        }
        fn integer(value: Option<u64>) -> String {
            value.map_or_else(|| "null".to_string(), |v| v.to_string())
        }

        let fields = [
            ("records".to_string(), self.count.to_string()),
            ("total_length".to_string(), self.total_length.to_string()),
            ("min_length".to_string(), integer(self.min_length)),
            ("max_length".to_string(), integer(self.max_length)),
            ("mean_length".to_string(), number(self.mean_length(), 2)),
            ("median_length".to_string(), integer(self.median_length())),
            ("n50".to_string(), integer(self.n50())),
            ("n90".to_string(), integer(self.n90())),
            ("l50".to_string(), integer(self.l50())),
            ("gc_content".to_string(), number(self.gc_content(), 6)),
            ("a".to_string(), self.bases.a.to_string()),
            ("c".to_string(), self.bases.c.to_string()),
            ("g".to_string(), self.bases.g.to_string()),
            ("t".to_string(), self.bases.t.to_string()),
            ("ambiguous".to_string(), self.bases.n.to_string()),
            ("other".to_string(), self.bases.other.to_string()),
            (
                "quality_bases".to_string(),
                self.quality_bases().to_string(),
            ),
            ("mean_quality".to_string(), number(self.mean_quality(), 4)),
            ("q20_fraction".to_string(), number(self.q20_fraction(), 6)),
            ("q30_fraction".to_string(), number(self.q30_fraction(), 6)),
        ];
        let body = fields
            .iter()
            .map(|(key, value)| format!("\"{key}\":{value}"))
            .collect::<Vec<_>>()
            .join(",");
        format!("{{{body}}}")
    }

    /// Drop the retained per-record lengths to cap memory on huge inputs.
    ///
    /// After this call [`SeqStats::n50`], [`SeqStats::median_length`] and
    /// [`SeqStats::l50`] return `None`, but counts and means stay correct.
    pub fn forget_lengths(&mut self) {
        self.lengths = Vec::new();
    }
}

fn min_option(a: Option<u64>, b: Option<u64>) -> Option<u64> {
    match (a, b) {
        (Some(a), Some(b)) => Some(a.min(b)),
        (a, b) => a.or(b),
    }
}

fn max_option(a: Option<u64>, b: Option<u64>) -> Option<u64> {
    match (a, b) {
        (Some(a), Some(b)) => Some(a.max(b)),
        (a, b) => a.or(b),
    }
}

impl fmt::Display for SeqStats {
    /// A block of aligned `label value` lines, without a trailing newline.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut lines: Vec<String> = Vec::with_capacity(14);
        let mut push = |label: &str, value: String| lines.push(format!("{label:<12} {value}"));

        push("records", self.count.to_string());
        push("total bases", self.total_length.to_string());
        if let Some(v) = self.min_length {
            push("min length", v.to_string());
        }
        if let Some(v) = self.max_length {
            push("max length", v.to_string());
        }
        if let Some(v) = self.mean_length() {
            push("avg length", format!("{v:.1}"));
        }
        if let Some(v) = self.median_length() {
            push("median len", v.to_string());
        }
        if let Some(v) = self.n50() {
            push("N50", v.to_string());
        }
        if let Some(v) = self.n90() {
            push("N90", v.to_string());
        }
        if let Some(v) = self.l50() {
            push("L50", v.to_string());
        }
        if let Some(v) = self.gc_content() {
            push("GC%", format!("{:.2}", v * 100.0));
        }
        push("N bases", self.bases.n.to_string());
        if let Some(v) = self.mean_quality() {
            push("avg quality", format!("Q{v:.1}"));
        }
        if let Some(v) = self.q20_fraction() {
            push("Q20%", format!("{:.2}", v * 100.0));
        }
        if let Some(v) = self.q30_fraction() {
            push("Q30%", format!("{:.2}", v * 100.0));
        }
        f.write_str(&lines.join("\n"))
    }
}

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

    fn stats_of(records: &[Sequence]) -> SeqStats {
        let mut stats = SeqStats::new();
        for record in records {
            stats.push(record);
        }
        stats
    }

    #[test]
    fn empty_stats_have_no_summaries() {
        let stats = SeqStats::new();
        assert!(stats.is_empty());
        assert_eq!(stats.mean_length(), None);
        assert_eq!(stats.n50(), None);
        assert_eq!(stats.l50(), None);
        assert_eq!(stats.gc_content(), None);
        assert_eq!(stats.mean_quality(), None);
    }

    #[test]
    fn length_statistics() {
        let stats = stats_of(&[
            Sequence::fasta("a", b"A".repeat(50)),
            Sequence::fasta("b", b"C".repeat(30)),
            Sequence::fasta("c", b"G".repeat(15)),
            Sequence::fasta("d", b"T".repeat(5)),
        ]);
        assert_eq!(stats.count, 4);
        assert_eq!(stats.total_length, 100);
        assert_eq!(stats.min_length, Some(5));
        assert_eq!(stats.max_length, Some(50));
        assert_eq!(stats.mean_length(), Some(25.0));
        assert_eq!(stats.median_length(), Some(15));
        assert_eq!(stats.n50(), Some(50));
        assert_eq!(stats.n90(), Some(15));
        assert_eq!(stats.l50(), Some(1));
        assert_eq!(stats.gc_content(), Some(0.45));
    }

    #[test]
    fn quality_statistics() {
        let stats = stats_of(&[
            Sequence::fastq("a", b"ACGT", b"IIII").unwrap(),
            Sequence::fastq("b", b"ACGT", b"!!!!").unwrap(),
        ]);
        assert_eq!(stats.q30_fraction(), Some(0.5));
        assert_eq!(stats.q20_fraction(), Some(0.5));
        let mean = stats.mean_quality().unwrap();
        assert!(mean > 2.0 && mean < 4.0, "{mean}");
    }

    #[test]
    fn quality_histogram_is_the_source_of_truth() {
        let stats = stats_of(&[
            Sequence::fastq("a", b"ACGT", b"IIII").unwrap(), // Q40
            Sequence::fastq("b", b"AC", b"!5").unwrap(),     // Q0 and Q20
        ]);
        let histogram = stats.quality_histogram();
        assert_eq!(histogram[40], 4);
        assert_eq!(histogram[20], 1);
        assert_eq!(histogram[0], 1);
        assert_eq!(stats.quality_bases(), 6);

        // Derived figures must agree with computing them the direct way.
        let direct: f64 = qual::expected_errors(b"IIII!5", PHRED33);
        assert!((stats.expected_errors() - direct).abs() < 1e-12);
        assert_eq!(stats.fraction_at_least(0), Some(1.0));
        assert_eq!(stats.fraction_at_least(20), Some(5.0 / 6.0));
        assert_eq!(stats.fraction_at_least(40), Some(4.0 / 6.0));
        assert_eq!(stats.fraction_at_least(93), Some(0.0));
        // Scores beyond the histogram are simply never reached.
        assert_eq!(stats.fraction_at_least(200), Some(0.0));
    }

    #[test]
    fn quality_scores_are_clamped_not_wrapped() {
        // A byte above '~' cannot appear in valid FASTQ, but it must not panic
        // or corrupt neighbouring buckets if it does.
        let mut record = Sequence::fastq("r", b"AC", b"II").unwrap();
        record.quality = Some(vec![255, 33]);
        let stats = stats_of(&[record]);
        assert_eq!(stats.quality_histogram()[93], 1);
        assert_eq!(stats.quality_histogram()[0], 1);
        assert_eq!(stats.quality_bases(), 2);
    }

    #[test]
    fn merging_matches_sequential() {
        let records: Vec<Sequence> = (1..20)
            .map(|i| Sequence::fasta(format!("s{i}"), b"ACGT".repeat(i)))
            .collect();
        let sequential = stats_of(&records);
        let (left, right) = records.split_at(7);
        let mut merged = stats_of(left);
        merged.merge(&stats_of(right));

        assert_eq!(merged.count, sequential.count);
        assert_eq!(merged.total_length, sequential.total_length);
        assert_eq!(merged.min_length, sequential.min_length);
        assert_eq!(merged.max_length, sequential.max_length);
        assert_eq!(merged.n50(), sequential.n50());
        assert_eq!(merged.gc_content(), sequential.gc_content());
    }

    #[test]
    fn forgetting_lengths_keeps_counts() {
        let mut stats = stats_of(&[Sequence::fasta("a", b"ACGT")]);
        stats.forget_lengths();
        assert_eq!(stats.count, 1);
        assert_eq!(stats.total_length, 4);
        assert_eq!(stats.n50(), None);
    }

    #[test]
    fn display_is_readable() {
        let text = stats_of(&[Sequence::fastq("a", b"ACGT", b"IIII").unwrap()]).to_string();
        assert!(text.contains("records      1"));
        assert!(text.contains("GC%"));
        assert!(text.contains("Q30%"));
    }
}