Merge pull request #308 from neherlab/refactor/indels

neherlab · Jan 2, 2025 · 56cd9ad · 56cd9ad
2 parents db3e1e8 + 7b5ad99
commit 56cd9ad
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 31 deletions.
diff --git a/packages/treetime/src/commands/ancestral/fitch.rs b/packages/treetime/src/commands/ancestral/fitch.rs
@@ -337,21 +337,13 @@ fn fitch_forward(graph: &SparseGraph, sparse_partitions: &[PartitionParsimony])
             if gap_in_parent == 0 {
               // If the gap is not in parent, add deletion.
               // the sequence that is deleted is the sequence of the parent
-              let indel = InDel {
-                range: *r,
-                seq: parent.sequence[r.0..r.1].into(),
-                deletion: true,
-              };
+              let indel = InDel::del(*r, &parent.sequence[r.0..r.1]);
               composition.add_indel(&indel);
               edge.indels.push(indel);
             }
           } else if gap_in_parent > 0 {
             // Add insertion if gap is present in parent.
-            let indel = InDel {
-              range: *r,
-              seq: sequence[r.0..r.1].into(),
-              deletion: false,
-            };
+            let indel = InDel::ins(*r, &sequence[r.0..r.1]);
             composition.add_indel(&indel);
             edge.indels.push(indel);
           }
@@ -363,11 +355,7 @@ fn fitch_forward(graph: &SparseGraph, sparse_partitions: &[PartitionParsimony])
             // all gaps in variable_indel are already processed
             continue;
           }
-          let indel = InDel {
-            range: r,
-            seq: sequence[r.0..r.1].into(),
-            deletion: true,
-          };
+          let indel = InDel::del(r, &sequence[r.0..r.1]);
           composition.add_indel(&indel);
           edge.indels.push(indel);
         }
@@ -378,11 +366,7 @@ fn fitch_forward(graph: &SparseGraph, sparse_partitions: &[PartitionParsimony])
             // all gaps in variable_indel are already processed
             continue;
           }
-          let indel = InDel {
-            range: r,
-            seq: sequence[r.0..r.1].into(),
-            deletion: false,
-          };
+          let indel = InDel::ins(r, &sequence[r.0..r.1]);
           composition.add_indel(&indel);
           edge.indels.push(indel);
         }

diff --git a/packages/treetime/src/seq/composition.rs b/packages/treetime/src/seq/composition.rs
@@ -90,7 +90,6 @@ mod tests {
   use super::*;
   use crate::alphabet::alphabet::{Alphabet, AlphabetName};
   use crate::representation::seq::Seq;
-  use crate::seq;
   use crate::seq::mutation::Sub;
   use maplit::btreemap;
   use pretty_assertions::assert_eq;
@@ -167,11 +166,7 @@ mod tests {
   #[test]
   fn test_composition_add_deletion() {
     let mut actual = Composition::with_sequence("AAAGCTTACGGGGTCAAGTCC".bytes(), "ACGT-".bytes(), b'-');
-    let indel = InDel {
-      range: (1, 5),
-      seq: seq!['A', 'A', 'G', 'C'],
-      deletion: true,
-    };
+    let indel = InDel::del((1, 5), "AAGC");
     actual.add_indel(&indel);
     let expected = Composition::from(btreemap! { b'-' => 4, b'A' => 4, b'C' => 4, b'G' => 5, b'T' => 4}, b'-');
     assert_eq!(expected, actual);
@@ -180,11 +175,7 @@ mod tests {
   #[test]
   fn test_composition_add_insertion() {
     let mut actual = Composition::with_sequence("AAAGCTTACGGGGTCAAGTCC".bytes(), "ACGT-".bytes(), b'-');
-    let indel = InDel {
-      range: (3, 6),
-      seq: seq!['A', 'T', 'C'],
-      deletion: false,
-    };
+    let indel = InDel::ins((3, 6), "ATC");
     actual.add_indel(&indel);
     let expected = Composition::from(btreemap! { b'-' => 0, b'A' => 7, b'C' => 6, b'G' => 6, b'T' => 5}, b'-');
     assert_eq!(expected, actual);

diff --git a/packages/treetime/src/seq/indel.rs b/packages/treetime/src/seq/indel.rs
@@ -9,6 +9,23 @@ pub struct InDel {
   pub deletion: bool, // deletion if True, insertion if False
 }
 
+impl InDel {
+  pub fn del(range: (usize, usize), seq: impl Into<Seq>) -> Self {
+    Self::new(range, seq, true)
+  }
+
+  pub fn ins(range: (usize, usize), seq: impl Into<Seq>) -> Self {
+    Self::new(range, seq, false)
+  }
+
+  pub fn new(range: (usize, usize), seq: impl Into<Seq>, deletion: bool) -> Self {
+    let seq = seq.into();
+    assert!(range.0 <= range.1);
+    assert_eq!(seq.len(), range.1 - range.0);
+    Self { range, seq, deletion }
+  }
+}
+
 impl fmt::Display for InDel {
   fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
     let delta_str = if self.deletion {