Fix linear search regression on invalid query k-mer lengths

Author: EricR86

newmap/unique_counts.py

@@ -463,7 +463,9 @@ def linear_search(index_filename: Path,
                       (~finished_search).sum()))
         verbose_print(verbose, "Counting {}-mers".format(kmer_length))
 
-        # Skip any kmers that contain an ambiguous base
+        max_kmer_query_lengths = []
+
+        # For every position that does not have an ambiguous base
         for i in np.nonzero(~finished_search)[0]:
             # Create the kmer from the sequence segment
             # NB: At the epilogue of the sequence, out of bounds
@@ -476,6 +478,10 @@ def linear_search(index_filename: Path,
                 # Ignore it for all longer kmer lengths (i.e. all
                 # future iterations)
                 finished_search[i] = True
+            # Otherwise:
+            else:
+                # Record the length of this k-mer
+                max_kmer_query_lengths.append(len(kmer))
 
         kmer_indices = np.nonzero(~finished_search)[0]
 
@@ -495,7 +501,7 @@ def linear_search(index_filename: Path,
         count_list = get_kmer_counts(index_filename,
                                      sequence_segment.data,
                                      kmer_indices.tolist(),
-                                     [kmer_length]*len(kmer_indices),
+                                     max_kmer_query_lengths,
                                      num_threads)
 
         # Assert that the number of indices to count and the number of counts

tests/test_count_kmers.py

@@ -3,7 +3,7 @@
 from tempfile import NamedTemporaryFile
 from util import TEST_DATA_PATH
 
-from newmap._c_newmap_count_kmers import count_kmers
+from newmap._c_newmap_count_kmers import count_kmers, count_kmers_from_sequence
 from newmap.main import (DEFAULT_SUFFIX_ARRAY_COMPRESSION_RATIO,
                          DEFAULT_KMER_LENGTH_IN_SEED_TABLE)
 from newmap.generate_index import generate_fm_index
@@ -18,21 +18,27 @@ def setUp(self):
                           DEFAULT_KMER_LENGTH_IN_SEED_TABLE)
         self.num_threads = 1
 
-    # @unittest.skip("Test relies on large data file not in respository")
     def test_count_kmers(self):
         counts = count_kmers(self.genome_index_filename,
                              [b'AAAA', b'AT', b'TAT', b'CCC', b'NNN', b'TCGT'],
                              self.num_threads)
         self.assertEqual(counts, [9, 3, 1, 8, 0, 0])
 
-    # @unittest.skip("Test relies on large data file not in respository")
     def test_count_wrong_type(self):
         with self.assertRaises(TypeError):
             count_kmers(self.genome_index_filename, ["AAAA"],
                         self.num_threads)
 
-    # @unittest.skip("Test relies on large data file not in respository")
     def test_empty_byte_string(self):
         with self.assertRaises(ValueError):
             count_kmers(self.genome_index_filename, [b'AAAA', b'', b'TAT'],
                         self.num_threads)
+
+    def test_sequence_counting(self):
+        counts = count_kmers_from_sequence(
+                    self.genome_index_filename,
+                    b'AAAAATTTTTATCGAATCGA',
+                    [0, 4, 9],
+                    [4, 2, 3],
+                    self.num_threads)
+        self.assertEqual(counts, [9, 3, 1])

tests/test_sequence_buffer_iter.py

@@ -187,3 +187,7 @@ def test_sizes_and_epilogues(self):
         self.assertTrue(sequence_buffer.epilogue)
 
         self.assertRaises(StopIteration, next, buffer_iter)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_unique_counts.py

@@ -22,16 +22,29 @@
 
 
 class TestCountKmers(unittest.TestCase):
-    def setUp(self):
-        self.genome_index_filename = NamedTemporaryFile(mode="w").name
-        self.fasta_filename = str(TEST_DATA_PATH / 'genome.fa')
-        generate_fm_index(self.fasta_filename,
-                          self.genome_index_filename,
+    genome_index_file = NamedTemporaryFile(mode="w")
+    genome_index_filename = genome_index_file.name
+    fasta_filename = str(TEST_DATA_PATH / 'genome.fa')
+    num_threads = 1
+
+    @classmethod
+    def setUpClass(cls):
+        generate_fm_index(cls.fasta_filename,
+                          cls.genome_index_filename,
                           DEFAULT_SUFFIX_ARRAY_COMPRESSION_RATIO,
                           DEFAULT_KMER_LENGTH_IN_SEED_TABLE)
-        self.num_threads = 1
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.genome_index_file.close()
 
     def test_binary_search(self):
+        self.search(use_binary_search=True)
+
+    def test_linear_search(self):
+        self.search(use_binary_search=False)
+
+    def search(self, use_binary_search):
         write_unique_counts(Path(self.fasta_filename),
                             Path(self.genome_index_filename),
                             15,  # Batch size
@@ -40,7 +53,7 @@ def test_binary_search(self):
                             [],  # Include chr ids
                             [],  # Exclude chr ids
                             self.num_threads,
-                            use_binary_search=True)
+                            use_binary_search)
 
         # Check the results in chr1.unique.uint8 and chr2.unique.uint8
         chr1_results = np.fromfile('chr1.unique.uint8', dtype=np.uint8)
@@ -51,3 +64,7 @@ def test_binary_search(self):
 
         self.assertTrue(np.array_equal(chr1_results, EXPECTED_CHR1_COUNTS))
         self.assertTrue(np.array_equal(chr2_results, EXPECTED_CHR2_COUNTS))
+
+
+if __name__ == '__main__':
+    unittest.main()