gwseq-io 0.2.0

//! Binning and reduction.
//!
//! The two stringly-typed parameter families of the Python API — `bin_mode` and
//! `reduce` — become enums with a `FromStr` whose error lists the accepted
//! spellings.
//!
//! There are **two** accumulator types, not one:
//! [`BinStats`] is what the binning hot loop carries (two `f64`s, one per bin
//! of a whole-genome walk), and [`ValueStats`] is what a reduction needs
//! (extremes and a sum of squares as well). Merging them would put five fields
//! where the inner loop wants two.

use crate::error::{Error, Result};

/// How the values covering a bin become the bin's value.
///
/// All three are per *base* of the bin, not per record of the file — a
/// distinction that matters for a bigWig whose intervals are wider than a bin,
/// and for a bigBed, where a bin's value is the depth of coverage its entries
/// make over that bin.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum BinMode {
    /// Base-weighted mean.
    #[default]
    Mean,
    /// Value summed over every base it covers.
    Sum,
    /// Bases of the bin carrying data — the bin's width where the file covers
    /// it fully.
    Count,
}

/// How a row of bins, or a column of loci, becomes one number.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Reduce {
    #[default]
    Mean,
    Sd,
    Sem,
    Sum,
    Count,
    Min,
    Max,
    L1Norm,
    L2Norm,
}

impl std::str::FromStr for BinMode {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "mean" => Ok(BinMode::Mean),
            "sum" => Ok(BinMode::Sum),
            "count" => Ok(BinMode::Count),
            // The wording is API: callers match on this message.
            o => Err(Error::invalid(format!("bin_mode {o} invalid"))),
        }
    }
}

impl std::str::FromStr for Reduce {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "mean" => Ok(Reduce::Mean),
            "sd" => Ok(Reduce::Sd),
            "sem" => Ok(Reduce::Sem),
            "sum" => Ok(Reduce::Sum),
            "count" => Ok(Reduce::Count),
            "min" => Ok(Reduce::Min),
            "max" => Ok(Reduce::Max),
            "l1norm" => Ok(Reduce::L1Norm),
            "l2norm" => Ok(Reduce::L2Norm),
            o => Err(Error::invalid(format!("reduce {o} invalid"))),
        }
    }
}

/// What one bin accumulates while it is being filled.
///
/// `f64` under an `f32` result: a bin of a whole-genome walk sums millions of
/// values, and an `f32` accumulator stops making progress once the running
/// total passes 2^24.
///
/// `count` is fractional rather than a record count. A bin is `span /
/// bin_count` bases wide and need not be a whole number of them, and a record
/// covers however many of those bases it overlaps.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct BinStats {
    pub sum: f64,
    pub count: f64,
}

impl BinStats {
    /// Fold in a value covering `bases` bases of this bin.
    #[inline]
    pub fn add(&mut self, value: f32, bases: f64) {
        self.sum += value as f64 * bases;
        self.count += bases;
    }

    #[inline]
    pub fn merge(&mut self, other: &BinStats) {
        self.sum += other.sum;
        self.count += other.count;
    }

