Repository: cogent3/cogent3
Pull request #2091 (2024-11-04 00:49:16)

ENH: add no_degenerates, sample, get_translation to new Alignment
## Summary by Sourcery

Add new methods 'no_degenerates', 'sample', and 'get_translation' to the Alignment class, enhancing its functionality for sequence manipulation and analysis. Refactor existing tests to improve coverage and add new tests for the introduced methods.

New Features:
- Introduce the 'no_degenerates' method to remove degenerate characters from alignments, with options for motif length and gap allowance.
- Add the 'sample' method to allow random sampling of alignment positions, supporting both with and without replacement options.
- Implement the 'get_translation' method to translate nucleotide sequences into protein sequences, with options for handling stop codons.

Enhancements:
- Refactor test cases to use parameterized tests for better coverage and flexibility in testing different sequence collection methods.

Tests:
- Add new tests for the 'sample' method to verify functionality with and without replacement, and for different motif sizes.
- Include tests for the 'no_degenerates' method to ensure correct exclusion of columns with degenerate characters.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 481f445ca..41b68d8c6 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -2499,7 +2499,7 @@ def _(
 
 @singledispatch
 def seq_to_gap_coords(
-    seq: StrORBytesORArray,
+    seq: typing.union[StrORBytesORArray, new_sequence.Sequence],
     *,
     alphabet: new_alphabet.AlphabetABC,
 ) -> tuple[numpy.ndarray, numpy.ndarray]:
@@ -4320,6 +4320,38 @@ class Alignment(SequenceCollection):
 
         return self.gapped_by_map(keep, info=self.info)
 
+    def no_degenerates(self, motif_length: int = 1, allow_gap: bool = False):
+        """returns new alignment without degenerate characters
+
+        Parameters
+        ----------
+        motif_length
+            sequences are segmented into units of this size and the segments are
+            excluded if they contain degenerate characters.
+        allow_gap
+            whether gaps are allowed or whether they are treated as a degenerate
+            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), "
+                "create the alignment using DNA, RNA or PROTEIN"
+            )
+            raise ValueError(msg) from e
+
+        if allow_gap:
+            chars.extend(self.moltype.gap)
+
+        def predicate(column):
+            data = set("".join([str(d) for d in column]))
+            return data <= set(chars)
+
+        return self.filtered(predicate, motif_length=motif_length)
+
     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.
 
@@ -4360,6 +4392,74 @@ class Alignment(SequenceCollection):
                 return True
         return False
 
+    def sample(
+        self,
+        *,
+        n: int = None,
+        with_replacement: bool = False,
+        motif_length: int = 1,
+        randint=numpy.random.randint,
+        permutation=numpy.random.permutation,
+    ):
+        """Returns random sample of positions from self, e.g. to bootstrap.
+
+        Parameters
+        ----------
+        n
+            number of positions to sample. If None, all positions are sampled.
+        with_replacement
+            if True, samples with replacement.
+        motif_length
+            number of positions to sample as a single motif.
+        randint
+            random number generator, default is numpy.randint
+        permutation
+            function to generate a random permutation of positions, default is
+            numpy.permutation
+
+        Notes
+        -----
+        By default (resampling all positions without replacement), generates
+        a permutation of the positions of the alignment.
+
+        Setting with_replacement to True and otherwise leaving parameters as
+        defaults generates a standard bootstrap resampling of the alignment.
+        """
+        # refactor: type hint for randint, permutation
+        # refactor: array
+        #  Given array_pos property, it will be efficient to generate random
+        # indices and use numpy.take() using that. In the case of motif_length
+        # != 1, the total number of positions is just len(self) // motif_length.
+        # Having produced that, those indices can be scaled back up, or the
+        # numpy array reshaped.
+
+        population_size = len(self) // motif_length
+        if not n:
+            n = population_size
+        if with_replacement:
+            locations = randint(0, population_size, n)
+        else:
+            assert n <= population_size, (n, population_size, motif_length)
+            locations = permutation(population_size)[:n]
+
+        positions = numpy.empty(n * motif_length, dtype=int)
+        for i, loc in enumerate(locations):
+            positions[i * motif_length : (i + 1) * motif_length] = range(
+                loc * motif_length, (loc + 1) * motif_length
+            )
+
+        new_seqs = {}
+        for seqid in self.names:
+            sampled = self._seqs_data.get_gapped_seq_array(seqid=seqid)[positions]
+            new_seqs[seqid] = sampled
+
+        new_seqs_data = self._seqs_data.from_seqs(
+            data=new_seqs, alphabet=self.moltype.most_degen_alphabet()
+        )
+        return self.__class__(
+            seqs_data=new_seqs_data, info=self.info, moltype=self.moltype
+        )
+
     def trim_stop_codons(self, gc: Any = None, strict: bool = False, **kwargs):
         # refactor: array
         if not self.has_terminal_stop(gc=gc, strict=strict):
@@ -4391,6 +4491,42 @@ class Alignment(SequenceCollection):
 
         return result
 
+    @extend_docstring_from(SequenceCollection.get_translation)
+    def get_translation(
+        self,
+        gc: int = None,
+        incomplete_ok: bool = False,
+        include_stop: bool = False,
+        trim_stop: bool = True,
+        **kwargs,
+    ):
+        if len(self.moltype.alphabet) != 4:
+            raise new_alphabet.AlphabetError(
+                f"Must be a DNA/RNA, not {self.moltype.label}"
+            )
+
+        translated = {}
+        if not trim_stop or include_stop:
+            seqs = self
+        else:
+            seqs = self.trim_stop_codons(gc=gc, strict=not incomplete_ok)
+        # do the translation
+        for seqname in seqs.names:
+            seq = seqs.get_gapped_seq(seqname)
+            pep = seq.get_translation(
+                gc,
+                incomplete_ok=incomplete_ok,
+                include_stop=include_stop,
+                trim_stop=trim_stop,
+            )
+            translated[seqname] = pep
+        kwargs["moltype"] = pep.moltype
+        seqs_data = self._seqs_data.from_seqs(
+            data=translated, alphabet=pep.moltype.most_degen_alphabet()
+        )
+
+        return self.__class__(seqs_data=seqs_data, info=self.info, **kwargs)
+
     def _get_seq_features(
         self,
         *,


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 9b5b7d621..9e98df901 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -1186,22 +1186,36 @@ def test_sequence_collection_has_terminal_stop_strict():
         seq_coll.has_terminal_stop(gc=1, strict=True)
 
 
-def test_sequence_collection_get_translation_trim_stop():
+@pytest.mark.parametrize(
+    "mk_cls,expect",
+    (
+        (new_alignment.make_unaligned_seqs, {"seq1": "DS", "seq2": "DSS"}),
+        (new_alignment.make_aligned_seqs, {"seq1": "DS-", "seq2": "DSS"}),
+    ),
+)
+def test_get_translation_trim_stop(mk_cls, expect):
     data = {"seq1": "GATTCCTAG", "seq2": "GATTCCTCC"}
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    got = seq_coll.get_translation(trim_stop=True)
-    expect = {"seq1": "DS", "seq2": "DSS"}
+    seqs = mk_cls(data, moltype="dna")
+    got = seqs.get_translation(trim_stop=True)
     assert got.to_dict() == expect
 
 
-def test_sequence_collection_get_translation_raises():
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+def test_get_translation_raises(mk_cls):
     """should raise error if self.moltype is not a nucleic acid"""
     data = {"seq1": "PAR", "seq2": "PQR"}
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="protein")
+    seqs = mk_cls(data, moltype="protein")
     with pytest.raises(new_alphabet.AlphabetError):
-        _ = seq_coll.get_translation(trim_stop=True)
+        _ = seqs.get_translation(trim_stop=True)
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
 @pytest.mark.parametrize(
     "seqs",
     (
@@ -1211,42 +1225,54 @@ def test_sequence_collection_get_translation_raises():
         {"seq1": "GATTTT", "seq2": "?GATCT"},
     ),
 )
-def test_sequence_collection_get_translation(seqs):
+def test_get_translation(seqs, mk_cls):
     """SequenceCollection.get_translation translates each seq"""
-    seq_coll = new_alignment.make_unaligned_seqs(seqs, moltype="dna")
-    got = seq_coll.get_translation(incomplete_ok=True)
+    seqs = mk_cls(seqs, moltype="dna")
+    got = seqs.get_translation(incomplete_ok=True)
     assert got.num_seqs == 2
     assert got.moltype == new_moltype.PROTEIN
 
 
-def test_sequence_collection_get_translation_with_stop():
-    data = {"seq1": "?GATCT", "seq2": "GATTAG"}
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    got = seq_coll.get_translation(
-        incomplete_ok=True, include_stop=True, trim_stop=False
-    )
-    assert got.to_dict() == {"seq1": "XS", "seq2": "D*"}
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+def test_get_translation_with_stop(mk_cls):
+    data = {"seq1": "?GATAG", "seq2": "GATTAG"}
+    seqs = mk_cls(data, moltype="dna")
+    got = seqs.get_translation(incomplete_ok=True, include_stop=True, trim_stop=False)
+    assert got.to_dict() == {"seq1": "X*", "seq2": "D*"}
     assert got.moltype == new_moltype.PROTEIN_WITH_STOP
 
 
-def test_sequence_collection_get_translation_non_div_3():
-    data = {"seq1": "?GATCTA", "seq2": "GATTAGGG"}
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    got = seq_coll.get_translation(incomplete_ok=True, include_stop=True)
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+def test_get_translation_non_div_3(mk_cls):
+    data = {"seq1": "?GATCTA", "seq2": "GATTAGG"}
+    seqs = mk_cls(data, moltype="dna")
+    got = seqs.get_translation(incomplete_ok=True, include_stop=True)
     assert got.to_dict() == {"seq1": "XS", "seq2": "D*"}
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
 @pytest.mark.parametrize(
     "data", ({"seq1": "GATTTT", "seq2": "GATC??"}, {"seq1": "GAT---", "seq2": "?GATCT"})
 )
-def test_sequence_collection_get_translation_error(data):
-    """SequenceCollection.get_translation translates each seq"""
-    # check for a failure when no moltype specified
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
+def test_get_translation_error(data, mk_cls):
+    seqs = mk_cls(data, moltype="dna")
     with pytest.raises(TypeError):
-        seq_coll.get_translation()
+        seqs.get_translation()
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
 @pytest.mark.parametrize(
     "data",
     (
@@ -1257,23 +1283,25 @@ def test_sequence_collection_get_translation_error(data):
         {"seq1": "GAT-T-", "seq2": "GATCTT"},
     ),
 )
-def test_sequence_collection_get_translation_info(data):
+def test_get_translation_info(data, mk_cls):
     """SequenceCollection.get_translation preserves info attribute"""
-    seq_coll = new_alignment.make_unaligned_seqs(
-        data, moltype="dna", info={"key": "value"}
-    )
-    got = seq_coll.get_translation(incomplete_ok=True)
+    seqs = mk_cls(data, moltype="dna", info={"key": "value"})
+    got = seqs.get_translation(incomplete_ok=True)
     assert got.info["key"] == "value"
 
 
-def test_sequence_collection_get_translation_incomplete():
+@pytest.mark.parametrize(
+    "mk_cls",
+    (new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs),
+)
+def test_get_translation_incomplete(mk_cls):
     """get translation works on incomplete codons"""
     data = {"seq1": "GATN--", "seq2": "?GATCT"}
-    seq_coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
-    got = seq_coll.get_translation(incomplete_ok=True)
+    seqs = mk_cls(data, moltype="dna")
+    got = seqs.get_translation(incomplete_ok=True)
     assert got.to_dict() == {"seq1": "DX", "seq2": "XS"}
     with pytest.raises(new_alphabet.AlphabetError):
-        _ = seq_coll.get_translation(incomplete_ok=False)
+        _ = seqs.get_translation(incomplete_ok=False)
 
 
 @pytest.mark.parametrize(
@@ -3848,6 +3876,43 @@ def test_get_feature():
     assert feat.get_slice().to_dict() == dict(x="AAA", y="CCT")
 
 
+def test_alignment_sample_with_replacement():
+    # test with replacement -- just verify that it rnus
+    alignment = new_alignment.make_aligned_seqs(
+        {"seq1": "GATC", "seq2": "GATC"}, moltype="dna"
+    )
+    sample = alignment.sample(n=100, with_replacement=True)
+    assert len(sample) == 100
+    # ensure that sampling with replacement works on single col alignment
+    alignment1 = new_alignment.make_aligned_seqs(
+        {"seq1": "A", "seq2": "A"}, moltype="dna"
+    )
+    result = alignment1.sample(with_replacement=True)
+    assert len(result) == 1
+
+
+def test_alignment_sample_tuples():
+    ##### test with motif size != 1 #####
+    alignment = new_alignment.make_aligned_seqs(
+        {
+            "seq1": "AACCDDEEFFGGHHIIKKLLMMNNPP",
+            "seq2": "AACCDDEEFFGGHHIIKKLLMMNNPP",
+        },
+        moltype="protein",
+    )
+    # ensure length correct
+    sample = alignment.sample(n=10, motif_length=2)
+    assert len(sample), 20
+    # test columns alignment preserved
+    seqs = list(sample.to_dict().values())
+    assert seqs[0] == seqs[1]
+    # ensure each char occurs twice as sampling dinucs without replacement
+    for char in seqs[0]:
+        assert seqs[0].count(char) == 2
+
+    assert all(len(set(seqs[0][i : i + 2])) == 1 for i in range(0, len(seqs[0]), 2))
+
+
 def test_alignment_take_positions():
     """SequenceCollection take_positions should return new alignment w/ specified pos"""
     gaps = new_alignment.make_aligned_seqs(
@@ -3964,6 +4029,52 @@ def test_filtered_drop_remainder():
     assert len(got) == 4
 
 
+def test_no_degenerates():
+    """no_degenerates correctly excludes columns containing IUPAC ambiguity codes"""
+    data = {
+        "s1": "AAA CCC GGG TTT".replace(" ", ""),
+        "s2": "CCC GGG T-T AAA".replace(" ", ""),
+        "s3": "GGR YTT AAA CCC".replace(" ", ""),
+    }
+    aln = new_alignment.make_aligned_seqs(data, moltype="dna")
+
+    # motif length of 1, defaults - no gaps allowed
+    result = aln.no_degenerates().to_dict()
+    expect = {
+        "s1": "AA CC GG TTT".replace(" ", ""),
+        "s2": "CC GG TT AAA".replace(" ", ""),
+        "s3": "GG TT AA CCC".replace(" ", ""),
+    }
+    assert result == expect
+
+    # allow gaps
+    result = aln.no_degenerates(allow_gap=True).to_dict()
+    expect = {
+        "s1": "AA CC GGG TTT".replace(" ", ""),
+        "s2": "CC GG T-T AAA".replace(" ", ""),
+        "s3": "GG TT AAA CCC".replace(" ", ""),
+    }
+    assert result == expect
+
+    # motif length of 3, defaults - no gaps allowed
+    result = aln.no_degenerates(motif_length=3, allow_gap=False).to_dict()
+    expect = {
+        "s1": "TTT".replace(" ", ""),
+        "s2": "AAA".replace(" ", ""),
+        "s3": "CCC".replace(" ", ""),
+    }
+    assert result == expect
+
+    # allow gaps
+    result = aln.no_degenerates(motif_length=3, allow_gap=True).to_dict()
+    expect = {
+        "s1": "GGG TTT".replace(" ", ""),
+        "s2": "T-T AAA".replace(" ", ""),
+        "s3": "AAA CCC".replace(" ", ""),
+    }
+    assert result == expect
+
+
 def test_omit_gap_pos_motif_length():
     """consistency with different motif_length values"""
     data = {