Finish codon-based trimming

Author: TJBIII

clipkit/api.py

@@ -39,6 +39,7 @@ def clipkit(
 
     # override some options not currently available through programmatic interface
     complement = False
+    codon = False
     use_log = False
     quiet = True
 
@@ -51,6 +52,7 @@ def clipkit(
         gaps,
         gap_characters,
         complement,
+        codon,
         TrimmingMode(mode),
         use_log,
         quiet,

clipkit/clipkit.py

@@ -47,6 +47,7 @@ class TrimRun:
     input_file_format: FileFormat
     output_file_format: FileFormat
     gaps: float
+    codon: bool
     version: str = current_version
 
     @property
@@ -67,6 +68,7 @@ def run(
     gaps: float,
     gap_characters: Union[list, None],
     complement: bool,
+    codon: bool,
     mode: TrimmingMode,
     use_log: bool,
     quiet: bool,
@@ -99,7 +101,7 @@ def run(
         gaps = smart_gap_threshold_determination(alignment, gap_characters)
 
     msa = create_msa(alignment, gap_characters)
-    msa.trim(mode, gap_threshold=gaps)
+    msa.trim(mode, gap_threshold=gaps, site_positions_to_trim=None, codon=codon)
 
     trim_run = TrimRun(
         alignment,
@@ -109,6 +111,7 @@ def run(
         input_file_format,
         output_file_format,
         gaps,
+        codon,
     )
 
     return trim_run, msa.stats
@@ -123,6 +126,7 @@ def execute(
     gaps: float,
     gap_characters: Union[list, None],
     complement: bool,
+    codon: bool,
     mode: TrimmingMode,
     use_log: bool,
     quiet: bool,
@@ -151,6 +155,7 @@ def execute(
         gaps,
         gap_characters,
         complement,
+        codon,
         mode,
         use_log,
         quiet,
@@ -168,6 +173,7 @@ def execute(
         trim_run.gap_characters,
         mode,
         complement,
+        codon,
         use_log,
     )
 

clipkit/msa.py

@@ -2,6 +2,7 @@
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 import numpy as np
+from itertools import chain
 from typing import Union
 
 from .modes import TrimmingMode
@@ -113,7 +114,12 @@ def trim(
                 site_positions_to_trim = np.array(site_positions_to_trim)
             if not isinstance(site_positions_to_trim, np.ndarray):
                 raise ValueError("site_positions_to_trim must be a list or np array")
-            self._site_positions_to_trim = site_positions_to_trim
+
+            self._site_positions_to_trim = (
+                self.determine_all_codon_sites_to_trim(site_positions_to_trim)
+                if codon is True
+                else site_positions_to_trim
+            )
         else:
             self._site_positions_to_trim = self.determine_site_positions_to_trim(
                 mode,
@@ -210,10 +216,8 @@ def determine_site_positions_to_trim(self, mode, gap_threshold, codon=False):
             Example:
                 [2, 9] -> [1, 2, 3, 7, 8, 9]
             """
-            sites_to_trim = map(
-                self.determine_codon_triplet_positions, sites_to_trim
-            ).flatten()
-            print(sites_to_trim)
+            return self.determine_all_codon_sites_to_trim(sites_to_trim)
+
         return sites_to_trim
 
     def generate_debug_log_info(self):
@@ -234,10 +238,31 @@ def generate_debug_log_info(self):
                 gappyness,
             )
 
-    def determine_codon_triplet_positions(alignment_position):
+    def determine_all_codon_sites_to_trim(self, sites_to_trim):
+        """
+        For each position in sites_to_trim we need the full triplet of codon positions.
+
+        Sites to trim -> all codon sites to trim
+        [2, 8] -> [0, 1, 2, 6, 7, 8]
+        """
+        sites_to_trim_codon = [
+            self.determine_codon_triplet_positions(site_pos)
+            for site_pos in sites_to_trim
+        ]
+        flattened_unique_sites = list(set(chain(*sites_to_trim_codon)))
+        return np.array(flattened_unique_sites)
+
+    def determine_codon_triplet_positions(self, alignment_position):
+        """
+        Block 0 -> [0,1,2], block 1 -> [3,4,5]
+
+        We filter to make sure we are not including any positions out of range
+        """
         block = alignment_position // 3
-        remainder = alignment_position % 3
-        if remainder:
-            block += 1
-        codon_triplet_index = block * 3
-        return [codon_triplet_index - 2, codon_triplet_index - 1, codon_triplet_index]
+        codon_triplet_index_start = block * 3
+        sites = [
+            codon_triplet_index_start,
+            codon_triplet_index_start + 1,
+            codon_triplet_index_start + 2,
+        ]
+        return list(filter(lambda x: x <= self._original_length - 1, sites))

clipkit/write.py

@@ -34,6 +34,7 @@ def write_user_args(
     gap_chars: list,
     mode: "TrimmingMode",
     complement: bool,
+    codon: bool,
     use_log: bool,
 ) -> None:
     if seq_type.value == "nt":
@@ -57,6 +58,7 @@ def write_user_args(
     Gap characters: {gap_chars}
     Trimming mode: {mode.value}
     Create complementary output: {complement}
+    Process as codons: {codon}
     Create log file: {use_log}
     """  # noqa
         )
@@ -77,9 +79,9 @@ def write_output_files_message(
         ------------------------
         | Writing output files |
         ------------------------
-        trimmed alignment: {out_file_name}
-        complement file: {out_file_name + '.complement' if complement else False}
-        log file: {out_file_name + '.log' if use_log else False}
+        Trimmed alignment: {out_file_name}
+        Complement file: {out_file_name + '.complement' if complement else False}
+        Log file: {out_file_name + '.log' if use_log else False}
     """
         )
     )

tests/integration/test_codon.py

@@ -26,7 +26,7 @@ def test_simple_codon(self):
             sequence_type=None,
             complement=False,
             codon=True,
-            gaps=0.2,
+            gaps=0.8,
             mode=TrimmingMode.gappy,
             use_log=False,
             gap_characters=DEFAULT_NT_GAP_CHARS,
@@ -42,3 +42,33 @@ def test_simple_codon(self):
             output_content = out_file.read()
 
         assert expected_content == output_content
+
+    def test_simple_codon_all_trimmed(self):
+        """
+        test codon
+        usage: clipkit simple.fa -co
+        """
+        output_file = "output/simple.fa_gappy_codon"
+        kwargs = dict(
+            input_file=f"{here.parent}/samples/simple.fa",
+            output_file=output_file,
+            input_file_format="fasta",
+            output_file_format="fasta",
+            sequence_type=None,
+            complement=False,
+            codon=True,
+            gaps=0.1,
+            mode=TrimmingMode.gappy,
+            use_log=False,
+            gap_characters=DEFAULT_NT_GAP_CHARS,
+            quiet=True,
+        )
+
+        execute(**kwargs)
+
+        expected_content = ">1\n>2\n>3\n>4\n>5\n"
+
+        with open(output_file, "r") as out_file:
+            output_content = out_file.read()
+
+        assert expected_content == output_content

tests/unit/test_msa.py

@@ -62,3 +62,75 @@ def test_trim_by_provided_site_positions_list(self):
             ]
         )
         np.testing.assert_equal(msa.sites_kept, expected_sites_kept)
+
+    @pytest.mark.parametrize(
+        "sites_to_trim, expected",
+        [
+            (
+                [0],
+                np.array(
+                    [
+                        ["T", "A", "T"],
+                        ["-", "A", "T"],
+                        ["-", "T", "A"],
+                        ["-", "T", "A"],
+                        ["-", "T", "-"],
+                    ]
+                ),
+            ),
+            (
+                [2],
+                np.array(
+                    [
+                        ["T", "A", "T"],
+                        ["-", "A", "T"],
+                        ["-", "T", "A"],
+                        ["-", "T", "A"],
+                        ["-", "T", "-"],
+                    ]
+                ),
+            ),
+            (
+                [3],
+                np.array(
+                    [
+                        ["A", "-", "G"],
+                        ["A", "-", "G"],
+                        ["A", "-", "G"],
+                        ["A", "G", "A"],
+                        ["A", "C", "a"],
+                    ]
+                ),
+            ),
+            (
+                [5],
+                np.array(
+                    [
+                        ["A", "-", "G"],
+                        ["A", "-", "G"],
+                        ["A", "-", "G"],
+                        ["A", "G", "A"],
+                        ["A", "C", "a"],
+                    ]
+                ),
+            ),
+            (
+                [0, 1, 2, 3, 4, 5],
+                np.array(
+                    [
+                        [],
+                        [],
+                        [],
+                        [],
+                        [],
+                    ],
+                    dtype=object,
+                ),
+            ),
+        ],
+    )
+    def test_trim_codons(self, sites_to_trim, expected):
+        bio_msa = get_biopython_msa("tests/unit/examples/simple.fa")
+        msa = MSA.from_bio_msa(bio_msa)
+        msa.trim(site_positions_to_trim=sites_to_trim, codon=True)
+        np.testing.assert_equal(msa.trimmed, expected)