Repository: cogent3/cogent3
Pull request #1945 (2024-07-10 22:55:00)

ENH: added CodonAlphabet


Patch:
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index de76c1305..731dad066 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -96,7 +96,6 @@ class convert_alphabet:
 
 
 class AlphabetABC(ABC):
-
     @property
     @abstractmethod
     def num_canonical(self) -> int: ...
@@ -427,7 +426,7 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
 
 
 @register_deserialiser(get_object_provenance(CharAlphabet))
-def deserialise_char_alphabet(data) -> CharAlphabet:
+def deserialise_char_alphabet(data: dict) -> CharAlphabet:
     return CharAlphabet.from_rich_dict(data)
 
 
@@ -694,7 +693,6 @@ class KmerAlphabet(tuple, AlphabetABC, KmerAlphabetABC):
     def from_indices(
         self, kmer_indices: numpy.ndarray, independent_kmer: bool = True
     ) -> numpy.ndarray:
-
         if independent_kmer:
             size = len(kmer_indices) * self.k
         else:
@@ -833,15 +831,164 @@ class KmerAlphabet(tuple, AlphabetABC, KmerAlphabetABC):
 
 
 @register_deserialiser(get_object_provenance(KmerAlphabet))
-def deserialise_kmer_alphabet(data) -> KmerAlphabet:
+def deserialise_kmer_alphabet(data: dict) -> KmerAlphabet:
     return KmerAlphabet.from_rich_dict(data)
 
 
+class CodonAlphabet(tuple):
+    """represents the sense-codons of a GeneticCode"""
+
+    def __new__(
+        cls,
+        words: tuple[StrORBytes, ...],
+        monomers: CharAlphabet,
+        gap: OptStr = None,
+    ):
+        if not words:
+            raise ValueError(f"cannot create empty {cls.__name__!r}")
+
+        if gap is not None:
+            assert _coerce_to_type(words[0], gap) in words
+
+        consistent_words(words, length=3)
+        return tuple.__new__(cls, words, monomers=monomers, gap=gap)
+
+    def __init__(
+        self,
+        words: tuple[StrORBytes, ...],
+        monomers: CharAlphabet,
+        gap: OptStr = None,
+    ):
+        """
+        Parameters
+        ----------
+        words
+            series of 3 character strings representing sense codons
+        monomers
+            CharAlphabet instance of DNA monomers, including degenerate
+            characters
+        gap
+            the gap state "---" if present
+        """
+        self._gap_char = gap
+        self.monomers = monomers
+        self.dtype = get_array_type(len(self))
+        self._words = set(self)  # for quick lookup
+        self._to_indices = {codon: i for i, codon in enumerate(self)}
+        self._from_indices = {i: codon for codon, i in self._to_indices.items()}
+        self.motif_length = 3
+
+    @property
+    def gap_char(self):
+        return self._gap_char
+
+    @property
+    def gap_index(self):
+        return self._to_indices(self.gap_char) if self._gap_char else None
+
+    @property
+    def missing_char(self):
+        """not supported on CodonAlphabet"""
+        return None
+
+    @property
+    def missing_index(self):
+        """not supported on CodonAlphabet"""
+        return None
+
+    @functools.singledispatchmethod
+    def to_indices(self, seq) -> numpy.ndarray:
+        """returns a sequence of codon indices"""
+        raise TypeError(f"{type(seq)} is not supported")
+
+    @to_indices.register
+    def _(self, seq: str) -> numpy.ndarray:
+        return self.to_indices([seq[i : i + 3] for i in range(0, len(seq), 3)])
+
+    @to_indices.register
+    def _(self, seq: list) -> numpy.ndarray:
+        return [self.to_index(c) for c in seq]
+
+    def to_index(self, codon: str) -> int:
+        if len(codon) != 3:
+            raise ValueError(f"{codon=!r} is not of length 3")
+        try:
+            return self._to_indices[codon]
+        except KeyError as e:
+            raise ValueError(f"{codon=!r} not in alphabet") from e
+
+    def from_index(self, index: int) -> str:
+        if index > len(self) or index < 0:
+            raise ValueError(f"{index=!r} is not within range")
+
+        try:
+            return self._from_indices[index]
+        except KeyError as e:
+            raise ValueError(f"invalid {index=}") from e
+
+    def from_indices(self, indices: numpy.ndarray) -> list[str]:
+        return [self.from_index(index) for index in indices]
+
+    @property
+    def num_canonical(self):
+        return len(self._words)
+
+    @functools.singledispatchmethod
+    def is_valid(self, seq) -> bool:
+        """seq is valid for alphabet"""
+        raise TypeError(f"{type(seq)} not supported")
+
+    @is_valid.register
+    def _(self, seq: str) -> bool:
+        try:
+            _ = self.to_indices(seq)
+            return True
+        except ValueError:
+            return False
+
+    @is_valid.register
+    def _(self, seq: numpy.ndarray) -> bool:
+        return seq.min() >= 0 and seq.max() < len(self)
+
+    def with_gap_motif(self):
+        if self.gap_char:
+            return self
+        monomers = self.monomers.with_gap_motif()
+        gap_char = monomers.gap_char * 3
+        words = tuple(self) + (gap_char,)
+        return self.__class__(words=words, monomers=monomers, gap=gap_char)
+
+    def to_rich_dict(self):
+        from cogent3._version import __version__
+
+        return {
+            "monomers": self.monomers.to_rich_dict(),
+            "type": get_object_provenance(self),
+            "version": __version__,
+            "words": list(self),
+        }
+
+    def to_json(self):
+        return json.dumps(self.to_rich_dict())
+
+    @classmethod
+    def from_rich_dict(cls, data):
+        data["monomers"] = deserialise_char_alphabet(data["monomers"])
+        data.pop("type", None)
+        data.pop("version", None)
+        return cls(**data)
+
+
+@register_deserialiser(get_object_provenance(CodonAlphabet))
+def deserialise_codon_alphabet(data: dict) -> CodonAlphabet:
+    return CodonAlphabet.from_rich_dict(data)
+
+
 _alphabet_moltype_map = {}
 
 
 def make_alphabet(*, chars, gap, missing, moltype):
