bixverse_core/core/base/
pert_dist.rs

1//! E-distance and related perturbation distances.
2//!
3//! E-distance compares two cell distributions in an embedding via
4//! `2 * mean(d(X, Y)) - mean(d(X, X')) - mean(d(Y, Y'))`.
5
6use crate::utils::gemm;
7use faer::Accum;
8use faer::linalg::matmul::triangular::BlockStructure;
9use faer::{Mat, MatRef, Par};
10use rayon::prelude::*;
11
12use crate::prelude::*;
13use crate::utils::faer_parallelism;
14
15////////////
16// Consts //
17////////////
18
19/// Row chunk of the within-group gram product
20const GRAM_ROW_TILE: usize = 512;
21
22///////////
23// Enums //
24///////////
25
26/// Cell-cell distance used in the E-distance formula.
27#[derive(Debug, Clone, Copy, Default)]
28pub enum PertDistance {
29    /// Energy distance with Euclidean cell-cell distance.
30    #[default]
31    Euclidean,
32    /// Energy distance with squared Euclidean cell-cell distance.
33    ///
34    /// Reduces algebraically to `2 * ||mu_X - mu_Y||^2`; within-group
35    /// dispersion does not affect the result.
36    SquaredEuclidean,
37}
38
39/// Parse the perturbation distance variant from a string.
40///
41/// ### Params
42///
43/// * `s` - The string to parse
44///
45/// ### Returns
46///
47/// The Option of [PertDistance]
48pub fn parse_perturbation_distance(s: &str) -> Option<PertDistance> {
49    match s.to_lowercase().as_str() {
50        "euclidean" => Some(PertDistance::Euclidean),
51        "squared_euclidean" | "sqeuclidean" => Some(PertDistance::SquaredEuclidean),
52        _ => None,
53    }
54}
55
56/////////////
57// Helpers //
58/////////////
59
60/// Calculate the squared norms of the rows
61///
62/// Accumulated feature by feature over contiguous column slices; the per-row
63/// summation order over features is unchanged.
64///
65/// ### Params
66///
67/// * `x` - The matrix for which to calculate the squared row norms
68///
69/// ### Returns
70///
71/// Vector of squared row norms
72fn row_norms_squared<T>(x: MatRef<T>) -> Vec<T>
73where
74    T: BixverseFloat,
75{
76    let mut out = vec![T::zero(); x.nrows()];
77    out.par_chunks_mut(GRAM_ROW_TILE)
78        .enumerate()
79        .for_each(|(c, chunk)| {
80            let start = c * GRAM_ROW_TILE;
81            for j in 0..x.ncols() {
82                let col = x.col(j).subrows(start, chunk.len());
83                for (o, &v) in chunk.iter_mut().zip(col.iter()) {
84                    *o += v * v;
85                }
86            }
87        });
88    out
89}
90
91/// Cell-cell distance from the squared distance
92#[inline(always)]
93fn finish_dist<T: BixverseFloat>(d_sq: T, dist: PertDistance) -> T {
94    match dist {
95        PertDistance::Euclidean => d_sq.sqrt(),
96        PertDistance::SquaredEuclidean => d_sq,
97    }
98}
99
100/// Mean pairwise distance within a group.
101///
102/// Returns zero if `x` has fewer than two rows.
103///
104/// ### Params
105///
106/// * `x` - The matrix for which to calculate the within distance
107/// * `dist` - Which [PertDistance] to calculate
108///
109/// ### Returns
110///
111/// The within pert distance
112pub fn mean_pairwise_within<T>(x: MatRef<T>, dist: PertDistance) -> T
113where
114    T: BixverseFloat + std::iter::Sum,
115{
116    mean_pairwise_within_impl(x, dist, true)
117}
118
119/// Within-group mean distance, optionally single-threaded
120///
121/// ### Params
122///
123/// * `x` - The matrix for which to calculate the within distance
124/// * `dist` - Which [PertDistance] to calculate
125/// * `parallel` - Use rayon; false for calls made inside an outer parallel loop
126///
127/// ### Returns
128///
129/// The within pert distance
130fn mean_pairwise_within_impl<T>(x: MatRef<T>, dist: PertDistance, parallel: bool) -> T
131where
132    T: BixverseFloat + std::iter::Sum,
133{
134    let n = x.nrows();
135    if n < 2 {
136        return T::zero();
137    }
138
139    // X * X^T, lower triangle: diagonal is ||x_i||^2, sub-diagonal is <x_i, x_j>.
140    let mut gram = Mat::<T>::zeros(n, n);
141    gemm::gram(
142        gram.as_mut(),
143        BlockStructure::TriangularLower,
144        Accum::Replace,
145        x,
146        x.transpose(),
147        T::one(),
148        if parallel {
149            faer_parallelism()
150        } else {
151            Par::Seq
152        },
153    );
154
155    let norms_sq: Vec<T> = (0..n).map(|i| *gram.get(i, i)).collect();
156    let two = T::from_f64(2.0).unwrap();
157
158    // Column j of the lower triangle is contiguous.
159    let col_sum = |j: usize| -> T {
160        let mut s = T::zero();
161        for (i, &g) in gram.col(j).iter().enumerate().skip(j + 1) {
162            // clamp: floating point can produce tiny negatives here.
163            let d_sq = (norms_sq[i] + norms_sq[j] - two * g).max(T::zero());
164            s += finish_dist(d_sq, dist);
165        }
166        s
167    };
168    let total: T = if parallel {
169        (0..n).into_par_iter().map(col_sum).sum()
170    } else {
171        (0..n).map(col_sum).sum()
172    };
173
174    let n_pairs = T::from_usize(n * (n - 1) / 2).unwrap();
175
176    total / n_pairs
177}
178
179/// Mean pairwise distance between two groups (rows are samples).
180///
181/// ### Params
182///
183/// * `x` - First matrix
184/// * `y` - Second matrix
185///
186/// ### Returns
187///
188/// The pairwise mean distance between the two matrices
189pub fn mean_pairwise_between<T>(
190    x: MatRef<T>,
191    y: MatRef<T>,
192    dist: PertDistance,
193) -> Result<T, CoreErrors>
194where
195    T: BixverseFloat + std::iter::Sum,
196{
197    mean_pairwise_between_impl(x, y, dist, true)
198}
199
200/// Between-group mean distance, optionally single-threaded
201///
202/// Tiled over row blocks of `x`, so the full `n_x x n_y` gram is never held.
203///
204/// ### Params
205///
206/// * `x` - First matrix
207/// * `y` - Second matrix
208/// * `dist` - Which [PertDistance] to calculate
209/// * `parallel` - Use rayon; false for calls made inside an outer parallel loop
210///
211/// ### Returns
212///
213/// The pairwise mean distance between the two matrices
214fn mean_pairwise_between_impl<T>(
215    x: MatRef<T>,
216    y: MatRef<T>,
217    dist: PertDistance,
218    parallel: bool,
219) -> Result<T, CoreErrors>
220where
221    T: BixverseFloat + std::iter::Sum,
222{
223    if x.ncols() != y.ncols() {
224        return Err(CoreErrors::NonMatchingFeatureDim {
225            dim_x: x.ncols(),
226            dim_y: y.ncols(),
227        });
228    }
229
230    let n_x = x.nrows();
231    let n_y = y.nrows();
232    // special case, return zero
233    if n_x == 0 || n_y == 0 {
234        return Ok(T::zero());
235    }
236
237    let x_norms = row_norms_squared(x);
238    let y_norms = row_norms_squared(y);
239    let two = T::from_f64(2.0).unwrap();
240    let par = if parallel {
241        faer_parallelism()
242    } else {
243        Par::Seq
244    };
245
246    let mut gram = Mat::<T>::zeros(GRAM_ROW_TILE.min(n_x), n_y);
247    let mut total = T::zero();
248    for start in (0..n_x).step_by(GRAM_ROW_TILE) {
249        let rows = GRAM_ROW_TILE.min(n_x - start);
250        let mut tile = gram.as_mut().submatrix_mut(0, 0, rows, n_y);
251        gemm::gemm(
252            tile.as_mut(),
253            Accum::Replace,
254            x.subrows(start, rows),
255            y.transpose(),
256            T::one(),
257            par,
258        );
259        let tile = tile.as_ref();
260        let x_tile = &x_norms[start..start + rows];
261
262        let col_sum = |j: usize| -> T {
263            let mut s = T::zero();
264            for (&g, &xn) in tile.col(j).iter().zip(x_tile) {
265                let d_sq = (xn + y_norms[j] - two * g).max(T::zero());
266                s += finish_dist(d_sq, dist);
267            }
268            s
269        };
270        let tile_total: T = if parallel {
271            (0..n_y).into_par_iter().map(col_sum).sum()
272        } else {
273            (0..n_y).map(col_sum).sum()
274        };
275        total += tile_total;
276    }
277
278    let n_pairs = T::from_usize(n_x * n_y).unwrap();
279
280    Ok(total / n_pairs)
281}
282
283/// Pack rows of a (samples x features) matrix into one `Mat<T>` per label.
284///
285/// ### Params
286///
287/// * `embedding` - The original embedding matrix of shape samples x features
288/// * `labels` - A grouping vector to use. Needs to have the same length as
289///   nrows embedding.
290///
291/// ### Returns
292///
293/// A vector of matrices per label
294fn group_rows_by_label<T>(embedding: MatRef<T>, labels: &[usize]) -> Result<Vec<Mat<T>>, CoreErrors>
295where
296    T: BixverseFloat,
297{
298    if embedding.nrows() != labels.len() {
299        return Err(CoreErrors::NumberLabelsNotEqualSampleNumber {
300            label_length: labels.len(),
301            n_samples: embedding.nrows(),
302        });
303    }
304
305    let d = embedding.ncols();
306    let n_groups = labels.iter().max().map(|&x| x + 1).unwrap_or(0);
307
308    let mut indices: Vec<Vec<usize>> = vec![Vec::new(); n_groups];
309    for (sample, &g) in labels.iter().enumerate() {
310        indices[g].push(sample);
311    }
312
313    let res = indices
314        .into_par_iter()
315        .map(|rows| Mat::from_fn(rows.len(), d, |i, j| *embedding.get(rows[i], j)))
316        .collect();
317
318    Ok(res)
319}
320
321////////////////////
322// Main functions //
323////////////////////
324
325/// E-distance between two groups, each given as (samples x features).
326///
327/// `2 * mean(d(X, Y)) - mean(d(X, X')) - mean(d(Y, Y'))`
328///
329/// ### Params
330///
331/// * `x` - First matrix
332/// * `y` - Second matrix
333/// * `dist` - The perturbation distance to calculate
334///
335/// ### Returns
336///
337/// The E-distance between the two groups
338pub fn edistance_two_matrices<T>(
339    x: MatRef<T>,
340    y: MatRef<T>,
341    dist: PertDistance,
342) -> Result<T, CoreErrors>
343where
344    T: BixverseFloat + std::iter::Sum,
345{
346    let within_x = mean_pairwise_within(x, dist);
347    let within_y = mean_pairwise_within(y, dist);
348    let between = mean_pairwise_between(x, y, dist)?;
349    let two = T::from_f64(2.0).unwrap();
350
351    Ok(two * between - within_x - within_y)
352}
353
354/// E-distance between two groups, single-threaded
355///
356/// For calls made from inside an outer parallel loop, where nested rayon or
357/// faer parallelism would only oversubscribe.
358///
359/// ### Params
360///
361/// * `x` - First matrix
362/// * `y` - Second matrix
363/// * `dist` - The perturbation distance to calculate
364///
365/// ### Returns
366///
367/// The E-distance between the two groups
368pub(crate) fn edistance_two_matrices_seq<T>(
369    x: MatRef<T>,
370    y: MatRef<T>,
371    dist: PertDistance,
372) -> Result<T, CoreErrors>
373where
374    T: BixverseFloat + std::iter::Sum,
375{
376    let within_x = mean_pairwise_within_impl(x, dist, false);
377    let within_y = mean_pairwise_within_impl(y, dist, false);
378    let between = mean_pairwise_between_impl(x, y, dist, false)?;
379    let two = T::from_f64(2.0).unwrap();
380
381    Ok(two * between - within_x - within_y)
382}
383
384/// Pairwise E-distances between all groups defined by `labels`.
385///
386/// ### Params
387///
388/// * `embedding` - Data matrix (samples x features)
389/// * `labels` - Per-sample group label (values in `0..n_groups`)
390/// * `dist` - Distance variant
391///
392/// ### Returns
393///
394/// Symmetric `n_groups x n_groups` E-distance matrix with a zero diagonal.
395pub fn pairwise_edistance<T>(
396    embedding: MatRef<T>,
397    labels: &[usize],
398    dist: PertDistance,
399) -> Result<Mat<T>, CoreErrors>
400where
401    T: BixverseFloat + std::iter::Sum,
402{
403    let groups = group_rows_by_label(embedding, labels)?;
404    let n_groups = groups.len();
405
406    // Only nest faer/rayon inside the outer loop when it cannot fill the pool.
407    let inner_par = groups.len() < rayon::current_num_threads();
408    let withins: Vec<T> = groups
409        .par_iter()
410        .map(|g| mean_pairwise_within_impl(g.as_ref(), dist, inner_par))
411        .collect();
412
413    let pairs: Vec<(usize, usize)> = (0..n_groups)
414        .flat_map(|i| (i + 1..n_groups).map(move |j| (i, j)))
415        .collect();
416
417    let two = T::from_f64(2.0).unwrap();
418    let results: Vec<(usize, usize, T)> = pairs
419        .par_iter()
420        .map(|&(i, j)| {
421            let between = mean_pairwise_between_impl(
422                groups[i].as_ref(),
423                groups[j].as_ref(),
424                dist,
425                inner_par,
426            )?;
427            let ed = two * between - withins[i] - withins[j];
428
429            Ok((i, j, ed))
430        })
431        .collect::<Result<Vec<_>, CoreErrors>>()?;
432
433    let mut result = Mat::<T>::zeros(n_groups, n_groups);
434    for (i, j, val) in results {
435        result[(i, j)] = val;
436        result[(j, i)] = val;
437    }
438
439    Ok(result)
440}
441
442/// One-sided E-distances from each group to a reference group.
443///
444/// ### Params
445///
446/// * `embedding` - Data matrix (samples x features)
447/// * `labels` - Per-sample group label
448/// * `reference` - Label of the reference group (e.g. control)
449/// * `dist` - Distance variant
450///
451/// ### Returns
452///
453/// Vector of length `n_groups`; the entry at `reference` is zero.
454pub fn onesided_edistance<T>(
455    embedding: MatRef<T>,
456    labels: &[usize],
457    reference: usize,
458    dist: PertDistance,
459) -> Result<Vec<T>, CoreErrors>
460where
461    T: BixverseFloat + std::iter::Sum,
462{
463    let groups = group_rows_by_label(embedding, labels)?;
464    let n_groups = groups.len();
465    assert!(reference < n_groups, "reference label out of range");
466
467    let inner_par = n_groups < rayon::current_num_threads();
468    let within_ref = mean_pairwise_within(groups[reference].as_ref(), dist);
469    let two = T::from_f64(2.0).unwrap();
470
471    let res = (0..n_groups)
472        .into_par_iter()
473        .map(|g| {
474            if g == reference {
475                Ok(T::zero())
476            } else {
477                let within_g = mean_pairwise_within_impl(groups[g].as_ref(), dist, inner_par);
478                let between = mean_pairwise_between_impl(
479                    groups[g].as_ref(),
480                    groups[reference].as_ref(),
481                    dist,
482                    inner_par,
483                )?;
484
485                Ok(two * between - within_g - within_ref)
486            }
487        })
488        .collect::<Result<Vec<_>, CoreErrors>>()?;
489
490    Ok(res)
491}
492
493///////////
494// Tests //
495///////////
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    const EPS: f64 = 1e-10;
502
503    fn close(a: f64, b: f64) -> bool {
504        (a - b).abs() < EPS
505    }
506
507    /// The R-facing parser is case insensitive, accepts both aliases and rejects the rest.
508    #[test]
509    fn parse_pert_distance_variants() {
510        assert!(matches!(
511            parse_perturbation_distance("euclidean"),
512            Some(PertDistance::Euclidean)
513        ));
514        assert!(matches!(
515            parse_perturbation_distance("Euclidean"),
516            Some(PertDistance::Euclidean)
517        ));
518        assert!(matches!(
519            parse_perturbation_distance("squared_euclidean"),
520            Some(PertDistance::SquaredEuclidean)
521        ));
522        assert!(matches!(
523            parse_perturbation_distance("sqeuclidean"),
524            Some(PertDistance::SquaredEuclidean)
525        ));
526        assert!(parse_perturbation_distance("manhattan").is_none());
527        assert!(parse_perturbation_distance("").is_none());
528    }
529
530    /// With no pair to average over the mean is defined as zero, not NaN from a zero divisor.
531    #[test]
532    fn within_lt_two_rows_is_zero() {
533        let single = Mat::<f64>::from_fn(1, 3, |_, j| j as f64);
534        assert_eq!(
535            mean_pairwise_within(single.as_ref(), PertDistance::Euclidean),
536            0.0
537        );
538        let empty = Mat::<f64>::zeros(0, 3);
539        assert_eq!(
540            mean_pairwise_within(empty.as_ref(), PertDistance::Euclidean),
541            0.0
542        );
543    }
544
545    /// Pins the within-group mean against a hand-computed value for both metrics.
546    #[test]
547    fn within_three_points_known() {
548        // Right triangle (0,0), (3,0), (0,4): pairwise 3, 4, 5; mean 4.
549        let m = Mat::<f64>::from_fn(3, 2, |i, j| match (i, j) {
550            (1, 0) => 3.0,
551            (2, 1) => 4.0,
552            _ => 0.0,
553        });
554        let d = mean_pairwise_within(m.as_ref(), PertDistance::Euclidean);
555        assert!(close(d, 4.0));
556        let d_sq = mean_pairwise_within(m.as_ref(), PertDistance::SquaredEuclidean);
557        assert!(close(d_sq, (9.0 + 16.0 + 25.0) / 3.0));
558    }
559
560    /// Mismatched feature counts error rather than reading past the shorter row.
561    #[test]
562    fn between_dim_mismatch_errors() {
563        let x = Mat::<f64>::zeros(2, 3);
564        let y = Mat::<f64>::zeros(2, 4);
565        assert!(matches!(
566            mean_pairwise_between(x.as_ref(), y.as_ref(), PertDistance::Euclidean),
567            Err(CoreErrors::NonMatchingFeatureDim { .. })
568        ));
569    }
570
571    /// An empty group on one side yields zero instead of dividing by an empty product.
572    #[test]
573    fn between_empty_returns_zero() {
574        let x = Mat::<f64>::zeros(0, 2);
575        let y = Mat::<f64>::from_fn(2, 2, |i, _| i as f64);
576        assert_eq!(
577            mean_pairwise_between(x.as_ref(), y.as_ref(), PertDistance::Euclidean).unwrap(),
578            0.0
579        );
580    }
581
582    /// Pins the between-group mean against a hand-computed value.
583    #[test]
584    fn between_known_value() {
585        // X = {(0,0)}, Y = {(3,4), (6,8)}. Distances 5, 10; mean 7.5.
586        let x = Mat::<f64>::zeros(1, 2);
587        let y = Mat::<f64>::from_fn(2, 2, |i, j| match (i, j) {
588            (0, 0) => 3.0,
589            (0, 1) => 4.0,
590            (1, 0) => 6.0,
591            (1, 1) => 8.0,
592            _ => 0.0,
593        });
594        let d = mean_pairwise_between(x.as_ref(), y.as_ref(), PertDistance::Euclidean).unwrap();
595        assert!(close(d, 7.5));
596    }
597
598    /// The energy distance is symmetric in its two groups despite the asymmetric argument order.
599    #[test]
600    fn edistance_symmetric() {
601        let x = Mat::<f64>::from_fn(4, 3, |i, j| (i + j) as f64);
602        let y = Mat::<f64>::from_fn(5, 3, |i, j| (i * 2 + j) as f64 - 1.0);
603        let exy = edistance_two_matrices(x.as_ref(), y.as_ref(), PertDistance::Euclidean).unwrap();
604        let eyx = edistance_two_matrices(y.as_ref(), x.as_ref(), PertDistance::Euclidean).unwrap();
605        assert!(close(exy, eyx));
606    }
607
608    /// The one closed form the energy distance has, so it pins the constant factor.
609    #[test]
610    fn edistance_singletons_squared_identity() {
611        // Singletons: within = 0, so E = 2*||mu_x - mu_y||^2 exactly.
612        let x = Mat::<f64>::zeros(1, 2);
613        let y = Mat::<f64>::from_fn(1, 2, |_, j| if j == 0 { 3.0 } else { 4.0 });
614        let e =
615            edistance_two_matrices(x.as_ref(), y.as_ref(), PertDistance::SquaredEuclidean).unwrap();
616        assert!(close(e, 50.0));
617    }
618
619    /// The pairwise matrix is group by group, symmetric, and zero on the diagonal.
620    #[test]
621    fn pairwise_edistance_shape_and_symmetry() {
622        let emb = Mat::<f64>::from_fn(6, 2, |i, j| (i + j) as f64);
623        let labels = vec![0, 0, 1, 1, 2, 2];
624        let m = pairwise_edistance(emb.as_ref(), &labels, PertDistance::Euclidean).unwrap();
625        assert_eq!(m.nrows(), 3);
626        assert_eq!(m.ncols(), 3);
627        for i in 0..3 {
628            assert_eq!(*m.get(i, i), 0.0);
629            for j in (i + 1)..3 {
630                assert!(close(*m.get(i, j), *m.get(j, i)));
631            }
632        }
633    }
634
635    /// Fewer labels than rows errors rather than silently dropping the tail.
636    #[test]
637    fn pairwise_edistance_label_count_mismatch() {
638        let emb = Mat::<f64>::zeros(4, 2);
639        let labels = vec![0_usize, 0, 1];
640        assert!(matches!(
641            pairwise_edistance(emb.as_ref(), &labels, PertDistance::Euclidean),
642            Err(CoreErrors::NumberLabelsNotEqualSampleNumber { .. })
643        ));
644    }
645
646    /// The reference group keeps its slot in the output and scores zero against itself.
647    #[test]
648    fn onesided_reference_entry_is_zero() {
649        let emb = Mat::<f64>::from_fn(6, 2, |i, j| (i + j) as f64);
650        let labels = vec![0_usize, 0, 1, 1, 2, 2];
651        let res = onesided_edistance(emb.as_ref(), &labels, 1, PertDistance::Euclidean).unwrap();
652        assert_eq!(res.len(), 3);
653        assert_eq!(res[1], 0.0);
654    }
655
656    /// The cheap one-sided path must agree with the column it replaces in the full matrix.
657    #[test]
658    fn onesided_matches_pairwise_column() {
659        let emb = Mat::<f64>::from_fn(8, 3, |i, j| ((i + 1) * (j + 1)) as f64);
660        let labels = vec![0_usize, 0, 0, 1, 1, 2, 2, 2];
661        let pw = pairwise_edistance(emb.as_ref(), &labels, PertDistance::Euclidean).unwrap();
662        let os = onesided_edistance(emb.as_ref(), &labels, 0, PertDistance::Euclidean).unwrap();
663        for g in 0..3 {
664            assert!(close(*pw.get(g, 0), os[g]));
665        }
666    }
667}
bixverse_core/core/base/pert_dist.rs

bixverse_core/core/base/
pert_dist.rs
