1use anyhow::{Context, Result};
2use clap::Args;
3use csv::WriterBuilder;
4use flate2::read::MultiGzDecoder;
5use mapping_info::MappingInfo;
6use serde::{Deserialize, Serialize};
7
8use std::{
9 collections::{HashMap, HashSet},
10 fmt,
11 fs::File,
12 io::{BufReader, Read, Write},
13 path::{Path, PathBuf},
14};
15
16#[derive(Debug, Clone, Args)]
17pub struct ReadTagTableCli {
18 #[arg(long = "read-tag-table", num_args = 1..)]
25 pub read_tag_table: Vec<PathBuf>,
26
27 #[arg(long = "rt-read-id-column", default_value = "read_id")]
28 pub rt_read_id_column: String,
29
30 #[arg(long = "rt-cell-column", default_value = "raw_cb")]
31 pub rt_cell_column: String,
32
33 #[arg(long = "rt-cell-qual-column", default_value = "quality_cb")]
34 pub rt_cell_qual_column: String,
35
36 #[arg(long = "rt-umi-column", default_value = "raw_umi")]
37 pub rt_umi_column: String,
38
39 #[arg(long = "rt-umi-qual-column", default_value = "quality_umi")]
40 pub rt_umi_qual_column: String,
41
42 #[arg(
43 long = "rt-original-read-id-column",
44 default_value = "original_read_id"
45 )]
46 pub rt_original_read_id_column: String,
47}
48
49impl ReadTagTableCli {
50 pub fn to_config_for_id(&self, id: usize) -> Result<ReadTagTableConfig> {
51 if self.read_tag_table.is_empty() {
52 anyhow::bail!("No --read-tag-table files supplied");
53 }
54
55 let path = self
56 .read_tag_table
57 .get(id)
58 .with_context(|| format!("No --read-tag-table for id {id}"))?
59 .clone();
60
61 Ok(self.config_for_path(path))
62 }
63
64 pub fn to_config(&self) -> Result<ReadTagTableConfig> {
65 if self.read_tag_table.is_empty() {
66 anyhow::bail!("No --read-tag-table files supplied");
67 }
68
69 if self.read_tag_table.len() > 1 {
70 anyhow::bail!(
71 "Number of --read-tag-table files (seen {}) > 1 is not supported by this function",
72 self.read_tag_table.len()
73 );
74 }
75
76 Ok(self.config_for_path(self.read_tag_table[0].clone()))
77 }
78
79 pub fn load(&self) -> Result<ReadTagTable> {
80 let config = self.to_config()?;
81 ReadTagTable::from_path_or_config(&config)
82 }
83
84 pub fn load_for_id(&self, id: usize) -> Result<ReadTagTable> {
85 let config = self.to_config_for_id(id)?;
86 ReadTagTable::from_path_or_config(&config)
87 }
88
89 fn config_for_path(&self, path: PathBuf) -> ReadTagTableConfig {
90 ReadTagTableConfig {
91 path,
92 read_id_column: self.rt_read_id_column.clone(),
93 original_read_id_column: self.rt_original_read_id_column.clone(),
94 cell_column: self.rt_cell_column.clone(),
95 cell_qual_column: self.rt_cell_qual_column.clone(),
96 umi_column: self.rt_umi_column.clone(),
97 umi_qual_column: self.rt_umi_qual_column.clone(),
98 }
99 }
100}
101
102#[derive(Debug, Clone)]
103pub struct ReadTagTableConfig {
104 pub path: PathBuf,
105 pub read_id_column: String,
106 pub original_read_id_column: String,
107 pub cell_column: String,
108 pub cell_qual_column: String,
109 pub umi_column: String,
110 pub umi_qual_column: String,
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct ReadTagRecord {
115 pub read_id: String,
116 pub original_read_id: Option<String>,
117 pub cell_seq: Vec<u8>,
118 pub cell_qual: Vec<u8>,
119 pub umi_seq: Vec<u8>,
120 pub umi_qual: Vec<u8>,
121}
122
123impl ReadTagRecord {
124 pub fn new(
125 read_id: String,
126 original_read_id: Option<String>,
127 cell_seq: impl AsRef<[u8]>,
128 cell_qual: impl AsRef<[u8]>,
129 umi_seq: impl AsRef<[u8]>,
130 umi_qual: impl AsRef<[u8]>,
131 ) -> Self {
132 Self {
133 read_id,
134 original_read_id,
135 cell_seq: cell_seq.as_ref().to_vec(),
136 cell_qual: cell_qual.as_ref().to_vec(),
137 umi_seq: umi_seq.as_ref().to_vec(),
138 umi_qual: umi_qual.as_ref().to_vec(),
139 }
140 }
141
142 pub fn from_slices(
143 read_id: impl Into<String>,
144 original_read_id: Option<String>,
145 cell_seq: &[u8],
146 cell_qual: &[u8],
147 umi_seq: &[u8],
148 umi_qual: &[u8],
149 ) -> Self {
150 Self::new(
151 read_id.into(),
152 original_read_id,
153 cell_seq,
154 cell_qual,
155 umi_seq,
156 umi_qual,
157 )
158 }
159
160 pub fn from_tsv_fields(
161 read_id: impl Into<String>,
162 original_read_id: Option<String>,
163 cell: &str,
164 cell_qual: Option<&str>,
165 umi: &str,
166 umi_qual: Option<&str>,
167 ) -> Self {
168 Self::new(
169 read_id.into(),
170 original_read_id,
171 cell.as_bytes(),
172 cell_qual.map(phred_from_ascii).unwrap_or_default(),
173 umi.as_bytes(),
174 umi_qual.map(phred_from_ascii).unwrap_or_default(),
175 )
176 }
177
178 pub fn cell_string(&self) -> String {
179 seq_to_string(&self.cell_seq)
180 }
181
182 pub fn umi_string(&self) -> String {
183 seq_to_string(&self.umi_seq)
184 }
185
186 pub fn cell_qual_string(&self) -> String {
187 phred_to_ascii(&self.cell_qual)
188 }
189
190 pub fn umi_qual_string(&self) -> String {
191 phred_to_ascii(&self.umi_qual)
192 }
193}
194
195#[derive(Debug, Clone, Default, Serialize, Deserialize)]
196pub struct ReadTagTable {
197 records: HashMap<String, ReadTagRecord>,
198}
199
200impl ReadTagTable {
201 pub fn new() -> Self {
202 Self {
203 records: HashMap::new(),
204 }
205 }
206
207 pub fn from_path_or_config(config: &ReadTagTableConfig) -> Result<Self> {
208 if is_binary_read_tag_table(&config.path) {
209 Self::load_binary(&config.path)
210 } else {
211 Self::from_config(config)
212 }
213 }
214
215 pub fn from_config(config: &ReadTagTableConfig) -> Result<Self> {
216 let mut mapping_info = MappingInfo::new(None, 0.0, 0);
217 let reader = open_maybe_gz(&config.path)?;
218
219 let mut rdr = csv::ReaderBuilder::new()
220 .delimiter(b'\t')
221 .has_headers(true)
222 .flexible(true)
223 .from_reader(reader);
224
225 let headers = rdr.headers()?.clone();
226
227 let read_id_ix = required_column_ix(&headers, &config.read_id_column)?;
228 let cell_ix = required_column_ix(&headers, &config.cell_column)?;
229 let umi_ix = required_column_ix(&headers, &config.umi_column)?;
230
231 let original_read_id_ix = optional_column_ix(&headers, &config.original_read_id_column);
232 let cell_qual_ix = optional_column_ix(&headers, &config.cell_qual_column);
233 let umi_qual_ix = optional_column_ix(&headers, &config.umi_qual_column);
234
235 let mut records = HashMap::new();
236
237 for rec in rdr.records() {
238 let rec = rec?;
239
240 let read_id = rec.get(read_id_ix).unwrap_or("").trim();
241 let cell = rec.get(cell_ix).unwrap_or("").trim();
242 let umi = rec.get(umi_ix).unwrap_or("").trim();
243
244 if read_id.is_empty() || cell.is_empty() || umi.is_empty() {
245 continue;
246 }
247
248 let record = ReadTagRecord::from_tsv_fields(
249 read_id.to_string(),
250 get_optional(&rec, original_read_id_ix),
251 cell,
252 get_optional_ref(&rec, cell_qual_ix),
253 umi,
254 get_optional_ref(&rec, umi_qual_ix),
255 );
256
257 records.insert(read_id.to_string(), record);
258 }
259
260 mapping_info.stop_file_io_time();
261
262 let (h, m, s, ms) = MappingInfo::split_duration(mapping_info.file_io_time);
263 println!(
264 "Read-tag table loaded from TSV: {} entries in {}:{:02}:{:02}.{:03}",
265 records.len(),
266 h,
267 m,
268 s,
269 ms,
270 );
271
272 Ok(Self { records })
273 }
274
275 pub fn save_binary<P: AsRef<Path>>(&self, path: P) -> Result<()> {
276 let path = path.as_ref();
277 let file = File::create(path).with_context(|| format!("creating {}", path.display()))?;
278
279 bincode::serialize_into(file, self)
280 .with_context(|| format!("writing binary read-tag table {}", path.display()))?;
281
282 Ok(())
283 }
284
285 pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<()> {
286 self.save_binary(path)
287 }
288
289 pub fn load_binary<P: AsRef<Path>>(path: P) -> Result<Self> {
290 let path = path.as_ref();
291 let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;
292
293 let table: Self = bincode::deserialize_from(file)
294 .with_context(|| format!("reading binary read-tag table {}", path.display()))?;
295
296 Ok(table)
297 }
298
299 pub fn get(&self, read_id: &str) -> Option<&ReadTagRecord> {
300 self.records.get(read_id)
301 }
302
303 pub fn cell_umi_for_read(&self, read_id: &str) -> Option<(String, String)> {
304 let rec = self.records.get(read_id)?;
305 Some((rec.cell_string(), rec.umi_string()))
306 }
307
308 pub fn cell_umi_bytes_for_read(&self, read_id: &str) -> Option<(&[u8], &[u8])> {
309 let rec = self.records.get(read_id)?;
310 Some((rec.cell_seq.as_slice(), rec.umi_seq.as_slice()))
311 }
312
313 pub fn len(&self) -> usize {
314 self.records.len()
315 }
316
317 pub fn is_empty(&self) -> bool {
318 self.records.is_empty()
319 }
320
321 pub fn pair_counts(&self) -> HashMap<(Vec<u8>, Vec<u8>), u64> {
322 let mut counts = HashMap::new();
323
324 for rec in self.records.values() {
325 if rec.cell_seq.is_empty() || rec.umi_seq.is_empty() {
326 continue;
327 }
328
329 *counts
330 .entry((rec.cell_seq.clone(), rec.umi_seq.clone()))
331 .or_insert(0) += 1;
332 }
333
334 counts
335 }
336
337 pub fn insert(&mut self, record: ReadTagRecord) -> Option<ReadTagRecord> {
338 self.records.insert(record.read_id.clone(), record)
339 }
340
341 pub fn merge(&mut self, other: Self) {
342 self.records.extend(other.records);
343 }
344
345 pub fn summarize_pairs(&self, min_pair_count: u64, min_cell_umis: u64) -> PairStats {
346 let counts = self.pair_counts();
347 PairStats::from_counts(&counts, min_pair_count, min_cell_umis)
348 }
349
350 pub fn write_tsv<W: Write>(&self, inner: W) -> Result<()> {
351 let mut writer = ReadTagTableWriter::new(inner)?;
352
353 for record in self.records.values() {
354 writer.write_record(record)?;
355 }
356
357 writer.flush()
358 }
359}
360
361fn is_binary_read_tag_table(path: &Path) -> bool {
362 path.extension()
363 .and_then(|x| x.to_str())
364 .is_some_and(|x| x.eq_ignore_ascii_case("bin"))
365}
366
367#[derive(Debug, Clone)]
368pub struct PairStats {
369 pub cell_entries: usize,
370 pub unique_cell_umi_combos: usize,
371 pub total_pair_observations: u64,
372 pub umis_per_cell: Summary,
373 pub detections_per_cell_umi: Summary,
374}
375
376impl PairStats {
377 pub fn from_counts(
378 cb_umi_counts: &HashMap<(Vec<u8>, Vec<u8>), u64>,
379 min_pair_count: u64,
380 min_cell_umis: u64,
381 ) -> Self {
382 let mut cell_to_umis: HashMap<&[u8], HashSet<&[u8]>> = HashMap::new();
383
384 for ((cell, umi), count) in cb_umi_counts {
385 if *count < min_pair_count {
386 continue;
387 }
388
389 cell_to_umis
390 .entry(cell.as_slice())
391 .or_default()
392 .insert(umi.as_slice());
393 }
394
395 let valid_cells: HashSet<&[u8]> = cell_to_umis
396 .iter()
397 .filter_map(|(cell, umis)| {
398 if umis.len() as u64 >= min_cell_umis {
399 Some(*cell)
400 } else {
401 None
402 }
403 })
404 .collect();
405
406 let mut final_cell_to_umi_count: HashMap<&[u8], u64> = HashMap::new();
407 let mut final_pair_counts = Vec::new();
408 let mut total_pair_observations = 0_u64;
409
410 for ((cell, _umi), count) in cb_umi_counts {
411 if *count < min_pair_count {
412 continue;
413 }
414
415 if !valid_cells.contains(cell.as_slice()) {
416 continue;
417 }
418
419 *final_cell_to_umi_count.entry(cell.as_slice()).or_insert(0) += 1;
420 final_pair_counts.push(*count);
421 total_pair_observations += *count;
422 }
423
424 let umis_per_cell: Vec<u64> = final_cell_to_umi_count.values().copied().collect();
425
426 Self {
427 cell_entries: final_cell_to_umi_count.len(),
428 unique_cell_umi_combos: final_pair_counts.len(),
429 total_pair_observations,
430 umis_per_cell: summarize(umis_per_cell),
431 detections_per_cell_umi: summarize(final_pair_counts),
432 }
433 }
434}
435
436#[derive(Debug, Clone)]
437pub struct Summary {
438 pub n: usize,
439 pub mean: f64,
440 pub median: f64,
441 pub min: u64,
442 pub max: u64,
443}
444
445fn summarize(mut values: Vec<u64>) -> Summary {
446 if values.is_empty() {
447 return Summary {
448 n: 0,
449 mean: 0.0,
450 median: 0.0,
451 min: 0,
452 max: 0,
453 };
454 }
455
456 values.sort_unstable();
457
458 let n = values.len();
459 let min = values[0];
460 let max = values[n - 1];
461
462 let sum: u128 = values.iter().map(|&x| x as u128).sum();
463 let mean = sum as f64 / n as f64;
464
465 let median = if n % 2 == 0 {
466 let a = values[n / 2 - 1];
467 let b = values[n / 2];
468 (a as f64 + b as f64) / 2.0
469 } else {
470 values[n / 2] as f64
471 };
472
473 Summary {
474 n,
475 mean,
476 median,
477 min,
478 max,
479 }
480}
481
482fn required_column_ix(headers: &csv::StringRecord, name: &str) -> Result<usize> {
483 headers
484 .iter()
485 .position(|h| h == name)
486 .with_context(|| format!("Could not find required column '{name}'"))
487}
488
489fn optional_column_ix(headers: &csv::StringRecord, name: &str) -> Option<usize> {
490 headers.iter().position(|h| h == name)
491}
492
493fn get_optional(rec: &csv::StringRecord, ix: Option<usize>) -> Option<String> {
494 let value = rec.get(ix?)?.trim();
495
496 if value.is_empty() {
497 None
498 } else {
499 Some(value.to_string())
500 }
501}
502
503fn get_optional_ref(rec: &csv::StringRecord, ix: Option<usize>) -> Option<&str> {
504 let value = rec.get(ix?)?.trim();
505
506 if value.is_empty() { None } else { Some(value) }
507}
508
509fn open_maybe_gz(path: &Path) -> Result<Box<dyn Read>> {
510 let file = File::open(path).with_context(|| format!("opening {}", path.display()))?;
511 let reader = BufReader::new(file);
512
513 if path.extension().is_some_and(|e| e == "gz") {
514 Ok(Box::new(MultiGzDecoder::new(reader)))
515 } else {
516 Ok(Box::new(reader))
517 }
518}
519
520fn seq_to_string(seq: &[u8]) -> String {
521 String::from_utf8_lossy(seq).to_string()
522}
523
524fn phred_from_ascii(text: &str) -> Vec<u8> {
525 text.as_bytes()
526 .iter()
527 .map(|q| q.saturating_sub(33))
528 .collect()
529}
530
531fn phred_to_ascii(qual: &[u8]) -> String {
532 qual.iter().map(|q| q.saturating_add(33) as char).collect()
533}
534
535fn display_qual(qual: &[u8]) -> String {
536 if qual.is_empty() {
537 "-".to_string()
538 } else {
539 phred_to_ascii(qual)
540 }
541}
542
543pub const READ_TAG_TABLE_COLUMNS: [&str; 8] = [
544 "read_id",
545 "original_read_id",
546 "orientation",
547 "raw_cb",
548 "quality_cb",
549 "raw_umi",
550 "quality_umi",
551 "status",
552];
553
554#[derive(Debug, Clone)]
555pub struct ReadTagWriteRecord<'a> {
556 pub read_id: &'a str,
557 pub original_read_id: Option<&'a str>,
558 pub orientation: Option<&'a str>,
559 pub raw_cb: String,
560 pub quality_cb: String,
561 pub raw_umi: String,
562 pub quality_umi: String,
563 pub status: &'a str,
564}
565
566pub struct ReadTagTableWriter<W: Write> {
567 writer: csv::Writer<W>,
568}
569
570impl<W: Write> ReadTagTableWriter<W> {
571 pub fn new(inner: W) -> Result<Self> {
572 let mut writer = WriterBuilder::new()
573 .delimiter(b'\t')
574 .has_headers(false)
575 .from_writer(inner);
576
577 writer
578 .write_record(READ_TAG_TABLE_COLUMNS)
579 .context("writing read-tag table header")?;
580
581 Ok(Self { writer })
582 }
583
584 pub fn write_record(&mut self, rec: &ReadTagRecord) -> Result<()> {
585 self.writer
586 .write_record([
587 rec.read_id.as_str(),
588 rec.original_read_id.as_deref().unwrap_or(""),
589 "",
590 rec.cell_string().as_str(),
591 rec.cell_qual_string().as_str(),
592 rec.umi_string().as_str(),
593 rec.umi_qual_string().as_str(),
594 "ok",
595 ])
596 .with_context(|| format!("writing read-tag table row for read_id '{}'", rec.read_id))?;
597
598 Ok(())
599 }
600
601 pub fn write_tsv_record(&mut self, rec: &ReadTagWriteRecord<'_>) -> Result<()> {
604 self.writer
605 .write_record([
606 rec.read_id,
607 rec.original_read_id.unwrap_or(""),
608 rec.orientation.unwrap_or(""),
609 &rec.raw_cb,
610 &rec.quality_cb,
611 &rec.raw_umi,
612 &rec.quality_umi,
613 rec.status,
614 ])
615 .with_context(|| format!("writing read-tag table row for read_id '{}'", rec.read_id))?;
616
617 Ok(())
618 }
619
620 pub fn flush(&mut self) -> Result<()> {
621 self.writer
622 .flush()
623 .context("flushing read-tag table writer")?;
624 Ok(())
625 }
626}
627
628impl fmt::Display for ReadTagRecord {
629 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
630 write!(
631 f,
632 "read_id={}, original_read_id={}, cell={}, cell_qual={}, umi={}, umi_qual={}",
633 self.read_id,
634 self.original_read_id.as_deref().unwrap_or("-"),
635 self.cell_string(),
636 display_qual(&self.cell_qual),
637 self.umi_string(),
638 display_qual(&self.umi_qual),
639 )
640 }
641}
642
643impl fmt::Display for ReadTagTable {
644 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
645 writeln!(f, "ReadTagTable: {} records", self.records.len())?;
646
647 for (i, record) in self.records.values().take(3).enumerate() {
648 writeln!(f, " [{}] {}", i, record)?;
649 }
650
651 Ok(())
652 }
653}
654
655#[cfg(test)]
656mod tests {
657 use super::*;
658 use flate2::{Compression, write::GzEncoder};
659 use std::fs;
660 use std::time::{SystemTime, UNIX_EPOCH};
661
662 fn tmp_path(name: &str) -> PathBuf {
663 let stamp = SystemTime::now()
664 .duration_since(UNIX_EPOCH)
665 .unwrap()
666 .as_nanos();
667
668 std::env::temp_dir().join(format!("read_tag_table_test_{stamp}_{name}"))
669 }
670
671 fn write_text(path: &Path, text: &str) {
672 fs::write(path, text).unwrap();
673 }
674
675 fn default_config(path: PathBuf) -> ReadTagTableConfig {
676 ReadTagTableConfig {
677 path,
678 read_id_column: "read_id".to_string(),
679 original_read_id_column: "original_read_id".to_string(),
680 cell_column: "raw_cb".to_string(),
681 cell_qual_column: "quality_cb".to_string(),
682 umi_column: "raw_umi".to_string(),
683 umi_qual_column: "quality_umi".to_string(),
684 }
685 }
686
687 #[test]
688 fn reads_plain_tsv_table() {
689 let path = tmp_path("table.tsv");
690
691 write_text(
692 &path,
693 concat!(
694 "read_id\toriginal_read_id\traw_cb\tquality_cb\traw_umi\tquality_umi\n",
695 "read1\torig1\tCELL1\tIIII\tUMI1\tJJJJ\n",
696 "read2\torig2\tCELL1\tIIII\tUMI2\tJJJJ\n",
697 "read3\t\tCELL2\t\tUMI3\t\n",
698 ),
699 );
700
701 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
702
703 assert_eq!(table.len(), 3);
704 assert!(!table.is_empty());
705
706 let rec = table.get("read1").unwrap();
707 assert_eq!(rec.read_id, "read1");
708 assert_eq!(rec.original_read_id.as_deref(), Some("orig1"));
709 assert_eq!(rec.cell_seq, b"CELL1");
710 assert_eq!(rec.cell_qual, vec![40, 40, 40, 40]);
711 assert_eq!(rec.umi_seq, b"UMI1");
712 assert_eq!(rec.umi_qual, vec![41, 41, 41, 41]);
713
714 let rec = table.get("read3").unwrap();
715 assert_eq!(rec.original_read_id, None);
716 assert!(rec.cell_qual.is_empty());
717 assert!(rec.umi_qual.is_empty());
718
719 fs::remove_file(path).ok();
720 }
721
722 #[test]
723 fn skips_rows_with_missing_required_values() {
724 let path = tmp_path("missing.tsv");
725
726 write_text(
727 &path,
728 concat!(
729 "read_id\traw_cb\traw_umi\n",
730 "read1\tCELL1\tUMI1\n",
731 "\tCELL2\tUMI2\n",
732 "read3\t\tUMI3\n",
733 "read4\tCELL4\t\n",
734 ),
735 );
736
737 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
738
739 assert_eq!(table.len(), 1);
740 assert!(table.get("read1").is_some());
741 assert!(table.get("read3").is_none());
742
743 fs::remove_file(path).ok();
744 }
745
746 #[test]
747 fn missing_required_column_returns_error() {
748 let path = tmp_path("bad_columns.tsv");
749
750 write_text(&path, concat!("read_id\traw_cb\n", "read1\tCELL1\n",));
751
752 let err = ReadTagTable::from_config(&default_config(path.clone()))
753 .unwrap_err()
754 .to_string();
755
756 assert!(
757 err.contains("Could not find required column 'raw_umi'"),
758 "unexpected error: {err}"
759 );
760
761 fs::remove_file(path).ok();
762 }
763
764 #[test]
765 fn reads_gzipped_tsv_table() {
766 let path = tmp_path("table.tsv.gz");
767
768 let file = File::create(&path).unwrap();
769 let mut gz = GzEncoder::new(file, Compression::default());
770
771 gz.write_all(
772 concat!(
773 "read_id\traw_cb\traw_umi\n",
774 "read1\tCELL1\tUMI1\n",
775 "read2\tCELL2\tUMI2\n",
776 )
777 .as_bytes(),
778 )
779 .unwrap();
780
781 gz.finish().unwrap();
782
783 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
784
785 assert_eq!(table.len(), 2);
786 assert_eq!(
787 table.cell_umi_for_read("read1"),
788 Some(("CELL1".to_string(), "UMI1".to_string()))
789 );
790 assert_eq!(table.cell_umi_for_read("missing"), None);
791
792 fs::remove_file(path).ok();
793 }
794
795 #[test]
796 fn loads_binary_when_extension_is_bin() {
797 let path = tmp_path("read_tags.bin");
798
799 let mut table = ReadTagTable::new();
800 table.insert(ReadTagRecord::new(
801 "read1".to_string(),
802 Some("orig1".to_string()),
803 b"CELL1",
804 vec![40, 40, 40, 40],
805 b"UMI1",
806 vec![41, 41, 41, 41],
807 ));
808
809 table.save_binary(&path).unwrap();
810
811 let loaded = ReadTagTable::from_path_or_config(&default_config(path.clone())).unwrap();
812
813 assert_eq!(loaded.len(), 1);
814 assert_eq!(
815 loaded.cell_umi_for_read("read1"),
816 Some(("CELL1".to_string(), "UMI1".to_string()))
817 );
818
819 fs::remove_file(path).ok();
820 }
821
822 #[test]
823 fn cli_load_uses_binary_for_bin_extension() {
824 let path = tmp_path("cli_read_tags.bin");
825
826 let mut table = ReadTagTable::new();
827 table.insert(ReadTagRecord::new(
828 "read1".to_string(),
829 None,
830 b"CELL1",
831 Vec::new(),
832 b"UMI1",
833 Vec::new(),
834 ));
835
836 table.save_binary(&path).unwrap();
837
838 let cli = ReadTagTableCli {
839 read_tag_table: vec![path.clone()],
840 rt_read_id_column: "read_id".to_string(),
841 rt_cell_column: "raw_cb".to_string(),
842 rt_cell_qual_column: "quality_cb".to_string(),
843 rt_umi_column: "raw_umi".to_string(),
844 rt_umi_qual_column: "quality_umi".to_string(),
845 rt_original_read_id_column: "original_read_id".to_string(),
846 };
847
848 let loaded = cli.load().unwrap();
849
850 assert_eq!(loaded.len(), 1);
851 assert_eq!(
852 loaded.cell_umi_for_read("read1"),
853 Some(("CELL1".to_string(), "UMI1".to_string()))
854 );
855
856 fs::remove_file(path).ok();
857 }
858
859 #[test]
860 fn pair_counts_counts_cell_umi_observations() {
861 let path = tmp_path("pairs.tsv");
862
863 write_text(
864 &path,
865 concat!(
866 "read_id\traw_cb\traw_umi\n",
867 "read1\tCELL1\tUMI1\n",
868 "read2\tCELL1\tUMI1\n",
869 "read3\tCELL1\tUMI2\n",
870 "read4\tCELL2\tUMI3\n",
871 ),
872 );
873
874 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
875 let counts = table.pair_counts();
876
877 assert_eq!(counts.get(&(b"CELL1".to_vec(), b"UMI1".to_vec())), Some(&2));
878 assert_eq!(counts.get(&(b"CELL1".to_vec(), b"UMI2".to_vec())), Some(&1));
879 assert_eq!(counts.get(&(b"CELL2".to_vec(), b"UMI3".to_vec())), Some(&1));
880
881 fs::remove_file(path).ok();
882 }
883
884 #[test]
885 fn summarize_pairs_applies_pair_and_cell_thresholds() {
886 let path = tmp_path("summary.tsv");
887
888 write_text(
889 &path,
890 concat!(
891 "read_id\traw_cb\traw_umi\n",
892 "r1\tCELL1\tUMI1\n",
893 "r2\tCELL1\tUMI1\n",
894 "r3\tCELL1\tUMI2\n",
895 "r4\tCELL1\tUMI2\n",
896 "r5\tCELL2\tUMI3\n",
897 "r6\tCELL2\tUMI3\n",
898 "r7\tCELL3\tUMI4\n",
899 ),
900 );
901
902 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
903
904 let stats = table.summarize_pairs(2, 2);
905
906 assert_eq!(stats.cell_entries, 1);
907 assert_eq!(stats.unique_cell_umi_combos, 2);
908 assert_eq!(stats.total_pair_observations, 4);
909
910 assert_eq!(stats.umis_per_cell.n, 1);
911 assert_eq!(stats.umis_per_cell.min, 2);
912 assert_eq!(stats.umis_per_cell.max, 2);
913
914 assert_eq!(stats.detections_per_cell_umi.n, 2);
915 assert_eq!(stats.detections_per_cell_umi.min, 2);
916 assert_eq!(stats.detections_per_cell_umi.max, 2);
917
918 fs::remove_file(path).ok();
919 }
920
921 #[test]
922 fn display_for_record_is_informative() {
923 let rec = ReadTagRecord::new(
924 "read1".to_string(),
925 Some("orig1".to_string()),
926 b"CELL1",
927 Vec::new(),
928 b"UMI1",
929 vec![41, 41, 41, 41],
930 );
931
932 let text = rec.to_string();
933
934 assert!(text.contains("read_id=read1"));
935 assert!(text.contains("original_read_id=orig1"));
936 assert!(text.contains("cell=CELL1"));
937 assert!(text.contains("cell_qual=-"));
938 assert!(text.contains("umi=UMI1"));
939 assert!(text.contains("umi_qual=JJJJ"));
940 }
941
942 #[test]
943 fn display_for_table_shows_count_and_at_most_three_records() {
944 let path = tmp_path("display.tsv");
945
946 write_text(
947 &path,
948 concat!(
949 "read_id\traw_cb\traw_umi\n",
950 "read1\tCELL1\tUMI1\n",
951 "read2\tCELL2\tUMI2\n",
952 "read3\tCELL3\tUMI3\n",
953 "read4\tCELL4\tUMI4\n",
954 ),
955 );
956
957 let table = ReadTagTable::from_config(&default_config(path.clone())).unwrap();
958 let text = table.to_string();
959
960 assert!(text.contains("ReadTagTable: 4 records"));
961
962 let shown_records = text
963 .lines()
964 .filter(|line| line.trim_start().starts_with('['))
965 .count();
966
967 assert_eq!(shown_records, 3);
968
969 fs::remove_file(path).ok();
970 }
971
972 #[test]
973 fn writer_writes_header_and_rows() {
974 let mut out = Vec::new();
975
976 {
977 let mut writer = ReadTagTableWriter::new(&mut out).unwrap();
978
979 writer
980 .write_record(&ReadTagRecord::new(
981 "read1".to_string(),
982 Some("orig1".to_string()),
983 b"CELL1",
984 vec![40, 40, 40, 40],
985 b"UMI1",
986 vec![41, 41, 41, 41],
987 ))
988 .unwrap();
989
990 writer.flush().unwrap();
991 }
992
993 let text = String::from_utf8(out).unwrap();
994
995 assert!(text.starts_with(&READ_TAG_TABLE_COLUMNS.join("\t")));
996 assert!(text.contains("read1\torig1\t\tCELL1\tIIII\tUMI1\tJJJJ\tok"));
997 }
998}