-    """constructs an alphabet and registers the associated moltype
+    """constructs a character alphabet and registers the associated moltype
 
     Notes
     -----
diff --git a/src/cogent3/core/new_genetic_code.py b/src/cogent3/core/new_genetic_code.py
index 66dbb5d85..c6361e986 100644
--- a/src/cogent3/core/new_genetic_code.py
+++ b/src/cogent3/core/new_genetic_code.py
@@ -257,6 +257,30 @@ class GeneticCode:
         """Returns True if codon is a stop codon, False otherwise."""
         return self[codon] == "*"
 
+    def get_alphabet(
+        self, include_gap: bool = False, include_stop: bool = False
+    ) -> new_alphabet.CodonAlphabet:
+        """returns a codon alphabet
+
+        Parameters
+        ----------
+        include_gap
+            alphabet includes the gap motif
+        """
+        if include_stop:
+            codons = tuple(self.moltype.alphabet.get_kmer_alphabet(k=3))
+        else:
+            codons = tuple(self.sense_codons)
+
+        gap = self.moltype.gapped_alphabet.gap_char * 3 if include_gap else None
+
+        if include_gap:
+            codons += (gap,)
+
+        return new_alphabet.CodonAlphabet(
+            words=codons, monomers=self.moltype.alphabet, gap=gap
+        )
+
 
 _mapping_cols = "ncbi_code_sequence", "ID", "name", "ncbi_start_codon_map"
 # code mappings are based on the product of bases in order TCAG


Test patch:
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index 99783a05f..3c8eacf2f 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -5,7 +5,7 @@ import pytest
 
 from numpy.testing import assert_allclose
 
