Repository: cogent3/cogent3
Pull request #2299 (2025-05-03 07:06:26)

Mixed patches
## Summary by Sourcery

Enhance sampling functionality for sequences and alignments with improved random sampling methods

New Features:
- Added new `sample()` method for Sequence class with support for motif-length sampling
- Enhanced sampling method for Alignment class with more robust sampling options

Enhancements:
- Improved type hints for sampling methods
- Added error handling for sampling scenarios
- Simplified sampling implementation in Alignment class

Tests:
- Added test cases for sequence and alignment sampling with different motif lengths
- Added test for sampling without replacement

Chores:
- Updated ruff configuration
- Updated noxfile for potential performance improvements


Patch:
diff --git a/changelog.d/20250406_125550_Gavin.Huttley.md b/changelog.d/20250406_125550_Gavin.Huttley.md
index 6cbec8c7a..2b81b3d44 100644
--- a/changelog.d/20250406_125550_Gavin.Huttley.md
+++ b/changelog.d/20250406_125550_Gavin.Huttley.md
@@ -35,6 +35,7 @@ Uncomment the section that is right (remove the HTML comment wrapper).
   Alignments, the gapped sequence is used. Added companion method
   `drop_duplicated_seqs()` which does as the name implies. The first name in each
   group of duplicates is retained.
+- Added new type Sequence.sample() method for generating random sequences
 
 <!--
 ### BUG
diff --git a/noxfile.py b/noxfile.py
index bc48d6fbb..ec3a1875f 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -1,7 +1,11 @@
+import os
 import pathlib
 
 import nox
 
+# on python >= 3.12 this will improve speed of test coverage a lot
+os.environ["COVERAGE_CORE"] = "sysmon"
+
 _py_versions = range(10, 14)
 
 
diff --git a/ruff.toml b/ruff.toml
index 0151300ab..da1a244cb 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -44,6 +44,8 @@ select = ["ALL"]
 # N801 CamelCase is not always appropriate
 # PLR0913 number of arguments sometimes needs to be > 5
 # FBT001 and FBT002, positional boolean arguments are allowed!
