Skip to main content

trx_rs/ops/
connectivity.rs

1use crate::dtype::TrxScalar;
2use crate::error::{Result, TrxError};
3use crate::trx_file::TrxFile;
4
5/// How to measure connectivity between groups.
6#[derive(Debug, Clone, Copy)]
7pub enum ConnectivityMeasure {
8    /// Count the number of streamlines connecting two groups.
9    Count,
10    /// Sum a DPS value for streamlines connecting two groups.
11    WeightedSum,
12}
13
14/// Compute a pairwise connectivity matrix between groups.
15///
16/// Returns a packed upper-triangle vector of size `n*(n+1)/2` where
17/// `n = group_names.len()`. Entry `(i,j)` with `i <= j` is at index
18/// `i*n - i*(i+1)/2 + j`.
19///
20/// Each streamline is assigned to a group if its index appears in that
21/// group's member list. A streamline connecting groups `i` and `j` increments
22/// the `(min(i,j), max(i,j))` entry.
23pub fn compute_group_connectivity<P: TrxScalar>(
24    trx: &TrxFile<P>,
25    group_names: &[&str],
26    measure: ConnectivityMeasure,
27    dps_weight_name: Option<&str>,
28) -> Result<Vec<f64>> {
29    let n = group_names.len();
30    let matrix_size = n * (n + 1) / 2;
31    let mut matrix = vec![0.0f64; matrix_size];
32
33    // Build streamline → group membership
34    let mut streamline_groups: Vec<Vec<usize>> = vec![Vec::new(); trx.nb_streamlines()];
35
36    for (gi, &gname) in group_names.iter().enumerate() {
37        let members = trx.group(gname)?;
38        for &m in members {
39            let idx = m as usize;
40            if idx < streamline_groups.len() {
41                streamline_groups[idx].push(gi);
42            }
43        }
44    }
45
46    // Optional DPS weights
47    let weights: Option<Vec<f64>> = match (measure, dps_weight_name) {
48        (ConnectivityMeasure::WeightedSum, Some(name)) => {
49            let view = trx.dps::<f32>(name)?;
50            Some(view.rows().map(|r| r[0] as f64).collect())
51        }
52        (ConnectivityMeasure::WeightedSum, None) => {
53            return Err(TrxError::Argument(
54                "WeightedSum requires a DPS weight name".into(),
55            ));
56        }
57        _ => None,
58    };
59
60    // Accumulate
61    for (si, groups) in streamline_groups.iter().enumerate() {
62        let val = weights.as_ref().map_or(1.0, |w| w[si]);
63        for &gi in groups {
64            for &gj in groups {
65                let (a, b) = if gi <= gj { (gi, gj) } else { (gj, gi) };
66                let idx = a * n - a * (a + 1) / 2 + b;
67                matrix[idx] += val;
68            }
69        }
70    }
71
72    Ok(matrix)
73}