Repository: cogent3/cogent3
Pull request #1970 (2024-07-31 00:26:59)

ENH: select_translatable frame argument
[NEW] users can now specify that all sequences in a collection are in
    the nominated frame.


Patch:
diff --git a/changelog.d/20240731_102115_Gavin.Huttley.md b/changelog.d/20240731_102115_Gavin.Huttley.md
new file mode 100644
index 000000000..a18db1e8d
--- /dev/null
+++ b/changelog.d/20240731_102115_Gavin.Huttley.md
@@ -0,0 +1,41 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+
+### Contributors
+
+- @GavinHuttley
+
+
+### ENH
+
+- The cogent3 select_translatable() app gets a `frame` argument. This allows
+  specifying the frame of reference for the translation for all sequences.
+
+<!--
+### BUG
+
+- A bullet item for the BUG category.
+
+-->
+<!--
+### DOC
+
+- A bullet item for the DOC category.
+
+-->
+<!--
+### Deprecations
+
+- A bullet item for the Deprecations category.
+
+-->
+<!--
+### Discontinued
+
+- A bullet item for the Discontinued category.
+
+-->
diff --git a/src/cogent3/app/translate.py b/src/cogent3/app/translate.py
index 13c5098ee..b2b9c6da5 100644
--- a/src/cogent3/app/translate.py
+++ b/src/cogent3/app/translate.py
@@ -6,7 +6,7 @@ from cogent3.core.genetic_code import GeneticCode, get_code
 from cogent3.core.moltype import Alphabet, MolType, get_moltype
 
 from .composable import NotCompleted, define_app
-from .typing import SeqsCollectionType, SerialisableType
+from .typing import SeqsCollectionType, SeqType, SerialisableType
 
 GeneticCodeTypes = Union[str, int, GeneticCode]
 MolTypes = Union[str, MolType]
@@ -166,6 +166,7 @@ class select_translatable:
         gc: GeneticCodeTypes = 1,
         allow_rc: bool = False,
         trim_terminal_stop: bool = True,
+        frame: Optional[int] = None,
     ):
         """
         Parameters
@@ -181,6 +182,9 @@ class select_translatable:
             best frame on rc, it will be negative
         trim_terminal_stop
             exclude terminal stop codon from seqs
+        frame
+            specify the coding frame (as an integer 1, 2, 3). If not specified
+            (default) uses best_frame.
 
         Returns
         -------
@@ -189,7 +193,6 @@ class select_translatable:
 
         Examples
         --------
-
         Create a sample sequence collection and an app that returns the sequences
         which are translatable.
 
@@ -202,19 +205,38 @@ class select_translatable:
         ... )
         >>> app = get_app("select_translatable")
         >>> result = app(aln)
-        >>> print(result.to_dict())
+        >>> result.to_dict()
         {'s1': 'AATATAAATGCCAGCTCATTACAGCATGAGAACAGCAGTTTATTACTTCATAAAGTCATA'}
 
         Use ``allow_rc=True`` to consider the reading frame for reverse strands.
 
         >>> app_rc = get_app("select_translatable", allow_rc=True)
         >>> result = app_rc(aln)
-        >>> print(result.to_dict())
+        >>> result.to_dict()
         {'s1': 'AATATAAATGCCAGCTCATTACAGCATGAGAACAGCAGTTTATTACTTCATAAAGTCATA', 's1_rc': 'AATATAAATGCCAGCTCATTACAGCATGAGAACAGCAGTTTATTACTTCATAAAGTCATA'}
+
+        If you know the sequences are all in a specific frame, you can specify
+        that using ``frame=<number>`` where the number is 1, 2 or 3. We make
+        two sequences in frame 1 and add an internal stop to one
+
+        >>> aln = make_unaligned_seqs(
+        ...     {
+        ...         "internal_stop": "AATTAAATGTGA",
+        ...         "s2": "TATGACTAA",
+        ...     }
+        ... )
+        >>> app = get_app("select_translatable", frame=1)
+        >>> result = app(aln)
+        >>> result.to_dict()
+        {'s2': 'TATGAC'}
         """
         moltype = get_moltype(moltype)
         assert moltype.label.lower() in ("dna", "rna"), "Invalid moltype"
 
+        if frame is not None:
+            assert 1 <= frame <= 3, f"{frame} not 1, 2 or 3"
+        self._frame = frame
+
         self._moltype = moltype
         self._gc = get_code(gc)
         self._allow_rc = allow_rc
@@ -222,6 +244,14 @@ class select_translatable:
 
     T = Union[SerialisableType, SeqsCollectionType]
 
