Repository: cogent3/cogent3
Pull request #2260: MAINT: new_alignment get_motif_probs() works with custom alphabets
Author: login: GavinHuttley
Created: 2025-03-22T02:14:24Z | Merged: 2025-03-22T02:31:00Z

[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.


Diff:
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 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 @@ def get_motif_probs(
             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 @@ def get_position_indices(
         # 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 @@ def take_positions(self, cols: list, negate: bool = False):
             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 @@ def _mapped(self, slicemap):
         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),
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)


Reviews:
user: login: sourcery-ai[bot]
body: Hey @GavinHuttley - I've reviewed your changes - here's some feedback:

**Overall Comments**:
- The use of `typing_extensions.Self` as a return annotation is a nice touch for improving type hinting.
- The deprecation warning is a good way to signal API changes.

<details>
<summary>Here's what I looked at during the review</summary>

- 🟡 **General issues**: 1 issue found
- 🟢 **Security**: all looks good
- 🟢 **Testing**: all looks good
- 🟡 **Complexity**: 1 issue found
- 🟢 **Documentation**: all looks good
</details>

***

<details>
<summary>Sourcery is free for open source - if you like our reviews please consider sharing them ✨</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20got%20an%20instant%20code%20review%20from%20%40SourceryAI%2C%20and%20it%20was%20brilliant%21%20It%27s%20free%20for%20open%20source%20and%20has%20a%20free%20trial%20for%20private%20code.%20Check%20it%20out%20https%3A//sourcery.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20got%20an%20instant%20code%20review%20from%20%40SourceryAI%2C%20and%20it%20was%20brilliant%21%20It%27s%20free%20for%20open%20source%20and%20has%20a%20free%20trial%20for%20private%20code.%20Check%20it%20out%20https%3A//sourcery.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https://sourcery.ai)
- [Facebook](https://www.facebook.com/sharer/sharer.php?u=https://sourcery.ai)

</details>

<sub>
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
</sub>
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-22T02:15:24Z

Review comments:
diff_hunk: @@ -6494,11 +6497,20 @@ def _mapped(self, slicemap):
         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,
-    ):
+    ) -> typing_extensions.Self:
         """applies gaps in self to unagpped sequences"""
path: src/cogent3/core/new_alignment.py
user: login: sourcery-ai[bot]
body: **nitpick (typo):** Typo noted in the docstring.

The word 'unagpped' appears to be a typo; it should likely be 'ungapped' to accurately describe the behavior.

```suggestion
        """applies gaps in self to ungapped sequences"""
```
created_at: 2025-03-22T02:15:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -6494,11 +6497,20 @@ def _mapped(self, slicemap):
         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:
path: src/cogent3/core/new_alignment.py
user: login: sourcery-ai[bot]
body: **issue (complexity):** Consider directly decorating the alias to simplify the deprecation and avoid an extra wrapper function, which will reduce indirection and abstraction while maintaining functionality and the deprecation warning.

You can simplify the deprecation alias by avoiding the extra wrapper function. Instead of defining a new function that just calls `apply_scaled_gaps`, you can decorate the alias directly. For example:

```python
apply_scaled_gaps = ...  # your existing implementation

replace_seqs = c3warn.deprecated_callable(
    "2025.6",
    reason="make intent clearer",
    new="apply_scaled_gaps"
)(apply_scaled_gaps)
```

This reduces the indirection / extra abstraction while keeping the same functionality and deprecation warning.
created_at: 2025-03-22T02:15:23Z
author_association: CONTRIBUTOR

PR comments:
user: login: sourcery-ai[bot]
created_at: 2025-03-22T02:14:28Z
body: <!-- Generated by sourcery-ai[bot]: start review_guide -->

## Reviewer's Guide by Sourcery

This pull request modifies the `get_motif_probs()` function to work with custom alphabets, including k-mer alphabets where k > 1. It also optimizes the `take_positions()` method in the `Alignment` class for performance and deprecates the `replace_seqs()` method in the `SequenceCollection` class, aliasing it to `apply_scaled_gaps()`.

#### Sequence diagram for get_motif_probs with k-mer alphabets

```mermaid
sequenceDiagram
    participant Alignment
    participant Alphabet
    participant Sequence

    Alignment->>Alphabet: get_motif_len()
    alt motif_len > 1
        Alignment->>Sequence: Iterate through sequence with step motif_len
        Sequence-->>Alignment: Returns motifs of length motif_len
    end
    Alignment->>Alignment: Count motifs
    Alignment->>Alignment: Calculate motif probabilities
```

