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;
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;
}
}