Repository: cogent3/cogent3
Pull request #2260 (2025-03-22 02:14:24)

MAINT: new_alignment get_motif_probs() works with custom alphabets
[CHANGED] now handles provision of k-mer alphabets with k > 1

## Summary by Sourcery

Fixes get_motif_probs() to work with custom alphabets, particularly k-mer alphabets with k > 1. Also includes minor optimisations and a deprecated function.

Bug Fixes:
- Fixes get_motif_probs() to handle provision of k-mer alphabets with k > 1.
- Fixes take_positions() to work with numpy arrays.

Enhancements:
- Optimises take_positions() by using numpy.array.take instead of creating a new array.

Tests:
- Adds a test case for get_motif_probs() with custom alphabets.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 9f17e11bc..28be7d575 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -8,7 +8,7 @@ import re
 import typing
 import warnings
 from abc import ABC, abstractmethod
-from collections import defaultdict
+from collections import Counter, defaultdict
 from collections.abc import Callable, Iterable, Iterator, Mapping
 from collections.abc import Sequence as PySeq
 from functools import singledispatch, singledispatchmethod
@@ -1749,21 +1749,20 @@ class SequenceCollection:
             if allow_gap:
                 alphabet = moltype.gapped_alphabet
 
-        counts = {}
+        motif_len = alphabet.motif_len
+        counts = Counter()
         for seq_name in self.names:
             sequence = self.seqs[seq_name]
-            motif_len = alphabet.motif_len
             if motif_len > 1:
-                posns = list(range(0, len(sequence) + 1 - motif_len, motif_len))
-                sequence = [sequence[i : i + motif_len] for i in posns]
+                sequence = [
+                    str(sequence[i : i + motif_len])
+                    for i in range(0, len(sequence) + 1 - motif_len, motif_len)
+                ]
             for motif in sequence:
                 if not allow_gap and self.moltype.gap in motif:
                     continue
 
-                if motif in counts:
-                    counts[motif] += 1
-                else:
-                    counts[motif] = 1
+                counts[motif] += 1
 
         probs = {}
         if not exclude_unobserved:
@@ -4556,7 +4555,11 @@ class Alignment(SequenceCollection):
         # use array_positions here
         return [i for i, col in enumerate(self.positions) if new_f(col)]
 
-    def take_positions(self, cols: list, negate: bool = False):
+    def take_positions(
+        self,
+        cols: list,
+        negate: bool = False,
+    ) -> typing_extensions.Self:
         """Returns new Alignment containing only specified positions.
 
         Parameters
@@ -4572,7 +4575,7 @@ class Alignment(SequenceCollection):
             cols = [i for i in range(len(self)) if i not in col_lookup]
 
         new_data = {
-            aligned.data.seqid: numpy.array(aligned.gapped_seq)[cols]
+            aligned.data.seqid: numpy.array(aligned.gapped_seq).take(cols)
             for aligned in self.seqs
         }
         seqs_data = self._seqs_data.from_seqs(
@@ -6494,12 +6497,21 @@ class Alignment(SequenceCollection):
         kwargs.pop("annotation_db", None)
         return self.__class__(**kwargs)
 
+    @c3warn.deprecated_callable(
+        "2025.6",
+        reason="make intent clearer",
+        new="apply_scaled_gaps",
+    )
+    def replace_seqs(self, *args, **kwargs) -> typing_extensions.Self:
+        """returns a new alignment with sequences replaced"""
+        return self.apply_scaled_gaps(*args, **kwargs)
+
     def apply_scaled_gaps(
         self,
         other: SequenceCollection,
-        aa_to_codon: bool | None = None,
-    ):
-        """applies gaps in self to unagpped sequences"""
+        aa_to_codon: bool = True,
+    ) -> typing_extensions.Self:
+        """applies gaps in self to ungapped sequences"""
         assert set(other.names) == set(self.names), "Names must match"
         if aa_to_codon and not all(
             (not self.moltype.is_nucleic, other.moltype.is_nucleic),


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 036c704a9..930c8e514 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -1790,6 +1790,31 @@ def test_sequence_collection_get_motif_probs():
     assert got == expect
 
 
+@pytest.mark.parametrize(
+    "mk_cls",
+    [new_alignment.make_unaligned_seqs, new_alignment.make_aligned_seqs],
+)
+def test_get_motif_probs_alpha(mk_cls):
+    data = {
+        "s1": "GT GT AG AA GT TC CA AA TG AA".replace(" ", ""),
+        "s2": "GT GT AG AA GT TC CA AA TG AA".replace(" ", ""),
+        "s3": "GT GT AG AA GT TC CA AA TG AA".replace(" ", ""),
+    }
+    seqs = mk_cls(data, moltype="dna")
+    dinucs = seqs.moltype.alphabet.get_kmer_alphabet(2)
+    mprobs = seqs.get_motif_probs(alphabet=seqs.moltype.alphabet.get_kmer_alphabet(2))
+    expect = {
+        "GT": 9 / 30,
+        "AG": 3 / 30,
+        "AA": 9 / 30,
+        "TC": 3 / 30,
+        "CA": 3 / 30,
+        "TG": 3 / 30,
+    }
+    expect |= {d: 0.0 for d in dinucs if d not in expect}
+    assert mprobs == expect
+
+
 def test_sequence_collection_get_motif_probs_protein():
     data = {"a": "MVSB", "b": "MVS", "c": "MVP"}
     aln = new_alignment.make_unaligned_seqs(data, moltype="protein")
@@ -3566,9 +3591,10 @@ def test_alignment_array_positions(simple_aln):
     assert numpy.array_equal(got, expect)
 
 
-def test_alignment_array_positions_take_positions(simple_aln):
+@pytest.mark.parametrize("pos_type", [list, tuple, numpy.array])
+def test_alignment_array_positions_take_positions(simple_aln, pos_type):
     """an alignment which has been subset should return the correct array_positions"""
-    subset = simple_aln.take_positions([0])
+    subset = simple_aln.take_positions(pos_type([0]))
     got = subset.array_positions
     expect = numpy.array([[0, 4, 2]])
     assert numpy.array_equal(got, expect)