Repository: cogent3/cogent3
Pull request #2092 (2024-11-04 05:48:34)

ENH: no_degenerates operates on array data
## Summary by Sourcery

Add a new method 'from_names_and_array' to the AlignedSeqsData class for constructing objects from names and array data, and refactor the 'no_degenerates' method to enhance performance by operating on array data. Update tests to cover new functionality and edge cases.

New Features:
- Introduce a new class method 'from_names_and_array' in the AlignedSeqsData class to construct an object from a list of names and a numpy array of aligned sequence data.

Enhancements:
- Refactor the 'no_degenerates' method to operate on array data, improving performance and clarity.

Tests:
- Add new tests for the 'from_names_and_array' method to ensure correct functionality, including cases for empty data and mismatched lengths.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 41b68d8c6..b8898be9d 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -2749,6 +2749,16 @@ class AlignedSeqsDataABC(SeqsDataABC):
         alphabet: new_alphabet.AlphabetABC,
     ): ...
 
+    @classmethod
+    @abstractmethod
+    def from_names_and_array(
+        cls,
+        *,
+        names: list[str],
+        data: numpy.ndarray,
+        alphabet: new_alphabet.AlphabetABC,
+    ): ...
+
     @property
     @abstractmethod
     def align_len(self) -> int: ...
@@ -2955,6 +2965,43 @@ class AlignedSeqsData(AlignedSeqsDataABC):
             **kwargs,
         )
 
+    @classmethod
+    def from_names_and_array(
+        cls,
+        *,
+        names: list[str],
+        data: numpy.ndarray,
+        alphabet: new_alphabet.AlphabetABC,
+    ):
+        """Construct an AlignedSeqsData object from a list of names and a numpy
+        array of aligned sequence data.
+
+        Parameters
+        ----------
+        names
+            list of sequence names
+        data
+            numpy array of aligned sequence data
+        alphabet
+            alphabet object for the sequences
+        """
+        if len(names) != data.shape[0]:
+            raise ValueError("Number of names must match number of rows in data.")
+
+        seqs = {}
+        gaps = {}
+        for name, seq in zip(names, data):
+            seq, gap_map = seq_to_gap_coords(seq, alphabet=alphabet)
+            seq.flags.writeable = False
+            gap_map.flags.writeable = False
+            seqs[name], gaps[name] = seq, gap_map
+
+        return cls(
+            seqs=seqs,
+            gaps=gaps,
+            alphabet=alphabet,
+        )
+
     @property
     def names(self) -> tuple[str]:
         return list(self._seqs.keys())
@@ -4333,24 +4380,38 @@ class Alignment(SequenceCollection):
             character (latter is default, as most evolutionary modelling treats
             gaps as N).
         """
-        # refactor: array
-        try:
-            chars = list(self.moltype.alphabet)
-        except AttributeError as e:
-            msg = (
-                "Invalid MolType (no degenerate characters), "
+        if self.moltype.degen_alphabet is None:
+            raise new_moltype.MolTypeError(
+                f"Invalid MolType={self.moltype.label} (no degenerate characters), "
                 "create the alignment using DNA, RNA or PROTEIN"
             )
-            raise ValueError(msg) from e
 
-        if allow_gap:
-            chars.extend(self.moltype.gap)
+        chars = len(self.moltype)
+
+        array_pos = self.array_positions
+        # by design, char alphabets are organised such that the canonical
+        # characters always occur first, followed by gap, then ambiguity
+        # characters. so we can define a cutoff as follows:
+        cutoff = chars + 1 if allow_gap else chars
+        indices = (array_pos < cutoff).all(axis=1)
 
-        def predicate(column):
-            data = set("".join([str(d) for d in column]))
-            return data <= set(chars)
+        if motif_length > 1:
+            num_motif = len(self) // motif_length
 
-        return self.filtered(predicate, motif_length=motif_length)
+            if remainder := len(self) % motif_length:
+                indices = indices[:-remainder]
+                array_pos = array_pos[:-remainder]
+
+            motif_valid = indices.reshape(num_motif, motif_length).all(axis=1).flatten()
+            indices = numpy.repeat(motif_valid, motif_length)
+        selected = array_pos[indices].T
+
+        aligned_seqs_data = self._seqs_data.from_names_and_array(
+            names=self.names, data=selected, alphabet=self.moltype.most_degen_alphabet()
+        )
+        return self.__class__(
+            seqs_data=aligned_seqs_data, info=self.info, moltype=self.moltype
+        )
 
     def omit_gap_pos(self, allowed_gap_frac: float = 1 - EPS, motif_length: int = 1):
         """Returns new alignment where all cols (motifs) have <= allowed_gap_frac gaps.
diff --git a/src/cogent3/core/new_moltype.py b/src/cogent3/core/new_moltype.py
index 9719830a1..f809af932 100644
--- a/src/cogent3/core/new_moltype.py
+++ b/src/cogent3/core/new_moltype.py
@@ -140,6 +140,9 @@ RNA_EXTENDED_PAIRS = {
 }
 
 
