Repository: cogent3/cogent3
Pull request #2129 (2024-11-22 04:09:41)

ENH: improved speed of variable_positions
[CHANGED] use numba.jit functions for different argument cases.
     These are now cached.

[NEW] added motif_length argument to variable_positions method

## Summary by Sourcery

Enhance the variable_positions method by integrating numba.jit for performance improvements and introducing a new motif_length argument. Update tests to validate the new functionality and performance enhancements.

New Features:
- Add a motif_length argument to the variable_positions method, enabling the identification of variable positions based on motif length.

Enhancements:
- Improve performance of the variable_positions method by using numba.jit for different argument cases, allowing for caching and faster execution.

Tests:
- Expand tests for the variable_positions method to cover new functionality, including the motif_length argument and different combinations of include_gap_motif and include_ambiguity parameters.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 15615837a..a0f5a3e1a 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -4268,6 +4268,11 @@ class Alignment(SequenceCollection):
             warns if motif_length > 1 and alignment trimmed to produce
             motif columns
         """
+        # refactor: performance
+        # use self.variable_positions and a numba decorated
+        # function for counting k-mers the latter should allow returning the
+        # first allowed state for when position is not variable
+
         align_len = len(self._slice_record)
         length = (align_len // motif_length) * motif_length
         if warn and align_len != length:
@@ -4541,27 +4546,53 @@ class Alignment(SequenceCollection):
 
         return result
 
-    def variable_positions(self, include_gap_motif: bool = True) -> list[int]:
+    def variable_positions(
+        self,
+        include_gap_motif: bool = True,
+        include_ambiguity: bool = False,
+        motif_length: int = 1,
+    ) -> tuple[int]:
         """Return a list of variable position indexes.
 
         Parameters
         ----------
         include_gap_motif
             if False, sequences with a gap motif in a column are ignored.
+        include_ambiguity
+            if True, all states are considered.
+        motif_length
+            if any position within a motif is variable, the entire motif is
+            considered variable.
+
+        Returns
+        -------
+        tuple of integers, if motif_length > 1, the returned positions are
+        motif_length long sequential indices.
 
+        Notes
+        -----
+        Truncates alignment to be modulo motif_length.
         """
-        # 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
+        align_len = len(self) // motif_length * motif_length
+        array_seqs = self.array_seqs[:, :align_len]
+        # both gap and missing data are treated as gaps
+        alpha = self.moltype.most_degen_alphabet()
+        gap_index = alpha.gap_index
+        missing_index = alpha.missing_index
+        # at the individual position level
+        if include_gap_motif and include_ambiguity:
+            indices = _var_pos_allow_all(array_seqs)
+        elif include_gap_motif:
+            indices = _var_pos_canonical_or_gap(array_seqs, gap_index, missing_index)
+        elif include_ambiguity:
+            indices = _var_pos_not_gap(array_seqs, gap_index)
+        else:
+            indices = _var_pos_canonical(array_seqs, gap_index)
+
+        if motif_length > 1:
+            indices = indices.reshape(-1, motif_length).any(axis=1).repeat(motif_length)
+
+        return tuple(numpy.where(indices)[0].tolist())
 
     def omit_bad_seqs(self, quantile: OptFloat = None):
         """Returns new alignment without sequences with a number of uniquely
@@ -6269,6 +6300,8 @@ def make_aligned_seqs(
     )
 
 
