Repository: cogent3/cogent3
Pull request #2171 (2024-12-16 05:03:35)

ENH: api compatbility and attrs for migration to new_type
## Summary by Sourcery

Enhance the genetic code translation to handle missing data and sequence ambiguities. Introduce a singledispatch method for detecting sequence ambiguities across various data types. Add comprehensive tests for ambiguity detection and translation handling of ambiguous and gapped sequences.

Enhancements:
- Add handling for missing data in genetic code translation by modifying codon indices to represent missing states.
- Introduce a singledispatch method for checking sequence ambiguity in different data types (str, bytes, numpy.ndarray).

Tests:
- Add tests for sequence ambiguity detection across different moltypes and data types.
- Add tests for translation handling of ambiguous and gapped sequences in genetic codes and alignments.


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 9b310bd66..8db8d45f8 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -1173,6 +1173,10 @@ class SenseCodonAlphabet(tuple, AlphabetABC):
     def to_index(self, codon: str) -> int:
         if len(codon) != 3:
             raise ValueError(f"{codon=!r} is not of length 3")
+        if not self.monomers.is_valid(codon):
+            raise AlphabetError(f"{codon=!r} elements not nucleotides")
+        if self.moltype.has_ambiguity(codon):
+            return len(self)
         try:
             return self._to_indices[codon]
         except KeyError as e:
diff --git a/src/cogent3/core/new_genetic_code.py b/src/cogent3/core/new_genetic_code.py
index 45c94fcee..225f446ab 100644
--- a/src/cogent3/core/new_genetic_code.py
+++ b/src/cogent3/core/new_genetic_code.py
@@ -21,6 +21,9 @@ import numpy
 from cogent3.core import new_alphabet, new_moltype
 from cogent3.util.table import Table
 
+if typing.TYPE_CHECKING:
+    from cogent3.core.sequence import Sequence
+
 OptStr = typing.Optional[str]
 SetStr = set[str]
 ConverterType = typing.Callable[[bytes, bytes], bytes]
@@ -253,9 +256,20 @@ class GeneticCode:
         # convert to indices and then bytes
         if incomplete_ok:
             seq = self.codons.to_indices(dna)
+            missing = seq >= len(self.codons)
         else:
-            codons = self.get_alphabet(include_stop=True, include_gap=True)
+            codons = self.get_alphabet(
+                include_stop=True,
+                include_gap=True,
+                include_missing=True,
+            )
             seq = codons.to_indices(dna)
+            missing = seq >= len(codons)
+        # any codon indices include missing data, they will have an index ==
+        # length of the alphabet employed, so we need to modify that index
+        # to represent the missing state which is X
+        if missing.any():
+            seq[missing] = len(self._codon_to_aa) - 1
 
         if rc:
             return self._translate_minus(seq.tobytes()).decode("utf8")[::-1]
@@ -280,7 +294,9 @@ class GeneticCode:
     @functools.cache
     def get_alphabet(
         self,
+        *,
         include_gap: bool = False,
+        include_missing: bool = False,
         include_stop: bool = False,
     ) -> new_alphabet.SenseCodonAlphabet:
         """returns a codon alphabet
@@ -289,19 +305,30 @@ class GeneticCode:
         ----------
         include_gap
             alphabet includes the gap motif
+        include_missing
+            alphabet includes the missing state as 3 * IUPAC_missing
+        include_stop
+            if True, this is just a kmer alphabet
         """
         if include_stop:
-            codons = tuple(self.moltype.alphabet.get_kmer_alphabet(k=3))
+            words = tuple(self.moltype.alphabet.get_kmer_alphabet(k=3))
         else:
-            codons = tuple(self.sense_codons)
+            words = tuple(self.sense_codons)
 
         gap = self.moltype.gapped_alphabet.gap_char * 3 if include_gap else None
-
+        missing = (
+            self.moltype.degen_gapped_alphabet.missing_char * 3
+            if include_missing
+            else None
+        )
         if include_gap:
