rsomics-popgen-core 0.1.1

use crate::{PopgenError, Result};

/// LD r² between two biallelic loci given the per-haplotype phased
/// genotype matrix `gens[hap_index][locus_index] ∈ {0, 1}` (0 = major,
/// 1 = minor). Returns NaN-free `r² ∈ [0, 1]`.
pub fn r_squared(genotypes: &[[u8; 2]]) -> Result<f64> {
    if genotypes.is_empty() {
        return Err(PopgenError::Empty);
    }
    let n = genotypes.len() as f64;
    let mut a_count = 0.0;
    let mut b_count = 0.0;
    let mut ab_count = 0.0;
    for hap in genotypes {
        let a = f64::from(hap[0]);
        let b = f64::from(hap[1]);
        a_count += a;
        b_count += b;
        ab_count += a * b;
    }
    let pa = a_count / n;
    let pb = b_count / n;
    let pab = ab_count / n;
    let d = pab - pa * pb;
    let denom = pa * (1.0 - pa) * pb * (1.0 - pb);
    if denom <= 0.0 {
        // One locus is monomorphic — r² is undefined; report 0 (matches
        // PLINK's behaviour for the no-variation case).
        return Ok(0.0);
    }
    Ok((d * d / denom).clamp(0.0, 1.0))
}

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

    fn approx(a: f64, b: f64, eps: f64) -> bool {
        (a - b).abs() < eps
    }

    #[test]
    fn perfect_ld_pair() {
        // Every haplotype carries both alleles together or neither →  r² = 1.
        let g = vec![[0, 0], [0, 0], [1, 1], [1, 1]];
        let r2 = r_squared(&g).unwrap();
        assert!(approx(r2, 1.0, 1e-9), "{r2}");
    }

    #[test]
    fn no_ld_independent_loci() {
        // 50/50 split, independent: r² ≈ 0.
        let g = vec![[0, 0], [0, 1], [1, 0], [1, 1]];
        let r2 = r_squared(&g).unwrap();
        assert!(approx(r2, 0.0, 1e-9), "{r2}");
    }

    #[test]
    fn monomorphic_locus_yields_zero() {
        // Locus A monomorphic.
        let g = vec![[0, 0], [0, 1], [0, 1], [0, 0]];
        let r2 = r_squared(&g).unwrap();
        assert_eq!(r2, 0.0);
    }

    #[test]
    fn empty_errors() {
        assert!(matches!(r_squared(&[]), Err(PopgenError::Empty)));
    }
}