    /// The bin's value under `mode`.
    ///
    /// `Mean` divides by `count` without guarding it: a bin nothing reached has
    /// `count == 0` and yields NaN, and the caller — which knows whether the
    /// bin was covered — writes `def_value` there instead. Moving the guard
    /// here would change which bins come back as `def_value`.
    #[inline]
    pub fn apply(&self, mode: BinMode) -> f32 {
        match mode {
            BinMode::Mean => (self.sum / self.count) as f32,
            BinMode::Sum => self.sum as f32,
            BinMode::Count => self.count as f32,
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.count == 0.0
    }
}

/// What a reduction accumulates over a locus or a profile bin.
///
/// `min` and `max` are data values and stay `f32`; the accumulators do not.
/// Both start NaN, so an untouched region reports no extremes rather than
/// `±inf`.
///
/// # Why the sums are shifted
///
/// The obvious accumulator holds `Σx` and `Σx²` and computes the variance as
/// `E[x²] − E[x]²`. That subtracts two nearly equal large numbers whenever the
/// values sit far from zero relative to their spread, which is the ordinary
/// case for coverage, CPM and log-ratio tracks — and it does not merely lose
/// a few bits. Measured against a float64 reference on `normal(1e4, 1e-2)`
/// data, the naive form reported a standard deviation of 1.9e-1 where the
/// truth was 9.8e-3: **wrong by a factor of twenty**, and in another case
/// wrong all the way to zero, the variance having gone negative and been
/// clamped.
///
/// So every sum here is taken relative to `shift`, the first value folded in:
/// `Σ(x − k)` and `Σ(x − k)²`. The variance is then
/// `[Σ(x−k)² − Σ(x−k)²/n] / n` over numbers of the size of the *spread*
/// rather than of the mean, and the cancellation goes with it. A constant
/// column gives exactly zero, because every difference is exactly zero.
/// `mean` and `sum` are recovered by adding `k` back, which costs one
/// multiply and no accuracy.
///
/// This is the standard shifted-data algorithm. Welford's would be equivalent
/// for values arriving one at a time, and cannot fold in a pre-aggregated zoom
/// record, which this has to do — see [`ValueStats::add_aggregate`].
///
/// **Accumulation order is behaviour.** These are summed in the order the
/// extraction visits intervals, which is block order within a batch and batch
/// order across the output. That order is deterministic on purpose: it is what
/// makes an answer independent of `parallel`, which `tests/roundtrip.rs` and
/// `tests/properties.rs` both check.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ValueStats {
    pub min: f32,
    pub max: f32,
    /// Every sum below is relative to this. Set by the first fold; meaningless
    /// while `count` is 0.
    shift: f64,
    /// `Σ(xᵢ − shift)`, each value weighted by the bases it covers.
    sum_shifted: f64,
    /// `Σ(xᵢ − shift)²`, same weighting.
    sum_sq_shifted: f64,
    /// `Σ|xᵢ|`, same weighting. Only `L1Norm` reads it.
    ///
    /// Exact everywhere a value arrives with its own sign, which is every path
    /// but one: a zoom record carries `Σx` and `Σx²` and no `Σ|x|`, so a
    /// record straddling zero cannot supply it. See [`Self::add_aggregate`].
    sum_abs: f64,
    pub count: i64,
}

impl Default for ValueStats {
    fn default() -> Self {
        Self {
            min: f32::NAN,
            max: f32::NAN,
            shift: 0.0,
            sum_shifted: 0.0,
            sum_sq_shifted: 0.0,
            sum_abs: 0.0,
            count: 0,
        }
    }
}

impl ValueStats {
    /// `Σx`, undoing the shift.
    #[inline]
    pub fn sum(&self) -> f64 {
        self.sum_shifted + self.shift * self.count as f64
    }

    /// `Σx²`, undoing the shift. Only `L2Norm` and the tests want this, and
    /// it is the one quantity the shift makes *less* accurate — which is the
    /// right trade, `l2norm` having no cancellation to suffer from.
    #[inline]
    pub fn sum_squared(&self) -> f64 {
        let n = self.count as f64;
        self.sum_sq_shifted + 2.0 * self.shift * self.sum_shifted + self.shift * self.shift * n
    }

    /// Fold in one value covering one base.
    #[inline]
    pub fn add(&mut self, value: f32) {
        self.add_repeated(value, 1);
    }

    /// Fold in `bases` bases all carrying `value`, as a wide interval does.
    ///
    /// The `f32::min`/`max` NaN rule does the seeding: `NAN.min(x) == x`, so
    /// the first value replaces the initial NaN without a branch.
    #[inline]
    pub fn add_repeated(&mut self, value: f32, bases: i64) {
        if bases <= 0 {
            return;
        }
        let v = value as f64;
        if self.count == 0 {
            self.shift = v;
        }
        self.min = self.min.min(value);
        self.max = self.max.max(value);
        let d = v - self.shift;
        let n = bases as f64;
        self.sum_shifted += d * n;
        self.sum_sq_shifted += d * d * n;
        self.sum_abs += v.abs() * n;
        self.count += bases;
    }

