Repository: cogent3/cogent3
Pull request #2112 (2024-11-15 00:23:08)

TST: porting testing to new alignment module
## Summary by Sourcery

Port tests to the new alignment module and introduce additional tests for alignment functionalities, including casting to DNA/RNA and variable position identification.

Tests:
- Port tests to the new alignment module, including tests for sequence collection initialization, counts per sequence, and alignment slicing.
- Add new tests for alignment casting to DNA and RNA, ensuring info attributes are preserved.
- Introduce tests for variable positions and gap filtering in alignments.


Patch:
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
index f1dd438ed..75d6fd974 100644
--- a/src/cogent3/core/alignment.py
+++ b/src/cogent3/core/alignment.py
@@ -3363,7 +3363,7 @@ class AlignmentI(object):
             counts[i] = c.tolist(motifs)
         return MotifCountsArray(counts, motifs, row_indices=self.names)
 
-    def variable_positions(self, include_gap_motif=True):
+    def variable_positions(self, include_gap_motif=True):  # ported
         """Return a list of variable position indexes.
 
         Parameters
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 6fda78d94..f4f4ab4d4 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -4190,7 +4190,13 @@ class Alignment(SequenceCollection):
 
         Parameters
         ----------
-
+        motif_length
+            number of elements per character.
+        include_ambiguity
+            if True, motifs containing ambiguous characters from the seq moltype
+            are included. No expansion of those is attempted.
+        allow_gap
+            if True, motifs containing a gap character are included.
         warn
             warns if motif_length > 1 and alignment trimmed to produce
             motif columns
@@ -4207,7 +4213,7 @@ class Alignment(SequenceCollection):
         if not allow_gap:
             exclude_chars.update(self.moltype.gap)
 
-        if not include_ambiguity:
+        if not include_ambiguity and self.moltype.degen_alphabet:
             ambigs = [c for c, v in self.moltype.ambiguities.items() if len(v) > 1]
             exclude_chars.update(ambigs)
 
@@ -4468,6 +4474,28 @@ class Alignment(SequenceCollection):
 
         return result
 
+    def variable_positions(self, include_gap_motif: bool = True) -> list[int]:
+        """Return a list of variable position indexes.
+
+        Parameters
+        ----------
+        include_gap_motif
+            if False, sequences with a gap motif in a column are ignored.
+
+        """
+        # refactor: use numba decorated function for returning the indices
+        indices = []
+        gap_index = self.moltype.most_degen_alphabet().gap_index
+        missing_index = self.moltype.most_degen_alphabet().missing_index
+
+        for i, pos in enumerate(self.array_positions):
+            unique = numpy.unique(pos)
+            if not include_gap_motif:
+                unique = unique[(unique != gap_index) & (unique != missing_index)]
+            if len(unique) > 1:
+                indices.append(i)
+        return indices
+
     def omit_bad_seqs(self, quantile: OptFloat = None):
         """Returns new alignment without sequences with a number of uniquely
         introduced gaps exceeding quantile


Test patch:
diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
index cc31ef70d..57700cd7b 100644
--- a/tests/test_core/test_alignment.py
+++ b/tests/test_core/test_alignment.py
@@ -539,7 +539,7 @@ class SequenceCollectionBaseTests(object):
             aln.with_modified_termini(), {"s1": "AATGR??", "s2": "?T-AG??"}
         )
 
-    def test_make_gap_filter(self):
+    def test_make_gap_filter(self):  # ported
         """make_gap_filter returns f(seq) -> True if aligned ok w/ query"""
         s1 = RNA.make_seq(seq="UC-----CU---C")
         s3 = RNA.make_seq(seq="UUCCUUCUU-UUC")
