fast-nnt 0.4.0

Rust implementation of the NeighborNet phylogenetic network algorithm.
Documentation
use rayon::prelude::*;

#[doc(hidden)]
pub const EPS: f64 = 1e-12;

#[derive(Debug)]
#[doc(hidden)]
pub struct Scratch {
    pub p: Vec<f64>,
    pub r: Vec<f64>,
    pub z: Vec<f64>,
    pub w: Vec<f64>,
}

impl Scratch {
    pub fn new(len: usize) -> Self {
        Self {
            p: vec![0.0; len],
            r: vec![0.0; len],
            z: vec![0.0; len],
            w: vec![0.0; len],
        }
    }
}

#[derive(Debug)]
#[doc(hidden)]
pub struct PackedTriIndex {
    pub n: usize,
    pub npairs: usize,
    pub row_offsets: Vec<usize>,
}

impl PackedTriIndex {
    pub fn new(n: usize) -> Self {
        let npairs = n * (n - 1) / 2;
        let mut row_offsets = vec![0usize; n + 1];
        let mut start = 0usize;
        for (i, off) in row_offsets.iter_mut().enumerate().take(n + 1).skip(1) {
            *off = start;
            start += n - i;
        }
        Self {
            n,
            npairs,
            row_offsets,
        }
    }
}

#[doc(hidden)]
pub fn calc_atx_indexed(x: &[f64], y: &mut [f64], tri_idx: &PackedTriIndex) {
    let n = tri_idx.n;
    debug_assert_eq!(x.len(), tri_idx.npairs);

    let mut s_index = 0usize;
    for i in 1..n {
        let mut y_s = 0.0;
        if i >= 2 {
            let mut d_idx = i - 2;
            for j in 1..i {
                y_s += x[d_idx];
                d_idx += n - j - 1;
            }
        }
        let end = s_index + (n - i);
        for xv in &x[s_index..end] {
            y_s += xv;
        }
        y[s_index] = y_s;
        s_index += n - i;
    }

    let mut s_index = 1usize;
    for i in 1..=(n - 2) {
        y[s_index] = y[s_index - 1] + y[s_index + n - i - 1] - 2.0 * x[s_index - 1];
        s_index += n - i;
    }

    for k in 3..=(n - 1) {
        s_index = k - 1;
        for i in 1..=(n - k) {
            let below = s_index + n - i - 1;
            y[s_index] = y[s_index - 1] + y[below] - y[below - 1] - 2.0 * x[s_index - 1];
            s_index += n - i;
        }
    }
}

pub fn calc_ainv_y(y: &[f64], x: &mut [f64], n: usize) {
    assert!(n >= 2, "n must be >= 2");
    let m = n * (n - 1) / 2;
    assert_eq!(y.len(), m, "y must have length n*(n-1)/2");
    assert_eq!(x.len(), m, "x must have length n*(n-1)/2");

    let tri_idx = PackedTriIndex::new(n);
    let offsets = &tri_idx.row_offsets;

    x[0] = (y[n - 2] + y[0] - y[2 * n - 4]) / 2.0;

    let mut d_index: usize = 2 * n - 4;
    for j in 3..=n {
        let t1 = y[d_index] + y[j - 2];
        let t2 = t1 - y[j - 3];
        let t3 = t2 - y[d_index + n - j];
        x[j - 2] = t3 / 2.0;
        d_index += n - j;
    }

    x[n - 2] = (y[n - 2] + y[m - 1] - y[n - 3]) / 2.0;

    for i in 2..=(n - 1) {
        let s_idx = offsets[i];
        let prev = offsets[i - 1];

        x[s_idx] = (y[prev] + y[s_idx] - y[prev + 1]) / 2.0;

        let row_len = n - i;
        for col in 1..row_len {
            x[s_idx + col] =
                (y[prev + col] + y[s_idx + col] - y[s_idx + col - 1] - y[prev + col + 1]) / 2.0;
        }
    }
}

#[doc(hidden)]
pub fn get_active_entries(x: &[f64], a: &mut [bool]) {
    a.par_iter_mut()
        .zip(x.par_iter())
        .for_each(|(ai, &xi)| *ai = xi <= 0.0);
}

#[doc(hidden)]
pub fn zero_negative_entries(x: &mut [f64]) {
    x.par_iter_mut().for_each(|xi| {
        if *xi <= 0.0 {
            *xi = 0.0
        }
    });
}

pub fn sum_array_squared(x: &[f64], n: usize) -> f64 {
    let tri_idx = PackedTriIndex::new(n);
    sum_array_squared_indexed(x, &tri_idx)
}

#[doc(hidden)]
pub fn sum_array_squared_indexed(x: &[f64], tri_idx: &PackedTriIndex) -> f64 {
    let expected = tri_idx.npairs;
    assert!(
        x.len() == expected,
        "x must have length n*(n-1)/2 (got {}, expected {})",
        x.len(),
        expected
    );
    let mut total = 0.0f64;

    // Per-row partial sums preserve the reference implementation's FP order.
    for i in 1..=tri_idx.n {
        let start = tri_idx.row_offsets[i];
        let len = tri_idx.n - i;
        let mut s_i = 0.0f64;
        for x_ij in &x[start..start + len] {
            s_i += x_ij * x_ij;
        }
        total += s_i;
    }

    total
}

pub fn mask_elements_branchless(r: &mut [f64], a: &[bool]) {
    assert_eq!(r.len(), a.len());
    for (x, &mask) in r.iter_mut().zip(a.iter()) {
        *x *= (!mask) as u8 as f64;
    }
}