    /// Fold in a pre-aggregated run: `bases` bases whose sum is `sum` and
    /// whose sum of squares is `sum_squared`, with known extremes.
    ///
    /// This is how a **zoom record** enters, prorated over the part of it a
    /// window covers. The record's own `Σx²` is used rather than its mean
    /// squared: squaring the mean would keep only the variance *between*
    /// records and drop the variance inside each, collapsing `sd` as the zoom
    /// level rises.
    ///
    /// `Σ|x|` cannot be recovered from a record in general — the format stores
    /// no such field — so it is derived where the record's sign is not in
    /// doubt (`min ≥ 0`, or `max ≤ 0`, which is every non-negative track) and
    /// approximated by `|Σx|` where the record straddles zero. That is a lower
    /// bound, it is the only thing the format allows, and the README says so
    /// under `quantify`'s `reduce`.
    #[inline]
    pub fn add_aggregate(&mut self, min: f32, max: f32, sum: f64, sum_squared: f64, bases: i64) {
        if bases <= 0 {
            return;
        }
        let n = bases as f64;
        if self.count == 0 {
            // The run's own mean: the shift that makes its internal spread the
            // thing being summed.
            self.shift = sum / n;
        }
        self.min = self.min.min(min);
        self.max = self.max.max(max);
        let k = self.shift;
        // Σ(x−k) = Σx − nk, and Σ(x−k)² = Σx² − 2kΣx + nk².
        self.sum_shifted += sum - n * k;
        self.sum_sq_shifted += sum_squared - 2.0 * k * sum + n * k * k;
        self.sum_abs += if min >= 0.0 {
            sum
        } else if max <= 0.0 {
            -sum
        } else {
            sum.abs()
        };
        self.count += bases;
    }

    /// Fold another accumulator in, rebasing it onto this one's shift.
    #[inline]
    pub fn merge(&mut self, other: &ValueStats) {
        if other.count == 0 {
            return;
        }
        if self.count == 0 {
            *self = *other;
            return;
        }
        self.min = self.min.min(other.min);
        self.max = self.max.max(other.max);
        // Rebase `other` from its shift to this one. With d = k_other − k_self:
        // Σ(x−k_self)  = Σ(x−k_other) + n·d
        // Σ(x−k_self)² = Σ(x−k_other)² + 2d·Σ(x−k_other) + n·d²
        let d = other.shift - self.shift;
        let n = other.count as f64;
        self.sum_shifted += other.sum_shifted + n * d;
        self.sum_sq_shifted += other.sum_sq_shifted + 2.0 * d * other.sum_shifted + n * d * d;
        self.sum_abs += other.sum_abs;
        self.count += other.count;
    }

    /// Population variance, from the shifted sums.
    ///
    /// Still clamped at zero, and now for a reason that is nearly theoretical
    /// rather than routine: with the shift the subtraction is between numbers
    /// of the size of the spread, so it goes negative only on a run that is
    /// constant to the last bit — where zero is the right answer anyway.
    fn variance(&self, count: f64) -> f64 {
        let mean_shifted = self.sum_shifted / count;
        ((self.sum_sq_shifted / count) - mean_shifted * mean_shifted).max(0.0)
    }

    /// Reduce to one number.
    ///
    /// A region no data reached has no mean, no extremes and no spread, so
    /// those keep `def_value`. Its count is not unknown but **zero**: leaving
    /// `def_value` there would make `count` the one reduction unable to say
    /// "nothing here".
    pub fn reduce(&self, reduce: Reduce, def_value: f32) -> f32 {
        if self.count == 0 {
            return match reduce {
                Reduce::Count => 0.0,
                _ => def_value,
            };
        }
        let count = self.count as f64;
        match reduce {
            Reduce::Mean => (self.shift + self.sum_shifted / count) as f32,
            Reduce::Sd => self.variance(count).sqrt() as f32,
            Reduce::Sem => (self.variance(count) / count).sqrt() as f32,
            Reduce::Sum => self.sum() as f32,
            Reduce::Count => count as f32,
            Reduce::Min => self.min,
            Reduce::Max => self.max,
            // Σ|x|, which is what an L1 norm is. Not Σx: those agree only on
            // data that never goes negative, and the reductions a caller
            // reaches for an L1 norm over are the ones that do.
            Reduce::L1Norm => self.sum_abs as f32,
            Reduce::L2Norm => self.sum_squared().max(0.0).sqrt() as f32,
        }
    }
}