-            codons += (gap,)
+            words += (gap,)
+
+        if include_missing:
+            words += (missing,)
 
         return new_alphabet.SenseCodonAlphabet(
-            words=codons,
+            words=words,
             monomers=self.moltype.degen_gapped_alphabet,
             gap=gap,
         )
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index 2c3eec936..5f443f1b5 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -715,6 +715,29 @@ class MolType:
         """
         return callable(self._complement)
 
+    @functools.singledispatchmethod
+    def has_ambiguity(self, seq: StrORBytesORArray) -> bool:
+        raise TypeError(f"{type(seq)} not supported")
+
+    @has_ambiguity.register
+    def _(self, seq: str) -> bool:
+        if not self.ambiguities:
+            return False
+        return any(self.is_ambiguity(c) for c in seq)
+
+    @has_ambiguity.register
+    def _(self, seq: bytes) -> bool:
+        return self.has_ambiguity(seq.decode("utf8"))
+
+    @has_ambiguity.register
+    def _(self, seq: numpy.ndarray) -> bool:
+        if not self.ambiguities:
+            return False
+
+        min_canonical = len(self.alphabet) - 1
+        max_degen = len(self.degen_alphabet)
+        return ((seq > min_canonical) & (seq < max_degen)).any()
+
     def rc(self, seq: str, validate: bool = True) -> str:
         """reverse reverse complement of a sequence"""
         if validate and not self.is_valid(seq):
@@ -874,8 +897,13 @@ class MolType:
         ----------
         query_motif
             the motif being queried.
-
+        validate
+            if True, checks the sequence is validated against the most
+            degenerate alphabet
         """
