Repository: cogent3/cogent3
Pull request #2166 (2024-12-12 21:27:59)

ENH: improve new seq / alignment moltype interconversion
## Summary by Sourcery

Enhance the sequence and alignment modules by introducing a method for converting sequence arrays between different alphabets, and refactor dtype determination logic. Update tests to cover new functionality and exception handling changes.

Enhancements:
- Add a method to convert sequence arrays between different alphabets, improving interconversion capabilities.
- Refactor the logic for determining the smallest numpy integer dtype to include boundary values.

Tests:
- Add tests for the new sequence array conversion method, including validation checks for input and output sequences.
- Update tests to reflect changes in exception handling, replacing ValueError with more specific AlphabetError and MolTypeError.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index e0743b107..dfbb3d348 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -502,12 +502,8 @@ class SeqsData(SeqsDataABC):
     def to_alphabet(
         self,
         alphabet: new_alphabet.AlphabetABC,
-        check_valid=True,
+        check_valid: bool = True,
     ) -> SeqsData:
-        # refactor: design -- map directly between arrays?
-        # if the length of the two alphabets are the same and the only difference
-        # between the sets of characters is U/T, then we have the special case of
-        # converting between DNA and RNA alphabets, which we can achieved easily.
         if (
             len(self.alphabet) == len(alphabet)
             and len(
@@ -519,6 +515,7 @@ class SeqsData(SeqsDataABC):
             )
             == 1
         ):
+            # rna <-> dna swap just replace alphabet
             return self.__class__(
                 data=self._data,
                 alphabet=alphabet,
@@ -526,20 +523,15 @@ class SeqsData(SeqsDataABC):
             )
 
         new_data = {}
-        old = self.alphabet.as_bytes()
-        new = alphabet.as_bytes()
-        convert_old_to_bytes = new_alphabet.array_to_bytes(old)
-        convert_bytes_to_new = new_alphabet.bytes_to_array(
-            new,
-            dtype=new_alphabet.get_array_type(len(new)),
-        )
-
         for seqid in self.names:
             seq_data = self.get_seq_array(seqid=seqid)
-            as_new_alpha = convert_bytes_to_new(convert_old_to_bytes(seq_data))
-
+            as_new_alpha = self.alphabet.convert_seq_array_to(
+                seq=seq_data,
+                alphabet=alphabet,
+                check_valid=False,
+            )
             if check_valid and not alphabet.is_valid(as_new_alpha):
