use std::sync::Arc;
use crate::sparse_io_vector::RowNameCanonicalizer;
use legume_numeric::matrix::membership::canon_locus;
use rustc_hash::FxHashMap as HashMap;
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum FeatureNameKind {
#[default]
Exact,
Gene { delim: char },
Locus { merge_overlapping: bool },
Mixed,
}
impl FeatureNameKind {
pub fn canonicalize(&self, name: &str) -> Box<str> {
match self {
FeatureNameKind::Exact => name.into(),
FeatureNameKind::Gene { delim } => gene_canonicalize(name, *delim),
FeatureNameKind::Locus { .. } => canon_locus(name),
FeatureNameKind::Mixed => {
if parse_locus(name).is_some() {
canon_locus(name)
} else if name.contains('_') {
gene_canonicalize(name, '_')
} else {
name.into()
}
}
}
}
pub fn is_exact(&self) -> bool {
matches!(self, FeatureNameKind::Exact)
}
pub fn needs_global_pass(&self) -> bool {
matches!(
self,
FeatureNameKind::Locus {
merge_overlapping: true
} | FeatureNameKind::Mixed
)
}
pub fn auto_detect(names: &[Box<str>]) -> Self {
let n = names.len();
if n == 0 {
return Self::Exact;
}
let mut n_locus = 0usize;
let mut n_gene_like = 0usize;
for name in names {
if parse_locus(name).is_some() {
n_locus += 1;
} else if name.contains('_') {
n_gene_like += 1;
}
}
let pct_locus = n_locus as f32 / n as f32;
let pct_gene = n_gene_like as f32 / n as f32;
if pct_locus >= 0.10 && pct_gene >= 0.10 {
Self::Mixed
} else if pct_locus >= 0.50 {
Self::Locus {
merge_overlapping: true,
}
} else if pct_gene >= 0.50 {
Self::Gene { delim: '_' }
} else {
Self::Exact
}
}
#[must_use]
pub fn reconcile(kinds: &[FeatureNameKind]) -> FeatureNameKind {
if kinds.iter().any(|k| matches!(k, FeatureNameKind::Mixed)) {
return FeatureNameKind::Mixed;
}
let gene = kinds
.iter()
.find(|k| matches!(k, FeatureNameKind::Gene { .. }));
let locus = kinds
.iter()
.find(|k| matches!(k, FeatureNameKind::Locus { .. }));
match (gene, locus) {
(Some(_), Some(_)) => FeatureNameKind::Mixed,
_ => gene.or(locus).cloned().unwrap_or(FeatureNameKind::Exact),
}
}
pub fn into_canonicalizer(self) -> Option<RowNameCanonicalizer> {
if self.is_exact() {
return None;
}
Some(Arc::new(move |name: &str| self.canonicalize(name)))
}
}
pub fn parse_locus(name: &str) -> Option<(Box<str>, u64, u64)> {
let lower = name.to_ascii_lowercase();
let stripped = lower
.strip_prefix("chr")
.map(str::to_string)
.unwrap_or(lower);
let parts: Vec<&str> = stripped.splitn(3, [':', '-', '_']).collect();
if parts.len() != 3 {
return None;
}
let chr = parts[0];
let start: u64 = parts[1].parse().ok()?;
let end: u64 = parts[2].parse().ok()?;
if end < start {
return None;
}
Some((chr.to_string().into_boxed_str(), start, end))
}
pub fn build_locus_overlap_canonical_map(names: &[Box<str>]) -> HashMap<Box<str>, Box<str>> {
let n = names.len();
let parsed: Vec<Option<(Box<str>, u64, u64)>> = names.iter().map(|n| parse_locus(n)).collect();
let mut by_chr: HashMap<Box<str>, Vec<usize>> = HashMap::default();
for (i, p) in parsed.iter().enumerate() {
if let Some((chr, _, _)) = p {
by_chr.entry(chr.clone()).or_default().push(i);
}
}
let mut parent: Vec<usize> = (0..n).collect();
fn find(p: &mut [usize], mut x: usize) -> usize {
while p[x] != x {
let g = p[p[x]];
p[x] = g;
x = g;
}
x
}
let mut cluster_extent: HashMap<usize, (u64, u64)> = HashMap::default();
for (_, mut idxs) in by_chr {
idxs.sort_by_key(|&i| parsed[i].as_ref().map(|p| p.1).unwrap_or(0));
let mut current_root: Option<usize> = None;
let mut current_min_start: u64 = 0;
let mut current_max_end: u64 = 0;
for i in idxs {
let (_, s, e) = parsed[i].as_ref().unwrap();
match current_root {
Some(root) if *s < current_max_end => {
let ra = find(&mut parent, root);
let rb = find(&mut parent, i);
if ra != rb {
parent[rb] = ra;
}
current_max_end = current_max_end.max(*e);
cluster_extent
.insert(find(&mut parent, i), (current_min_start, current_max_end));
}
_ => {
current_root = Some(i);
current_min_start = *s;
current_max_end = *e;
cluster_extent.insert(i, (*s, *e));
}
}
}
}
let mut out: HashMap<Box<str>, Box<str>> = HashMap::default();
for (i, p) in parsed.iter().enumerate() {
if let Some((chr, _, _)) = p {
let root = find(&mut parent, i);
let (mn, mx) = cluster_extent.get(&root).copied().unwrap_or((0, 0));
let canonical = format!("{}_{}_{}", chr, mn, mx).into_boxed_str();
out.insert(names[i].clone(), canonical);
}
}
out
}
pub fn build_locus_overlap_canonicalizer(names: &[Box<str>]) -> RowNameCanonicalizer {
let map = Arc::new(build_locus_overlap_canonical_map(names));
Arc::new(move |name: &str| map.get(name).cloned().unwrap_or_else(|| canon_locus(name)))
}
pub fn build_mixed_kind_canonicalizer(names: &[Box<str>]) -> RowNameCanonicalizer {
let map = Arc::new(build_locus_overlap_canonical_map(names));
Arc::new(move |name: &str| {
if let Some(c) = map.get(name) {
c.clone()
} else if parse_locus(name).is_some() {
canon_locus(name)
} else if name.contains('_') {
gene_canonicalize(name, '_')
} else {
name.into()
}
})
}
fn gene_canonicalize(name: &str, delim: char) -> Box<str> {
let stripped = strip_feature_type_suffix(name, delim);
stripped.rsplit(delim).next().unwrap_or(stripped).into()
}
fn strip_feature_type_suffix(name: &str, delim: char) -> &str {
const TAGS: &[&str] = &[
"Gene_Expression",
"Gene",
"Antibody_Capture",
"CRISPR_Guide_Capture",
"Multiplexing_Capture",
"Custom",
"Peaks",
];
for tag in TAGS {
let mut suffix = String::with_capacity(tag.len() + 1);
suffix.push(delim);
suffix.push_str(tag);
if let Some(rest) = name.strip_suffix(suffix.as_str()) {
return rest;
}
}
name
}
#[derive(clap::ValueEnum, Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum FeatureNameKindArg {
#[default]
Auto,
Exact,
Gene,
Locus,
LocusOverlap,
Mixed,
}
impl FeatureNameKindArg {
pub fn resolve_or_gene(&self) -> FeatureNameKind {
Option::<FeatureNameKind>::from(self.clone())
.unwrap_or(FeatureNameKind::Gene { delim: '_' })
}
}
impl From<FeatureNameKindArg> for Option<FeatureNameKind> {
fn from(arg: FeatureNameKindArg) -> Self {
match arg {
FeatureNameKindArg::Auto => None,
FeatureNameKindArg::Exact => Some(FeatureNameKind::Exact),
FeatureNameKindArg::Gene => Some(FeatureNameKind::Gene { delim: '_' }),
FeatureNameKindArg::Locus => Some(FeatureNameKind::Locus {
merge_overlapping: false,
}),
FeatureNameKindArg::LocusOverlap => Some(FeatureNameKind::Locus {
merge_overlapping: true,
}),
FeatureNameKindArg::Mixed => Some(FeatureNameKind::Mixed),
}
}
}
#[cfg(test)]
#[path = "feature_names_tests.rs"]
mod feature_names_tests;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn exact_passthrough() {
let k = FeatureNameKind::Exact;
assert_eq!(
k.canonicalize("ENSG00000000003_TSPAN6").as_ref(),
"ENSG00000000003_TSPAN6"
);
assert!(k.is_exact());
assert!(k.into_canonicalizer().is_none());
}
#[test]
fn gene_takes_last_underscore_component() {
let k = FeatureNameKind::Gene { delim: '_' };
assert_eq!(k.canonicalize("ENSG00000000003_TSPAN6").as_ref(), "TSPAN6");
assert_eq!(k.canonicalize("TSPAN6").as_ref(), "TSPAN6");
assert_eq!(k.canonicalize("A_B_C").as_ref(), "C");
assert!(!k.is_exact());
assert!(k.into_canonicalizer().is_some());
}
#[test]
fn gene_strips_cell_ranger_feature_type_suffix() {
let k = FeatureNameKind::Gene { delim: '_' };
assert_eq!(
k.canonicalize("ENSG00000187634_SAMD11_Gene").as_ref(),
"SAMD11"
);
assert_eq!(
k.canonicalize("ENSG00000187634_SAMD11_Gene_Expression")
.as_ref(),
"SAMD11"
);
assert_eq!(k.canonicalize("FakeGene").as_ref(), "FakeGene");
}
#[test]
fn locus_strips_chr_and_folds_separators() {
let k = FeatureNameKind::Locus {
merge_overlapping: false,
};
assert_eq!(k.canonicalize("chr1:1000-2000").as_ref(), "1_1000_2000");
assert_eq!(k.canonicalize("1_1000_2000").as_ref(), "1_1000_2000");
assert_eq!(k.canonicalize("ChrX:5000-6000").as_ref(), "X_5000_6000");
}
#[test]
fn parse_locus_accepts_common_formats() {
assert_eq!(
parse_locus("chr1:1000-2000"),
Some(("1".into(), 1000, 2000))
);
assert_eq!(parse_locus("1:1000-2000"), Some(("1".into(), 1000, 2000)));
assert_eq!(
parse_locus("chr1_1000_2000"),
Some(("1".into(), 1000, 2000))
);
assert_eq!(
parse_locus("CHR1:1000-2000"),
Some(("1".into(), 1000, 2000))
);
assert_eq!(
parse_locus("chrX:5000-6000"),
Some(("x".into(), 5000, 6000))
);
assert_eq!(parse_locus("chrMT:1-100"), Some(("mt".into(), 1, 100)));
}
#[test]
fn parse_locus_rejects_non_loci() {
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()); }
#[test]
fn overlap_map_merges_two_overlapping_intervals() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"chr1:15-30".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
let c0 = map.get(&names[0]).unwrap();
let c1 = map.get(&names[1]).unwrap();
assert_eq!(c0, c1, "both inputs should map to the same canonical");
assert_eq!(c0.as_ref(), "1_1_30"); }
#[test]
fn overlap_map_keeps_non_overlapping_separate() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"chr1:100-200".to_string().into_boxed_str(),
"chr2:1-20".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
assert_eq!(map.get(&names[0]).unwrap().as_ref(), "1_1_20");
assert_eq!(map.get(&names[1]).unwrap().as_ref(), "1_100_200");
assert_eq!(map.get(&names[2]).unwrap().as_ref(), "2_1_20");
}
#[test]
fn overlap_map_handles_transitive_chain() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"chr1:15-30".to_string().into_boxed_str(),
"chr1:25-40".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
let c0 = map.get(&names[0]).unwrap();
let c1 = map.get(&names[1]).unwrap();
let c2 = map.get(&names[2]).unwrap();
assert_eq!(c0, c1);
assert_eq!(c1, c2);
assert_eq!(c0.as_ref(), "1_1_40"); }
#[test]
fn overlap_map_handles_full_containment() {
let names = vec![
"chr1:1-100".to_string().into_boxed_str(),
"chr1:30-50".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
let c0 = map.get(&names[0]).unwrap();
let c1 = map.get(&names[1]).unwrap();
assert_eq!(c0, c1);
assert_eq!(c0.as_ref(), "1_1_100");
}
#[test]
fn overlap_map_treats_adjacent_as_separate() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"chr1:20-30".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
assert_ne!(map.get(&names[0]).unwrap(), map.get(&names[1]).unwrap());
}
#[test]
fn overlap_map_normalizes_chr_prefix_within_cluster() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"1:15-30".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
let c0 = map.get(&names[0]).unwrap();
let c1 = map.get(&names[1]).unwrap();
assert_eq!(c0, c1);
assert_eq!(c0.as_ref(), "1_1_30");
}
#[test]
fn overlap_map_normalizes_separators_within_cluster() {
let names = vec![
"chr1:1-20".to_string().into_boxed_str(),
"chr1_15_30".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
assert_eq!(map.get(&names[0]).unwrap(), map.get(&names[1]).unwrap());
}
#[test]
fn overlap_map_ignores_non_locus_names() {
let names = vec![
"TGFB1".to_string().into_boxed_str(),
"chr1:1-20".to_string().into_boxed_str(),
];
let map = build_locus_overlap_canonical_map(&names);
assert!(!map.contains_key(&names[0]));
assert!(map.contains_key(&names[1]));
}
#[test]
fn overlap_map_handles_zero_length_interval() {
let names = vec!["chr1:1000-1000".to_string().into_boxed_str()];
let map = build_locus_overlap_canonical_map(&names);
assert_eq!(map.get(&names[0]).unwrap().as_ref(), "1_1000_1000");
}
#[test]
fn overlap_canonicalizer_falls_back_to_canon_locus_for_unmatched() {
let names = vec!["chr1:1-20".to_string().into_boxed_str()];
let canon = build_locus_overlap_canonicalizer(&names);
assert_eq!(canon("chr1:1-20").as_ref(), "1_1_20");
assert_eq!(canon("chr2:500-600").as_ref(), "2_500_600");
let g = canon("TGFB1");
assert!(!g.is_empty());
}
#[test]
fn auto_detect_pure_locus_axis() {
let names: Vec<Box<str>> = (0..100)
.map(|i| format!("chr1:{}-{}", i * 100, i * 100 + 50).into_boxed_str())
.collect();
assert!(matches!(
FeatureNameKind::auto_detect(&names),
FeatureNameKind::Locus {
merge_overlapping: true
}
));
}
#[test]
fn auto_detect_pure_gene_axis() {
let names: Vec<Box<str>> = (0..100)
.map(|i| format!("ENSG000_GENE{}", i).into_boxed_str())
.collect();
assert!(matches!(
FeatureNameKind::auto_detect(&names),
FeatureNameKind::Gene { delim: '_' }
));
}
#[test]
fn auto_detect_mixed_axis() {
let mut names: Vec<Box<str>> = (0..80)
.map(|i| format!("chr1:{}-{}", i * 1000, i * 1000 + 500).into_boxed_str())
.collect();
names.extend((0..20).map(|i| format!("ENSG000_GENE{}", i).into_boxed_str()));
assert!(matches!(
FeatureNameKind::auto_detect(&names),
FeatureNameKind::Mixed
));
}
#[test]
fn auto_detect_empty_or_exact() {
assert!(matches!(
FeatureNameKind::auto_detect(&[]),
FeatureNameKind::Exact
));
let names = vec!["TGFB1".into(), "CD4".into(), "IL2".into(), "GAPDH".into()];
assert!(matches!(
FeatureNameKind::auto_detect(&names),
FeatureNameKind::Exact
));
}
#[test]
fn mixed_dispatcher_canonicalizes_each_name_by_kind() {
let names: Vec<Box<str>> = vec![
"chr1:1-20".into(), "chr1:15-30".into(), "ENSG000_TGFB1".into(), "CD4".into(), ];
let canon = build_mixed_kind_canonicalizer(&names);
assert_eq!(canon("chr1:1-20").as_ref(), "1_1_30");
assert_eq!(canon("chr1:15-30").as_ref(), "1_1_30");
assert_eq!(canon("ENSG000_TGFB1").as_ref(), "TGFB1");
assert_eq!(canon("CD4").as_ref(), "CD4");
}
}