+    def _get_frame(self, seq: SeqType) -> int:
+        if self._frame is not None:
+            tr = self._gc.translate(str(seq), start=self._frame - 1)
+            if "*" in tr[:-1]:
+                raise ValueError(f"internal stop in {seq.name}")
+            return self._frame
+        return best_frame(seq, self._gc, allow_rc=self._allow_rc)
+
     def main(self, seqs: SeqsCollectionType) -> T:
         """returns the translatable sequences from seqs.
 
@@ -234,7 +264,7 @@ class select_translatable:
         error_log = []
         for seq in seqs.seqs:
             try:
-                frame = best_frame(seq, self._gc, allow_rc=self._allow_rc)
+                frame = self._get_frame(seq)
                 if frame < 0:
                     seq = seq.rc()
                     frame *= -1
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 5bf9ec86f..a2984e56f 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -474,7 +474,7 @@ def index_to_coord(index: int, coeffs: numpy.ndarray):  # pragma: no cover
     return coord
 
 
-@numba.jit()
+@numba.jit(nopython=True)
 def seq_to_kmer_indices(
     seq: numpy.ndarray,
     result: numpy.ndarray,
@@ -538,7 +538,7 @@ def seq_to_kmer_indices(
 
 
 # todo: profile this against pure python
-@numba.jit()
+@numba.jit(nopython=True)
 def kmer_indices_to_seq(
     kmer_indices: numpy.ndarray,
     result: numpy.ndarray,


Test patch:
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index 7df4c6d98..9d83518b0 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -498,14 +498,13 @@ def test_str():
     got = str(func)
     assert (
         got
-        == "select_translatable(moltype='dna', gc=1, allow_rc=False,\ntrim_terminal_stop=True)"
+        == "select_translatable(moltype='dna', gc=1, allow_rc=False,\ntrim_terminal_stop=True, frame=None)"
     )
 
     func = select_translatable(allow_rc=True)
     got = str(func)
-    assert (
-        got
-        == "select_translatable(moltype='dna', gc=1, allow_rc=True, trim_terminal_stop=True)"
+    assert got.startswith(
+        "select_translatable(moltype='dna', gc=1, allow_rc=True, trim_terminal_stop=True"
     )
 
     nodegen = omit_degenerates()
diff --git a/tests/test_app/test_translate.py b/tests/test_app/test_translate.py
index 113960b13..f73d40ff5 100644
--- a/tests/test_app/test_translate.py
+++ b/tests/test_app/test_translate.py
@@ -1,5 +1,6 @@
 from unittest import TestCase
 
+import pytest
 from cogent3 import DNA, make_aligned_seqs, make_unaligned_seqs
 from cogent3.app.composable import NotCompleted
 from cogent3.app.translate import (
@@ -48,17 +49,17 @@ class TestTranslatable(TestCase):
     def test_select_translatable(self):
         """correctly get translatable seqs"""
         data = {
-            "a": "AATATAAATGCCAGCTCATTACAGCATGAGAACA" "GCAGTTTATTACTTCATAAAGTCATA",
-            "rc": "TATGACTTTATGAAGTAATAAACTGCTGTTCTCA" "TGCTGTAATGAGCTGGCATTTATATT",
+            "a": "AATATAAATGCCAGCTCATTACAGCATGAGAACAGCAGTTTATTACTTCATAAAGTCATA",
+            "rc": "TATGACTTTATGAAGTAATAAACTGCTGTTCTCATGCTGTAATGAGCTGGCATTTATATT",
         }
         seqs = make_unaligned_seqs(data=data, moltype=DNA)
         trans = select_translatable(allow_rc=False)
-        tr = trans(seqs)
+        tr = trans(seqs)  # pylint: disable=not-callable
         ex = data.copy()
         ex.pop("rc")
         self.assertEqual(tr.to_dict(), ex)
         trans = select_translatable(allow_rc=True)
-        tr = trans(seqs)
+        tr = trans(seqs)  # pylint: disable=not-callable
         ex = data.copy()
         ex["rc"] = data["a"]
         self.assertEqual(tr.to_dict(), ex)
@@ -66,7 +67,7 @@ class TestTranslatable(TestCase):
         # if seqs not translatable returns NotCompletedResult
         data = dict(a="TAATTGATTAA", b="GCAGTTTATTA")
         seqs = make_unaligned_seqs(data=data, moltype=DNA)
-        got = select_translatable(allow_rc=False)
+        got = select_translatable(allow_rc=False)(seqs)  # pylint: disable=not-callable
         self.assertTrue(type(got), NotCompleted)
 
     def test_translate_frames(self):
@@ -86,12 +87,12 @@ class TestTranslate(TestCase):
         seqs = make_unaligned_seqs(seqs, moltype="dna")
         # trim terminal stops
         translater = translate_seqs()
-        aa = translater(seqs)
+        aa = translater(seqs)  # pylint: disable=not-callable
         self.assertEqual(aa.to_dict(), dict(a="MR", b="M"))
         self.assertEqual(aa.moltype.label, "protein")
         # don't trim terminal stops, returns NotCompleted
         translater = translate_seqs(trim_terminal_stop=False)
-        aa = translater(seqs)
+        aa = translater(seqs)  # pylint: disable=not-callable
         self.assertIsInstance(aa, NotCompleted)
 
     def test_translate_aln(self):
@@ -100,13 +101,13 @@ class TestTranslate(TestCase):
         # an array alignment
         aln = make_aligned_seqs(data)
         translater = translate_seqs()
-        aa = translater(aln)
+        aa = translater(aln)  # pylint: disable=not-callable
         self.assertEqual(aa.to_dict(), dict(a="MRP", b="MF-"))
         self.assertEqual(aa.moltype.label, "protein")
         self.assertIsInstance(aa, type(aln))
         # Alignment
         aln = aln.to_type(array_align=True)
-        aa = translater(aln)
+        aa = translater(aln)  # pylint: disable=not-callable
         self.assertEqual(aa.to_dict(), dict(a="MRP", b="MF-"))
         self.assertEqual(aa.moltype.label, "protein")
         self.assertIsInstance(aa, type(aln))
@@ -143,3 +144,58 @@ class TestFourFoldDegen(TestCase):
                 get_code(i), alphabet=DNA.alphabet, as_indices=True
             )
             self.assertEqual(got, expect)
+
+
+@pytest.fixture(params=(None, 0, 1, 2))
+def framed_seqs(DATA_DIR, request):
+    # sample sequences with terminating stop codon
+    # using valid values for frame
+    data = {
+        "NineBande": "GCAAGGCGCCAACAGAGCAGATGGGCTGAAAGTAAGGAAACATGTAATGATAGGCAGACTTAA",
+        "Mouse": "GCAGTGAGCCAGCAGAGCAGATGGGCTGCAAGTAAAGGAACATGTAACGACAGGCAGGTTTAA",
+        "Human": "GCAAGGAGCCAACATAACAGATGGGCTGGAAGTAAGGAAACATGTAATGATAGGCGGACTTAA",
+        "HowlerMon": "GCAAGGAGCCAACATAACAGATGGGCTGAAAGTGAGGAAACATGTAATGATAGGCAGACTTAA",
+        "DogFaced": "GCAAGGAGCCAGCAGAACAGATGGGTTGAAACTAAGGAAACATGTAATGATAGGCAGACTTAA",
+    }
+    prefix = "A" * (request.param or 0)
+    frame = None if request.param is None else request.param + 1
+    for k, s in data.items():
+        data[k] = prefix + s
+    return make_unaligned_seqs(data=data, moltype="dna", info={"frame": frame})
+
+
+def test_select_translatable_with_frame_terminal_stop(framed_seqs):
+    frame = framed_seqs.info.frame
+    sl = slice(None, None) if frame is None else slice(frame - 1, None)
+    expect = {s.name: str(s[sl]) for s in framed_seqs.seqs}
+    app = select_translatable(frame=frame, trim_terminal_stop=False)
+    got = app(framed_seqs)  # pylint: disable=not-callable
+    assert got.to_dict() == expect
+
+
+def test_select_translatable_with_frame_no_stop(framed_seqs):
+    frame = framed_seqs.info.frame
+    sl = slice(None, -3) if frame is None else slice(frame - 1, -3)
+    expect = {s.name: str(s[sl]) for s in framed_seqs.seqs}
+    app = select_translatable(frame=frame, trim_terminal_stop=True)
+    got = app(framed_seqs)  # pylint: disable=not-callable
+    assert got.to_dict() == expect
+
+
+def test_select_trabnslatable_exclude_internal_stop():
+    aln = make_unaligned_seqs(
+        {
+            "internal_stop": "AATTAAATGTGA",
+            "s2": "TATGACTAA",
+        }
+    )
+    app = select_translatable(frame=1)
+    result = app(aln)  # pylint: disable=not-callable
+    expect = {"s2": "TATGAC"}
+    assert result.to_dict() == expect
+
+
+@pytest.mark.parametrize("frame", (-1, 0, 4))
+def test_select_translatable_invalid_frame(frame):
+    with pytest.raises(AssertionError):
+        _ = select_translatable(frame=frame)