#### Updated class diagram for Alignment

```mermaid
classDiagram
    class Alignment {
        +take_positions(cols: list, negate: bool = False)
    }
```

### File-Level Changes

| Change | Details | Files |
| ------ | ------- | ----- |
| Modified `get_motif_probs()` to handle k-mer alphabets with k > 1. | <ul><li>Modified `get_motif_probs()` to accept a custom alphabet.</li><li>Modified `get_motif_probs()` to correctly count motifs when using a k-mer alphabet.</li><li>Added a test case for `get_motif_probs()` with a custom alphabet.</li></ul> | `src/cogent3/core/new_alignment.py`<br/>`tests/test_core/test_new_alignment.py` |
| Optimized `take_positions()` method in `Alignment` class. | <ul><li>Replaced list comprehension with `numpy.array.take()` for improved performance.</li><li>Added type hint for return value.</li><li>Added a parameterize decorator to test with different position types.</li></ul> | `src/cogent3/core/new_alignment.py`<br/>`tests/test_core/test_new_alignment.py` |
| Deprecated `replace_seqs()` method in `SequenceCollection` class. | <ul><li>Deprecated `replace_seqs()` and aliased to `apply_scaled_gaps()` with a deprecation warning.</li></ul> | `src/cogent3/core/new_alignment.py` |

---

<details>
<summary>Tips and commands</summary>

#### Interacting with Sourcery

- **Trigger a new review:** Comment `@sourcery-ai review` on the pull request.
- **Continue discussions:** Reply directly to Sourcery's review comments.
- **Generate a GitHub issue from a review comment:** Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with `@sourcery-ai issue` to create an issue from it.
- **Generate a pull request title:** Write `@sourcery-ai` anywhere in the pull
  request title to generate a title at any time. You can also comment
  `@sourcery-ai title` on the pull request to (re-)generate the title at any time.
- **Generate a pull request summary:** Write `@sourcery-ai summary` anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment `@sourcery-ai summary` on the pull request to
  (re-)generate the summary at any time.
- **Generate reviewer's guide:** Comment `@sourcery-ai guide` on the pull
  request to (re-)generate the reviewer's guide at any time.
- **Resolve all Sourcery comments:** Comment `@sourcery-ai resolve` on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
- **Dismiss all Sourcery reviews:** Comment `@sourcery-ai dismiss` on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  `@sourcery-ai review` to trigger a new review!
- **Generate a plan of action for an issue:** Comment `@sourcery-ai plan` on
  an issue to generate a plan of action for it.

#### Customizing Your Experience

Access your [dashboard](https://app.sourcery.ai) to:
- Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
- Change the review language.
- Add, remove or edit custom review instructions.
- Adjust other review settings.

#### Getting Help

- [Contact our support team](mailto:support@sourcery.ai) for questions or feedback.
- Visit our [documentation](https://docs.sourcery.ai) for detailed guides and information.
- Keep in touch with the Sourcery team by following us on [X/Twitter](https://x.com/SourceryAI), [LinkedIn](https://www.linkedin.com/company/sourcery-ai/) or [GitHub](https://github.com/sourcery-ai).

</details>

<!-- Generated by sourcery-ai[bot]: end review_guide -->
author_association: CONTRIBUTOR
user: login: coveralls
created_at: 2025-03-22T02:26:03Z
body: ## Pull Request Test Coverage Report for [Build 14004221843](https://coveralls.io/builds/72867234)


### Details

* **8** of **9**   **(88.89%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.003%**) to **90.996%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [src/cogent3/core/new_alignment.py](https://coveralls.io/builds/72867234/source?filename=src%2Fcogent3%2Fcore%2Fnew_alignment.py#L6507) | 8 | 9 | 88.89%
<!-- | **Total:** | **8** | **9** | **88.89%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/72867234/badge)](https://coveralls.io/builds/72867234) |
| :-- | --: |
| Change from base [Build 14003345304](https://coveralls.io/builds/72866790): |  0.003% |
| Covered Lines: | 29814 |
| Relevant Lines: | 32764 |

---
##### 💛  - [Coveralls](https://coveralls.io)

author_association: COLLABORATOR