bixverse_methods/methods/
rbh.rs

1//! Implementations of the reciprocal best hit (RBH) methods for intersections
2//! or correlations, based on the works of Cantini, et al., Bioinformatics,
3//! 2019
4
5use faer::{Mat, unzip, zip};
6use rayon::prelude::*;
7use rustc_hash::FxHashSet;
8use std::collections::BTreeMap;
9
10use crate::prelude::*;
11use bixverse_core::core::base::cors_similarity::{cor_two_matrices, set_similarity};
12use bixverse_core::core::mat_struct::NamedMatrix;
13use bixverse_core::utils::vec_utils::{array_max, flatten_vector};
14
15/////////////
16// Helpers //
17/////////////
18
19/// Identify the k-th largest values
20///
21/// ### Params
22///
23/// * `values` - The slice of numerics to look at
24/// * `k` - The k-th values to check
25///
26/// ### Returns
27///
28/// The k-th largest value
29fn kth_largest<T: BixverseFloat>(values: &[T], k: usize) -> T {
30    let mut sorted: Vec<T> = values.to_vec();
31    sorted.sort_by(|a, b| b.partial_cmp(a).unwrap());
32    let idx = k.saturating_sub(1).min(sorted.len().saturating_sub(1));
33    sorted[idx]
34}
35
36/// Helper to get the names by index order
37///
38/// ### Params
39///
40/// * `names` - The BTreeMap containing string to index information
41///
42/// ### Returns
43///
44/// The names by index order
45fn names_by_index(names: &BTreeMap<String, usize>) -> Vec<&str> {
46    let mut pairs: Vec<(&str, usize)> = names.iter().map(|(k, &v)| (k.as_str(), v)).collect();
47    pairs.sort_by_key(|(_, idx)| *idx);
48    pairs.into_iter().map(|(k, _)| k).collect()
49}
50
51////////////////
52// Structures //
53////////////////
54
55/// Structure for an Rbh triplet Result
56#[derive(Clone, Debug)]
57pub struct RbhTripletStruc<'a, T> {
58    /// Name of module 1 of the RBH hit
59    pub t1: &'a str,
60    /// Name of module 2 of the RBH hit
61    pub t2: &'a str,
62    /// Similarity value between the two
63    pub sim: T,
64}
65
66/// Structure to store the RBH results.
67#[derive(Clone, Debug)]
68#[allow(dead_code)]
69pub struct RbhResult<T> {
70    /// Name of the origin data set
71    pub origin: String,
72    /// Name of the target data set
73    pub target: String,
74    /// Names of the origin modules/gene sets
75    pub origin_modules: Vec<String>,
76    /// Names of the target modules/gene sets
77    pub target_modules: Vec<String>,
78    /// Similarities between the modules/gene sets
79    pub similarities: Vec<T>,
80}
81
82////////////////////
83// Set similarity //
84////////////////////
85
86/// Calculates the reciprocal best hits based on set similarities.
87///
88/// Function will calculate the set similarities (Jaccard or Overlap
89/// coefficient) between all of the gene sets between the two data sets and
90/// calculate the reciprocal best hits based on this similarity matrix.
91/// Option to return reciprocal k-best hits.
92///
93/// ### Params
94///
95/// * `origin_modules` - A BTreeMap containing the identified modules of the
96///   the origin data set.
97/// * `target_modules` - A BTreeMap containing the identified modules of the
98///   the target data set.
99/// * `k` - Top k hits to consider. If `k=1` this becomes the reciprocal best
100///   hit.
101/// * `overlap_coefficient` - Shall the overlap coefficient be used instead of
102///   Jaccard similarity.
103/// * `min_similarity` - Minimum similarity to be returned
104///
105/// ### Returns
106///
107/// A vector of `RbhTripletStruc`.
108pub fn calculate_rbh_set<'a, T>(
109    origin_modules: &'a BTreeMap<String, FxHashSet<String>>,
110    target_modules: &'a BTreeMap<String, FxHashSet<String>>,
111    k: usize,
112    overlap_coefficient: bool,
113    min_similarity: T,
114) -> Vec<RbhTripletStruc<'a, T>>
115where
116    T: BixverseFloat,
117{
118    let names_targets: Vec<&String> = target_modules.keys().collect();
119    let names_origin: Vec<&String> = origin_modules.keys().collect();
120
121    let origin_sets: Vec<&FxHashSet<String>> = origin_modules.values().collect();
122    let target_sets: Vec<&FxHashSet<String>> = target_modules.values().collect();
123
124    let similarities_flat: Vec<Vec<T>> = origin_sets
125        .into_par_iter()
126        .map(|v1| {
127            target_sets
128                .iter()
129                .map(|v2| set_similarity(v1, v2, overlap_coefficient))
130                .collect()
131        })
132        .collect();
133
134    let mat_data: Vec<T> = flatten_vector(similarities_flat);
135
136    let max_sim = array_max(&mat_data);
137
138    if max_sim < min_similarity {
139        vec![RbhTripletStruc {
140            t1: "NA",
141            t2: "NA",
142            sim: T::zero(),
143        }]
144    } else {
145        let nrow = names_origin.len();
146        let ncol = names_targets.len();
147
148        let sim_mat = Mat::from_fn(nrow, ncol, |i, j| mat_data[j + i * ncol]);
149
150        let row_thresholds: Vec<T> = sim_mat
151            .row_iter()
152            .map(|x| {
153                let row: Vec<T> = x.iter().copied().collect();
154                kth_largest(&row, k)
155            })
156            .collect();
157
158        let col_thresholds: Vec<T> = sim_mat
159            .col_iter()
160            .map(|x| {
161                let row: Vec<T> = x.iter().copied().collect();
162                kth_largest(&row, k)
163            })
164            .collect();
165
166        let mut matching_pairs: Vec<RbhTripletStruc<T>> = Vec::new();
167
168        for r in 0..nrow {
169            for c in 0..ncol {
170                let value = sim_mat[(r, c)];
171                if value >= row_thresholds[r] && value >= col_thresholds[c] {
172                    let triplet = RbhTripletStruc {
173                        t1: names_origin[r],
174                        t2: names_targets[c],
175                        sim: value,
176                    };
177
178                    matching_pairs.push(triplet)
179                }
180            }
181        }
182
183        if !matching_pairs.is_empty() {
184            matching_pairs
185        } else {
186            vec![RbhTripletStruc {
187                t1: "NA",
188                t2: "NA",
189                sim: T::zero(),
190            }]
191        }
192    }
193}
194
195///////////////////////
196// Correlation based //
197///////////////////////
198
199/// Calculate the RBH based on correlation of two NamedMatrices
200///
201/// The function will intersect into shared features and calculate the
202/// correlation matrix and subsequently reciprocal best hits based on the
203/// absolute correlation. Option to return reciprocal k-best hits.
204///
205/// ### Params
206///
207/// * `x1` - `NamedMatrix` of the origin data
208/// * `x2` - `NamedMatrix` of the target data
209/// * `k` - Top k hits to consider. If `k=1` this becomes the reciprocal best
210///   hit.
211/// * `spearman` - Shall Spearman correlations be used.
212///
213/// ### Returns
214///
215/// A vector of `RbhTripletStruc`.
216pub fn calculate_rbh_cor<'a, T>(
217    x1: &'a NamedMatrix<'a, T>,
218    x2: &'a NamedMatrix<'a, T>,
219    k: usize,
220    spearman: bool,
221) -> Vec<RbhTripletStruc<'a, T>>
222where
223    T: BixverseFloat,
224{
225    let row_names_1: FxHashSet<String> = x1.row_names.keys().cloned().collect();
226    let row_names_2: FxHashSet<String> = x2.row_names.keys().cloned().collect();
227
228    let names_targets: Vec<&str> = names_by_index(&x1.col_names);
229    let names_origin: Vec<&str> = names_by_index(&x2.col_names);
230
231    let intersecting_rows: Vec<String> = row_names_1.intersection(&row_names_2).cloned().collect();
232
233    // early return if there are no intersecting rows
234    if intersecting_rows.is_empty() {
235        return vec![RbhTripletStruc {
236            t1: "NA",
237            t2: "NA",
238            sim: T::zero(),
239        }];
240    }
241
242    let row_refs: Vec<&str> = intersecting_rows.iter().map(|s| s.as_str()).collect();
243    let sub_x1 = x1.get_rows(&row_refs).unwrap();
244    let sub_x2 = x2.get_rows(&row_refs).unwrap();
245
246    let mut correlations = cor_two_matrices(&sub_x1.as_ref(), &sub_x2.as_ref(), spearman);
247
248    zip!(correlations.as_mut()).for_each(|unzip!(x)| *x = x.abs());
249
250    let row_thresholds: Vec<T> = correlations
251        .row_iter()
252        .map(|x| {
253            let row: Vec<T> = x.iter().copied().collect();
254            kth_largest(&row, k)
255        })
256        .collect();
257
258    let col_thresholds: Vec<T> = correlations
259        .col_iter()
260        .map(|x| {
261            let col: Vec<T> = x.iter().copied().collect();
262            kth_largest(&col, k)
263        })
264        .collect();
265
266    let mut matching_pairs: Vec<RbhTripletStruc<'a, T>> = Vec::new();
267    let nrow = names_targets.len();
268    let ncol = names_origin.len();
269
270    for r in 0..nrow {
271        for c in 0..ncol {
272            let value = correlations[(r, c)];
273            if value >= row_thresholds[r] && value >= col_thresholds[c] {
274                let triplet = RbhTripletStruc {
275                    t1: names_targets[r],
276                    t2: names_origin[c],
277                    sim: value,
278                };
279                matching_pairs.push(triplet);
280            }
281        }
282    }
283
284    if !matching_pairs.is_empty() {
285        matching_pairs
286    } else {
287        vec![RbhTripletStruc {
288            t1: "NA",
289            t2: "NA",
290            sim: T::zero(),
291        }]
292    }
293}
294
295///////////
296// Tests //
297///////////
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    /////////
304    // Set //
305    /////////
306
307    fn module(items: &[&str]) -> FxHashSet<String> {
308        items.iter().map(|s| s.to_string()).collect()
309    }
310
311    fn modules(pairs: &[(&str, &[&str])]) -> BTreeMap<String, FxHashSet<String>> {
312        pairs
313            .iter()
314            .map(|(n, it)| (n.to_string(), module(it)))
315            .collect()
316    }
317
318    /// Helper used by the k-best tests below.
319    ///
320    /// Built so that the resulting Jaccard matrix is:
321    /// ```text
322    ///        X       Y      Z
323    /// A    1.0    0.091   0.0
324    /// B   0.333    0.2   0.333
325    /// C    0.0    0.333   1.0
326    /// ```
327    /// At k = 1 only (A,X) and (C,Z) are reciprocal; at k = 2 the set grows
328    /// to include (B,X), (B,Z) and (C,Y).
329    fn three_by_three_modules() -> (
330        BTreeMap<String, FxHashSet<String>>,
331        BTreeMap<String, FxHashSet<String>>,
332    ) {
333        let origin = modules(&[
334            ("A", &["1", "2", "3", "4", "5", "6"]),
335            ("B", &["4", "5", "6", "7", "8", "9"]),
336            ("C", &["7", "8", "9", "10", "11", "12"]),
337        ]);
338        let target = modules(&[
339            ("X", &["1", "2", "3", "4", "5", "6"]),
340            ("Y", &["6", "7", "10", "11", "13", "14"]),
341            ("Z", &["7", "8", "9", "10", "11", "12"]),
342        ]);
343        (origin, target)
344    }
345
346    fn pair_set<T>(result: &[RbhTripletStruc<T>]) -> FxHashSet<(String, String)> {
347        result
348            .iter()
349            .map(|r| (r.t1.to_string(), r.t2.to_string()))
350            .collect()
351    }
352
353    /// Perfectly matching module pairs come back as reciprocal best hits at similarity 1.
354    #[test]
355    fn k1_identifies_clear_reciprocal_hits() {
356        let origin = modules(&[("A", &["g1", "g2", "g3"]), ("B", &["g4", "g5", "g6"])]);
357        let target = modules(&[("X", &["g1", "g2", "g3"]), ("Y", &["g4", "g5", "g6"])]);
358
359        let result = calculate_rbh_set(&origin, &target, 1, false, 0.0_f64);
360
361        assert_eq!(result.len(), 2);
362        let pairs = pair_set(&result);
363        assert!(pairs.contains(&("A".into(), "X".into())));
364        assert!(pairs.contains(&("B".into(), "Y".into())));
365        for r in &result {
366            assert!((r.sim - 1.0).abs() < 1e-12);
367        }
368    }
369
370    /// A hit that is best in its row but not in its column must not survive k = 1.
371    #[test]
372    fn k1_excludes_non_reciprocal_hits() {
373        let (origin, target) = three_by_three_modules();
374
375        let result = calculate_rbh_set(&origin, &target, 1, false, 0.0_f64);
376
377        let pairs = pair_set(&result);
378        let expected: FxHashSet<(String, String)> =
379            [("A".into(), "X".into()), ("C".into(), "Z".into())]
380                .into_iter()
381                .collect();
382        assert_eq!(pairs, expected);
383    }
384
385    /// Raising k to 2 admits the second-best hits that k = 1 discards.
386    #[test]
387    fn k2_includes_pairs_missed_by_k1() {
388        let (origin, target) = three_by_three_modules();
389
390        let result = calculate_rbh_set(&origin, &target, 2, false, 0.0_f64);
391
392        let pairs = pair_set(&result);
393        let expected: FxHashSet<(String, String)> = [
394            ("A".into(), "X".into()),
395            ("B".into(), "X".into()),
396            ("B".into(), "Z".into()),
397            ("C".into(), "Y".into()),
398            ("C".into(), "Z".into()),
399        ]
400        .into_iter()
401        .collect();
402        assert_eq!(pairs, expected);
403    }
404
405    /// A k at or above both dimensions thresholds nothing, so every pair is returned.
406    #[test]
407    fn k_at_least_as_large_as_dim_returns_all_pairs() {
408        let (origin, target) = three_by_three_modules();
409
410        let result = calculate_rbh_set(&origin, &target, 5, false, 0.0_f64);
411
412        assert_eq!(result.len(), 9);
413    }
414
415    /// Nothing clearing `min_similarity` yields the single NA placeholder, not an empty vector.
416    #[test]
417    fn returns_na_when_max_below_threshold() {
418        let origin = modules(&[("A", &["g1", "g2"])]);
419        let target = modules(&[("X", &["g3", "g4"])]);
420
421        let result = calculate_rbh_set(&origin, &target, 1, false, 0.5_f64);
422
423        assert_eq!(result.len(), 1);
424        assert_eq!(result[0].t1, "NA");
425        assert_eq!(result[0].t2, "NA");
426        assert_eq!(result[0].sim, 0.0);
427    }
428
429    /// The overlap coefficient flag switches metric: a strict subset scores 1.0, not 0.25.
430    #[test]
431    fn overlap_coefficient_differs_from_jaccard() {
432        let origin = modules(&[("A", &["g1"])]);
433        let target = modules(&[("X", &["g1", "g2", "g3", "g4"])]);
434
435        let jaccard = calculate_rbh_set(&origin, &target, 1, false, 0.0_f64);
436        let overlap = calculate_rbh_set(&origin, &target, 1, true, 0.0_f64);
437
438        assert!((jaccard[0].sim - 0.25).abs() < 1e-12);
439        assert!((overlap[0].sim - 1.0).abs() < 1e-12);
440    }
441
442    //////////////////
443    // Correlations //
444    //////////////////
445
446    fn mat_from_rows(rows: &[&[f64]]) -> Mat<f64> {
447        let nrow = rows.len();
448        let ncol = rows[0].len();
449        Mat::from_fn(nrow, ncol, |i, j| rows[i][j])
450    }
451
452    fn build_named<'a>(
453        values: &'a Mat<f64>,
454        row_names: &[&str],
455        col_names: &[&str],
456    ) -> NamedMatrix<'a, f64> {
457        NamedMatrix {
458            col_names: col_names
459                .iter()
460                .enumerate()
461                .map(|(i, s)| (s.to_string(), i))
462                .collect(),
463            row_names: row_names
464                .iter()
465                .enumerate()
466                .map(|(i, s)| (s.to_string(), i))
467                .collect(),
468            values: values.as_ref(),
469        }
470    }
471
472    /// Two named matrices designed so the Pearson correlation matrix after `abs()` is:
473    /// ```text
474    ///         X       Y
475    /// A     1.000   0.944
476    /// B     0.944   1.000
477    /// ```
478    /// At k = 1 only (A,X) and (B,Y) are reciprocal; at k = 2 every entry passes both
479    /// thresholds (the 2nd-largest in each row and column is 0.944), so all 4 pairs return.
480    fn two_by_two_cor_setup() -> (Mat<f64>, Mat<f64>) {
481        let m1 = mat_from_rows(&[&[1.0, 1.0], &[2.0, 1.0], &[3.0, 2.0], &[4.0, 3.0]]);
482        let m2 = m1.clone();
483        (m1, m2)
484    }
485
486    /// Identical inputs give a diagonal of 1.0, so k = 1 returns exactly the matched column pairs.
487    #[test]
488    fn cor_k1_identifies_perfect_pairs() {
489        let (m1, m2) = two_by_two_cor_setup();
490        let x1 = build_named(&m1, &["g1", "g2", "g3", "g4"], &["A", "B"]);
491        let x2 = build_named(&m2, &["g1", "g2", "g3", "g4"], &["X", "Y"]);
492
493        let result = calculate_rbh_cor(&x1, &x2, 1, false);
494
495        let pairs: FxHashSet<(String, String)> = result
496            .iter()
497            .map(|r| (r.t1.to_string(), r.t2.to_string()))
498            .collect();
499        let expected: FxHashSet<(String, String)> =
500            [("A".into(), "X".into()), ("B".into(), "Y".into())]
501                .into_iter()
502                .collect();
503        assert_eq!(pairs, expected);
504
505        for r in &result {
506            assert!((r.sim - 1.0).abs() < 1e-10);
507        }
508    }
509
510    /// At k = 2 every entry of a 2x2 correlation matrix clears both thresholds.
511    #[test]
512    fn cor_k2_returns_all_pairs_in_2x2() {
513        let (m1, m2) = two_by_two_cor_setup();
514        let x1 = build_named(&m1, &["g1", "g2", "g3", "g4"], &["A", "B"]);
515        let x2 = build_named(&m2, &["g1", "g2", "g3", "g4"], &["X", "Y"]);
516
517        let result = calculate_rbh_cor(&x1, &x2, 2, false);
518
519        assert_eq!(result.len(), 4);
520    }
521
522    /// Perfectly anticorrelated columns still pair, since the correlation matrix is taken absolute.
523    #[test]
524    fn cor_uses_absolute_correlation() {
525        // m2 = -m1, so cor(A,X) = cor(B,Y) = -1; after abs they should still pair.
526        let m1 = mat_from_rows(&[&[1.0, 1.0], &[2.0, 1.0], &[3.0, 2.0], &[4.0, 3.0]]);
527        let m2 = mat_from_rows(&[&[-1.0, -1.0], &[-2.0, -1.0], &[-3.0, -2.0], &[-4.0, -3.0]]);
528        let x1 = build_named(&m1, &["g1", "g2", "g3", "g4"], &["A", "B"]);
529        let x2 = build_named(&m2, &["g1", "g2", "g3", "g4"], &["X", "Y"]);
530
531        let result = calculate_rbh_cor(&x1, &x2, 1, false);
532
533        let pairs: FxHashSet<(String, String)> = result
534            .iter()
535            .map(|r| (r.t1.to_string(), r.t2.to_string()))
536            .collect();
537        assert!(pairs.contains(&("A".into(), "X".into())));
538        assert!(pairs.contains(&("B".into(), "Y".into())));
539
540        for r in &result {
541            let p = (r.t1, r.t2);
542            if p == ("A", "X") || p == ("B", "Y") {
543                assert!((r.sim - 1.0).abs() < 1e-10);
544            }
545        }
546    }
547
548    /// Matrices with no shared row names return the NA placeholder rather than correlating nothing.
549    #[test]
550    fn cor_no_row_intersection_returns_na() {
551        let m1 = mat_from_rows(&[&[1.0], &[2.0]]);
552        let m2 = mat_from_rows(&[&[1.0], &[2.0]]);
553        let x1 = build_named(&m1, &["g1", "g2"], &["A"]);
554        let x2 = build_named(&m2, &["g3", "g4"], &["X"]);
555
556        let result = calculate_rbh_cor(&x1, &x2, 1, false);
557
558        assert_eq!(result.len(), 1);
559        assert_eq!(result[0].t1, "NA");
560        assert_eq!(result[0].t2, "NA");
561        assert_eq!(result[0].sim, 0.0);
562    }
563
564    /// Rows in only one matrix are dropped before correlating, so their values cannot leak in.
565    #[test]
566    fn cor_uses_only_intersecting_rows() {
567        // x1 has g1..g4; x2 has g2..g5. Shared rows are {g2, g3, g4}.
568        // On those shared rows: A=[2,3,4], B=[4,2,3], X=A, Y=B.
569        // After abs: cor matrix is [[1.0, 0.5],[0.5, 1.0]], so k=1 yields (A,X),(B,Y).
570        // g1 and g5 carry junk values that would distort the cors if they leaked in.
571        let m1 = mat_from_rows(&[&[1.0, 1.0], &[2.0, 4.0], &[3.0, 2.0], &[4.0, 3.0]]);
572        let m2 = mat_from_rows(&[&[2.0, 4.0], &[3.0, 2.0], &[4.0, 3.0], &[99.0, -99.0]]);
573        let x1 = build_named(&m1, &["g1", "g2", "g3", "g4"], &["A", "B"]);
574        let x2 = build_named(&m2, &["g2", "g3", "g4", "g5"], &["X", "Y"]);
575
576        let result = calculate_rbh_cor(&x1, &x2, 1, false);
577
578        let pairs: FxHashSet<(String, String)> = result
579            .iter()
580            .map(|r| (r.t1.to_string(), r.t2.to_string()))
581            .collect();
582        let expected: FxHashSet<(String, String)> =
583            [("A".into(), "X".into()), ("B".into(), "Y".into())]
584                .into_iter()
585                .collect();
586        assert_eq!(pairs, expected);
587    }
588
589    /// Columns are resolved by stored index, not by the BTreeMap's alphabetical key order.
590    #[test]
591    fn cor_respects_non_alphabetical_column_order() {
592        let m1 = mat_from_rows(&[&[1.0, 1.0], &[2.0, 1.0], &[3.0, 2.0], &[4.0, 3.0]]);
593        let m2 = mat_from_rows(&[&[1.0, 1.0], &[1.0, 2.0], &[2.0, 3.0], &[3.0, 4.0]]);
594
595        let x1 = NamedMatrix {
596            col_names: [("Z".to_string(), 0), ("A".to_string(), 1)]
597                .into_iter()
598                .collect(),
599            row_names: ["g1", "g2", "g3", "g4"]
600                .iter()
601                .enumerate()
602                .map(|(i, s)| (s.to_string(), i))
603                .collect(),
604            values: m1.as_ref(),
605        };
606        let x2 = NamedMatrix {
607            col_names: [("A".to_string(), 0), ("Z".to_string(), 1)]
608                .into_iter()
609                .collect(),
610            row_names: ["g1", "g2", "g3", "g4"]
611                .iter()
612                .enumerate()
613                .map(|(i, s)| (s.to_string(), i))
614                .collect(),
615            values: m2.as_ref(),
616        };
617
618        let result = calculate_rbh_cor(&x1, &x2, 1, false);
619
620        // correct alignment yields Z<->Z and A<->A, both at correlation 1.0.
621        let pairs: FxHashSet<(String, String)> = result
622            .iter()
623            .map(|r| (r.t1.to_string(), r.t2.to_string()))
624            .collect();
625        assert!(pairs.contains(&("Z".into(), "Z".into())));
626        assert!(pairs.contains(&("A".into(), "A".into())));
627        for r in &result {
628            assert!((r.sim - 1.0).abs() < 1e-10);
629        }
630    }
631}
bixverse_methods/methods/rbh.rs

bixverse_methods/methods/
rbh.rs
