use ann_search_rs::{build_exhaustive_index, query_exhaustive_self};
use faer::linalg::matmul::matmul;
use faer::linalg::solvers::Solve;
use faer::{Accum, Mat, MatRef};
use rayon::prelude::*;
use crate::core::math::vector_helpers::{mad, median};
use crate::graph::graph_components::induced_components;
use crate::prelude::*;
use crate::single_cell::sc_trajectory::diffusion::{boundary_cells, nearest_candidate};
use crate::utils::faer_parallelism;
const ADAPTIVE_K_CAP: usize = 30;
const MIN_KNN: usize = 6;
const TERMINAL_CUTOFF_Z: f64 = 3.719_016_485_455_680_4;
const POWER_ITER_MAX: usize = 5_000;
const POWER_ITER_TOL: f64 = 1e-12;
const MAX_DENSE_TRANSIENT: usize = 20_000;
const ABSORPTION_ROW_SUM_TOL: f64 = 1e-3;
const ABSORPTION_RESIDUAL_TOL: f64 = 1e-5;
pub fn build_waypoint_transitions(
wp_data: &[Vec<f32>],
wp_pseudotime: &[f32],
knn: usize,
verbose: bool,
) -> Result<CompressedSparseData2<f32>, BixverseErrors> {
let n_wp = wp_data.len();
if n_wp == 0 {
return Err(BixverseErrors::InvalidArgument(
"Palantir: the waypoint set is empty".to_string(),
));
}
if wp_pseudotime.len() != n_wp {
return Err(BixverseErrors::PalantirWaypointDataMismatch {
n_waypoints: n_wp,
n_pseudotime: wp_pseudotime.len(),
});
}
if knn < MIN_KNN {
return Err(BixverseErrors::PalantirKnnTooSmall {
knn,
minimum: MIN_KNN,
});
}
if n_wp < MIN_KNN {
return Err(BixverseErrors::PalantirTooFewWaypoints {
n_waypoints: n_wp,
minimum: MIN_KNN,
});
}
let knn = knn.min(n_wp);
let (indices, distances) = waypoint_knn(wp_data, knn, verbose)?;
let adaptive_std = adaptive_bandwidths(&distances, knn);
let mut indptr: Vec<u32> = vec![0];
let mut col_idx: Vec<u32> = Vec::new();
let mut data: Vec<f32> = Vec::new();
for i in 0..n_wp {
let sigma_i = adaptive_std[i];
let cut = wp_pseudotime[i] - sigma_i;
let start = data.len();
for (slot, &j) in indices[i].iter().enumerate() {
let d = distances[i][slot];
if d <= 0.0 {
continue;
}
if wp_pseudotime[j] < cut {
continue;
}
let sigma_j = adaptive_std[j];
let exponent = -0.5 * d * d * (1.0 / (sigma_i * sigma_i) + 1.0 / (sigma_j * sigma_j));
col_idx.push(j as u32);
data.push(exponent.exp());
}
if data.len() == start {
let mut best: Option<usize> = None;
let mut best_pt = f32::NEG_INFINITY;
for (slot, &j) in indices[i].iter().enumerate() {
if j != i && distances[i][slot] > 0.0 && wp_pseudotime[j] > best_pt {
best_pt = wp_pseudotime[j];
best = Some(j);
}
}
let target = match best {
Some(j) if wp_pseudotime[j] >= wp_pseudotime[i] => j,
_ => i,
};
col_idx.push(target as u32);
data.push(1.0);
}
let row_sum: f32 = data[start..].iter().sum();
if row_sum > 0.0 {
for v in data[start..].iter_mut() {
*v /= row_sum;
}
}
indptr.push(col_idx.len() as u32);
}
Ok(CompressedSparseData2::new_csr(
&data,
&col_idx,
&indptr,
None,
(n_wp, n_wp),
))
}
fn waypoint_knn(wp_data: &[Vec<f32>], k: usize, verbose: bool) -> ScKnnResults {
let n_wp = wp_data.len();
let n_dims = wp_data[0].len();
let mat = Mat::<f32>::from_fn(n_wp, n_dims, |i, j| wp_data[i][j]);
let index = build_exhaustive_index(mat.as_ref(), "euclidean");
let (indices, distances) = query_exhaustive_self(&index, k, true, verbose)?;
let distances = distances
.map(|d| {
d.into_iter()
.map(|row| row.into_iter().map(|v| v.sqrt()).collect())
.collect::<Vec<Vec<f32>>>()
})
.ok_or(BixverseErrors::InvalidArgument(
"Palantir: the exhaustive waypoint search returned no distances".to_string(),
))?;
Ok((indices, distances))
}
fn adaptive_bandwidths(distances: &[Vec<f32>], knn: usize) -> Vec<f32> {
let adaptive_k = (knn / 3).saturating_sub(1).min(ADAPTIVE_K_CAP);
let mut sigma: Vec<f32> = distances
.iter()
.map(|row| {
let idx = adaptive_k.min(row.len().saturating_sub(1));
row.get(idx).copied().unwrap_or(0.0)
})
.collect();
let smallest_positive = sigma
.iter()
.copied()
.filter(|s| *s > 0.0)
.fold(f32::INFINITY, f32::min);
let fallback = if smallest_positive.is_finite() {
smallest_positive
} else {
1.0
};
for s in sigma.iter_mut() {
if !s.is_finite() || *s <= 0.0 {
*s = fallback;
}
}
sigma
}
fn row_rescale_factors(transitions: &CompressedSparseData2<f32>) -> Vec<f64> {
(0..transitions.shape.0)
.map(|i| {
let mass: f64 = (transitions.indptr[i] as usize..transitions.indptr[i + 1] as usize)
.map(|idx| transitions.data[idx] as f64)
.sum();
if mass > 0.0 { 1.0 / mass } else { 1.0 }
})
.collect()
}
pub fn stationary_distribution(
transitions: &CompressedSparseData2<f32>,
) -> Result<Vec<f64>, BixverseErrors> {
if !transitions.cs_type.is_csr() {
return Err(BixverseErrors::SparseMatrixMustBeCsr);
}
let n = transitions.shape.0;
let scale = row_rescale_factors(transitions);
let mut v = vec![1.0 / n as f64; n];
let mut next = vec![0.0f64; n];
for iteration in 1..=POWER_ITER_MAX {
next.iter_mut()
.zip(v.iter())
.for_each(|(o, &x)| *o = 0.5 * x);
for i in 0..n {
let vi = v[i];
if vi == 0.0 {
continue;
}
let scaled = 0.5 * scale[i] * vi;
for idx in transitions.indptr[i] as usize..transitions.indptr[i + 1] as usize {
let j = transitions.indices[idx] as usize;
next[j] += transitions.data[idx] as f64 * scaled;
}
}
let total: f64 = next.iter().sum();
if total <= 0.0 || !total.is_finite() {
return Err(BixverseErrors::PalantirStationaryMassLost { iteration, total });
}
for x in next.iter_mut() {
*x /= total;
}
let residual: f64 = next
.iter()
.zip(v.iter())
.map(|(&a, &b)| (a - b).abs())
.sum();
std::mem::swap(&mut v, &mut next);
if residual < POWER_ITER_TOL {
break;
}
}
Ok(v)
}
pub fn detect_terminal_states(
transitions: &CompressedSparseData2<f32>,
wp_data: &[Vec<f32>],
wp_pseudotime: &[f32],
) -> Result<Vec<usize>, BixverseErrors> {
let ranks: Vec<f64> = stationary_distribution(transitions)?
.into_iter()
.map(f64::abs)
.collect();
let centre = median(&ranks).ok_or(BixverseErrors::PalantirNoTerminalStates)?;
let spread = mad(&ranks, None).ok_or(BixverseErrors::PalantirNoTerminalStates)?;
if spread <= 0.0 || !spread.is_finite() {
return Err(BixverseErrors::PalantirStationaryRanksDegenerate {
n_waypoints: ranks.len(),
});
}
let cutoff = centre + TERMINAL_CUTOFF_Z * spread;
let above: Vec<usize> = (0..ranks.len()).filter(|&i| ranks[i] > cutoff).collect();
if above.is_empty() {
return Err(BixverseErrors::PalantirNoTerminalStates);
}
let (n_comps, labels) = induced_components(transitions, &above)?;
let mut candidates: Vec<usize> = vec![usize::MAX; n_comps];
let mut best: Vec<f32> = vec![f32::NEG_INFINITY; n_comps];
for (slot, &wp) in above.iter().enumerate() {
let comp = labels[slot];
if wp_pseudotime[wp] > best[comp] {
best[comp] = wp_pseudotime[wp];
candidates[comp] = wp;
}
}
let boundaries = boundary_cells(wp_data);
if boundaries.is_empty() {
return Err(BixverseErrors::PalantirNoBoundaryCells);
}
let mut terminal: Vec<usize> = candidates
.into_iter()
.filter(|&c| c != usize::MAX)
.map(|c| nearest_candidate(wp_data, &boundaries, c))
.collect::<Result<Vec<usize>, BixverseErrors>>()?;
terminal.sort_unstable();
terminal.dedup();
if terminal.is_empty() {
return Err(BixverseErrors::PalantirNoTerminalStates);
}
Ok(terminal)
}
pub fn absorption_probabilities(
transitions: &CompressedSparseData2<f32>,
absorbing: &[usize],
) -> Result<(Mat<f32>, usize), BixverseErrors> {
if !transitions.cs_type.is_csr() {
return Err(BixverseErrors::SparseMatrixMustBeCsr);
}
let n = transitions.shape.0;
let n_abs = absorbing.len();
if n_abs == 0 {
return Err(BixverseErrors::PalantirNoTerminalStates);
}
let mut pos_abs = vec![-1i64; n];
for (slot, &a) in absorbing.iter().enumerate() {
if a >= n {
return Err(BixverseErrors::SliceIndexOutOfBounds { index: a, len: n });
}
pos_abs[a] = slot as i64;
}
let reachable = reaches_absorbing(transitions, &pos_abs);
let transient: Vec<usize> = (0..n).filter(|&i| pos_abs[i] < 0 && reachable[i]).collect();
let stranded = (0..n).filter(|&i| pos_abs[i] < 0 && !reachable[i]).count();
let n_trans = transient.len();
if n_trans > MAX_DENSE_TRANSIENT {
return Err(BixverseErrors::PalantirTooManyWaypointsForDenseSolve {
n_transient: n_trans,
limit: MAX_DENSE_TRANSIENT,
});
}
let mut pos_trans = vec![-1i64; n];
for (slot, &t) in transient.iter().enumerate() {
pos_trans[t] = slot as i64;
}
let mut out = Mat::<f32>::zeros(n, n_abs);
for (slot, &a) in absorbing.iter().enumerate() {
out[(a, slot)] = 1.0;
}
if n_trans == 0 {
return Ok((out, stranded));
}
let scale = row_rescale_factors(transitions);
let mut lhs = Mat::<f64>::zeros(n_trans, n_trans);
let mut rhs = Mat::<f64>::zeros(n_trans, n_abs);
for i in 0..n {
let ti = pos_trans[i];
if ti < 0 {
continue;
}
let ti = ti as usize;
for idx in transitions.indptr[i] as usize..transitions.indptr[i + 1] as usize {
let j = transitions.indices[idx] as usize;
let v = transitions.data[idx] as f64 * scale[i];
if pos_trans[j] >= 0 {
lhs[(ti, pos_trans[j] as usize)] -= v;
} else if pos_abs[j] >= 0 {
rhs[(ti, pos_abs[j] as usize)] += v;
}
}
lhs[(ti, ti)] += 1.0;
}
let solved = lhs.partial_piv_lu().solve(&rhs);
let mut residual = Mat::<f64>::zeros(n_trans, n_abs);
matmul(
residual.as_mut(),
Accum::Replace,
lhs.as_ref(),
solved.as_ref(),
1.0f64,
faer_parallelism(),
);
let mut max_residual = 0.0f64;
for ti in 0..n_trans {
for k in 0..n_abs {
let r = (residual[(ti, k)] - rhs[(ti, k)]).abs();
if r.is_nan() || r > max_residual {
max_residual = r;
}
}
}
if !max_residual.is_finite() || max_residual > ABSORPTION_RESIDUAL_TOL {
return Err(BixverseErrors::PalantirAbsorbingSolveFailed {
residual: max_residual,
n_transient: n_trans,
});
}
for ti in 0..n_trans {
let mut row_sum = 0.0f64;
for k in 0..n_abs {
let p = solved[(ti, k)].max(0.0);
row_sum += p;
out[(transient[ti], k)] = p as f32;
}
if !row_sum.is_finite() {
return Err(BixverseErrors::PalantirAbsorbingRowNotFinite {
row_sum,
waypoint: transient[ti],
n_transient: n_trans,
});
}
if row_sum > 1.0 + ABSORPTION_ROW_SUM_TOL {
return Err(BixverseErrors::PalantirAbsorbingRowSum {
row_sum,
waypoint: transient[ti],
n_transient: n_trans,
});
}
}
Ok((out, stranded))
}
fn reaches_absorbing(transitions: &CompressedSparseData2<f32>, pos_abs: &[i64]) -> Vec<bool> {
let n = transitions.shape.0;
let reversed = transitions.transpose_and_convert();
let mut reached = vec![false; n];
let mut stack: Vec<usize> = (0..n).filter(|&i| pos_abs[i] >= 0).collect();
for &s in &stack {
reached[s] = true;
}
while let Some(u) = stack.pop() {
for idx in reversed.indptr[u] as usize..reversed.indptr[u + 1] as usize {
let v = reversed.indices[idx] as usize;
if reversed.data[idx] > 0.0 && !reached[v] {
reached[v] = true;
stack.push(v);
}
}
}
reached
}
pub fn project_to_cells(
weights: &[f32],
n_waypoints: usize,
n_cells: usize,
branch_probs_wp: MatRef<f32>,
) -> Result<Mat<f32>, BixverseErrors> {
if weights.len() != n_waypoints * n_cells {
return Err(BixverseErrors::DimensionMisMatchSparse {
indices_len: weights.len(),
data_len: n_waypoints * n_cells,
});
}
if branch_probs_wp.nrows() != n_waypoints {
return Err(BixverseErrors::ShapeMismatch {
expected: (n_waypoints, branch_probs_wp.ncols()),
got: (branch_probs_wp.nrows(), branch_probs_wp.ncols()),
});
}
let w = MatRef::<f32>::from_row_major_slice(weights, n_waypoints, n_cells);
let mut out = Mat::<f32>::zeros(n_cells, branch_probs_wp.ncols());
matmul(
out.as_mut(),
Accum::Replace,
w.transpose(),
branch_probs_wp,
1.0f32,
faer_parallelism(),
);
Ok(out)
}
pub fn branch_entropy(branch_probs: MatRef<f32>) -> Vec<f32> {
let n_states = branch_probs.ncols();
(0..branch_probs.nrows())
.into_par_iter()
.map(|i| {
let total: f64 = (0..n_states).map(|k| branch_probs[(i, k)] as f64).sum();
if total <= 0.0 || !total.is_finite() {
return 0.0;
}
let h: f64 = (0..n_states)
.map(|k| {
let p = branch_probs[(i, k)] as f64 / total;
if p > 0.0 { -p * p.ln() } else { 0.0 }
})
.sum();
h as f32
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_relative_eq;
fn chain_from_rows(rows: &[Vec<f32>]) -> CompressedSparseData2<f32> {
let n = rows.len();
let mut indptr = vec![0u32];
let mut indices = Vec::new();
let mut data = Vec::new();
for row in rows {
for (j, &v) in row.iter().enumerate() {
if v != 0.0 {
indices.push(j as u32);
data.push(v);
}
}
indptr.push(indices.len() as u32);
}
CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (n, n))
}
#[test]
fn test_transitions_rows_sum_to_one_and_have_no_self_loops() {
let wp_data: Vec<Vec<f32>> = (0..12).map(|i| vec![i as f32, 0.0]).collect();
let wp_pt: Vec<f32> = (0..12).map(|i| i as f32 / 11.0).collect();
let t = build_waypoint_transitions(&wp_data, &wp_pt, 6, false).unwrap();
for i in 0..12 {
let lo = t.indptr[i] as usize;
let hi = t.indptr[i + 1] as usize;
let sum: f32 = t.data[lo..hi].iter().sum();
assert_relative_eq!(sum, 1.0, epsilon = 1e-5);
assert!(
!t.indices[lo..hi].contains(&(i as u32)),
"self loop on waypoint {i}"
);
}
}
#[test]
fn test_transitions_prune_far_back_edges() {
let wp_data: Vec<Vec<f32>> = (0..12).map(|i| vec![i as f32 * 0.01]).collect();
let wp_pt: Vec<f32> = (0..12).map(|i| i as f32).collect();
let t = build_waypoint_transitions(&wp_data, &wp_pt, 6, false).unwrap();
for i in 0..12 {
for idx in t.indptr[i] as usize..t.indptr[i + 1] as usize {
let j = t.indices[idx] as usize;
assert!(
wp_pt[j] >= wp_pt[i] - 1.0,
"edge {i} -> {j} runs too far backwards"
);
}
}
}
#[test]
fn test_transitions_reject_small_knn() {
let wp_data: Vec<Vec<f32>> = (0..12).map(|i| vec![i as f32]).collect();
let wp_pt: Vec<f32> = (0..12).map(|i| i as f32).collect();
assert!(matches!(
build_waypoint_transitions(&wp_data, &wp_pt, 4, false),
Err(BixverseErrors::PalantirKnnTooSmall { .. })
));
}
#[test]
fn test_transitions_clamp_knn_to_the_waypoint_count() {
let wp_data: Vec<Vec<f32>> = (0..12).map(|i| vec![i as f32 * 0.01]).collect();
let wp_pt: Vec<f32> = (0..12).map(|i| i as f32 * 0.02).collect();
let oversized = build_waypoint_transitions(&wp_data, &wp_pt, 40, false).unwrap();
let exact = build_waypoint_transitions(&wp_data, &wp_pt, 12, false).unwrap();
assert_eq!(oversized.indptr, exact.indptr);
assert_eq!(oversized.indices, exact.indices);
}
#[test]
fn test_transitions_reject_too_few_waypoints_for_the_bandwidth() {
let wp_data: Vec<Vec<f32>> = (0..4).map(|i| vec![i as f32]).collect();
let wp_pt: Vec<f32> = (0..4).map(|i| i as f32).collect();
assert!(matches!(
build_waypoint_transitions(&wp_data, &wp_pt, 30, false),
Err(BixverseErrors::PalantirTooFewWaypoints {
n_waypoints: 4,
minimum: 6
})
));
}
#[test]
fn test_transitions_reject_a_waypoint_pseudotime_mismatch() {
let wp_data: Vec<Vec<f32>> = (0..8).map(|i| vec![i as f32]).collect();
let wp_pt: Vec<f32> = (0..3).map(|i| i as f32).collect();
assert!(matches!(
build_waypoint_transitions(&wp_data, &wp_pt, 6, false),
Err(BixverseErrors::PalantirWaypointDataMismatch {
n_waypoints: 8,
n_pseudotime: 3
})
));
}
#[test]
fn test_fully_pruned_row_never_runs_backwards() {
let wp_data: Vec<Vec<f32>> = (0..12).map(|i| vec![i as f32 * 0.01]).collect();
let wp_pt: Vec<f32> = (0..12).map(|i| i as f32).collect();
let t = build_waypoint_transitions(&wp_data, &wp_pt, 6, false).unwrap();
for i in 0..12 {
for idx in t.indptr[i] as usize..t.indptr[i + 1] as usize {
let j = t.indices[idx] as usize;
assert!(
wp_pt[j] >= wp_pt[i],
"edge {i} -> {j} runs backwards in pseudotime"
);
}
}
let last = 11usize;
let lo = t.indptr[last] as usize;
let hi = t.indptr[last + 1] as usize;
assert_eq!(&t.indices[lo..hi], &[last as u32]);
}
#[test]
fn test_stationary_of_two_state_chain() {
let t = chain_from_rows(&[vec![0.75, 0.25], vec![0.5, 0.5]]);
let pi = stationary_distribution(&t).unwrap();
assert_relative_eq!(pi[0], 2.0 / 3.0, epsilon = 1e-9);
assert_relative_eq!(pi[1], 1.0 / 3.0, epsilon = 1e-9);
}
#[test]
fn test_stationary_ignores_per_row_scale_drift() {
let drift = 1.05f32;
let t = chain_from_rows(&[vec![0.75 * drift, 0.25 * drift], vec![0.5, 0.5]]);
let pi = stationary_distribution(&t).unwrap();
assert_relative_eq!(pi[0], 2.0 / 3.0, epsilon = 1e-9);
assert_relative_eq!(pi[1], 1.0 / 3.0, epsilon = 1e-9);
}
#[test]
fn test_stationary_survives_a_periodic_chain() {
let t = chain_from_rows(&[vec![0.0, 1.0], vec![1.0, 0.0]]);
let pi = stationary_distribution(&t).unwrap();
assert_relative_eq!(pi[0], 0.5, epsilon = 1e-9);
assert_relative_eq!(pi[1], 0.5, epsilon = 1e-9);
}
#[test]
fn test_absorption_on_a_three_state_chain() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0],
vec![0.3, 0.0, 0.7],
vec![0.0, 0.0, 1.0],
]);
let (b, stranded) = absorption_probabilities(&t, &[0, 2]).unwrap();
assert_eq!(stranded, 0);
assert_relative_eq!(b[(1, 0)], 0.3, epsilon = 1e-6);
assert_relative_eq!(b[(1, 1)], 0.7, epsilon = 1e-6);
assert_relative_eq!(b[(0, 0)], 1.0, epsilon = 1e-6);
assert_relative_eq!(b[(2, 1)], 1.0, epsilon = 1e-6);
}
#[test]
fn test_absorption_through_a_transient_cascade() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0],
vec![0.25, 0.0, 0.0, 0.75],
vec![0.0, 0.0, 0.0, 1.0],
]);
let (b, _) = absorption_probabilities(&t, &[0, 3]).unwrap();
assert_relative_eq!(b[(1, 0)], 0.25, epsilon = 1e-6);
assert_relative_eq!(b[(1, 1)], 0.75, epsilon = 1e-6);
assert_relative_eq!(b[(2, 0)], 0.25, epsilon = 1e-6);
}
#[test]
fn test_absorption_rows_sum_to_one() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0, 0.0],
vec![0.1, 0.4, 0.5, 0.0],
vec![0.0, 0.2, 0.3, 0.5],
vec![0.0, 0.0, 0.0, 1.0],
]);
let (b, _) = absorption_probabilities(&t, &[0, 3]).unwrap();
for i in 0..4 {
let sum: f32 = (0..2).map(|k| b[(i, k)]).sum();
assert_relative_eq!(sum, 1.0, epsilon = 1e-5);
}
}
#[test]
fn test_absorption_reports_a_stranded_transient() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0],
vec![0.5, 0.0, 0.5],
]);
let (b, stranded) = absorption_probabilities(&t, &[0]).unwrap();
assert_eq!(stranded, 1);
assert_relative_eq!(b[(1, 0)], 0.0, epsilon = 1e-6);
assert_relative_eq!(b[(2, 0)], 1.0, epsilon = 1e-6);
assert_relative_eq!(b[(0, 0)], 1.0, epsilon = 1e-6);
}
#[test]
fn test_absorption_does_not_redistribute_mass_lost_to_a_strand() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
vec![0.1, 0.0, 0.9, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 0.0, 0.0, 0.5, 0.5, 0.0],
]);
let (b, stranded) = absorption_probabilities(&t, &[0, 1]).unwrap();
assert_eq!(stranded, 1);
assert_relative_eq!(b[(5, 0)], 0.05, epsilon = 1e-6);
assert_relative_eq!(b[(5, 1)], 0.5, epsilon = 1e-6);
assert_relative_eq!(b[(3, 0)], 0.1, epsilon = 1e-6);
assert_relative_eq!(b[(3, 1)], 0.0, epsilon = 1e-6);
assert_relative_eq!(b[(4, 1)], 1.0, epsilon = 1e-6);
assert_relative_eq!(b[(2, 0)], 0.0, epsilon = 1e-6);
assert_relative_eq!(b[(2, 1)], 0.0, epsilon = 1e-6);
}
#[test]
fn test_absorption_rows_never_sum_above_one() {
let t = chain_from_rows(&[
vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
vec![0.1, 0.0, 0.9, 0.0, 0.0, 0.0],
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
vec![0.0, 0.0, 0.0, 0.5, 0.5, 0.0],
]);
let (b, _) = absorption_probabilities(&t, &[0, 1]).unwrap();
for i in 0..6 {
let sum: f32 = (0..2).map(|k| b[(i, k)]).sum();
assert!((0.0..=1.0 + 1e-5).contains(&sum), "row {i} sums to {sum}");
}
}
#[test]
fn test_absorption_survives_a_chain_with_a_huge_absorption_time() {
const N: usize = 31;
let mut rows = vec![vec![0.0f32; N]; N];
rows[0][1] = 1.0;
for i in 1..N - 1 {
rows[i][i + 1] = 0.4;
rows[i][i - 1] = 0.6;
}
rows[N - 1][N - 1] = 1.0;
let t = chain_from_rows(&rows);
let (b, stranded) = absorption_probabilities(&t, &[N - 1]).unwrap();
assert_eq!(stranded, 0);
for i in 0..N {
assert_relative_eq!(b[(i, 0)], 1.0, epsilon = 1e-4);
}
}
#[test]
fn test_absorption_rejects_an_empty_absorbing_set() {
let t = chain_from_rows(&[vec![1.0, 0.0], vec![0.0, 1.0]]);
assert!(matches!(
absorption_probabilities(&t, &[]),
Err(BixverseErrors::PalantirNoTerminalStates)
));
}
#[test]
fn test_projection_is_a_convex_combination() {
let weights = vec![0.25, 0.25, 0.25, 0.75, 0.75, 0.75];
let b_wp = Mat::<f32>::from_fn(2, 2, |i, j| if i == j { 1.0 } else { 0.0 });
let out = project_to_cells(&weights, 2, 3, b_wp.as_ref()).unwrap();
for c in 0..3 {
assert_relative_eq!(out[(c, 0)], 0.25, epsilon = 1e-6);
assert_relative_eq!(out[(c, 1)], 0.75, epsilon = 1e-6);
}
}
#[test]
fn test_entropy_matches_known_values() {
let bp = Mat::<f32>::from_fn(3, 2, |i, j| match (i, j) {
(0, 0) => 1.0,
(0, 1) => 0.0,
(1, _) => 0.5,
(2, _) => 4.0,
_ => 0.0,
});
let h = branch_entropy(bp.as_ref());
assert_relative_eq!(h[0], 0.0, epsilon = 1e-6);
assert_relative_eq!(h[1], std::f32::consts::LN_2, epsilon = 1e-6);
assert_relative_eq!(h[2], std::f32::consts::LN_2, epsilon = 1e-6);
}
#[test]
fn test_entropy_of_an_all_zero_row_is_zero_not_nan() {
let bp = Mat::<f32>::zeros(2, 3);
let h = branch_entropy(bp.as_ref());
assert!(h.iter().all(|v| v.is_finite()));
assert_relative_eq!(h[0], 0.0, epsilon = 1e-6);
}
#[test]
fn test_detect_terminal_states_finds_the_late_end_of_a_line() {
let wp_data: Vec<Vec<f32>> = (0..30).map(|i| vec![i as f32 / 29.0]).collect();
let wp_pt: Vec<f32> = (0..30).map(|i| i as f32 / 29.0).collect();
let t = build_waypoint_transitions(&wp_data, &wp_pt, 6, false).unwrap();
let terminal = detect_terminal_states(&t, &wp_data, &wp_pt).unwrap();
assert_eq!(terminal, vec![29]);
}
#[test]
fn test_detect_terminal_states_rejects_a_flat_stationary_spread() {
let n = 8usize;
let mut rows = vec![vec![0.0f32; n]; n];
for (i, row) in rows.iter_mut().enumerate() {
row[(i + 1) % n] = 1.0;
}
let t = chain_from_rows(&rows);
let wp_data: Vec<Vec<f32>> = (0..n).map(|i| vec![i as f32]).collect();
let wp_pt: Vec<f32> = (0..n).map(|i| i as f32).collect();
assert!(matches!(
detect_terminal_states(&t, &wp_data, &wp_pt),
Err(BixverseErrors::PalantirStationaryRanksDegenerate { n_waypoints: 8 })
));
}
#[test]
fn test_stranded_waypoints_count_a_zero_weight_route_out() {
let data = vec![1.0f32, 0.0, 1.0, 0.5, 0.5];
let indices = vec![0u32, 0, 1, 0, 2];
let indptr = vec![0u32, 1, 3, 5];
let t = CompressedSparseData2::new_csr(&data, &indices, &indptr, None, (3, 3));
let (b, stranded) = absorption_probabilities(&t, &[0]).unwrap();
assert_eq!(stranded, 1);
assert_relative_eq!(b[(1, 0)], 0.0, epsilon = 1e-6);
}
}