+# PT011, I disagree about checking error messages as well
+# as exception types as it doubles the "cost" of fixing typos
 ignore = [
     "COM812",
     "EXE002", 
@@ -55,6 +57,7 @@ ignore = [
     "PLR0913",
     "FBT001",
     "FBT002",
+    "PT011",
     ]
 
 # Allow fix for all enabled rules (when `--fix`) is provided.
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 2f8a39b09..d4dd8516f 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -4781,7 +4781,7 @@ class Alignment(SequenceCollection):
 
     def take_positions(
         self,
-        cols: list,
+        cols: list[int] | numpy.ndarray[int],
         negate: bool = False,
     ) -> typing_extensions.Self:
         """Returns new Alignment containing only specified positions.
@@ -4799,8 +4799,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).take(cols)
-            for aligned in self.seqs
+            aligned.data.seqid: numpy.array(aligned).take(cols) for aligned in self.seqs
         }
         seqs_data = self._seqs_data.from_seqs(
             data=new_data,
@@ -5577,9 +5576,13 @@ class Alignment(SequenceCollection):
         n: int | None = None,
         with_replacement: bool = False,
         motif_length: int = 1,
-        randint=numpy.random.randint,
-        permutation=numpy.random.permutation,
-    ):
+        randint: typing.Callable[
+            [int, int | None, int | None], numpy.ndarray
+        ] = numpy.random.randint,
+        permutation: typing.Callable[
+            [numpy.ndarray], numpy.ndarray
+        ] = numpy.random.permutation,
+    ) -> typing_extensions.Self:
         """Returns random sample of positions from self, e.g. to bootstrap.
 
         Parameters
@@ -5613,35 +5616,28 @@ class Alignment(SequenceCollection):
         # numpy array reshaped.
 
         population_size = len(self) // motif_length
-        if not n:
-            n = population_size
+        if not with_replacement and n and n > population_size:
+            msg = f"cannot sample without replacement when {n=} > {population_size=}"
+            raise ValueError(msg)
+
+        n = n or population_size
+
         if with_replacement:
             locations = randint(0, population_size, n)
         else:
-            assert n <= population_size, (n, population_size, motif_length)
             locations = permutation(population_size)[:n]
 
-        positions = numpy.empty(n * motif_length, dtype=int)
-        for i, loc in enumerate(locations):
-            positions[i * motif_length : (i + 1) * motif_length] = range(
-                loc * motif_length,
-                (loc + 1) * motif_length,
-            )
-
-        new_seqs = {}
-        for aligned in self.seqs:
-            sampled = numpy.array(aligned)[positions]
-            new_seqs[aligned.data.seqid] = sampled
+        if motif_length == 1:
+            positions = locations
+        else:
+            positions = numpy.empty(n * motif_length, dtype=int)
+            for i, loc in enumerate(locations):
+                positions[i * motif_length : (i + 1) * motif_length] = range(
+                    loc * motif_length,
+                    (loc + 1) * motif_length,
+                )
 
-        new_seqs_data = self._seqs_data.from_seqs(
-            data=new_seqs,
-            alphabet=self.moltype.most_degen_alphabet(),
-        )
-        kwargs = self._get_init_kwargs()
-        kwargs["seqs_data"] = new_seqs_data
-        kwargs.pop("annotation_db", None)
-        kwargs.pop("slice_record", None)
-        return self.__class__(**kwargs)
+        return self.take_positions(positions)
 
     def distance_matrix(
         self,
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index a58108fd8..1d4a7ad90 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -1845,9 +1845,71 @@ class Sequence:
             strand,
         )
 
+    def sample(
+        self,
+        *,
+        n: int | None = None,
+        with_replacement: bool = False,
+        motif_length: int = 1,
+        randint: typing.Callable[
+            [int, int | None, int | None], numpy.ndarray
+        ] = numpy.random.randint,
+        permutation: typing.Callable[
+            [numpy.ndarray], numpy.ndarray
+        ] = numpy.random.permutation,
+    ) -> typing_extensions.Self:
+        """Returns random sample of positions from self, e.g. to bootstrap.
+
+        Parameters
+        ----------
+        n
+            number of positions to sample. If None, all positions are sampled.
+        with_replacement
+            if True, samples with replacement.
+        motif_length
+            number of positions to sample as a single motif. Starting point
+            of each sampled motif is modulo motif_length in the original sequence.
+        randint
+            random number generator, default is numpy.randint
+        permutation
+            function to generate a random permutation of positions, default is
+            numpy.permutation
+
+        Notes
+        -----
+        By default (resampling all positions without replacement), generates
+        a permutation of the positions of the alignment.
+        """
+        population_size = len(self) // motif_length
+        if not with_replacement and n and n > population_size:
+            msg = f"cannot sample without replacement when {n=} > {population_size=}"
+            raise ValueError(msg)
+
+        n = n or population_size
+
+        if with_replacement:
+            locations = randint(0, population_size, n)
+        else:
+            locations = permutation(population_size)[:n]
+
+        if motif_length == 1:
+            positions = locations
+        else:
+            positions = numpy.empty(n * motif_length, dtype=int)
+            for i, loc in enumerate(locations):
+                positions[i * motif_length : (i + 1) * motif_length] = range(
+                    loc * motif_length,
+                    (loc + 1) * motif_length,
+                )
+
+        sampled = numpy.array(self).take(positions)
+        return self.moltype.make_seq(
+            seq=sampled, name=f"{self.name}-randomised", check_seq=False
+        )
+
 
 @register_deserialiser(get_object_provenance(Sequence))
-def deserialise_sequence(data) -> Sequence:
+def deserialise_sequence(data: dict) -> Sequence:
     return Sequence.from_rich_dict(data)
 
 


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index fda5050a9..32569b37c 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -4559,6 +4559,16 @@ def test_alignment_sample_with_replacement():
     assert len(result) == 1
 
 
+def test_alignment_sample_without_replacement():
+    # test with replacement -- just verify that it rnus
+    alignment = new_alignment.make_aligned_seqs(
+        {"seq1": "GATC", "seq2": "GATC"},
+        moltype="dna",
+    )
+    with pytest.raises(ValueError):
+        alignment.sample(n=100, with_replacement=False)
+
+
 def test_alignment_sample_tuples():
     ##### test with motif size != 1 #####
     alignment = new_alignment.make_aligned_seqs(
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 3fcd67cdf..41842f343 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -3055,3 +3055,31 @@ def test_coerce_moltype(moltype, seq):
     assert seq.moltype.name == moltype
     expect = "ATGC" if moltype == "dna" else "AUGC"
     assert str(seq) == expect
+
+
+def test_sample_motif_length_3():
+    seq = cogent3.get_dataset("brca1")["Human"].seq
+    got = seq.sample(motif_length=3)
+    assert len(got) == len(seq)
+    assert got != seq
+    # following will fail if motif length is 1 due
+    # to random creation of stops
+    aa = got.get_translation()
+    assert len(aa) == len(seq) // 3
+
+
+def test_sample_motif_length_1():
+    seq = cogent3.get_dataset("brca1")["Human"].seq
+    got = seq.sample(motif_length=1)
+    assert len(got) == len(seq)
+    assert got != seq
+    # following will fail if motif length is 1 due
+    # to random creation of stop codons
+    with pytest.raises(new_alphabet.AlphabetError):
+        got.get_translation()
+
+
+def test_sample_without_replacement():
+    seq = cogent3.get_dataset("brca1")["Human"].seq
+    with pytest.raises(ValueError):
+        seq[:20].sample(n=100, with_replacement=False)