Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
527 lines
16 KiB
Rust
527 lines
16 KiB
Rust
//! Cluster-based permutation tests for spatio-temporal data
|
|
//!
|
|
//! Implements cluster-based permutation tests as described in:
|
|
//! Maris & Oostenveld (2007) "Nonparametric statistical testing of EEG- and MEG-data"
|
|
//!
|
|
//! These tests control for multiple comparisons by forming clusters of
|
|
//! adjacent significant points and testing cluster-level statistics.
|
|
|
|
use crate::{Result, StatsError, permutation::Tail, utils};
|
|
use rand::prelude::*;
|
|
use rayon::prelude::*;
|
|
use std::collections::{HashSet, VecDeque};
|
|
|
|
/// Sparse adjacency matrix for spatial connectivity
|
|
#[derive(Debug, Clone)]
|
|
pub struct Adjacency {
|
|
/// Number of nodes (e.g., channels or sources)
|
|
pub n_nodes: usize,
|
|
/// Adjacency list: for each node, list of neighbor indices
|
|
pub neighbors: Vec<Vec<usize>>,
|
|
}
|
|
|
|
impl Adjacency {
|
|
/// Create adjacency from neighbor lists
|
|
pub fn new(neighbors: Vec<Vec<usize>>) -> Self {
|
|
let n_nodes = neighbors.len();
|
|
Self { n_nodes, neighbors }
|
|
}
|
|
|
|
/// Create fully connected adjacency (all nodes are neighbors)
|
|
pub fn fully_connected(n_nodes: usize) -> Self {
|
|
let neighbors: Vec<Vec<usize>> = (0..n_nodes)
|
|
.map(|i| (0..n_nodes).filter(|&j| j != i).collect())
|
|
.collect();
|
|
Self { n_nodes, neighbors }
|
|
}
|
|
|
|
/// Create grid adjacency (4-connectivity for 2D grid)
|
|
pub fn grid_2d(n_rows: usize, n_cols: usize) -> Self {
|
|
let n_nodes = n_rows * n_cols;
|
|
let mut neighbors = vec![Vec::new(); n_nodes];
|
|
|
|
for row in 0..n_rows {
|
|
for col in 0..n_cols {
|
|
let idx = row * n_cols + col;
|
|
// Up
|
|
if row > 0 {
|
|
neighbors[idx].push((row - 1) * n_cols + col);
|
|
}
|
|
// Down
|
|
if row < n_rows - 1 {
|
|
neighbors[idx].push((row + 1) * n_cols + col);
|
|
}
|
|
// Left
|
|
if col > 0 {
|
|
neighbors[idx].push(row * n_cols + col - 1);
|
|
}
|
|
// Right
|
|
if col < n_cols - 1 {
|
|
neighbors[idx].push(row * n_cols + col + 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
Self { n_nodes, neighbors }
|
|
}
|
|
|
|
/// Create adjacency from distance matrix (connect nodes within threshold)
|
|
pub fn from_distance_matrix(distances: &[Vec<f64>], threshold: f64) -> Self {
|
|
let n_nodes = distances.len();
|
|
let neighbors: Vec<Vec<usize>> = (0..n_nodes)
|
|
.map(|i| {
|
|
(0..n_nodes)
|
|
.filter(|&j| i != j && distances[i][j] <= threshold)
|
|
.collect()
|
|
})
|
|
.collect();
|
|
Self { n_nodes, neighbors }
|
|
}
|
|
}
|
|
|
|
/// A cluster of connected significant points
|
|
#[derive(Debug, Clone)]
|
|
pub struct Cluster {
|
|
/// Indices of points in the cluster (flattened: node * n_times + time)
|
|
pub indices: Vec<usize>,
|
|
/// Cluster statistic (sum of t-values)
|
|
pub statistic: f64,
|
|
/// Node indices in the cluster
|
|
pub nodes: Vec<usize>,
|
|
/// Time indices in the cluster
|
|
pub times: Vec<usize>,
|
|
}
|
|
|
|
/// Result of cluster-based permutation test
|
|
#[derive(Debug, Clone)]
|
|
pub struct ClusterTestResult {
|
|
/// Observed t-statistics [n_nodes x n_times]
|
|
pub t_obs: Vec<Vec<f64>>,
|
|
/// Clusters found in observed data
|
|
pub clusters: Vec<Cluster>,
|
|
/// P-values for each cluster
|
|
pub cluster_pvalues: Vec<f64>,
|
|
/// Null distribution of maximum cluster statistics
|
|
pub h0: Vec<f64>,
|
|
/// Number of permutations
|
|
pub n_permutations: usize,
|
|
/// Threshold used for clustering
|
|
pub threshold: f64,
|
|
}
|
|
|
|
/// Find clusters in a 2D array of statistics
|
|
///
|
|
/// # Arguments
|
|
/// * `stats` - 2D array of test statistics [n_nodes x n_times]
|
|
/// * `threshold` - Cluster-forming threshold (absolute value)
|
|
/// * `adjacency` - Spatial adjacency matrix (None = no spatial clustering)
|
|
/// * `tail` - Which tail to consider for clustering
|
|
///
|
|
/// # Returns
|
|
/// Vector of clusters found
|
|
pub fn find_clusters(
|
|
stats: &[Vec<f64>],
|
|
threshold: f64,
|
|
adjacency: Option<&Adjacency>,
|
|
tail: Tail,
|
|
) -> Vec<Cluster> {
|
|
if stats.is_empty() || stats[0].is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
let n_nodes = stats.len();
|
|
let n_times = stats[0].len();
|
|
|
|
// Create mask of significant points
|
|
let mask: Vec<Vec<bool>> = stats
|
|
.iter()
|
|
.map(|row| {
|
|
row.iter()
|
|
.map(|&v| match tail {
|
|
Tail::TwoSided => v.abs() >= threshold,
|
|
Tail::Greater => v >= threshold,
|
|
Tail::Less => v <= -threshold,
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
// Find connected components
|
|
let mut visited = vec![vec![false; n_times]; n_nodes];
|
|
let mut clusters = Vec::new();
|
|
|
|
for node in 0..n_nodes {
|
|
for time in 0..n_times {
|
|
if mask[node][time] && !visited[node][time] {
|
|
// BFS to find cluster
|
|
let cluster = bfs_cluster(
|
|
&mask,
|
|
&mut visited,
|
|
node,
|
|
time,
|
|
n_nodes,
|
|
n_times,
|
|
adjacency,
|
|
stats,
|
|
);
|
|
if !cluster.indices.is_empty() {
|
|
clusters.push(cluster);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Sort clusters by absolute statistic (largest first)
|
|
clusters.sort_by(|a, b| {
|
|
b.statistic
|
|
.abs()
|
|
.partial_cmp(&a.statistic.abs())
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
|
|
clusters
|
|
}
|
|
|
|
/// BFS to find connected cluster
|
|
fn bfs_cluster(
|
|
mask: &[Vec<bool>],
|
|
visited: &mut [Vec<bool>],
|
|
start_node: usize,
|
|
start_time: usize,
|
|
n_nodes: usize,
|
|
n_times: usize,
|
|
adjacency: Option<&Adjacency>,
|
|
stats: &[Vec<f64>],
|
|
) -> Cluster {
|
|
let mut indices = Vec::new();
|
|
let mut nodes_set = HashSet::new();
|
|
let mut times_set = HashSet::new();
|
|
let mut statistic = 0.0;
|
|
|
|
let mut queue = VecDeque::new();
|
|
queue.push_back((start_node, start_time));
|
|
visited[start_node][start_time] = true;
|
|
|
|
while let Some((node, time)) = queue.pop_front() {
|
|
indices.push(node * n_times + time);
|
|
nodes_set.insert(node);
|
|
times_set.insert(time);
|
|
statistic += stats[node][time];
|
|
|
|
// Check temporal neighbors
|
|
if time > 0 && mask[node][time - 1] && !visited[node][time - 1] {
|
|
visited[node][time - 1] = true;
|
|
queue.push_back((node, time - 1));
|
|
}
|
|
if time < n_times - 1 && mask[node][time + 1] && !visited[node][time + 1] {
|
|
visited[node][time + 1] = true;
|
|
queue.push_back((node, time + 1));
|
|
}
|
|
|
|
// Check spatial neighbors
|
|
if let Some(adj) = adjacency {
|
|
for &neighbor in &adj.neighbors[node] {
|
|
if neighbor < n_nodes && mask[neighbor][time] && !visited[neighbor][time] {
|
|
visited[neighbor][time] = true;
|
|
queue.push_back((neighbor, time));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut nodes: Vec<usize> = nodes_set.into_iter().collect();
|
|
let mut times: Vec<usize> = times_set.into_iter().collect();
|
|
nodes.sort_unstable();
|
|
times.sort_unstable();
|
|
|
|
Cluster {
|
|
indices,
|
|
statistic,
|
|
nodes,
|
|
times,
|
|
}
|
|
}
|
|
|
|
/// Spatio-temporal cluster-based permutation test
|
|
///
|
|
/// # Arguments
|
|
/// * `x` - Condition A data [n_observations x n_nodes x n_times]
|
|
/// * `y` - Condition B data (optional, for two-sample test)
|
|
/// * `threshold` - Cluster-forming threshold (t-value)
|
|
/// * `n_permutations` - Number of permutations
|
|
/// * `adjacency` - Spatial adjacency (None = temporal clustering only)
|
|
/// * `tail` - Test tail
|
|
/// * `seed` - Optional random seed
|
|
///
|
|
/// # Returns
|
|
/// ClusterTestResult with clusters and p-values
|
|
pub fn spatio_temporal_cluster_test(
|
|
x: &[Vec<Vec<f64>>],
|
|
y: Option<&[Vec<Vec<f64>>]>,
|
|
threshold: f64,
|
|
n_permutations: usize,
|
|
adjacency: Option<&Adjacency>,
|
|
tail: Tail,
|
|
seed: Option<u64>,
|
|
) -> Result<ClusterTestResult> {
|
|
if x.is_empty() {
|
|
return Err(StatsError::InvalidInput("Empty data".to_string()));
|
|
}
|
|
|
|
let n_obs_x = x.len();
|
|
let n_nodes = x[0].len();
|
|
let n_times = x[0][0].len();
|
|
|
|
// Validate dimensions
|
|
for obs in x {
|
|
if obs.len() != n_nodes || obs.iter().any(|row| row.len() != n_times) {
|
|
return Err(StatsError::DimensionMismatch(
|
|
"All observations must have same dimensions".to_string(),
|
|
));
|
|
}
|
|
}
|
|
|
|
// Determine test type
|
|
let (t_obs, _compute_stat): (
|
|
Vec<Vec<f64>>,
|
|
Box<dyn Fn(&[Vec<Vec<f64>>], Option<&[Vec<Vec<f64>>]>) -> Vec<Vec<f64>> + Send + Sync>,
|
|
) = if let Some(y_data) = y {
|
|
// Two-sample test
|
|
let _n_obs_y = y_data.len();
|
|
for obs in y_data {
|
|
if obs.len() != n_nodes || obs.iter().any(|row| row.len() != n_times) {
|
|
return Err(StatsError::DimensionMismatch(
|
|
"Y dimensions must match X".to_string(),
|
|
));
|
|
}
|
|
}
|
|
|
|
let t = compute_t_stats_ind(x, y_data, n_nodes, n_times);
|
|
let func = Box::new(move |a: &[Vec<Vec<f64>>], b: Option<&[Vec<Vec<f64>>]>| {
|
|
compute_t_stats_ind(a, b.unwrap(), n_nodes, n_times)
|
|
});
|
|
(t, func)
|
|
} else {
|
|
// One-sample test
|
|
let t = compute_t_stats_1samp(x, n_nodes, n_times);
|
|
let func = Box::new(move |a: &[Vec<Vec<f64>>], _: Option<&[Vec<Vec<f64>>]>| {
|
|
compute_t_stats_1samp(a, n_nodes, n_times)
|
|
});
|
|
(t, func)
|
|
};
|
|
|
|
// Find clusters in observed data
|
|
let clusters = find_clusters(&t_obs, threshold, adjacency, tail);
|
|
|
|
// Generate null distribution
|
|
let h0: Vec<f64> = if y.is_some() {
|
|
// Two-sample: permute group labels
|
|
let _n_obs_y = y.unwrap().len();
|
|
let mut combined: Vec<Vec<Vec<f64>>> = x.to_vec();
|
|
combined.extend(y.unwrap().iter().cloned());
|
|
|
|
(0..n_permutations)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let mut rng = match seed {
|
|
Some(s) => StdRng::seed_from_u64(s.wrapping_add(i as u64)),
|
|
None => StdRng::from_entropy(),
|
|
};
|
|
|
|
let mut shuffled = combined.clone();
|
|
shuffled.shuffle(&mut rng);
|
|
|
|
let perm_x = &shuffled[..n_obs_x];
|
|
let perm_y = &shuffled[n_obs_x..];
|
|
|
|
let t_perm = compute_t_stats_ind(perm_x, perm_y, n_nodes, n_times);
|
|
let perm_clusters = find_clusters(&t_perm, threshold, adjacency, tail);
|
|
|
|
perm_clusters.first().map_or(0.0, |c| c.statistic.abs())
|
|
})
|
|
.collect()
|
|
} else {
|
|
// One-sample: sign flipping
|
|
(0..n_permutations)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let mut rng = match seed {
|
|
Some(s) => StdRng::seed_from_u64(s.wrapping_add(i as u64)),
|
|
None => StdRng::from_entropy(),
|
|
};
|
|
|
|
// Generate sign flips
|
|
let signs: Vec<f64> = (0..n_obs_x)
|
|
.map(|_| if rng.r#gen::<bool>() { 1.0 } else { -1.0 })
|
|
.collect();
|
|
|
|
// Apply signs
|
|
let flipped: Vec<Vec<Vec<f64>>> = x
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(k, obs)| {
|
|
obs.iter()
|
|
.map(|row| row.iter().map(|&v| v * signs[k]).collect())
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let t_perm = compute_t_stats_1samp(&flipped, n_nodes, n_times);
|
|
let perm_clusters = find_clusters(&t_perm, threshold, adjacency, tail);
|
|
|
|
perm_clusters.first().map_or(0.0, |c| c.statistic.abs())
|
|
})
|
|
.collect()
|
|
};
|
|
|
|
// Compute cluster p-values
|
|
let cluster_pvalues: Vec<f64> = clusters
|
|
.iter()
|
|
.map(|cluster| {
|
|
let count = h0.iter().filter(|&&v| v >= cluster.statistic.abs()).count();
|
|
(count as f64 + 1.0) / (n_permutations as f64 + 1.0)
|
|
})
|
|
.collect();
|
|
|
|
Ok(ClusterTestResult {
|
|
t_obs,
|
|
clusters,
|
|
cluster_pvalues,
|
|
h0,
|
|
n_permutations,
|
|
threshold,
|
|
})
|
|
}
|
|
|
|
/// Compute one-sample t-statistics for each node and time
|
|
fn compute_t_stats_1samp(data: &[Vec<Vec<f64>>], n_nodes: usize, n_times: usize) -> Vec<Vec<f64>> {
|
|
let _n_obs = data.len() as f64;
|
|
|
|
(0..n_nodes)
|
|
.map(|node| {
|
|
(0..n_times)
|
|
.map(|time| {
|
|
let values: Vec<f64> = data.iter().map(|obs| obs[node][time]).collect();
|
|
let mean = utils::mean(&values);
|
|
let se = utils::sem(&values);
|
|
if se < 1e-10 { 0.0 } else { mean / se }
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Compute independent two-sample t-statistics for each node and time
|
|
fn compute_t_stats_ind(
|
|
x: &[Vec<Vec<f64>>],
|
|
y: &[Vec<Vec<f64>>],
|
|
n_nodes: usize,
|
|
n_times: usize,
|
|
) -> Vec<Vec<f64>> {
|
|
let n_x = x.len() as f64;
|
|
let n_y = y.len() as f64;
|
|
|
|
(0..n_nodes)
|
|
.map(|node| {
|
|
(0..n_times)
|
|
.map(|time| {
|
|
let vals_x: Vec<f64> = x.iter().map(|obs| obs[node][time]).collect();
|
|
let vals_y: Vec<f64> = y.iter().map(|obs| obs[node][time]).collect();
|
|
|
|
let mean_x = utils::mean(&vals_x);
|
|
let mean_y = utils::mean(&vals_y);
|
|
let var_x = utils::variance(&vals_x, 1);
|
|
let var_y = utils::variance(&vals_y, 1);
|
|
|
|
let se = (var_x / n_x + var_y / n_y).sqrt();
|
|
if se < 1e-10 {
|
|
0.0
|
|
} else {
|
|
(mean_x - mean_y) / se
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_adjacency_grid() {
|
|
let adj = Adjacency::grid_2d(3, 3);
|
|
assert_eq!(adj.n_nodes, 9);
|
|
// Corner has 2 neighbors
|
|
assert_eq!(adj.neighbors[0].len(), 2);
|
|
// Center has 4 neighbors
|
|
assert_eq!(adj.neighbors[4].len(), 4);
|
|
// Edge has 3 neighbors
|
|
assert_eq!(adj.neighbors[1].len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_find_clusters() {
|
|
// Simple 2x5 array with clusters
|
|
// Without spatial adjacency, each row forms separate temporal clusters
|
|
let stats = vec![vec![0.5, 2.5, 2.8, 0.3, 0.1], vec![0.2, 2.6, 2.4, 0.4, 3.0]];
|
|
|
|
let clusters = find_clusters(&stats, 2.0, None, Tail::Greater);
|
|
|
|
// Should find clusters (row0 cols 1-2, row1 cols 1-2, row1 col 4)
|
|
assert!(!clusters.is_empty());
|
|
// With spatial adjacency, the two middle clusters merge
|
|
let adj = Adjacency::fully_connected(2);
|
|
let clusters_spatial = find_clusters(&stats, 2.0, Some(&adj), Tail::Greater);
|
|
assert!(clusters_spatial.len() <= clusters.len());
|
|
}
|
|
|
|
#[test]
|
|
fn test_cluster_test_1samp() {
|
|
// Create data with a significant cluster
|
|
let mut rng = StdRng::seed_from_u64(42);
|
|
|
|
let n_obs = 20;
|
|
let n_nodes = 4;
|
|
let n_times = 10;
|
|
|
|
// Generate data with effect in nodes 1-2, times 3-6
|
|
let data: Vec<Vec<Vec<f64>>> = (0..n_obs)
|
|
.map(|_| {
|
|
(0..n_nodes)
|
|
.map(|node| {
|
|
(0..n_times)
|
|
.map(|time| {
|
|
let noise: f64 = rng.r#gen::<f64>() * 0.5 - 0.25;
|
|
if (1..=2).contains(&node) && (3..=6).contains(&time) {
|
|
1.5 + noise // Effect
|
|
} else {
|
|
noise // No effect
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
let adj = Adjacency::fully_connected(n_nodes);
|
|
let result = spatio_temporal_cluster_test(
|
|
&data,
|
|
None,
|
|
2.0,
|
|
500,
|
|
Some(&adj),
|
|
Tail::Greater,
|
|
Some(42),
|
|
)
|
|
.unwrap();
|
|
|
|
// Should find at least one significant cluster
|
|
assert!(!result.clusters.is_empty());
|
|
// First cluster should be significant
|
|
assert!(result.cluster_pvalues[0] < 0.05);
|
|
}
|
|
}
|