use rustc_hash::FxHashSet;
use crate::prelude::*;
use crate::single_cell::mc_generation::seacells::{
compute_diffusion_kernel, determine_multiscale_space, diffusion_map_from_kernel_diag,
max_min_sampling,
};
pub struct MultiscaleSpace {
pub components: Vec<Vec<f32>>,
pub converged: bool,
pub residual: f64,
}
const MAX_MULTISCALE_EIGENVALUE: f64 = 1.0 - 1e-10;
const MIN_MULTISCALE_EIGS: usize = 3;
const MIN_MULTISCALE_EIGENVALUES: usize = 2;
pub fn multiscale_components(
knn_indices: &[Vec<usize>],
knn_distances: &[Vec<f32>],
squared_dist: bool,
n_dcs: usize,
n_eigs: Option<usize>,
seed: u64,
lanczos_params: Option<LanczosParams>,
) -> Result<MultiscaleSpace, BixverseErrors> {
if n_dcs == 0 {
return Err(BixverseErrors::InvalidArgument(
"Palantir: n_dcs must be positive".to_string(),
));
}
let mut kernel = compute_diffusion_kernel(knn_indices, knn_distances, squared_dist)?;
let degrees = kernel_row_sums(&kernel);
let solve = diffusion_map_from_kernel_diag(&mut kernel, n_dcs + 1, seed, lanczos_params)?;
if solve.eigenvalues.len() < MIN_MULTISCALE_EIGENVALUES {
return Err(BixverseErrors::InvalidArgument(
"Palantir: fewer than two diffusion components could be resolved".to_string(),
));
}
let eigenvectors = back_transform_eigenvectors(&solve.eigenvectors, °rees);
let n_keep = resolve_n_eigs(&solve.eigenvalues, n_eigs)?;
let clamped = clamp_eigenvalues(&solve.eigenvalues)?;
Ok(MultiscaleSpace {
components: determine_multiscale_space(&clamped, &eigenvectors, Some(n_keep)),
converged: solve.converged,
residual: solve.residual,
})
}
fn kernel_row_sums(kernel: &CompressedSparseData2<f32>) -> Vec<f32> {
(0..kernel.shape.0)
.map(|i| {
(kernel.indptr[i]..kernel.indptr[i + 1])
.map(|idx| kernel.data[idx as usize])
.sum()
})
.collect()
}
fn back_transform_eigenvectors(eigenvectors: &[Vec<f32>], degrees: &[f32]) -> Vec<Vec<f32>> {
let n = eigenvectors.len();
if n == 0 {
return Vec::new();
}
let n_comps = eigenvectors[0].len();
let inv_sqrt_degree: Vec<f32> = degrees
.iter()
.map(|&d| if d > 0.0 { 1.0 / d.sqrt() } else { 0.0 })
.collect();
let mut out = vec![vec![0.0f32; n_comps]; n];
for c in 0..n_comps {
let mut norm_sq = 0.0f64;
for i in 0..n {
let v = eigenvectors[i][c] * inv_sqrt_degree[i];
out[i][c] = v;
norm_sq += (v as f64) * (v as f64);
}
let norm = norm_sq.sqrt() as f32;
if norm > 0.0 {
for row in out.iter_mut() {
row[c] /= norm;
}
}
}
out
}
fn clamp_eigenvalues(eigenvalues: &[f64]) -> Result<Vec<f64>, BixverseErrors> {
eigenvalues
.iter()
.map(|&l| {
if l.is_finite() {
Ok(l.min(MAX_MULTISCALE_EIGENVALUE))
} else {
Err(BixverseErrors::PalantirNonFiniteEigenvalue { eigenvalue: l })
}
})
.collect()
}
fn resolve_n_eigs(eigenvalues: &[f64], requested: Option<usize>) -> Result<usize, BixverseErrors> {
let available = eigenvalues.len();
if let Some(n) = requested {
if n < MIN_MULTISCALE_EIGENVALUES {
return Err(BixverseErrors::PalantirNEigsTooSmall {
n_eigs: n,
minimum: MIN_MULTISCALE_EIGENVALUES,
});
}
return Ok(n.min(available));
}
let gaps: Vec<f64> = eigenvalues.windows(2).map(|w| w[0] - w[1]).collect();
if gaps.is_empty() {
return Ok(MIN_MULTISCALE_EIGS.min(available));
}
let mut order: Vec<usize> = (0..gaps.len()).collect();
order.sort_by(|&a, &b| gaps[a].total_cmp(&gaps[b]));
let mut n_eigs = order[order.len() - 1] + 1;
if n_eigs < MIN_MULTISCALE_EIGS && order.len() >= 2 {
n_eigs = order[order.len() - 2] + 1;
}
if n_eigs < MIN_MULTISCALE_EIGS {
n_eigs = MIN_MULTISCALE_EIGS;
}
Ok(n_eigs.min(available))
}
pub fn minmax_scale_columns(data: &mut [Vec<f32>]) {
if data.is_empty() {
return;
}
let n_cols = data[0].len();
for c in 0..n_cols {
let mut lo = f32::INFINITY;
let mut hi = f32::NEG_INFINITY;
for row in data.iter() {
lo = lo.min(row[c]);
hi = hi.max(row[c]);
}
let range = hi - lo;
if range > 0.0 {
for row in data.iter_mut() {
row[c] = (row[c] - lo) / range;
}
} else {
for row in data.iter_mut() {
row[c] = 0.0;
}
}
}
}
pub fn boundary_cells(data: &[Vec<f32>]) -> Vec<usize> {
if data.is_empty() {
return Vec::new();
}
let n_cols = data[0].len();
let mut found: FxHashSet<usize> = FxHashSet::default();
for c in 0..n_cols {
let (mut arg_lo, mut arg_hi) = (0usize, 0usize);
let (mut lo, mut hi) = (data[0][c], data[0][c]);
for (i, row) in data.iter().enumerate().skip(1) {
if row[c] < lo {
lo = row[c];
arg_lo = i;
}
if row[c] > hi {
hi = row[c];
arg_hi = i;
}
}
found.insert(arg_lo);
found.insert(arg_hi);
}
let mut out: Vec<usize> = found.into_iter().collect();
out.sort_unstable();
out
}
pub fn nearest_candidate(
data: &[Vec<f32>],
candidates: &[usize],
query: usize,
) -> Result<usize, BixverseErrors> {
let first = *candidates
.first()
.ok_or(BixverseErrors::PalantirNoBoundaryCells)?;
let mut best = first;
let mut best_dist = f32::INFINITY;
for &c in candidates {
let dist: f32 = data[c]
.iter()
.zip(data[query].iter())
.map(|(&a, &b)| (a - b) * (a - b))
.sum();
if dist < best_dist {
best_dist = dist;
best = c;
}
}
Ok(best)
}
pub fn assemble_waypoints(
data: &[Vec<f32>],
num_waypoints: usize,
boundaries: &[usize],
terminal_states: Option<&[usize]>,
start_cell: usize,
seed: u64,
) -> Vec<usize> {
if data.is_empty() {
return Vec::new();
}
let n_dims = data[0].len();
let target = num_waypoints.max(MIN_MULTISCALE_EIGS.max(n_dims));
let mut set: FxHashSet<usize> = max_min_sampling(data, target, seed).into_iter().collect();
set.extend(boundaries.iter().copied());
if let Some(states) = terminal_states {
set.extend(states.iter().copied());
}
set.remove(&start_cell);
let mut rest: Vec<usize> = set.into_iter().collect();
rest.sort_unstable();
let mut out = Vec::with_capacity(rest.len() + 1);
out.push(start_cell);
out.extend(rest);
out
}
#[cfg(test)]
mod tests {
use super::*;
use approx::assert_relative_eq;
#[test]
fn test_minmax_scale_maps_to_unit_range() {
let mut data = vec![vec![1.0, 10.0], vec![3.0, 20.0], vec![5.0, 30.0]];
minmax_scale_columns(&mut data);
assert_relative_eq!(data[0][0], 0.0, epsilon = 1e-6);
assert_relative_eq!(data[1][0], 0.5, epsilon = 1e-6);
assert_relative_eq!(data[2][0], 1.0, epsilon = 1e-6);
assert_relative_eq!(data[1][1], 0.5, epsilon = 1e-6);
}
#[test]
fn test_minmax_scale_constant_column_is_zero_not_nan() {
let mut data = vec![vec![7.0, 1.0], vec![7.0, 2.0]];
minmax_scale_columns(&mut data);
assert!(data[0][0].is_finite() && data[1][0].is_finite());
assert_relative_eq!(data[0][0], 0.0, epsilon = 1e-6);
assert_relative_eq!(data[1][0], 0.0, epsilon = 1e-6);
}
#[test]
fn test_boundary_cells_are_column_extremes() {
let data = vec![
vec![0.0, 0.5],
vec![1.0, 0.9],
vec![2.0, 0.1],
vec![3.0, 0.4],
];
assert_eq!(boundary_cells(&data), vec![0, 1, 2, 3]);
}
#[test]
fn test_boundary_cells_take_first_index_on_ties() {
let data = vec![vec![1.0], vec![1.0], vec![1.0]];
assert_eq!(boundary_cells(&data), vec![0]);
}
#[test]
fn test_nearest_candidate_picks_closest() {
let data = vec![vec![0.0, 0.0], vec![10.0, 0.0], vec![1.0, 1.0]];
assert_eq!(nearest_candidate(&data, &[0, 1], 2).unwrap(), 0);
assert!(nearest_candidate(&data, &[], 0).is_err());
}
#[test]
fn test_waypoints_start_first_then_ascending() {
let data: Vec<Vec<f32>> = (0..50)
.map(|i| vec![i as f32, (i % 7) as f32, (i % 3) as f32])
.collect();
let boundaries = boundary_cells(&data);
let waypoints = assemble_waypoints(&data, 20, &boundaries, Some(&[42]), 5, 42);
assert_eq!(waypoints[0], 5);
assert!(waypoints.windows(2).skip(1).all(|w| w[0] < w[1]));
assert!(!waypoints[1..].contains(&5));
assert!(waypoints.contains(&42));
for b in &boundaries {
if *b != 5 {
assert!(waypoints.contains(b), "boundary {b} missing");
}
}
}
#[test]
fn test_resolve_n_eigs_takes_largest_gap() {
let eigenvalues = [1.0, 0.9, 0.4, 0.35];
assert_eq!(resolve_n_eigs(&eigenvalues, None).unwrap(), 3);
}
#[test]
fn test_resolve_n_eigs_uses_a_late_gap_when_it_is_largest() {
let eigenvalues = [1.0, 0.98, 0.95, 0.93, 0.30];
assert_eq!(resolve_n_eigs(&eigenvalues, None).unwrap(), 4);
}
#[test]
fn test_resolve_n_eigs_honours_explicit_requests_below_the_floor() {
let eigenvalues = [1.0, 0.9, 0.5, 0.2];
assert_eq!(resolve_n_eigs(&eigenvalues, Some(2)).unwrap(), 2);
assert_eq!(resolve_n_eigs(&eigenvalues, Some(4)).unwrap(), 4);
assert_eq!(resolve_n_eigs(&eigenvalues, Some(99)).unwrap(), 4);
}
#[test]
fn test_resolve_n_eigs_rejects_a_request_below_two() {
let eigenvalues = [1.0, 0.9, 0.5, 0.2];
assert!(matches!(
resolve_n_eigs(&eigenvalues, Some(1)),
Err(BixverseErrors::PalantirNEigsTooSmall { n_eigs: 1, .. })
));
assert!(matches!(
resolve_n_eigs(&eigenvalues, Some(0)),
Err(BixverseErrors::PalantirNEigsTooSmall { n_eigs: 0, .. })
));
}
#[test]
fn test_resolve_n_eigs_never_exceeds_what_is_available() {
let two = [1.0, 0.5];
let one = [1.0];
assert_eq!(resolve_n_eigs(&two, Some(2)).unwrap(), 2);
assert_eq!(resolve_n_eigs(&two, Some(99)).unwrap(), 2);
assert_eq!(resolve_n_eigs(&two, None).unwrap(), 2);
assert_eq!(resolve_n_eigs(&one, Some(3)).unwrap(), 1);
assert_eq!(resolve_n_eigs(&one, None).unwrap(), 1);
assert_eq!(resolve_n_eigs(&[], None).unwrap(), 0);
}
#[test]
fn test_clamp_eigenvalues_pins_the_ceiling() {
let clamped = clamp_eigenvalues(&[1.0, 1.5, 0.999_999, 0.5]).unwrap();
assert_eq!(clamped[0], MAX_MULTISCALE_EIGENVALUE);
assert_eq!(clamped[1], MAX_MULTISCALE_EIGENVALUE);
assert_eq!(clamped[2], 0.999_999);
assert_eq!(clamped[3], 0.5);
const {
assert!(1.0 - MAX_MULTISCALE_EIGENVALUE < 4e-5);
}
}
#[test]
fn test_clamp_eigenvalues_rejects_non_finite_input() {
assert!(matches!(
clamp_eigenvalues(&[1.0, f64::NAN]),
Err(BixverseErrors::PalantirNonFiniteEigenvalue { .. })
));
assert!(clamp_eigenvalues(&[f64::INFINITY]).is_err());
}
#[test]
fn test_back_transform_renormalises_to_unit_l2() {
let eigenvectors = vec![vec![0.6, 1.0], vec![0.8, 2.0]];
let degrees = vec![4.0, 9.0];
let out = back_transform_eigenvectors(&eigenvectors, °rees);
for c in 0..2 {
let norm: f32 = out.iter().map(|r| r[c] * r[c]).sum::<f32>().sqrt();
assert_relative_eq!(norm, 1.0, epsilon = 1e-5);
}
assert!(out[0][0] > out[1][0]);
}
#[test]
fn test_back_transform_survives_zero_degree_rows() {
let eigenvectors = vec![vec![1.0], vec![2.0]];
let degrees = vec![0.0, 4.0];
let out = back_transform_eigenvectors(&eigenvectors, °rees);
assert!(out.iter().all(|r| r[0].is_finite()));
assert_relative_eq!(out[0][0], 0.0, epsilon = 1e-6);
}
}