Repository: cogent3/cogent3 (python)

Issue:
ENH: add a SequenceCollection.count_ambiguous_per_seq method
This is really valuable for identifying low quality sequences. It should also be added to the `drop_bad_seqs()` method.

This should be implemented with `cogent3.core.new_sequence.Sequence.count_ambiguous()` and `cogent3.core.new_alignment.SequenceCollection.count_ambiguous_per_seq()`

Pull request:
ENH: add a SequenceCollection.count_ambiguous_per_seq method, fixes #2005
## Summary by Sourcery

Add a new method to count ambiguous characters per sequence and enhance sequence filtering by introducing an ambig_fraction parameter. Update tests to cover these new functionalities.

New Features:
- Introduce a method count_ambiguous_per_seq in SequenceCollection to count ambiguous characters per sequence.

Enhancements:
- Add ambig_fraction parameter to omit_bad_seqs function to allow filtering sequences based on ambiguous character content.

Tests:
- Add tests for the new count_ambiguous_per_seq method in SequenceCollection.
- Add tests for the ambig_fraction parameter in omit_bad_seqs function.

Patch:
diff --git a/src/cogent3/app/align.py b/src/cogent3/app/align.py
index 35d379900..0ba0df9e7 100644
--- a/src/cogent3/app/align.py
+++ b/src/cogent3/app/align.py
@@ -828,7 +828,7 @@ def cogent3_score(aln: AlignedSeqsType) -> float:
     if aln.num_seqs == 1 or len(aln) == 0:
         msg = "zero length" if len(aln) == 0 else "single sequence"
         return NotCompleted(
-            "FAL",
+            "FAIL",
             "cogent3_score",
             f"cannot compute alignment quality because {msg}",
             source=aln.info,
@@ -838,7 +838,7 @@ def cogent3_score(aln: AlignedSeqsType) -> float:
     return align_params.get(
         "lnL",
         NotCompleted(
-            "FAL", "cogent3_score", "no alignment quality score", source=aln.info
+            "FAIL", "cogent3_score", "no alignment quality score", source=aln.info
         ),
     )
 
diff --git a/src/cogent3/app/sample.py b/src/cogent3/app/sample.py
index 5ba807e4f..bd7a81869 100644
--- a/src/cogent3/app/sample.py
+++ b/src/cogent3/app/sample.py
@@ -16,6 +16,7 @@ from .typing import AlignedSeqsType, SeqsCollectionType, SerialisableType
 # set.
 
 MolTypes = Union[str, MolType]
+OptInt = Optional[int]
 
 
 def intersection(groups):
@@ -892,6 +893,7 @@ class omit_bad_seqs:
         self,
         quantile: Optional[float] = None,
         gap_fraction: int = 1,
+        ambig_fraction: OptInt = None,  # refactor: set default to 1 when support for old style aln is dropped
         moltype: MolTypes = "dna",
     ):
         """
@@ -905,6 +907,10 @@ class omit_bad_seqs:
         gap_fraction
             sequences whose proportion of gaps is >= this value are excluded, the
             default excludes sequences that are just gaps.
+        ambig_fraction
+            sequences whose proportion of ambiguous characters is >= this value
+            are excluded. Only supported for new style alignments. By default,
+            not applied.
         moltype
             molecular type, can be string or instance
 
@@ -954,8 +960,10 @@ class omit_bad_seqs:
         assert (
             moltype.label.lower() in "dna rna protein protein_with_stop"
         ), "moltype must be one of DNA, RNA or PROTEIN"
+        # refactor: design, this should raise a MolTypeError
         self._quantile = quantile
         self._gap_fraction = gap_fraction
+        self._ambig_fraction = ambig_fraction
         self._moltype = moltype
 
     T = Union[SerialisableType, AlignedSeqsType]
@@ -967,6 +975,24 @@ class omit_bad_seqs:
         gaps_per_seq = aln.count_gaps_per_seq()
         length = len(aln)
         keep = [n for n, c in gaps_per_seq.items() if c / length < self._gap_fraction]
+
+        from cogent3.core import new_alignment
+
+        if self._ambig_fraction is not None:
+            if isinstance(aln, new_alignment.Alignment):
+                ambigs_per_seq = aln.count_ambiguous_per_seq()
+                keep = [
+                    n for n in keep if ambigs_per_seq[n] / length < self._ambig_fraction
+                ]
+            else:
+                msg = (
+                    "ambig_fraction cannot be applied to old style alignments, set "
+                    "new_type=True in make_aligned_seqs to create an Alignment which "
+                    "supports this feature"
+                )
+
+                return NotCompleted("FAIL", self, msg, aln.info.source)
+
         result = aln.take_seqs(keep)
         if self._quantile is not None:
             result = result.omit_bad_seqs(quantile=self._quantile)
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 4df43ebc1..66b4ed777 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -43,7 +43,7 @@ from cogent3.format.phylip import alignment_to_phylip
 from cogent3.maths.stats.number import CategoryCounter
 from cogent3.util import progress_display as UI
 from cogent3.util import warning as c3warn
-from cogent3.util.deserialise import deserialise_object, register_deserialiser
+from cogent3.util.deserialise import register_deserialiser
 from cogent3.util.dict_array import DictArray, DictArrayTemplate
 from cogent3.util.io import atomic_write, get_format_suffixes
 from cogent3.util.misc import (
@@ -1592,6 +1592,17 @@ class SequenceCollection:
         )
         return per_seq.motif_totals()
 
+    def count_ambiguous_per_seq(self) -> DictArray:
+        """Counts of ambiguous characters per sequence."""
+
+        darr = DictArrayTemplate(self.names)
+        counts = numpy.array(
+            [self.seqs[name].count_ambiguous() for name in self.names],
+            dtype=numpy.uint32,
+        )
+
+        return darr.wrap(counts)
+
     def get_motif_probs(
         self,
         alphabet: new_alphabet.AlphabetABC = None,
@@ -4496,7 +4507,7 @@ class Alignment(SequenceCollection):
         unique: bool = False,
         include_ambiguity: bool = True,
         drawable: bool = False,
-    ):
+    ) -> DictArray:
         """return counts of gaps per sequence as a DictArray
 
         Parameters