+class MolTypeError(TypeError): ...
+
+
 def make_pairs(
     *,
     pairs: dict[tuple[str, str], bool] = None,


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 9e98df901..13da959ee 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -2773,7 +2773,7 @@ def test_aligned_seqs_data_unequal_seqlens_raises(data_type, dna_alphabet, dna_m
         )
 
 
-def test_from_seqs_and_maps(dna_alphabet):
+def test_from_seqs_and_gaps(dna_alphabet):
     # AlignedSeqsData should be able to be constructed from sequences and gap maps
     seqs = {"seq1": "ACCTA", "seq2": ""}
     gaps = {"seq1": numpy.array([[0, 1]]), "seq2": numpy.array([[0, 6]])}
@@ -2786,7 +2786,7 @@ def test_from_seqs_and_maps(dna_alphabet):
     assert asd.get_gapped_seq_str(seqid="seq2") == "------"
 
 
-def test_from_seqs_and_maps_diff_seq_lens_raises(dna_alphabet):
+def test_from_seqs_and_gaps_diff_seq_lens_raises(dna_alphabet):
     seqs = {"seq1": "ACCTA", "seq2": "A"}
     gaps = {"seq1": numpy.array([[0, 1]]), "seq2": numpy.array([[0, 1]])}
     with pytest.raises(ValueError):
@@ -2795,7 +2795,7 @@ def test_from_seqs_and_maps_diff_seq_lens_raises(dna_alphabet):
         )
 
 
-def test_from_seqs_and_maps_diff_keys_raises(dna_alphabet):
+def test_from_seqs_and_gaps_diff_keys_raises(dna_alphabet):
     seqs = {"seq1": "ACCTA", "seq2": "A"}
     gaps = {"seq1": numpy.array([[0, 1]]), "seq3": numpy.array([[0, 1]])}
     with pytest.raises(ValueError):
@@ -2804,6 +2804,39 @@ def test_from_seqs_and_maps_diff_keys_raises(dna_alphabet):
         )
 
 
+@pytest.mark.parametrize("seqid, i", (("seq1", 0), ("seq2", 1), ("seq3", 2)))
+def test_from_names_and_array(dna_alphabet, seqid, i):
+    names = ["seq1", "seq2", "seq3"]
+    data = numpy.array([[0, 1, 2, 3], [3, 2, 1, 0], [4, 4, 4, 4]])
+    asd = new_alignment.AlignedSeqsData.from_names_and_array(
+        names=names, data=data, alphabet=dna_alphabet
+    )
+    assert asd.names == ["seq1", "seq2", "seq3"]
+    got = asd.get_gapped_seq_array(seqid=seqid)
+    expect = data[i]
+    assert numpy.array_equal(got, expect)
+
+
+def test_from_names_and_array_empty_raises(dna_alphabet):
+    names = []
+    data = numpy.array([]).reshape(0, 0)
+
+    with pytest.raises(ValueError):
+        _ = new_alignment.AlignedSeqsData.from_names_and_array(
+            names=names, data=data, alphabet=dna_alphabet
+        )
+
+
+def test_from_names_and_array_mismatched_length(dna_alphabet):
+    names = ["seq1", "seq2"]
+    data = numpy.array([[1, 0, 1], [0, 1, 0], [1, 1, 1]])
+
+    with pytest.raises(ValueError):
+        _ = new_alignment.AlignedSeqsData.from_names_and_array(
+            names=names, data=data, alphabet=dna_alphabet
+        )
+
+
 def test_aligned_seqs_data_diff_keys_raises(dna_alphabet):
     """AlignedSeqsData expect identical keys in seqs and gaps"""
     seqs = dict(seq1=numpy.array([2, 1]), seq2=numpy.array([2, 0, 3, 1]))
@@ -4074,6 +4107,27 @@ def test_no_degenerates():
     }
     assert result == expect
 
+    # for length non-divisible by motif_length
+    data = {
+        "s1": "AAA CCC GGG TTT T".replace(" ", ""),
+        "s2": "CCC GGG T-T AAA A".replace(" ", ""),
+        "s3": "GGR YTT AAA CCC C".replace(" ", ""),
+    }
+    aln = new_alignment.make_aligned_seqs(data, moltype="dna")
+    result = aln.no_degenerates(motif_length=3, allow_gap=False).to_dict()
+    expect = {
+        "s1": "TTT".replace(" ", ""),
+        "s2": "AAA".replace(" ", ""),
+        "s3": "CCC".replace(" ", ""),
+    }
+
+
+def test_no_degenerates_bad_moltype_raises():
+    """no_degenerates should raise an error if moltype has no degenerate symbols"""
+    aln = new_alignment.make_aligned_seqs({"s1": "ACGT", "s2": "ACGT"}, moltype="text")
+    with pytest.raises(new_moltype.MolTypeError):
+        _ = aln.no_degenerates()
+
 
 def test_omit_gap_pos_motif_length():
     """consistency with different motif_length values"""