-from cogent3.core import new_alphabet, new_moltype
+from cogent3.core import new_alphabet, new_genetic_code, new_moltype
 
 
 @pytest.mark.parametrize(
@@ -546,3 +546,62 @@ def test_deserialise_alphas(alpha):
     ser = alpha.to_json()
     got = deserialise_object(ser)
     assert got == alpha
+
+
+@pytest.fixture
+def calpha():
+    gc = new_genetic_code.get_code(1)
+    return new_alphabet.CodonAlphabet(
+        words=gc.sense_codons, monomers=gc.moltype.alphabet
+    )
+
+
+def test_codon_alphabet_init(calpha):
+    assert len(calpha) == 61
+
+
+def test_codon_alphabet_index(calpha):
+    assert calpha.to_index("TTT") == 0
+
+
+@pytest.mark.parametrize("seq", (["TTT", "TTC"], "TTTTTC"))
+def test_to_indices(calpha, seq):
+    got = calpha.to_indices(seq)
+    assert_allclose(got, numpy.array([0, 1], dtype=numpy.uint8))
+
+
+def test_codon_alphabet_invalid_codon(calpha):
+    with pytest.raises(ValueError):
+        calpha.to_index("TGA")
+
+
+def test_codon_alphabet_from_index(calpha):
+    assert calpha.from_index(0) == "TTT"
+
+
+def test_codon_alphabet_from_indices(calpha):
+    assert calpha.from_indices(numpy.array([1, 0], dtype=numpy.uint8)) == ["TTC", "TTT"]
+
+
+@pytest.mark.parametrize("seq", ("GGGAGA", numpy.array([1, 2]), numpy.array([2, 60])))
+def test_codon_alphabet_is_valid(seq, calpha):
+    assert calpha.is_valid(seq)
+
+
+@pytest.mark.parametrize("seq", ("GGGTGA", numpy.array([-1, 2]), numpy.array([2, 66])))
+def test_codon_alphabet_not_is_valid(seq, calpha):
+    assert not calpha.is_valid(seq)
+
+
+def test_codon_alphabet_with_gap_motif(calpha):
+    with_gap = calpha.with_gap_motif()
+    assert len(with_gap) == len(calpha) + 1
+
+
+def test_codon_alphabet_serlialise_round_trip(calpha):
+    from cogent3.util.deserialise import deserialise_object
+
+    got = deserialise_object(calpha.to_json())
+    assert isinstance(got, new_alphabet.CodonAlphabet)
+    assert list(got) == list(calpha)
+    assert calpha.to_index("TTC") == 1
diff --git a/tests/test_core/test_new_genetic_code.py b/tests/test_core/test_new_genetic_code.py
index 78f7c1ef9..d76aa8d87 100644
--- a/tests/test_core/test_new_genetic_code.py
+++ b/tests/test_core/test_new_genetic_code.py
@@ -127,3 +127,22 @@ def test_translate_incomplete_rc(seq, expect):
     gc = new_genetic_code.DEFAULT
     got = gc.translate(seq, rc=True)
     assert got == expect
+
+
+def test_get_alphabet():
+    gc = new_genetic_code.get_code(1)
+    alpha = gc.get_alphabet()
+    assert len(alpha) == 61
+    assert alpha.to_index("TTT") == 0
+
+    alpha_gap = gc.get_alphabet(include_gap=True)
+    assert len(alpha_gap) == 62
+
+
+def test_get_alphabet_with_stop():
+    gc = new_genetic_code.get_code(1)
+    alpha = gc.get_alphabet(include_stop=True)
+    assert len(alpha) == 64
+
+    alpha_gap = gc.get_alphabet(include_gap=True, include_stop=True)
+    assert len(alpha_gap) == 65
diff --git a/tests/test_util/test_deserialise.py b/tests/test_util/test_deserialise.py
index d0f46d0e0..4ee5adbcf 100644
--- a/tests/test_util/test_deserialise.py
+++ b/tests/test_util/test_deserialise.py
@@ -11,6 +11,7 @@ from numpy.testing import assert_allclose
 
 from cogent3 import (
     get_app,
+    get_code,
     load_aligned_seqs,
     load_tree,
     make_aligned_seqs,
@@ -34,10 +35,11 @@ from cogent3.util.deserialise import (
 class TestDeserialising(TestCase):
     def test_roundtrip_codon_alphabet(self):
         """codon alphabet to_json enables roundtrip"""
-        data = moltype.STANDARD_CODON.to_json()
+        STANDARD_CODON = get_code(1).get_alphabet()
+        data = STANDARD_CODON.to_json()
         got = deserialise_object(data)
-        self.assertEqual(type(got), type(moltype.STANDARD_CODON))
-        self.assertEqual(list(got), list(moltype.STANDARD_CODON))
+        self.assertEqual(type(got), type(STANDARD_CODON))
+        self.assertEqual(list(got), list(STANDARD_CODON))
 
     def test_roundtrip_alphabet(self):
         """alphabet to_json enables roundtrip"""