-                raise new_moltype.MolTypeError(
+                raise new_alphabet.AlphabetError(
                     f"Changing from old alphabet={self.alphabet} to new "
                     f"{alphabet=} is not valid for this data",
                 )
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 3454cfddc..39df257b5 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -155,18 +155,27 @@ class MonomerAlphabetABC(ABC):
     @abstractmethod
     def as_bytes(self) -> bytes: ...
 
+    @abstractmethod
+    def convert_seq_array_to(
+        self,
+        *,
+        alphabet: typing_extensions.Self,
+        seq: numpy.ndarray,
+        check_valid: bool = True,
+    ) -> numpy.ndarray: ...
+
 
 def get_array_type(num_elements: int):
     """Returns the smallest numpy integer dtype that can contain elements
     within num_elements.
     """
-    if num_elements < 2**8:
+    if num_elements <= 2**8:
         dtype = numpy.uint8
-    elif num_elements < 2**16:
+    elif num_elements <= 2**16:
         dtype = numpy.uint16
-    elif num_elements < 2**32:
+    elif num_elements <= 2**32:
         dtype = numpy.uint32
-    elif num_elements < 2**64:
+    elif num_elements <= 2**64:
         dtype = numpy.uint64
     else:
         raise NotImplementedError(f"{num_elements} is too big for 64-bit integer")
@@ -484,7 +493,7 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         return json.dumps(self.to_rich_dict())
 
     @classmethod
-    def from_rich_dict(cls, data: dict) -> "CharAlphabet":
+    def from_rich_dict(cls, data: dict) -> typing_extensions.Self:
         """returns an instance from a serialised dictionary"""
         return cls(data["chars"], gap=data["gap"], missing=data["missing"])
 
@@ -519,6 +528,80 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         missing = self.missing_char if self.missing_char in motif_subset else None
         return self.__class__(tuple(motif_subset), gap=gap, missing=missing)
 
+    def convert_seq_array_to(
+        self,
+        *,
+        alphabet: typing_extensions.Self,
+        seq: numpy.ndarray,
+        check_valid: bool = True,
+    ) -> numpy.ndarray:
+        """converts a numpy array with indices from self to other
+
+        Parameters
+        ----------
+        alphabet
+            alphabet to convert to
+        seq
+            ndarray of uint8 integers
+        check_valid
+            validates both input and out sequences are valid for self and
+            other respectively. Validation failure raises an AlphabetError.
+
+        Returns
+        -------
+        the indices of characters in common between self and other
+        are swapped
+        """
+        if check_valid and not self.is_valid(seq):
+            raise AlphabetError(f"input sequence not valid for {self}")
+
+        converter = make_converter(self, alphabet)
+        output = numpy.frombuffer(converter(seq.tobytes()), dtype=alphabet.dtype)
+        if check_valid and not alphabet.is_valid(output):
+            raise AlphabetError(f"output sequence not valid for {alphabet}")
+
+        return output
+
+
+@functools.cache
+def make_converter(
+    src_alpha: CharAlphabet,
+    dest_alpha: CharAlphabet,
+) -> convert_alphabet:
+    """makes a convert_alphabet instance between two CharAlphabet alphabets
+
+    Parameters
+    ----------
+    src_alpha, dest_alpha
+        the two CharAlphabet instances
+
+    Notes
+    -----
+    The common characters between the two alphabets are used to create the
+    convert_alphabet() instance. The result of applying the convertor to a
+    sequence should be validated by the caller to only shared characters
+    were present.
+    """
+    # we need to make a converter that maps the indices of common characters
+    # to new their values in dest_alpha and maps chars unique to src_alpha
+    # to a single value that exceeds the length of dest_alpha so it's
+    # easy to validate
+    src_bytes = {bytes([c]) for c in src_alpha.as_bytes()}
+    other_bytes = {bytes([c]) for c in dest_alpha.as_bytes()}
+    common = b"".join(src_bytes & other_bytes)
+
+    diff = b"".join(src_bytes - other_bytes)
+    src_indices = numpy.concatenate(
+        [src_alpha.to_indices(common), src_alpha.to_indices(diff)],
+    )
+    dest_indices = numpy.concatenate(
+        [
+            dest_alpha.to_indices(common),
+            numpy.array([len(dest_alpha)] * len(diff), dtype=dest_alpha.dtype),
+        ],
+    )
+    return convert_alphabet(src_indices.tobytes(), dest_indices.tobytes())
+
 
 @register_deserialiser(get_object_provenance(CharAlphabet))
 def deserialise_char_alphabet(data: dict) -> CharAlphabet:
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index 07d1d6397..36aa3fd87 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -1244,20 +1244,19 @@ class Sequence:
         if moltype is self.moltype:
             return self
 
-        if len(self.moltype.alphabet) == len(moltype.alphabet):
-            # converting between dna and rna
-            seq = (
-                moltype.most_degen_alphabet()
-                .array_to_bytes(array(self))
-                .decode("utf-8")
+        seq = array(self)
+        self_alpha = self.moltype.most_degen_alphabet()
+        other_alpha = moltype.most_degen_alphabet()
+        if len(self.moltype.alphabet) != len(moltype.alphabet):
+            # use alphabet converter
+            seq = self_alpha.convert_seq_array_to(
+                seq=seq,
+                alphabet=other_alpha,
+                check_valid=False,
             )
 
-        else:
-            # assume converting between ASCII/BYTES and nucleic acid
-            seq = str(self)
-
-        if not moltype.most_degen_alphabet().is_valid(seq):
-            raise ValueError(
+        if not other_alpha.is_valid(seq):
+            raise new_moltype.MolTypeError(
                 f"Changing from old moltype={self.moltype.label!r} to new "
                 f"moltype={moltype.label!r} is not valid for this data",
             )


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index c995822b4..bca93ca43 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -337,7 +337,7 @@ def test_seqs_data_to_alphabet_invalid():
         data={"a": "AAA", "b": "TTT", "c": "LLL"},
         alphabet=ASCII,
     )
-    with pytest.raises(new_moltype.MolTypeError):
+    with pytest.raises(new_alphabet.AlphabetError):
         _ = seqs.to_alphabet(DNA)
 
 
@@ -2081,10 +2081,14 @@ def test_sequence_collection_to_moltype_with_gaps(mk_cls):
     # we can convert from text to dna to rna
     rna_seqs = dna_seqs.to_moltype("rna")
     assert rna_seqs.moltype.label == "rna"
-
-    # but not from text to rna directly
-    with pytest.raises(new_moltype.MolTypeError):
-        seqs.to_moltype("rna")
+    assert rna_seqs.to_dict() != dna_seqs.to_dict()
+
+    # and reverse this back to text
+    text_seqs = rna_seqs.to_moltype("text")
+    assert text_seqs.to_dict() == rna_seqs.to_dict()
+    # but this does not equal the original text seqs
+    # because of changing T to U
+    assert text_seqs.to_dict() != seqs.to_dict()
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_core/test_new_alphabet.py b/tests/test_core/test_new_alphabet.py
index c049f86ef..7efd6be64 100644
--- a/tests/test_core/test_new_alphabet.py
+++ b/tests/test_core/test_new_alphabet.py
@@ -22,7 +22,7 @@ def test_get_array_type(num, expect):
 
 def test_get_array_type_invalid():
     with pytest.raises(NotImplementedError):
-        new_alphabet.get_array_type(2**64)
+        new_alphabet.get_array_type(2**64 + 1)
 
 
 def test_make_char_alphabet():
@@ -740,3 +740,17 @@ def test_get_subset_invalid():
     dna_alpha = new_moltype.get_moltype("dna").alphabet
     with pytest.raises(new_alphabet.AlphabetError):
         dna_alpha.get_subset(("A", "G", "-"))
+
+
+def test_convert_seq_array_to_check_valid():
+    bytes_alpha = new_moltype.BYTES.most_degen_alphabet()
+    dna_alpha = new_moltype.DNA.most_degen_alphabet()
+    # seq is invalid for DNA
+    seq = numpy.array([65, 67, 71, 84], dtype=numpy.uint8)
+    with pytest.raises(new_alphabet.AlphabetError):
+        dna_alpha.convert_seq_array_to(alphabet=bytes_alpha, seq=seq, check_valid=True)
+
+    # bytes alphabet seq with these low ordinals invalid for DNA output
+    seq = numpy.array([0, 1, 2, 3], dtype=numpy.uint8)
+    with pytest.raises(new_alphabet.AlphabetError):
+        bytes_alpha.convert_seq_array_to(alphabet=dna_alpha, seq=seq, check_valid=True)
diff --git a/tests/test_core/test_new_moltype.py b/tests/test_core/test_new_moltype.py
index fc8f569fa..7f8df9efa 100644
--- a/tests/test_core/test_new_moltype.py
+++ b/tests/test_core/test_new_moltype.py
@@ -615,8 +615,8 @@ def test_make_seq_with_seq(seq):
 
 
 def test_make_seq_with_seq_invalid_moltype():
-    seq = new_moltype.PROTEIN.make_seq(seq="ACQK")
-    with pytest.raises(ValueError):
+    seq = new_moltype.BYTES.make_seq(seq=b"".join([bytes([i]) for i in range(4)]))
+    with pytest.raises(new_moltype.MolTypeError):
         _ = new_moltype.DNA.make_seq(seq=seq)
 
 
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 5b1a87edd..20eaecc7f 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -155,7 +155,7 @@ def test_sequence_to_moltype():
     assert str(got[trev]) == "AAAA"
 
     # should raise exception if moltype not compatible with sequence data
-    with pytest.raises(ValueError):
+    with pytest.raises(new_moltype.MolTypeError):
         s.to_moltype("rna")
 
     # calling with a null object should raise an exception