@@ -944,7 +944,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
     as a constructor.
     """
 
-    def test_alignment_quality(self):
+    def test_alignment_quality(self):  # ported
         """check alignment method correctly invokes underlying app"""
         aln = self.Class(["AAAC", "ACGC", "AGCC", "A-TC"], moltype="dna")
         got = aln.alignment_quality(equifreq_mprobs=False)
@@ -1128,7 +1128,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
             alignmentUpper.iupac_consensus(PROTEIN), "ACBEFGHIKLMNPZRS?UVWX-"
         )
 
-    def test_is_ragged(self):
+    def test_is_ragged(self):  # ported
         """SequenceCollection is_ragged should return true if ragged alignment"""
         assert not self.identical.is_ragged()
         assert not self.gaps.is_ragged()
@@ -1269,7 +1269,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         # order now changes
         self.assertTrue(got.find(ref_row) < got.find(other_row))
 
-    def test_variable_positions(self):
+    def test_variable_positions(self):  # ported
         """correctly identify variable positions"""
         new_seqs = {"A": "-CG-C", "B": "ACAA?", "C": "GCGAC"}
         aln = self.Class(data=new_seqs, moltype=DNA)
@@ -1280,7 +1280,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         self.assertEqual(aln.variable_positions(include_gap_motif=True), [])
         self.assertEqual(aln.variable_positions(include_gap_motif=False), [])
 
-    def test_to_type_info(self):
+    def test_to_type_info(self):  # will not port
         """interconverting between alignment types preserves info attribute"""
         new_seqs = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
         array_align = self.Class == ArrayAlignment
@@ -1291,7 +1291,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         self.assertEqual(id(aln), id(new))
         self.assertEqual(new.info["key"], "value")
 
-    def test_to_dna(self):
+    def test_to_dna(self):  # ported
         """alignment cast to DNA works"""
         data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
         aln = self.Class(data=data)
@@ -1302,14 +1302,14 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         paln = dna.get_translation()
         self.assertRaises(AlphabetError, paln.to_dna)
 
-    def test_to_dna_info(self):
+    def test_to_dna_info(self):  # ported
         """alignment cast to DNA preserves info attribute"""
         data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
         aln = self.Class(data=data, info={"key": "value"})
         dna = aln.to_dna()
         self.assertEqual(dna.info["key"], "value")
 
-    def test_to_rna(self):
+    def test_to_rna(self):  # ported
         """alignment cast to RNA works"""
         data = {"seq1": "ACGUACGUA", "seq2": "ACCGAA---", "seq3": "ACGUACGUU"}
         aln = self.Class(data=data)
@@ -1317,14 +1317,14 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         self.assertEqual(set(rna.names), set(aln.names))
         self.assertTrue(rna.moltype == RNA)
 
-    def test_to_rna_info(self):
+    def test_to_rna_info(self):  # ported
         """alignment cast to RNA preserves info attribute"""
         data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
         aln = self.Class(data=data, info={"key": "value"})
         rna = aln.to_rna()
         self.assertEqual(rna.info["key"], "value")
 
-    def test_to_protein(self):
+    def test_to_protein(self):  # will not port (use get_translation)
         """alignment cast to protein works"""
         data = {"seq1": "TYV", "seq3": "TYV", "seq2": "TE-"}
         aln = self.Class(data=data)
@@ -1334,7 +1334,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         # should fail if invalid character set
         self.assertRaises(AlphabetError, paln.to_dna)
 
-    def test_to_protein_info(self):
+    def test_to_protein_info(self):  # will not port (use get_translation)
         """alignment cast to protein preserves info attribute"""
         data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
         aln = self.Class(data=data, info={"key": "value"})
@@ -1418,7 +1418,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         c = aln.counts_per_pos(include_ambiguity=False, allow_gap=True)
         assert_equal(set(c.motifs), set("ACGT-"))
 
-    def test_counts_per_seq_default_moltype(self):
+    def test_counts_per_seq_default_moltype(self):  # ported
         """produce correct counts per seq with default moltypes"""
         data = {"a": "AAAA??????", "b": "CCCGGG--NN", "c": "CCGGTTCCAA"}
         coll = self.Class(data=data)
@@ -1430,7 +1430,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         got = coll.counts_per_seq(include_ambiguity=True, allow_gap=True)
         self.assertEqual(got.col_sum()["-"], 2)
 
-    def test_counts_per_pos_default_moltype(self):
+    def test_counts_per_pos_default_moltype(self):  # ported
         """produce correct counts per pos with default moltypes"""
         data = {"a": "AAAA??????", "b": "CCCGGG--NN", "c": "CCGGTTCCAA"}
         coll = self.Class(data=data)
@@ -1453,7 +1453,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         self.assertTrue("-" not in found_motifs)
         self.assertEqual(lengths, {2})
 
-    def test_entropy_per_pos_just_gaps(self):
+    def test_entropy_per_pos_just_gaps(self):  # ported
         """pos with just gaps have nan"""
         a = self.Class(dict(a="A---", b="C---", c="C---"), moltype=DNA)
         entropy = a.entropy_per_pos()
@@ -1529,7 +1529,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
         assert_equal(got.array, [2, 1, 2])
         assert_equal(got["b"], 1)
 
-    def test_info_source(self):
+    def test_info_source(self):  # ported
         """info.source exists if load_aligned_seqs given a filename"""
         array_align = self.Class == ArrayAlignment
         seqs = load_aligned_seqs("data/brca1.fasta", array_align=array_align)
@@ -1601,7 +1601,7 @@ class AlignmentBaseTests(SequenceCollectionBaseTests):
 class ArrayAlignmentTests(AlignmentBaseTests, TestCase):
     Class = ArrayAlignment
 
-    def test_slice_align(self):
+    def test_slice_align(self):  # ported
         """slicing alignment should work correctly"""
         data = {"seq1": "ACGACGACG", "seq2": "ACGACGACG", "seq3": "ACGACGACG"}
         alignment = self.Class(data=data)
@@ -1614,7 +1614,7 @@ class ArrayAlignmentTests(AlignmentBaseTests, TestCase):
         expect = {"seq1": "GGG", "seq2": "GGG", "seq3": "GGG"}
         self.assertEqual(sub_align.to_dict(), expect)
 
-    def test_slice_align_info(self):
+    def test_slice_align_info(self):  # ported
         """slicing alignment preserves info attribute"""
         data = {"seq1": "ACGACGACG", "seq2": "ACGACGACG", "seq3": "ACGACGACG"}
         alignment = self.Class(data=data, info={"key": "value"})
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 4ac50d158..8b980d299 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -675,10 +675,11 @@ def test_sequence_collection_init_ordered(ordered1, ordered2):
     assert sec.names == ["c", "a"]
 
 
-def test_sequence_collection_info_source():
+@pytest.mark.parametrize("load_cls", [load_aligned_seqs, load_unaligned_seqs])
+def test_sequence_collection_info_source(load_cls):
     """info.source exists if load seqs given a filename"""
     path = pathlib.Path("data/brca1.fasta")
-    seqs = load_unaligned_seqs(path, moltype="dna", new_type=True)
+    seqs = load_cls(path, moltype="dna", new_type=True)
     assert seqs.info.source == str(path)
 
 
@@ -1699,10 +1700,13 @@ def test_sequence_collection_get_motif_probs_protein():
     expect = {"M": 3 / 10, "V": 3 / 10, "S": 2 / 10, "P": 1 / 10, "B": 1 / 10}
 
 
-def test_sequence_collection_counts_per_seq():
+@pytest.mark.parametrize(
+    "mk_cls", (new_alignment.make_aligned_seqs, new_alignment.make_unaligned_seqs)
+)
+def test_sequence_collection_counts_per_seq(mk_cls):
     """SequenceCollection.counts_per_seq handles motif length, allow_gaps etc.."""
     data = {"a": "AAAA??????", "b": "CCCGGG--NN", "c": "CCGGTTCCAA"}
-    coll = new_alignment.make_unaligned_seqs(data, moltype="dna")
+    coll = mk_cls(data, moltype="dna")
     mtype = coll.moltype
     got = coll.counts_per_seq()
     assert got["a", "A"] == 4
@@ -1726,12 +1730,46 @@ def test_sequence_collection_counts_per_seq():
         assert c[k] == expect[k]
 
 
+def test_counts_per_seq_text_moltype():
+    """produce correct counts per seq with text moltypes"""
+    data = {"a": "AAAA??????", "b": "CCCGGG--NN", "c": "CCGGTTCCAA"}
+    coll = new_alignment.make_aligned_seqs(data, moltype="text")
+    got = coll.counts_per_seq(include_ambiguity=True, allow_gap=True)
+    assert got.col_sum()["-"] == 2
+    assert got.col_sum()["?"] == 6
+    assert got.col_sum()["T"] == 2
+
+
+def test_counts_per_pos_text_moltype():
+    """produce correct counts per pos with default moltypes"""
+    data = {"a": "AAAA??????", "b": "CCCGGG--NN", "c": "CCGGTTCCAA"}
+    coll = new_alignment.make_aligned_seqs(data, moltype="text")
+    got = coll.counts_per_pos()
+    # should not include gap character
+    assert "-" not in got.motifs
+    # allowing gaps
+    got = coll.counts_per_pos(allow_gap=True)
+    # should include gap character
+    assert got[5, "-"] == 0
+    assert got[6, "-"] == 1
+
+    # now with motif-length 2
+    got = coll.counts_per_pos(motif_length=2)
+    found_motifs = set()
+    lengths = set()
+    for m in got.motifs:
+        lengths.add(len(m))
+        found_motifs.update(m)
+    assert "-" not in found_motifs
+    assert lengths == {2}
+
+
 @pytest.mark.parametrize(
-    "coll_maker", (new_alignment.make_aligned_seqs, new_alignment.make_unaligned_seqs)
+    "mk_cls", (new_alignment.make_aligned_seqs, new_alignment.make_unaligned_seqs)
 )
-def test_sequence_collection_probs_per_seq(coll_maker):
+def test_sequence_collection_probs_per_seq(mk_cls):
     data = {"seq1": "AA??", "seq2": "CG-N", "seq3": "CGAA"}
-    coll = coll_maker(data, moltype="dna")
+    coll = mk_cls(data, moltype="dna")
     got = coll.probs_per_seq()
     assert got["seq1", "A"] == 1.0
     assert got["seq2", "C"] == 0.5
@@ -2003,6 +2041,44 @@ def test_sequence_collection_to_dna(mk_cls):
     }
 
 
+def test_to_dna():
+    """alignment cast to DNA works"""
+    data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
+    aln = new_alignment.make_aligned_seqs(data, moltype="text")
+    dna = aln.to_dna()
+    assert set(dna.names) == set(aln.names)
+    assert dna.moltype.label == "dna"
+    # should fail if invalid character set
+    paln = dna.get_translation()
+    with pytest.raises(new_moltype.MolTypeError):
+        _ = paln.to_dna()
+
+
+def test_to_dna_info():
+    """alignment cast to DNA preserves info attribute"""
+    data = {"seq1": "ACGTACGTA", "seq2": "ACCGAA---", "seq3": "ACGTACGTT"}
+    aln = new_alignment.make_aligned_seqs(data, info={"key": "value"}, moltype="text")
+    dna = aln.to_dna()
+    assert dna.info["key"] == "value"
+
+
+def test_to_rna():
+    """alignment cast to RNA works"""
+    data = {"seq1": "ACGUACGUA", "seq2": "ACCGAA---", "seq3": "ACGUACGUU"}
+    aln = new_alignment.make_aligned_seqs(data, moltype="text")
+    rna = aln.to_rna()
+    assert set(rna.names) == set(aln.names)
+    assert rna.moltype.label == "rna"
+
+
+def test_to_rna_info():
+    """alignment cast to RNA preserves info attribute"""
+    data = {"seq1": "ACGUACGUA", "seq2": "ACCGAA---", "seq3": "ACGUACGUU"}
+    aln = new_alignment.make_aligned_seqs(data, info={"key": "value"}, moltype="text")
+    rna = aln.to_rna()
+    assert rna.info["key"] == "value"
+
+
 @pytest.mark.parametrize(
     "mk_cls",
     [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs],
@@ -3686,6 +3762,31 @@ def test_alignment_slice_neg_step_ungapped(aligned_dict, start, stop, step, seqi
     assert got == expect
 
 
+def test_slice_align():
+    """slicing alignment should work correctly"""
+    data = {"seq1": "ACGACGACG", "seq2": "ACGACGACG", "seq3": "ACGACGACG"}
+    alignment = new_alignment.make_aligned_seqs(data, moltype="dna")
+    sub_align = alignment[2:5]
+    assert isinstance(sub_align, new_alignment.Alignment)
+    expect = {"seq1": "GAC", "seq2": "GAC", "seq3": "GAC"}
+    assert sub_align.to_dict() == expect
+    # slice third positions
+    sub_align = alignment[2::3]
+    expect = {"seq1": "GGG", "seq2": "GGG", "seq3": "GGG"}
+    assert sub_align.to_dict() == expect
+
+
+def test_slice_align_info():
+    """slicing alignment preserves info attribute"""
+    data = {"seq1": "ACGACGACG", "seq2": "ACGACGACG", "seq3": "ACGACGACG"}
+    alignment = new_alignment.make_aligned_seqs(
+        data, info={"key": "value"}, moltype="dna"
+    )
+    sub_align = alignment[2:5]
+    assert len(sub_align) == 3
+    assert sub_align.info["key"] == "value"
+
+
 @pytest.mark.parametrize("method", ("ic_score", "cogent3_score", "sp_score"))
 def test_alignment_quality_methods(method):
     data = {
@@ -3706,6 +3807,33 @@ def test_alignment_quality_methods(method):
     assert numpy.allclose(score, expected_score)
 
 
+def test_alignment_quality():
+    """check alignment method correctly invokes underlying app"""
+    aln = new_alignment.make_aligned_seqs(
+        ["AAAC", "ACGC", "AGCC", "A-TC"], moltype="dna"
+    )
+    got = aln.alignment_quality(equifreq_mprobs=False)
+    expect = (
+        2 * numpy.log2(1 / 0.4)
+        + numpy.log2(1 / (4 * 0.4))
+        + (1 / 2) * numpy.log2(1 / (8 / 15))
+        + (1 / 4) * numpy.log2(1 / (4 / 15))
+    )
+    assert numpy.all(got == expect)
+
+
+def test_variable_positions():
+    """correctly identify variable positions"""
+    new_seqs = {"A": "-CG-C", "B": "ACAA?", "C": "GCGAC"}
+    aln = new_alignment.make_aligned_seqs(new_seqs, moltype="dna")
+    assert aln.variable_positions(include_gap_motif=True) == [0, 2, 3, 4]
+    assert aln.variable_positions(include_gap_motif=False) == [0, 2]
+    new_seqs = {"A": "GCGAC", "B": "GCGAC", "C": "GCGAC"}
+    aln = new_alignment.make_aligned_seqs(new_seqs, moltype="dna")
+    assert aln.variable_positions(include_gap_motif=True) == []
+    assert aln.variable_positions(include_gap_motif=False) == []
+
+
 def test_alignment_counts_per_pos():
     """correctly count motifs"""
     exp = numpy.array(
@@ -4736,3 +4864,30 @@ def test_asd_get_gapped_seq(aligned_dict, dna_alphabet):
     numpy.testing.assert_array_equal(fwd_gapped, orig_array)
     rev_gapped = rev.gapped_array_value
     numpy.testing.assert_array_equal(rev_gapped, orig_revd)
+
+
+def test_make_gap_filter():
+    """make_gap_filter returns f(seq) -> True if aligned ok w/ query"""
+    RNA = new_moltype.get_moltype("rna")
+    s1 = RNA.make_seq(seq="UC-----CU---C")
+    s3 = RNA.make_seq(seq="UUCCUUCUU-UUC")
+    s4 = RNA.make_seq(seq="UU-UUUU-UUUUC")
+    # check that the behavior is ok for gap runs
+    f1 = new_alignment.make_gap_filter(s1, 0.9, 5)
+    f3 = new_alignment.make_gap_filter(s3, 0.9, 5)
+    # Should return False since s1 has gap run >= 5 with respect to s3
+    assert f3(s1) == False
+    # Should return False since s3 has an insertion run >= 5 to s1
+    assert f1(s3) == False
+    # Should retun True since s4 does not have a long enough gap or ins run
+    assert f3(s4) == True
+    f3 = new_alignment.make_gap_filter(s3, 0.9, 6)
+    assert f3(s1) == True
+
+    # Check that behavior is ok for gap_fractions
+    f1 = new_alignment.make_gap_filter(s1, 0.5, 6)
+    f3 = new_alignment.make_gap_filter(s3, 0.5, 6)
+    # Should return False since 0.53% of positions are diff for gaps
+    assert f3(s1) == False
+    assert f1(s3) == False
+    assert f3(s4) == True