/// The bin grid a request resolves to: how wide a bin is, how many there are,
/// and whether a partial trailing bin is walked.
///
/// `bin_size` is `f64` because the API accepts a fractional one — it snaps the
/// window to a grid of that width, so the edges no longer fall on whole bases.
/// The `iter_all_*` paths reject a fractional size, the windows there having to
/// tile the genome on a whole-base grid.
#[derive(Debug, Clone, Copy)]
pub struct BinPlan {
    pub bin_size: f64,
    pub bin_count: Option<usize>,
    pub full_bin: bool,
}

impl BinPlan {
    pub fn new(bin_size: f64, bin_count: Option<usize>, full_bin: bool) -> Result<Self> {
        // is_finite first, so NaN is rejected before it reaches a comparison
        // that would answer false either way.
        if !bin_size.is_finite() || bin_size <= 0.0 {
            return Err(Error::invalid(format!(
                "bin_size must be a positive finite number, got {bin_size}"
            )));
        }
        // A bin is a run of bases, so its width is a whole number of them.
        // A fractional one used to be accepted and snapped the window to a
        // grid finer than a base, which put bin edges between bases and made
        // every value a weighted split of two — an answer no caller asked for
        // and one the whole-file iterators refused outright, so the API
        // disagreed with itself. `bin_count` is how to ask for a window
        // divided into a number of parts.
        if bin_size.fract() != 0.0 {
            return Err(Error::invalid(format!(
                "bin_size must be a whole number of base pairs, got {bin_size}. \
                 Use bin_count to divide a window into a fixed number of bins."
            )));
        }
        if bin_count == Some(0) {
            return Err(Error::invalid("bin_count must be at least 1"));
        }
        Ok(Self {
            bin_size,
            bin_count,
            full_bin,
        })
    }

    /// The bin size as the whole number of base pairs it is.
    ///
    /// Infallible: [`BinPlan::new`] is the only constructor and it refuses a
    /// fractional one.
    pub fn whole_bin_size(&self) -> i64 {
        self.bin_size as i64
    }
}

