1use std::sync::Arc;
30
31use crate::sparse_io_vector::RowNameCanonicalizer;
32use legume_numeric::matrix::membership::canon_locus;
33use rustc_hash::FxHashMap as HashMap;
34
35#[derive(Clone, Debug, Default, PartialEq, Eq)]
40pub enum FeatureNameKind {
41 #[default]
43 Exact,
44 Gene { delim: char },
48 Locus { merge_overlapping: bool },
55 Mixed,
61}
62
63impl FeatureNameKind {
64 pub fn canonicalize(&self, name: &str) -> Box<str> {
71 match self {
72 FeatureNameKind::Exact => name.into(),
73 FeatureNameKind::Gene { delim } => gene_canonicalize(name, *delim),
74 FeatureNameKind::Locus { .. } => canon_locus(name),
75 FeatureNameKind::Mixed => {
76 if parse_locus(name).is_some() {
77 canon_locus(name)
78 } else if name.contains('_') {
79 gene_canonicalize(name, '_')
80 } else {
81 name.into()
82 }
83 }
84 }
85 }
86
87 pub fn is_exact(&self) -> bool {
89 matches!(self, FeatureNameKind::Exact)
90 }
91
92 pub fn needs_global_pass(&self) -> bool {
96 matches!(
97 self,
98 FeatureNameKind::Locus {
99 merge_overlapping: true
100 } | FeatureNameKind::Mixed
101 )
102 }
103
104 pub fn auto_detect(names: &[Box<str>]) -> Self {
112 let n = names.len();
113 if n == 0 {
114 return Self::Exact;
115 }
116 let mut n_locus = 0usize;
117 let mut n_gene_like = 0usize;
118 for name in names {
119 if parse_locus(name).is_some() {
120 n_locus += 1;
121 } else if name.contains('_') {
122 n_gene_like += 1;
123 }
124 }
125 let pct_locus = n_locus as f32 / n as f32;
126 let pct_gene = n_gene_like as f32 / n as f32;
127 if pct_locus >= 0.10 && pct_gene >= 0.10 {
128 Self::Mixed
129 } else if pct_locus >= 0.50 {
130 Self::Locus {
131 merge_overlapping: true,
132 }
133 } else if pct_gene >= 0.50 {
134 Self::Gene { delim: '_' }
135 } else {
136 Self::Exact
137 }
138 }
139
140 #[must_use]
155 pub fn reconcile(kinds: &[FeatureNameKind]) -> FeatureNameKind {
156 if kinds.iter().any(|k| matches!(k, FeatureNameKind::Mixed)) {
157 return FeatureNameKind::Mixed;
158 }
159 let gene = kinds
160 .iter()
161 .find(|k| matches!(k, FeatureNameKind::Gene { .. }));
162 let locus = kinds
163 .iter()
164 .find(|k| matches!(k, FeatureNameKind::Locus { .. }));
165 match (gene, locus) {
166 (Some(_), Some(_)) => FeatureNameKind::Mixed,
167 _ => gene.or(locus).cloned().unwrap_or(FeatureNameKind::Exact),
168 }
169 }
170
171 pub fn into_canonicalizer(self) -> Option<RowNameCanonicalizer> {
177 if self.is_exact() {
178 return None;
179 }
180 Some(Arc::new(move |name: &str| self.canonicalize(name)))
181 }
182}
183
184pub fn parse_locus(name: &str) -> Option<(Box<str>, u64, u64)> {
189 let lower = name.to_ascii_lowercase();
190 let stripped = lower
191 .strip_prefix("chr")
192 .map(str::to_string)
193 .unwrap_or(lower);
194 let parts: Vec<&str> = stripped.splitn(3, [':', '-', '_']).collect();
196 if parts.len() != 3 {
197 return None;
198 }
199 let chr = parts[0];
200 let start: u64 = parts[1].parse().ok()?;
201 let end: u64 = parts[2].parse().ok()?;
202 if end < start {
203 return None;
204 }
205 Some((chr.to_string().into_boxed_str(), start, end))
206}
207
208pub fn build_locus_overlap_canonical_map(names: &[Box<str>]) -> HashMap<Box<str>, Box<str>> {
219 let n = names.len();
220 let parsed: Vec<Option<(Box<str>, u64, u64)>> = names.iter().map(|n| parse_locus(n)).collect();
221
222 let mut by_chr: HashMap<Box<str>, Vec<usize>> = HashMap::default();
224 for (i, p) in parsed.iter().enumerate() {
225 if let Some((chr, _, _)) = p {
226 by_chr.entry(chr.clone()).or_default().push(i);
227 }
228 }
229
230 let mut parent: Vec<usize> = (0..n).collect();
232 fn find(p: &mut [usize], mut x: usize) -> usize {
233 while p[x] != x {
234 let g = p[p[x]];
235 p[x] = g;
236 x = g;
237 }
238 x
239 }
240
241 let mut cluster_extent: HashMap<usize, (u64, u64)> = HashMap::default();
243 for (_, mut idxs) in by_chr {
244 idxs.sort_by_key(|&i| parsed[i].as_ref().map(|p| p.1).unwrap_or(0));
245 let mut current_root: Option<usize> = None;
246 let mut current_min_start: u64 = 0;
247 let mut current_max_end: u64 = 0;
248 for i in idxs {
249 let (_, s, e) = parsed[i].as_ref().unwrap();
250 match current_root {
251 Some(root) if *s < current_max_end => {
252 let ra = find(&mut parent, root);
253 let rb = find(&mut parent, i);
254 if ra != rb {
255 parent[rb] = ra;
256 }
257 current_max_end = current_max_end.max(*e);
258 cluster_extent
259 .insert(find(&mut parent, i), (current_min_start, current_max_end));
260 }
261 _ => {
262 current_root = Some(i);
263 current_min_start = *s;
264 current_max_end = *e;
265 cluster_extent.insert(i, (*s, *e));
266 }
267 }
268 }
269 }
270
271 let mut out: HashMap<Box<str>, Box<str>> = HashMap::default();
274 for (i, p) in parsed.iter().enumerate() {
275 if let Some((chr, _, _)) = p {
276 let root = find(&mut parent, i);
277 let (mn, mx) = cluster_extent.get(&root).copied().unwrap_or((0, 0));
278 let canonical = format!("{}_{}_{}", chr, mn, mx).into_boxed_str();
279 out.insert(names[i].clone(), canonical);
280 }
281 }
282 out
283}
284
285pub fn build_locus_overlap_canonicalizer(names: &[Box<str>]) -> RowNameCanonicalizer {
291 let map = Arc::new(build_locus_overlap_canonical_map(names));
292 Arc::new(move |name: &str| map.get(name).cloned().unwrap_or_else(|| canon_locus(name)))
293}
294
295pub fn build_mixed_kind_canonicalizer(names: &[Box<str>]) -> RowNameCanonicalizer {
305 let map = Arc::new(build_locus_overlap_canonical_map(names));
306 Arc::new(move |name: &str| {
307 if let Some(c) = map.get(name) {
308 c.clone()
309 } else if parse_locus(name).is_some() {
310 canon_locus(name)
311 } else if name.contains('_') {
312 gene_canonicalize(name, '_')
313 } else {
314 name.into()
315 }
316 })
317}
318
319fn gene_canonicalize(name: &str, delim: char) -> Box<str> {
328 let stripped = strip_feature_type_suffix(name, delim);
329 stripped.rsplit(delim).next().unwrap_or(stripped).into()
330}
331
332fn strip_feature_type_suffix(name: &str, delim: char) -> &str {
338 const TAGS: &[&str] = &[
343 "Gene_Expression",
344 "Gene",
345 "Antibody_Capture",
346 "CRISPR_Guide_Capture",
347 "Multiplexing_Capture",
348 "Custom",
349 "Peaks",
350 ];
351 for tag in TAGS {
352 let mut suffix = String::with_capacity(tag.len() + 1);
355 suffix.push(delim);
356 suffix.push_str(tag);
357 if let Some(rest) = name.strip_suffix(suffix.as_str()) {
358 return rest;
359 }
360 }
361 name
362}
363
364#[derive(clap::ValueEnum, Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
371#[serde(rename_all = "kebab-case")]
372pub enum FeatureNameKindArg {
373 #[default]
374 Auto,
375 Exact,
376 Gene,
377 Locus,
378 LocusOverlap,
379 Mixed,
380}
381
382impl FeatureNameKindArg {
383 pub fn resolve_or_gene(&self) -> FeatureNameKind {
387 Option::<FeatureNameKind>::from(self.clone())
388 .unwrap_or(FeatureNameKind::Gene { delim: '_' })
389 }
390}
391
392impl From<FeatureNameKindArg> for Option<FeatureNameKind> {
393 fn from(arg: FeatureNameKindArg) -> Self {
394 match arg {
395 FeatureNameKindArg::Auto => None,
396 FeatureNameKindArg::Exact => Some(FeatureNameKind::Exact),
397 FeatureNameKindArg::Gene => Some(FeatureNameKind::Gene { delim: '_' }),
398 FeatureNameKindArg::Locus => Some(FeatureNameKind::Locus {
399 merge_overlapping: false,
400 }),
401 FeatureNameKindArg::LocusOverlap => Some(FeatureNameKind::Locus {
402 merge_overlapping: true,
403 }),
404 FeatureNameKindArg::Mixed => Some(FeatureNameKind::Mixed),
405 }
406 }
407}
408
409#[cfg(test)]
410#[path = "feature_names_tests.rs"]
411mod feature_names_tests;
412
413#[cfg(test)]
414mod tests {
415 use super::*;
416
417 #[test]
418 fn exact_passthrough() {
419 let k = FeatureNameKind::Exact;
420 assert_eq!(
421 k.canonicalize("ENSG00000000003_TSPAN6").as_ref(),
422 "ENSG00000000003_TSPAN6"
423 );
424 assert!(k.is_exact());
425 assert!(k.into_canonicalizer().is_none());
426 }
427
428 #[test]
429 fn gene_takes_last_underscore_component() {
430 let k = FeatureNameKind::Gene { delim: '_' };
431 assert_eq!(k.canonicalize("ENSG00000000003_TSPAN6").as_ref(), "TSPAN6");
432 assert_eq!(k.canonicalize("TSPAN6").as_ref(), "TSPAN6");
434 assert_eq!(k.canonicalize("A_B_C").as_ref(), "C");
437 assert!(!k.is_exact());
438 assert!(k.into_canonicalizer().is_some());
439 }
440
441 #[test]
442 fn gene_strips_cell_ranger_feature_type_suffix() {
443 let k = FeatureNameKind::Gene { delim: '_' };
444 assert_eq!(
447 k.canonicalize("ENSG00000187634_SAMD11_Gene").as_ref(),
448 "SAMD11"
449 );
450 assert_eq!(
452 k.canonicalize("ENSG00000187634_SAMD11_Gene_Expression")
453 .as_ref(),
454 "SAMD11"
455 );
456 assert_eq!(k.canonicalize("FakeGene").as_ref(), "FakeGene");
459 }
460
461 #[test]
462 fn locus_strips_chr_and_folds_separators() {
463 let k = FeatureNameKind::Locus {
464 merge_overlapping: false,
465 };
466 assert_eq!(k.canonicalize("chr1:1000-2000").as_ref(), "1_1000_2000");
467 assert_eq!(k.canonicalize("1_1000_2000").as_ref(), "1_1000_2000");
468 assert_eq!(k.canonicalize("ChrX:5000-6000").as_ref(), "X_5000_6000");
469 }
470
471 #[test]
474 fn parse_locus_accepts_common_formats() {
475 assert_eq!(
477 parse_locus("chr1:1000-2000"),
478 Some(("1".into(), 1000, 2000))
479 );
480 assert_eq!(parse_locus("1:1000-2000"), Some(("1".into(), 1000, 2000)));
481 assert_eq!(
482 parse_locus("chr1_1000_2000"),
483 Some(("1".into(), 1000, 2000))
484 );
485 assert_eq!(
486 parse_locus("CHR1:1000-2000"),
487 Some(("1".into(), 1000, 2000))
488 );
489 assert_eq!(
490 parse_locus("chrX:5000-6000"),
491 Some(("x".into(), 5000, 6000))
492 );
493 assert_eq!(parse_locus("chrMT:1-100"), Some(("mt".into(), 1, 100)));
494 }
495
496 #[test]
497 fn parse_locus_rejects_non_loci() {
498 assert!(parse_locus("TGFB1").is_none()); assert!(parse_locus("ENSG00000105329").is_none()); assert!(parse_locus("chr1:bad-2000").is_none()); assert!(parse_locus("chr1:1000").is_none()); assert!(parse_locus("chr1:2000-1000").is_none()); assert!(parse_locus("").is_none()); assert!(parse_locus("chr1").is_none()); }
506
507 #[test]
508 fn overlap_map_merges_two_overlapping_intervals() {
509 let names = vec![
511 "chr1:1-20".to_string().into_boxed_str(),
512 "chr1:15-30".to_string().into_boxed_str(),
513 ];
514 let map = build_locus_overlap_canonical_map(&names);
515 let c0 = map.get(&names[0]).unwrap();
516 let c1 = map.get(&names[1]).unwrap();
517 assert_eq!(c0, c1, "both inputs should map to the same canonical");
518 assert_eq!(c0.as_ref(), "1_1_30"); }
520
521 #[test]
522 fn overlap_map_keeps_non_overlapping_separate() {
523 let names = vec![
524 "chr1:1-20".to_string().into_boxed_str(),
525 "chr1:100-200".to_string().into_boxed_str(),
526 "chr2:1-20".to_string().into_boxed_str(),
527 ];
528 let map = build_locus_overlap_canonical_map(&names);
529 assert_eq!(map.get(&names[0]).unwrap().as_ref(), "1_1_20");
530 assert_eq!(map.get(&names[1]).unwrap().as_ref(), "1_100_200");
531 assert_eq!(map.get(&names[2]).unwrap().as_ref(), "2_1_20");
533 }
534
535 #[test]
536 fn overlap_map_handles_transitive_chain() {
537 let names = vec![
541 "chr1:1-20".to_string().into_boxed_str(),
542 "chr1:15-30".to_string().into_boxed_str(),
543 "chr1:25-40".to_string().into_boxed_str(),
544 ];
545 let map = build_locus_overlap_canonical_map(&names);
546 let c0 = map.get(&names[0]).unwrap();
547 let c1 = map.get(&names[1]).unwrap();
548 let c2 = map.get(&names[2]).unwrap();
549 assert_eq!(c0, c1);
550 assert_eq!(c1, c2);
551 assert_eq!(c0.as_ref(), "1_1_40"); }
553
554 #[test]
555 fn overlap_map_handles_full_containment() {
556 let names = vec![
558 "chr1:1-100".to_string().into_boxed_str(),
559 "chr1:30-50".to_string().into_boxed_str(),
560 ];
561 let map = build_locus_overlap_canonical_map(&names);
562 let c0 = map.get(&names[0]).unwrap();
563 let c1 = map.get(&names[1]).unwrap();
564 assert_eq!(c0, c1);
565 assert_eq!(c0.as_ref(), "1_1_100");
566 }
567
568 #[test]
569 fn overlap_map_treats_adjacent_as_separate() {
570 let names = vec![
573 "chr1:1-20".to_string().into_boxed_str(),
574 "chr1:20-30".to_string().into_boxed_str(),
575 ];
576 let map = build_locus_overlap_canonical_map(&names);
577 assert_ne!(map.get(&names[0]).unwrap(), map.get(&names[1]).unwrap());
578 }
579
580 #[test]
581 fn overlap_map_normalizes_chr_prefix_within_cluster() {
582 let names = vec![
585 "chr1:1-20".to_string().into_boxed_str(),
586 "1:15-30".to_string().into_boxed_str(),
587 ];
588 let map = build_locus_overlap_canonical_map(&names);
589 let c0 = map.get(&names[0]).unwrap();
590 let c1 = map.get(&names[1]).unwrap();
591 assert_eq!(c0, c1);
592 assert_eq!(c0.as_ref(), "1_1_30");
593 }
594
595 #[test]
596 fn overlap_map_normalizes_separators_within_cluster() {
597 let names = vec![
599 "chr1:1-20".to_string().into_boxed_str(),
600 "chr1_15_30".to_string().into_boxed_str(),
601 ];
602 let map = build_locus_overlap_canonical_map(&names);
603 assert_eq!(map.get(&names[0]).unwrap(), map.get(&names[1]).unwrap());
604 }
605
606 #[test]
607 fn overlap_map_ignores_non_locus_names() {
608 let names = vec![
611 "TGFB1".to_string().into_boxed_str(),
612 "chr1:1-20".to_string().into_boxed_str(),
613 ];
614 let map = build_locus_overlap_canonical_map(&names);
615 assert!(!map.contains_key(&names[0]));
616 assert!(map.contains_key(&names[1]));
617 }
618
619 #[test]
620 fn overlap_map_handles_zero_length_interval() {
621 let names = vec!["chr1:1000-1000".to_string().into_boxed_str()];
623 let map = build_locus_overlap_canonical_map(&names);
624 assert_eq!(map.get(&names[0]).unwrap().as_ref(), "1_1000_1000");
625 }
626
627 #[test]
628 fn overlap_canonicalizer_falls_back_to_canon_locus_for_unmatched() {
629 let names = vec!["chr1:1-20".to_string().into_boxed_str()];
630 let canon = build_locus_overlap_canonicalizer(&names);
631 assert_eq!(canon("chr1:1-20").as_ref(), "1_1_20");
633 assert_eq!(canon("chr2:500-600").as_ref(), "2_500_600");
635 let g = canon("TGFB1");
639 assert!(!g.is_empty());
640 }
641
642 #[test]
645 fn auto_detect_pure_locus_axis() {
646 let names: Vec<Box<str>> = (0..100)
647 .map(|i| format!("chr1:{}-{}", i * 100, i * 100 + 50).into_boxed_str())
648 .collect();
649 assert!(matches!(
650 FeatureNameKind::auto_detect(&names),
651 FeatureNameKind::Locus {
652 merge_overlapping: true
653 }
654 ));
655 }
656
657 #[test]
658 fn auto_detect_pure_gene_axis() {
659 let names: Vec<Box<str>> = (0..100)
660 .map(|i| format!("ENSG000_GENE{}", i).into_boxed_str())
661 .collect();
662 assert!(matches!(
663 FeatureNameKind::auto_detect(&names),
664 FeatureNameKind::Gene { delim: '_' }
665 ));
666 }
667
668 #[test]
669 fn auto_detect_mixed_axis() {
670 let mut names: Vec<Box<str>> = (0..80)
672 .map(|i| format!("chr1:{}-{}", i * 1000, i * 1000 + 500).into_boxed_str())
673 .collect();
674 names.extend((0..20).map(|i| format!("ENSG000_GENE{}", i).into_boxed_str()));
675 assert!(matches!(
676 FeatureNameKind::auto_detect(&names),
677 FeatureNameKind::Mixed
678 ));
679 }
680
681 #[test]
682 fn auto_detect_empty_or_exact() {
683 assert!(matches!(
684 FeatureNameKind::auto_detect(&[]),
685 FeatureNameKind::Exact
686 ));
687 let names = vec!["TGFB1".into(), "CD4".into(), "IL2".into(), "GAPDH".into()];
688 assert!(matches!(
689 FeatureNameKind::auto_detect(&names),
690 FeatureNameKind::Exact
691 ));
692 }
693
694 #[test]
695 fn mixed_dispatcher_canonicalizes_each_name_by_kind() {
696 let names: Vec<Box<str>> = vec![
697 "chr1:1-20".into(), "chr1:15-30".into(), "ENSG000_TGFB1".into(), "CD4".into(), ];
702 let canon = build_mixed_kind_canonicalizer(&names);
703 assert_eq!(canon("chr1:1-20").as_ref(), "1_1_30");
704 assert_eq!(canon("chr1:15-30").as_ref(), "1_1_30");
705 assert_eq!(canon("ENSG000_TGFB1").as_ref(), "TGFB1");
706 assert_eq!(canon("CD4").as_ref(), "CD4");
707 }
708}