+# I'm explicitly encoding the alternate variants of the following as numba
+# can then cache the byte-compiiled functions.
 @make_aligned_seqs.register
 def _(
     data: AlignedSeqsDataABC,
@@ -6313,3 +6346,136 @@ def _(
     if label_to_name:
         aln = aln.rename_seqs(label_to_name)
     return aln
+
+
+@numba.jit(cache=True)
+def _var_pos_canonical_or_gap(
+    arr: numpy.ndarray, gap_index: int, missing_index: int
+) -> numpy.ndarray:  # pragma: no cover
+    """return boolean array indicating columns with more than one value below threshold
+
+    Parameters
+    ----------
+    arr
+        a 2D array with rows being sequences and columns positions
+    gap_index
+        value of gap state
+    missing_index
+        value of missing state
+
+    Returns
+    ------
+    a boolean array
+    """
+    # relying on consistent ordering of gap as num canonical + 1
+    m, n = arr.shape
+    result = numpy.zeros(n, dtype=numpy.bool_)
+    for pos in numpy.arange(n):
+        last = -1
+        for seq in numpy.arange(m):
+            state = arr[seq, pos]
+            if state <= gap_index or state == missing_index:
+                if last == -1:
+                    last = state
+                elif state != last:
+                    result[pos] = True
+                    break
+
+    return result
+
+
+@numba.jit(cache=True)
+def _var_pos_canonical(
+    arr: numpy.ndarray, gap_index: int
+) -> numpy.ndarray:  # pragma: no cover
+    """return boolean array indicating columns with more than one value below threshold
+
+    Parameters
+    ----------
+    arr
+        a 2D array with rows being sequences and columns positions
+    gap_index
+        value of gap state
+
+    Returns
+    ------
+    a boolean array
+    """
+    # relying on consistent ordering of gap as num canonical + 1
+    m, n = arr.shape
+    result = numpy.zeros(n, dtype=numpy.bool_)
+    for pos in numpy.arange(n):
+        last = -1
+        for seq in numpy.arange(m):
+            state = arr[seq, pos]
+            if state < gap_index:
+                if last == -1:
+                    last = state
+                elif state != last:
+                    result[pos] = True
+                    break
+
+    return result
+
+
+@numba.jit(cache=True)
+def _var_pos_allow_all(
+    arr: numpy.ndarray,
+) -> numpy.ndarray:  # pragma: no cover
+    """return boolean array indicating columns with more than one value below threshold
+
+    Parameters
+    ----------
+    arr
+        a 2D array with rows being sequences and columns positions
+
+    Returns
+    ------
+    a boolean array
+    """
+    m, n = arr.shape
+    result = numpy.zeros(n, dtype=numpy.bool_)
+    for pos in numpy.arange(n):
+        last = -1
+        for seq in numpy.arange(m):
+            state = arr[seq, pos]
+            if last == -1:
+                last = state
+            elif state != last:
+                result[pos] = True
+                break
+    return result
+
+
+@numba.jit(cache=True)
+def _var_pos_not_gap(
+    arr: numpy.ndarray,
+    gap_index: int,
+) -> numpy.ndarray:  # pragma: no cover
+    """return boolean array indicating columns with more than one value below threshold
+
+    Parameters
+    ----------
+    arr
+        a 2D array with rows being sequences and columns positions
+    gap_index
+        value of gap state
+
+    Returns
+    ------
+    a boolean array
+    """
+    m, n = arr.shape
+    result = numpy.zeros(n, dtype=numpy.bool_)
+    for pos in numpy.arange(n):
+        last = -1
+        for seq in numpy.arange(m):
+            state = arr[seq, pos]
+            if state != gap_index:
+                if last == -1:
+                    last = state
+                elif state != last:
+                    result[pos] = True
+                    break
+
+    return result
diff --git a/src/cogent3/core/new_alphabet.py b/src/cogent3/core/new_alphabet.py
index 832dd782f..a29b19f8d 100644
--- a/src/cogent3/core/new_alphabet.py
+++ b/src/cogent3/core/new_alphabet.py
@@ -260,17 +260,17 @@ class CharAlphabet(tuple, AlphabetABC, MonomerAlphabetABC):
         gap: OptStr = None,
         missing: OptStr = None,
     ):
+        self.dtype = get_array_type(len(self))
         self._gap_char = gap
-        self._gap_index = self.index(gap) if gap else None
+        self._gap_index = self.dtype(self.index(gap)) if gap else None
         self._missing_char = missing
-        self._missing_index = self.index(missing) if missing else None
+        self._missing_index = self.dtype(self.index(missing)) if missing else None
 
         # the number of canonical states are non-gap, non-missing states
         adj = (1 if gap else 0) + (1 if missing else 0)
 
         self._num_canonical = len(self) - adj
 
-        self.dtype = get_array_type(len(self))
         self._chars = set(self)  # for quick lookup
         byte_chars = self.as_bytes()
         self._bytes2arr = bytes_to_array(byte_chars, dtype=self.dtype)


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index d07bbb1d5..2c9c80c7d 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -3816,12 +3816,56 @@ 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]
+    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) == []
+    assert aln.variable_positions(include_gap_motif=True) == ()
+    assert aln.variable_positions(include_gap_motif=False) == ()
+    new_seqs = {"A": "-CG?C-", "B": "ACAAYA", "C": "GCGACA"}
+    aln = new_alignment.make_aligned_seqs(new_seqs, moltype="dna")
+    assert aln.variable_positions(include_gap_motif=False, include_ambiguity=True) == (
+        0,
+        2,
+        3,
+        4,
+    )
+    assert aln.variable_positions(include_gap_motif=True, include_ambiguity=True) == (
+        0,
+        2,
+        3,
+        4,
+        5,
+    )
+
+
+def test_variable_positions_motif_length():
+    """correctly identify variable positions"""
+    #                 *  - ?
+    new_seqs = {"A": "-CG-CA", "B": "ACGACA", "C": "GCGACY"}
+    aln = new_alignment.make_aligned_seqs(new_seqs, moltype="dna")
+    # only the first dinucleotide is variable if gaps and ambigs disallowed
+    assert aln.variable_positions(
+        motif_length=2, include_gap_motif=False, include_ambiguity=False
+    ) == (0, 1)
+    # if gaps are allowed, the second dinucleotide is also variable
+    assert aln.variable_positions(
+        motif_length=2, include_gap_motif=True, include_ambiguity=False
+    ) == (0, 1, 2, 3)
+    # if no gaps but ambig allowed, the third dinucleotide is also variable
+    assert aln.variable_positions(
+        motif_length=2, include_gap_motif=False, include_ambiguity=True
+    ) == (0, 1, 4, 5)
+    # if gaps and ambig allowed all are variable
+    assert aln.variable_positions(
+        motif_length=2, include_gap_motif=True, include_ambiguity=True
+    ) == (0, 1, 2, 3, 4, 5)
+    # variable remains modulo motif_length
+    new_seqs = {"A": "-CG-CAA", "B": "ACGACAC", "C": "GCGACYG"}
+    aln = new_alignment.make_aligned_seqs(new_seqs, moltype="dna")
+    assert aln.variable_positions(
+        motif_length=2, include_gap_motif=True, include_ambiguity=True
+    ) == (0, 1, 2, 3, 4, 5)
 
 
 def test_alignment_counts_per_pos():