impl Default for BinPlan {
    fn default() -> Self {
        Self {
            bin_size: 1.0,
            bin_count: None,
            full_bin: false,
        }
    }
}

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

    #[test]
    fn bin_stats_weight_by_bases() {
        let mut s = BinStats::default();
        s.add(2.0, 10.0);
        s.add(4.0, 30.0);
        assert_eq!(s.count, 40.0);
        assert_eq!(s.sum, 140.0);
        assert_eq!(s.apply(BinMode::Mean), 3.5);
        assert_eq!(s.apply(BinMode::Sum), 140.0);
        assert_eq!(s.apply(BinMode::Count), 40.0);
    }

    #[test]
    fn an_untouched_bin_means_nan_not_zero() {
        // The caller substitutes def_value; this must not do it silently.
        assert!(BinStats::default().apply(BinMode::Mean).is_nan());
        assert_eq!(BinStats::default().apply(BinMode::Count), 0.0);
    }

    #[test]
    fn value_stats_seed_extremes_from_nan() {
        let mut s = ValueStats::default();
        assert!(s.min.is_nan() && s.max.is_nan());
        s.add(3.0);
        s.add(-1.0);
        s.add(7.0);
        assert_eq!((s.min, s.max, s.count), (-1.0, 7.0, 3));
        assert_eq!(s.sum(), 9.0);
        assert_eq!(s.sum_squared(), 59.0);
    }

    /// The L1 norm is `Σ|x|`. It agrees with the sum only on data that never
    /// goes negative, and the data a caller reaches for an L1 norm over is
    /// exactly the data that does.
    #[test]
    fn l1norm_is_the_sum_of_absolute_values() {
        let mut s = ValueStats::default();
        for v in [3.0f32, -4.0, 5.0, -6.0] {
            s.add(v);
        }
        assert_eq!(s.reduce(Reduce::Sum, 0.0), -2.0);
        assert_eq!(s.reduce(Reduce::L1Norm, 0.0), 18.0);
        assert_eq!(s.reduce(Reduce::L2Norm, 0.0), (86.0f32).sqrt());
    }

    /// The whole point of the shift. `E[x²] − E[x]²` on this data subtracts
    /// two numbers that agree to fifteen digits; the shifted form subtracts
    /// two that agree to none.
    #[test]
    fn a_tiny_spread_on_a_large_mean_survives() {
        let mean = 1.0e4f32;
        let values: Vec<f32> = (0..2000)
            .map(|i| mean + (i % 7) as f32 * 1.0e-2 - 0.03)
            .collect();
        let mut s = ValueStats::default();
        for v in &values {
            s.add(*v);
        }
        // The reference, in f64 and two passes.
        let n = values.len() as f64;
        let m: f64 = values.iter().map(|v| *v as f64).sum::<f64>() / n;
        let want = (values.iter().map(|v| (*v as f64 - m).powi(2)).sum::<f64>() / n).sqrt();
        let got = s.reduce(Reduce::Sd, 0.0) as f64;
        assert!((got - want).abs() <= want * 1e-3, "sd {got} against {want}");
    }

    /// A column that never varies has a standard deviation of exactly zero,
    /// not "nearly zero" and not a clamped negative.
    #[test]
    fn a_constant_column_has_no_spread_at_all() {
        for value in [1.0f32, 12345.678, -9876.5, 1.0e-7] {
            let mut s = ValueStats::default();
            for _ in 0..500 {
                s.add(value);
            }
            assert_eq!(s.reduce(Reduce::Sd, -1.0), 0.0, "value {value}");
            assert_eq!(s.reduce(Reduce::Sem, -1.0), 0.0, "value {value}");
            assert_eq!(s.reduce(Reduce::Mean, -1.0), value, "value {value}");
        }
    }

    /// Merging two accumulators rebases one onto the other's shift, and has to
    /// give what folding everything into one would have given.
    #[test]
    fn merging_is_folding_by_another_route() {
        let a_values = [1000.0f32, 1000.5, 999.5, 1001.0];
        let b_values = [-3.0f32, 2000.25, 7.5];
        let (mut a, mut b, mut whole) = (
            ValueStats::default(),
            ValueStats::default(),
            ValueStats::default(),
        );
        for v in a_values {
            a.add(v);
            whole.add(v);
        }
        for v in b_values {
            b.add(v);
            whole.add(v);
        }
        a.merge(&b);
        assert_eq!(a.count, whole.count);
        for r in [
            Reduce::Mean,
            Reduce::Sum,
            Reduce::L1Norm,
            Reduce::Min,
            Reduce::Max,
        ] {
            assert_eq!(a.reduce(r, 0.0), whole.reduce(r, 0.0), "{r:?}");
        }
        let (got, want) = (a.reduce(Reduce::Sd, 0.0), whole.reduce(Reduce::Sd, 0.0));
        assert!((got - want).abs() <= want * 1e-5, "sd {got} against {want}");
    }

    /// A pre-aggregated run — a zoom record — carries no `Σ|x|`, so the L1
    /// norm is derived from its sign where the sign is not in doubt.
    #[test]
    fn an_aggregate_run_folds_in_with_its_own_spread() {
        // 100 bases summing to 500 with a sum of squares of 3000: mean 5,
        // variance 3000/100 - 25 = 5.
        let mut s = ValueStats::default();
        s.add_aggregate(1.0, 9.0, 500.0, 3000.0, 100);
        assert_eq!(s.count, 100);
        assert_eq!(s.reduce(Reduce::Mean, 0.0), 5.0);
        assert_eq!(s.reduce(Reduce::Sum, 0.0), 500.0);
        assert!((s.reduce(Reduce::Sd, 0.0) - 5.0f32.sqrt()).abs() < 1e-4);
        // min >= 0, so every value in the run is positive and Σ|x| == Σx.
        assert_eq!(s.reduce(Reduce::L1Norm, 0.0), 500.0);

        // All-negative: Σ|x| == -Σx, equally certain.
        let mut s = ValueStats::default();
        s.add_aggregate(-9.0, -1.0, -500.0, 3000.0, 100);
        assert_eq!(s.reduce(Reduce::L1Norm, 0.0), 500.0);
    }

    #[test]
    fn empty_reduces_to_def_value_except_count() {
        let s = ValueStats::default();
        for r in [
            Reduce::Mean,
            Reduce::Sd,
            Reduce::Sum,
            Reduce::Min,
            Reduce::Max,
        ] {
            assert_eq!(s.reduce(r, -5.0), -5.0, "{r:?}");
        }
        assert_eq!(s.reduce(Reduce::Count, -5.0), 0.0);
    }

    #[test]
    fn reductions_match_their_definitions() {
        let mut s = ValueStats::default();
        for v in [2.0f32, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0] {
            s.add(v);
        }
        assert_eq!(s.reduce(Reduce::Mean, 0.0), 5.0);
        assert_eq!(s.reduce(Reduce::Sd, 0.0), 2.0); // textbook population sd
        assert_eq!(s.reduce(Reduce::Sum, 0.0), 40.0);
        assert_eq!(s.reduce(Reduce::Count, 0.0), 8.0);
        assert_eq!(s.reduce(Reduce::Min, 0.0), 2.0);
        assert_eq!(s.reduce(Reduce::Max, 0.0), 9.0);
        // Every value here is positive, so the L1 norm and the sum agree.
        assert_eq!(s.reduce(Reduce::L1Norm, 0.0), 40.0);
        assert_eq!(s.reduce(Reduce::L2Norm, 0.0), 232.0f32.sqrt());
    }

    #[test]
    fn variance_of_a_constant_column_is_not_negative() {
        let mut s = ValueStats::default();
        for _ in 0..1000 {
            s.add(1e7);
        }
        assert_eq!(s.reduce(Reduce::Sd, 0.0), 0.0);
    }

    #[test]
    fn add_repeated_matches_repeated_add() {
        let mut a = ValueStats::default();
        for _ in 0..5 {
            a.add(3.5);
        }
        let mut b = ValueStats::default();
        b.add_repeated(3.5, 5);
        assert_eq!(a, b);
    }

    #[test]
    fn bin_plan_rejects_nonsense() {
        assert!(BinPlan::new(0.0, None, false).is_err());
        assert!(BinPlan::new(-1.0, None, false).is_err());
        assert!(BinPlan::new(f64::NAN, None, false).is_err());
        assert!(BinPlan::new(1.0, Some(0), false).is_err());
        assert_eq!(
            BinPlan::new(10.0, None, false).unwrap().whole_bin_size(),
            10
        );
    }

    /// A bin is a run of bases. A fractional width used to be accepted here
    /// and refused by the whole-file iterators, so the same argument was legal
    /// through one door and not the other.
    #[test]
    fn a_fractional_bin_size_is_refused_everywhere() {
        for bad in [0.5f64, 2.5, 1.000001, 99.9] {
            let err = BinPlan::new(bad, None, false).unwrap_err().to_string();
            assert!(err.contains("whole number of base pairs"), "{bad}: {err}");
            assert!(err.contains("bin_count"), "{bad}: {err}");
        }
        // A whole number written as a float is fine; that is what the Python
        // layer hands over for `bin_size=100`.
        for good in [1.0f64, 100.0, 1e6] {
            assert!(BinPlan::new(good, None, false).is_ok(), "{good}");
        }
    }

    #[test]
    fn mode_and_reduce_parse_and_reject() {
        use std::str::FromStr;
        assert_eq!(BinMode::from_str("sum").unwrap(), BinMode::Sum);
        assert_eq!(Reduce::from_str("l2norm").unwrap(), Reduce::L2Norm);
        let err = BinMode::from_str("median").unwrap_err().to_string();
        assert_eq!(err, "bin_mode median invalid");
        assert_eq!(
            Reduce::from_str("median").unwrap_err().to_string(),
            "reduce median invalid"
        );
    }
}