@@ -4551,6 +4562,15 @@ class Alignment(SequenceCollection):
 
         return result
 
+    def count_ambiguous_per_seq(self) -> DictArray:
+        """Return the counts of ambiguous characters per sequence as a DictArray."""
+
+        gap_index = self.moltype.most_degen_alphabet().gap_index
+        ambigs_pos = self.array_seqs > gap_index
+        ambigs = ambigs_pos.sum(axis=1)
+
+        return DictArray.from_array_names(ambigs, self.names)
+
     def variable_positions(
         self,
         include_gap_motif: bool = True,
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index b859978f9..8d74a9602 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -363,6 +363,12 @@ class Sequence:
 
         return counts
 
+    def count_ambiguous(self) -> int:
+        """Returns the number of ambiguous characters in the sequence."""
+        data = numpy.array(self)
+        gap_index = self.moltype.most_degen_alphabet().gap_index
+        return numpy.sum(data > gap_index)
+
     def __lt__(self, other):
         """compares based on the sequence string."""
         return str(self) < str(other)


Test patch:
diff --git a/tests/test_app/test_sample.py b/tests/test_app/test_sample.py
index eb1ca31dc..f0768b699 100644
--- a/tests/test_app/test_sample.py
+++ b/tests/test_app/test_sample.py
@@ -293,40 +293,6 @@ class TranslateTests(TestCase):
         for name, seq in got.to_dict().items():
             self.assertIn(seq, expect[name])
 
-    def test_omit_bad_seqs(self):
-        """correctly omit bad sequences from an alignment"""
-        data = {
-            "s1": "---ACC---TT-",
-            "s2": "---ACC---TT-",
-            "s3": "---ACC---TT-",
-            "s4": "--AACCG-GTT-",
-            "s5": "--AACCGGGTTT",
-            "s6": "AGAACCGGGTT-",
-            "s7": "------------",
-        }
-        aln = make_aligned_seqs(data=data, moltype=DNA)
-        # default just eliminates strict gap sequences
-        dropbad = sample.omit_bad_seqs()
-        got = dropbad(aln)
-        expect = data.copy()
-        del expect["s7"]
-        self.assertEqual(got.to_dict(), expect)
-        # providing a more stringent gap_frac
-        dropbad = sample.omit_bad_seqs(gap_fraction=0.5)
-        got = dropbad(aln)
-        expect = data.copy()
-        for n in ("s1", "s2", "s3", "s7"):
-            del expect[n]
-        self.assertEqual(got.to_dict(), expect)
-
-        # setting quantile drops additional sequences
-        dropbad = sample.omit_bad_seqs(quantile=6 / 7)
-        got = dropbad(aln)
-        expect = data.copy()
-        for n in ("s6", "s7"):
-            del expect[n]
-        self.assertEqual(got.to_dict(), expect)
-
     def test_omit_duplicated(self):
         """correctly drop duplicated sequences"""
         # strict omit_duplicated
@@ -565,3 +531,92 @@ def test_concat_empty(data):
     ccat = sample.concat()
     got = ccat(data)
     assert isinstance(got, NotCompleted)
+
+
+@pytest.fixture
+def bad_gap_data():
+    return {
+        "s1": "---ACC---TT-",
+        "s2": "---ACC---TT-",
+        "s3": "---ACC---TT-",
+        "s4": "--AACCG-GTT-",
+        "s5": "--AACCGGGTTT",
+        "s6": "AGAACCGGGTT-",
+        "s7": "------------",
+    }
+
+
+@pytest.mark.parametrize(
+    "gap_fraction, quantile, expected_keys",
+    [
+        (
+            1,
+            None,
+            {"s1", "s2", "s3", "s4", "s5", "s6"},
+        ),  # default just eliminates strict gap sequences
+        (0.5, None, {"s4", "s5", "s6"}),  # providing a more stringent gap_frac
+        (
+            1,
+            6 / 7,
+            {"s1", "s2", "s3", "s4", "s5"},
+        ),  # setting quantile drops additional sequences
+    ],
+)
+def test_omit_bad_seqs(bad_gap_data, gap_fraction, quantile, expected_keys):
+    """correctly omit bad sequences from an alignment"""
+
+    dropbad = sample.omit_bad_seqs(gap_fraction=gap_fraction, quantile=quantile)
+    aln = make_aligned_seqs(bad_gap_data, moltype=DNA)
+    got = dropbad(aln)
+    expected = {k: bad_gap_data[k] for k in expected_keys}
+    assert got.to_dict() == expected
+
+
+@pytest.fixture
+def bad_ambig_gap_data():
+    return {
+        "s1": "---ACC---TT-",
+        "s2": "---ACC---TT-",
+        "s3": "SSSACCSSSTTS",
+        "s4": "SSAACCGSGTTS",
+        "s5": "--AACCGGGTTT",
+        "s6": "AGAACCGGGTT-",
+        "s7": "SSSSSSSSSSSS",
+    }
+
+
+def test_omit_bad_seqs_ambigs(bad_ambig_gap_data):
+    """correctly omit sequences from a new style alignment via fraction of ambiguous data"""
+    from cogent3.core import new_alignment
+
+    aln = new_alignment.make_aligned_seqs(bad_ambig_gap_data, moltype="dna")
+
+    # drop sequences with all ambiguous data
+    dropbad = sample.omit_bad_seqs(ambig_fraction=1)
+    got = dropbad(aln)
+    assert set(got.to_dict().keys()) == {"s1", "s2", "s3", "s4", "s5", "s6"}
+
+    # drop sequences with more than 50% ambiguous data
+    dropbad = sample.omit_bad_seqs(ambig_fraction=0.5)
+    got = dropbad(aln)
+    assert set(got.to_dict().keys()) == {"s1", "s2", "s4", "s5", "s6"}
+
+    # drop sequences with any ambiguous data
+    dropbad = sample.omit_bad_seqs(ambig_fraction=0.01)
+    got = dropbad(aln)
+    assert set(got.to_dict().keys()) == {"s1", "s2", "s5", "s6"}
+
+    # drop sequences with more than 50% ambiguous data and 50% gaps
+    dropbad = sample.omit_bad_seqs(gap_fraction=0.5, ambig_fraction=0.5)
+    got = dropbad(aln)
+    assert set(got.to_dict().keys()) == {"s4", "s5", "s6"}
+
+
+def test_omit_bad_seqs_ambigs_old_aln(bad_ambig_gap_data):
+    # ambig_fraction should be ignored if using old style alignment
+
+    aln = make_aligned_seqs(bad_ambig_gap_data, moltype=DNA)
+    dropbad = sample.omit_bad_seqs(gap_fraction=0.5, ambig_fraction=0.5)
+
+    got = dropbad(aln)
+    assert isinstance(got, NotCompleted)
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index e8bd69342..a0d9ad405 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -1839,6 +1839,22 @@ def test_sequence_collection_counts():
     assert all(got[k] == v for k, v in expect.items())
 
 
+@pytest.mark.parametrize(
+    "mk_cls", [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs]
+)
+def test_sequence_collection_count_ambiguous_per_seq(mk_cls):
+    data = {
+        "a": "AGGGT",
+        "b": "AGGN?",
+        "c": "?????",
+        "d": "-----",
+    }
+    coll = mk_cls(data, moltype="dna")
+    got = coll.count_ambiguous_per_seq()
+    expect = {"a": 0, "b": 2, "c": 5, "d": 0}
+    assert got == expect
+
+
 def test_sequence_collection_make_feature(seqs):
     data = {"seqid": "seq1", "biotype": "xyz", "name": "abc", "spans": [(1, 2)]}
     feat = seqs.make_feature(feature=data)
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 7a6883a35..e19e793b8 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -831,6 +831,20 @@ def test_counts():
     assert dict(got) == expect
 
 
+@pytest.mark.parametrize(
+    "data, expect",
+    [
+        ("CCTTNN", 2),
+        ("???MMM", 6),
+        ("TGCAGG", 0),
+        ("", 0),
+    ],
+)
+def test_count_ambiguous(data, expect):
+    seq = new_moltype.DNA.make_seq(seq=data)
+    assert seq.count_ambiguous() == expect
+
+
 def test_strand_symmetry():
     """correctly compute test of strand symmetry"""
     from cogent3 import get_moltype