+        if not self.ambiguities:
+            return False
+
         if validate and not self.is_valid(query_motif):
             raise new_alphabet.AlphabetError(
                 f"{query_motif[:4]!r} not valid for moltype {self.name!r}",
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index 36aa3fd87..b4e50cc3d 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -2049,7 +2049,7 @@ class NucleicAcidSequenceMixin:
 
         if not incomplete_ok and "X" in pep:
             raise new_alphabet.AlphabetError(
-                "Incomplete codon in translation, set incomplete_ok=True to "
+                "A codon is incomplete or contains an ambiguity, set incomplete_ok=True to "
                 "allow translation",
             )
         return protein.make_seq(seq=pep, name=self.name)


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 34e3bc669..825f314b4 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -1370,6 +1370,22 @@ def test_get_translation_trim_stop(mk_cls, expect):
     assert got.to_dict() == expect
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+@pytest.mark.parametrize("seq", ("ATG---NTT", "ATGCAY"))
+def test_get_translation_ambigs(mk_cls, seq):
+    data = {"s1": seq}
+    expect = "M-X" if "-" in seq else "MX"
+    seqs = mk_cls(data, moltype="dna")
+    got = seqs.get_translation(incomplete_ok=True)
+    assert str(got.seqs["s1"]) == expect
+    with pytest.raises(new_alphabet.AlphabetError):
+        # ambiguity codes raise an exception unless explicitly allowed
+        _ = seqs.get_translation(incomplete_ok=False)
+
+
 @pytest.mark.parametrize(
     "mk_cls",
     (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
diff --git a/tests/test_core/test_new_genetic_code.py b/tests/test_core/test_new_genetic_code.py
index cc946db1a..97831a003 100644
--- a/tests/test_core/test_new_genetic_code.py
+++ b/tests/test_core/test_new_genetic_code.py
@@ -55,6 +55,14 @@ def test_translate():
     assert code.translate("ATGGGG") == code["ATG"] + code["GGG"]
 
 
+@pytest.mark.parametrize("incomplete_ok", (False, True))
+@pytest.mark.parametrize("seq", ("ATGNTT", "ATGCAY"))
+def test_translation_ambig(seq, incomplete_ok):
+    code = new_genetic_code.DEFAULT
+    aa = code.translate(seq, incomplete_ok=incomplete_ok)
+    assert aa == "MX"
+
+
 def test_translate_rc():
     code = new_genetic_code.DEFAULT
     seq = "ATGGGG"
@@ -189,3 +197,24 @@ def test_to_regex():
     aa = make_seq(aa, moltype="protein", new_type=True)
     pattern = new_genetic_code.DEFAULT.to_regex(aa)
     assert "".join(re.findall(pattern, dna)) == dna
+
+
+@pytest.mark.parametrize("include_stop", [True, False])
+@pytest.mark.parametrize("include_missing", [True, False])
+def test_get_alphabet_missing(include_stop, include_missing):
+    # handles include_missing argument
+    gc = new_genetic_code.get_code(1)
+    alpha = gc.get_alphabet(include_stop=include_stop, include_missing=include_missing)
+
+    base_size = 61  # standard genetic code has 61 sense codons
+    if include_stop:
+        base_size += 3  # standard genetic code has 3 stop codons
+    if include_missing:
+        base_size += 1  # adds '???' when include_missing is True, hence + 1
+
+    assert len(alpha) == base_size
+
+    # Check presence/absence of missing character
+    missing = "???"
+    result = (missing in alpha) if include_missing else (missing not in alpha)
+    assert result
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index 7f8df9efa..5c52fc548 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -667,3 +667,52 @@ def test_json_roundtrip_moltype(moltype):
     # because our moltype instances are singletons,
     # we expect the same object
     assert got is mt
+
+
+@pytest.mark.parametrize(
+    "moltype, seq, expect",
+    [
+        (new_moltype.DNA, "ACGT", False),  # No ambiguity
+        (new_moltype.DNA, "ACGTN", True),  # N is ambiguous
+        (new_moltype.DNA, "ACGTY", True),  # Y is ambiguous
+        (new_moltype.RNA, "ACGU", False),  # No ambiguity
+        (new_moltype.RNA, "ACGUN", True),  # N is ambiguous
+        (new_moltype.RNA, "ACGUY", True),  # Y is ambiguous
+        (new_moltype.PROTEIN, "ACDEFGHIKLMNPQRSTVWY", False),  # No ambiguity
+        (new_moltype.PROTEIN, "ACDEFGHIKLMNPQRSTVWYX", True),  # X is ambiguous
+        (
+            new_moltype.PROTEIN_WITH_STOP,
+            "ACDEFGHIKLM*NPQRSTVWY",
+            False,
+        ),  # * is not ambiguous (stop)
+    ],
+)
+@pytest.mark.parametrize("cast", (str, bytes, numpy.ndarray))
+def test_has_ambiguity(moltype, seq, expect, cast):
+    seq = make_typed(seq, cast, moltype)
+    assert moltype.has_ambiguity(seq) == expect
+
+
+@pytest.mark.parametrize(
+    "seq",
+    ["ACGT", "ACGTN"],
+)
+@pytest.mark.parametrize("moltype", ("text", "bytes"))
+@pytest.mark.parametrize("cast", (str, bytes, numpy.ndarray))
+def test_has_ambiguity_all_false(moltype, seq, cast):
+    moltype = new_moltype.get_moltype(moltype)
+    # these moltypes don't have an ambuity code
+    seq = make_typed(seq, cast, moltype)
+    assert not moltype.has_ambiguity(seq)
+
+
+def test_has_ambiguity_validation():
+    """raises TypeError if input not a string"""
+    with pytest.raises(TypeError):
+        new_moltype.DNA.has_ambiguity(None)
+
+    with pytest.raises(TypeError):
+        new_moltype.DNA.has_ambiguity([])
+
+    with pytest.raises(TypeError):
+        new_moltype.DNA.has_ambiguity(1)
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 04f32b740..2ebc04ccc 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -2913,6 +2913,19 @@ def test_translation(seq, rc):
     assert str(aa) == "MF"
 
 
+@pytest.mark.parametrize("seq", ("ATGNTT", "ATGCAY"))
+def test_translation_ambig(seq):
+    seq = new_moltype.DNA.make_seq(seq=seq)
+    aa = seq.get_translation(incomplete_ok=True)
+    assert str(aa) == "MX"
+
+
+def test_translation_gapped():
+    seq = new_moltype.DNA.make_seq(seq="ATG---")
+    aa = seq.get_translation()
+    assert str(aa) == "M-"
+
+
 def test_get_translation_include_stop():
     s = new_moltype.DNA.make_seq(seq="ATTTAACTT", name="s1")
     aa = s.get_translation(include_stop=True)