1use crate::error::{Error, Result};
11use crate::genomic::ChrMap;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
14pub enum Strand {
15 #[default]
16 Forward,
17 Reverse,
18}
19
20impl Strand {
21 #[inline]
22 pub fn is_reverse(self) -> bool {
23 matches!(self, Strand::Reverse)
24 }
25}
26
27impl std::str::FromStr for Strand {
28 type Err = Error;
29
30 fn from_str(s: &str) -> Result<Self> {
33 match s {
34 "+" | "." | "" => Ok(Strand::Forward),
35 "-" => Ok(Strand::Reverse),
36 other => Err(Error::invalid(format!("Strand {other} invalid (+ or -)"))),
37 }
38 }
39}
40
41#[derive(Debug, Clone, Default, PartialEq)]
43pub struct Locs {
44 pub chr_ids: Vec<String>,
45 pub starts: Vec<i64>,
46 pub ends: Vec<i64>,
47 pub strands: Vec<Strand>,
49}
50
51impl Locs {
52 pub fn parse(
65 chr_ids: &[String],
66 starts: &[i64],
67 ends: &[i64],
68 centers: &[i64],
69 span: Option<i64>,
70 ) -> Result<Self> {
71 if chr_ids.is_empty() && starts.is_empty() && ends.is_empty() && centers.is_empty() {
72 return Ok(Self::default());
73 }
74 if span.is_some_and(|span| span < -1) {
75 return Err(Error::invalid(format!(
76 "span {} is negative (-1, or no span at all, means none was given)",
77 span.expect("checked just above")
78 )));
79 }
80
81 let (parsed_starts, parsed_ends) = match span {
82 Some(span) if span >= 0 => {
83 let given = [!starts.is_empty(), !ends.is_empty(), !centers.is_empty()];
84 if given.iter().filter(|g| **g).count() != 1 {
85 return Err(Error::invalid(
86 "Exactly one of starts, ends or centers must be specified when using span.",
87 ));
88 }
89 if !starts.is_empty() {
90 (starts.to_vec(), starts.iter().map(|s| s + span).collect())
91 } else if !ends.is_empty() {
92 (ends.iter().map(|e| e - span).collect(), ends.to_vec())
93 } else {
94 (
98 centers.iter().map(|c| c - span / 2).collect(),
99 centers.iter().map(|c| c + (span + 1) / 2).collect(),
100 )
101 }
102 }
103 _ => {
104 if starts.is_empty() || ends.is_empty() {
105 return Err(Error::invalid(
106 "Either starts+ends or exactly one of starts/ends/centers \
107 together with span must be specified.",
108 ));
109 }
110 (starts.to_vec(), ends.to_vec())
111 }
112 };
113
114 if chr_ids.len() != parsed_starts.len() || chr_ids.len() != parsed_ends.len() {
115 return Err(Error::invalid(format!(
116 "Length mismatch between chr_ids ({}) and starts/ends/centers ({}/{})",
117 chr_ids.len(),
118 parsed_starts.len(),
119 parsed_ends.len()
120 )));
121 }
122
123 Ok(Self {
124 chr_ids: chr_ids.to_vec(),
125 starts: parsed_starts,
126 ends: parsed_ends,
127 strands: Vec::new(),
128 })
129 }
130
131 pub fn spans(chr_ids: &[String], starts: &[i64], ends: &[i64]) -> Result<Self> {
133 Self::parse(chr_ids, starts, ends, &[], None)
134 }
135
136 pub fn from_starts(chr_ids: &[String], starts: &[i64], span: i64) -> Result<Self> {
138 Self::parse(chr_ids, starts, &[], &[], Some(span))
139 }
140
141 pub fn from_ends(chr_ids: &[String], ends: &[i64], span: i64) -> Result<Self> {
143 Self::parse(chr_ids, &[], ends, &[], Some(span))
144 }
145
146 pub fn centered(chr_ids: &[String], centers: &[i64], span: i64) -> Result<Self> {
148 Self::parse(chr_ids, &[], &[], centers, Some(span))
149 }
150
151 pub fn chromosomes(chr_ids: Vec<String>) -> Self {
159 Self {
160 chr_ids,
161 starts: Vec::new(),
162 ends: Vec::new(),
163 strands: Vec::new(),
164 }
165 }
166
167 pub fn whole_chromosomes(map: &ChrMap, chr_ids: &[String]) -> Result<Self> {
170 let selected = map.select(chr_ids)?;
171 Ok(Self {
172 chr_ids: selected.iter().map(|e| e.id.clone()).collect(),
173 starts: vec![0; selected.len()],
174 ends: selected.iter().map(|e| e.size).collect(),
175 strands: Vec::new(),
176 })
177 }
178
179 pub fn with_strands(mut self, strands: &[String]) -> Result<Self> {
183 if strands.is_empty() {
184 self.strands = Vec::new();
185 return Ok(self);
186 }
187 if strands.len() != self.chr_ids.len() {
188 return Err(Error::invalid(format!(
189 "Length mismatch between strands ({}) and the {} loci of the request",
190 strands.len(),
191 self.chr_ids.len()
192 )));
193 }
194 self.strands = strands
195 .iter()
196 .enumerate()
197 .map(|(i, s)| {
198 s.parse::<Strand>().map_err(|_| {
199 Error::invalid(format!("Strand {s} at index {i} invalid (+ or -)"))
200 })
201 })
202 .collect::<Result<_>>()?;
203 Ok(self)
204 }
205
206 #[inline]
208 pub fn strand(&self, i: usize) -> Strand {
209 self.strands.get(i).copied().unwrap_or_default()
210 }
211
212 pub fn len(&self) -> usize {
213 self.chr_ids.len()
214 }
215
216 pub fn is_empty(&self) -> bool {
217 self.chr_ids.is_empty()
218 }
219}
220
221#[derive(Debug, Clone)]
225pub struct IndexedLoc {
226 pub chr_index: usize,
229 pub start: i64,
230 pub end: i64,
231 pub binned_start: i64,
235 pub binned_end: i64,
236 pub bin_size: f64,
240 pub reverse: bool,
241 pub output_start: usize,
245 pub output_end: usize,
246}
247
248impl IndexedLoc {
249 #[inline]
251 pub fn span(&self) -> i64 {
252 self.binned_end - self.binned_start
253 }
254
255 #[inline]
257 pub fn bin_count(&self) -> i64 {
258 (self.output_end - self.output_start) as i64
259 }
260
261 #[inline]
269 pub fn bin_at(&self, pos: i64) -> i64 {
270 ((pos - self.binned_start) * self.bin_count()) / self.span()
271 }
272
273 #[inline]
275 pub fn bin_after(&self, pos: i64) -> i64 {
276 let span = self.span();
277 ((pos - self.binned_start) * self.bin_count() + span - 1) / span
278 }
279
280 #[inline]
290 pub fn bin_coverage(&self, b: i64, from: i64, to: i64) -> f64 {
291 let n = self.bin_count();
292 let s = self.span();
293 let lo = ((from - self.binned_start) * n).max(b * s);
294 let hi = ((to - self.binned_start) * n).min((b + 1) * s);
295 if hi > lo {
296 (hi - lo) as f64 / n as f64
297 } else {
298 0.0
299 }
300 }
301
302 #[inline]
309 pub fn bin_fraction(&self, b: i64, from: i64, to: i64) -> f64 {
310 let n = self.bin_count();
311 let s = self.span();
312 let lo = ((from - self.binned_start) * n).max(b * s);
313 let hi = ((to - self.binned_start) * n).min((b + 1) * s);
314 if hi > lo {
315 (hi - lo) as f64 / s as f64
316 } else {
317 0.0
318 }
319 }
320
321 #[inline]
329 pub fn row(&self, bin_count: usize) -> usize {
330 self.output_start.checked_div(bin_count).unwrap_or(0)
331 }
332}
333
334#[derive(Debug, Clone, Default)]
336pub struct IndexedLocs {
337 pub locs: Vec<IndexedLoc>,
341 pub bin_count: usize,
344 pub output_len: usize,
346}
347
348#[derive(Debug, Clone, Copy, PartialEq, Eq)]
352pub struct LocBatch {
353 pub start: usize,
354 pub end: usize,
355}
356
357impl LocBatch {
358 #[inline]
359 pub fn is_empty(&self) -> bool {
360 self.start >= self.end
361 }
362 #[inline]
363 pub fn len(&self) -> usize {
364 self.end.saturating_sub(self.start)
365 }
366}
367
368impl IndexedLocs {
369 pub fn build(
374 map: &ChrMap,
375 locs: &Locs,
376 bin_size: f64,
377 bin_count: Option<usize>,
378 full_bin: bool,
379 ) -> Result<Self> {
380 let bin = crate::genomic::BinPlan::new(bin_size, bin_count, full_bin)?.whole_bin_size();
384 if locs.is_empty() {
385 return Ok(Self {
386 locs: Vec::new(),
387 bin_count: bin_count.unwrap_or(0),
388 output_len: 0,
389 });
390 }
391
392 let mut built = Vec::with_capacity(locs.len());
393 let mut max_binned_span = 0i64;
394 for i in 0..locs.len() {
395 let entry = map.resolve(&locs.chr_ids[i])?;
396 let (start, end) = (locs.starts[i], locs.ends[i]);
397 if start > end {
398 return Err(Error::invalid(format!(
399 "Locus {}:{start}-{end} at index {i} ends before it starts",
400 locs.chr_ids[i]
401 )));
402 }
403 let binned_start = start.div_euclid(bin) * bin;
415 let binned_end = if full_bin {
416 end.saturating_add(bin - 1).div_euclid(bin) * bin
417 } else {
418 end.div_euclid(bin) * bin
419 };
420 max_binned_span = max_binned_span.max(binned_end - binned_start);
421 built.push(IndexedLoc {
422 chr_index: entry.index,
423 start,
424 end,
425 binned_start,
426 binned_end,
427 bin_size: 0.0,
428 reverse: locs.strand(i).is_reverse(),
429 output_start: 0,
430 output_end: 0,
431 });
432 }
433
434 let bin_count = match bin_count {
435 Some(n) => n,
436 None => (max_binned_span / bin) as usize,
437 };
438 if bin_count < 1 {
439 return Err(Error::invalid(format!(
440 "No locus is a full bin wide (widest binned locus: {max_binned_span} bp, \
441 bin size: {bin_size})"
442 )));
443 }
444
445 for (i, loc) in built.iter_mut().enumerate() {
446 loc.bin_size = (loc.binned_end - loc.binned_start) as f64 / bin_count as f64;
447 loc.output_start = i * bin_count;
448 loc.output_end = loc.output_start + bin_count;
449 }
450 let output_len = built.len() * bin_count;
451
452 built.sort_by(|a, b| {
456 (a.chr_index, a.binned_start, a.binned_end).cmp(&(
457 b.chr_index,
458 b.binned_start,
459 b.binned_end,
460 ))
461 });
462
463 Ok(Self {
464 locs: built,
465 bin_count,
466 output_len,
467 })
468 }
469
470 pub fn coverage(&self) -> u64 {
476 self.locs
477 .iter()
478 .map(|l| (l.binned_end - l.binned_start).max(0) as u64)
479 .sum()
480 }
481
482 pub fn effective_bin_size(&self, fallback: f64) -> f64 {
493 let narrowest = self
494 .locs
495 .iter()
496 .map(|l| l.bin_size)
497 .filter(|s| *s > 0.0)
498 .fold(f64::INFINITY, f64::min);
499 if narrowest.is_finite() {
500 narrowest
501 } else {
502 fallback
503 }
504 }
505
506 pub fn batches(&self, parallel: usize) -> (Vec<LocBatch>, u64) {
519 let total = self.coverage();
520 if self.locs.is_empty() {
521 return (Vec::new(), total);
522 }
523 let parallel = parallel.max(1) as u64;
524 let per_batch = total.div_ceil(parallel).max(1);
526
527 let mut batches = vec![LocBatch { start: 0, end: 0 }];
528 let mut coverage = 0u64;
529 for i in 0..self.locs.len().saturating_sub(1) {
530 let loc = &self.locs[i];
531 coverage += (loc.binned_end - loc.binned_start).max(0) as u64;
532 if coverage < per_batch {
533 continue;
534 }
535 if batches.len() as u64 >= parallel {
541 continue;
542 }
543 batches.last_mut().unwrap().end = i + 1;
544 batches.push(LocBatch {
545 start: i + 1,
546 end: i + 1,
547 });
548 coverage = 0;
549 }
550 batches.last_mut().unwrap().end = self.locs.len();
551 (batches, total)
552 }
553
554 pub fn reverse_output_rows(&self, output: &mut [f32]) {
561 for loc in &self.locs {
562 if loc.reverse {
563 output[loc.output_start..loc.output_end].reverse();
564 }
565 }
566 }
567}
568
569#[cfg(test)]
570mod tests {
571 use super::*;
572
573 fn ids(n: usize) -> Vec<String> {
574 (0..n).map(|_| "chr1".to_string()).collect()
575 }
576
577 #[test]
578 fn starts_and_ends_pass_through() {
579 let l = Locs::spans(&ids(2), &[10, 20], &[15, 25]).unwrap();
580 assert_eq!(l.starts, [10, 20]);
581 assert_eq!(l.ends, [15, 25]);
582 }
583
584 #[test]
585 fn span_extends_from_starts_and_back_from_ends() {
586 assert_eq!(
587 Locs::from_starts(&ids(2), &[10, 20], 5).unwrap().ends,
588 [15, 25]
589 );
590 assert_eq!(
591 Locs::from_ends(&ids(2), &[15, 25], 5).unwrap().starts,
592 [10, 20]
593 );
594 }
595
596 #[test]
597 fn centered_windows_are_asymmetric_on_an_odd_span() {
598 let l = Locs::centered(&ids(1), &[100], 5).unwrap();
599 assert_eq!((l.starts[0], l.ends[0]), (98, 103));
600 let l = Locs::centered(&ids(1), &[100], 4).unwrap();
601 assert_eq!((l.starts[0], l.ends[0]), (98, 102));
602 }
603
604 #[test]
605 fn no_loci_at_all_is_a_request_for_nothing() {
606 let l = Locs::parse(&[], &[], &[], &[], None).unwrap();
607 assert!(l.is_empty());
608 }
609
610 #[test]
611 fn ambiguous_and_missing_combinations_are_refused() {
612 let e = Locs::parse(&ids(1), &[1], &[2], &[], Some(5))
614 .unwrap_err()
615 .to_string();
616 assert!(e.contains("Exactly one of starts, ends or centers"), "{e}");
617 let e = Locs::parse(&ids(1), &[], &[], &[5], None)
619 .unwrap_err()
620 .to_string();
621 assert!(e.contains("Either starts+ends"), "{e}");
622 assert!(Locs::parse(&ids(1), &[1], &[], &[], None).is_err());
624 }
625
626 #[test]
627 fn length_mismatch_names_both_counts() {
628 let e = Locs::spans(&ids(3), &[1, 2], &[2, 3])
629 .unwrap_err()
630 .to_string();
631 assert!(e.contains("chr_ids (3)") && e.contains("(2/2)"), "{e}");
632 }
633
634 #[test]
635 fn strands_default_forward_and_accept_dot_and_empty() {
636 let l = Locs::spans(&ids(3), &[0, 0, 0], &[1, 1, 1]).unwrap();
637 assert_eq!(l.strand(0), Strand::Forward);
638 let l = l
639 .with_strands(&["+".into(), "-".into(), ".".into()])
640 .unwrap();
641 assert_eq!(l.strand(1), Strand::Reverse);
642 assert_eq!(l.strand(2), Strand::Forward);
643 let l = Locs::spans(&ids(1), &[0], &[1])
644 .unwrap()
645 .with_strands(&["".into()])
646 .unwrap();
647 assert_eq!(l.strand(0), Strand::Forward);
648 }
649
650 #[test]
651 fn a_bad_strand_names_its_index() {
652 let e = Locs::spans(&ids(2), &[0, 0], &[1, 1])
653 .unwrap()
654 .with_strands(&["+".into(), "?".into()])
655 .unwrap_err()
656 .to_string();
657 assert!(e.contains("Strand ? at index 1"), "{e}");
658 }
659
660 #[test]
661 fn strand_count_must_match_loci() {
662 let e = Locs::spans(&ids(2), &[0, 0], &[1, 1])
663 .unwrap()
664 .with_strands(&["+".into()])
665 .unwrap_err()
666 .to_string();
667 assert!(e.contains("strands (1)") && e.contains("2 loci"), "{e}");
668 }
669
670 #[test]
671 fn chromosomes_names_without_coordinates() {
672 let l = Locs::chromosomes(vec!["chr1".into(), "chr2".into()]);
673 assert_eq!(l.chr_ids, ["chr1", "chr2"]);
674 assert!(l.starts.is_empty() && l.ends.is_empty());
675 assert!(Locs::parse(&l.chr_ids, &[], &[], &[], None).is_err());
677 }
678
679 #[test]
680 fn whole_chromosomes_span_each_one() {
681 let map = ChrMap::from_entries([("chr1".into(), 100), ("chr2".into(), 50)]);
682 let l = Locs::whole_chromosomes(&map, &[]).unwrap();
683 assert_eq!(l.chr_ids, ["chr1", "chr2"]);
684 assert_eq!(l.starts, [0, 0]);
685 assert_eq!(l.ends, [100, 50]);
686 let l = Locs::whole_chromosomes(&map, &["2".into()]).unwrap();
687 assert_eq!(l.chr_ids, ["chr2"]);
688 assert_eq!(l.ends, [50]);
689 }
690
691 #[test]
692 fn reverse_output_rows_mirrors_only_reverse_loci() {
693 let locs = IndexedLocs {
694 locs: vec![
695 IndexedLoc {
696 chr_index: 0,
697 start: 0,
698 end: 3,
699 binned_start: 0,
700 binned_end: 3,
701 bin_size: 1.0,
702 reverse: false,
703 output_start: 0,
704 output_end: 3,
705 },
706 IndexedLoc {
707 chr_index: 0,
708 start: 0,
709 end: 3,
710 binned_start: 0,
711 binned_end: 3,
712 bin_size: 1.0,
713 reverse: true,
714 output_start: 3,
715 output_end: 6,
716 },
717 ],
718 bin_count: 3,
719 output_len: 6,
720 };
721 let mut out = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
722 locs.reverse_output_rows(&mut out);
723 assert_eq!(out, [1.0, 2.0, 3.0, 6.0, 5.0, 4.0]);
724 }
725}