Repository: iqbal-lab-org/make_prg

Issue:
Repeats can cause underclustering
This is an interesting issue (but broke me for one day).

Consider the following dummy alignment:
```
>1
TTTTTTTGGGGGGGAAAAAAATTTTTTT-------AAAAAAATTTTTTTAAAAAAA-------
>2
-------TTTTTTTAAAAAAATTTTTTTGGGGGGGAAAAAAATTTTTTT-------AAAAAAA
>3
TTTTTTTAAAAAAATTTTTTTAAAAAAATTTTTTT-------GGGGGGG-------AAAAAAA
```

If you take the Hamming distance between any pair, it's high, but if you count all the 7-mers for each, the kmer counts are identical.
(They are: ```{'TTTTTTT': 3.0, 'TTTTTTG': 1.0, 'TTTTTGG': 1.0, 'TTTTGGG': 1.0, 'TTTGGGG': 1.0, 'TTGGGGG': 1.0, 'TGGGGGG': 1.0, 'GGGGGGG': 1.0, 'GGGGGGA': 1.0, 'GGGGGAA': 1.0, 'GGGGAAA': 1.0, 'GGGAAAA': 1.0, 'GGAAAAA': 1.0, 'GAAAAAA': 1.0, 'AAAAAAA': 3.0, 'AAAAAAT': 2.0, 'AAAAATT': 2.0, 'AAAATTT': 2.0, 'AAATTTT': 2.0, 'AATTTTT': 2.0, 'ATTTTTT': 2.0, 'TTTTTTA': 2.0, 'TTTTTAA': 2.0, 'TTTTAAA': 2.0, 'TTTAAAA': 2.0, 'TTAAAAA': 2.0, 'TAAAAAA': 2.0}```)

The alignment wouldn't actually look like this, it's the best simple example I could come up with. In practice, I saw this happen in MSAs of 5,800 sequences of 23kb each around EBA175 gene in P falciparum. Ie, looking at some repeat-rich (Plasmo ATs) subalignments the algorithm says they are not 'one-reference-like', so asks kmeans to cluster with increasing number of clusters; but kmeans when looking at the kmers, finds a smaller number of distinct data points due to same kmer counts, and cannot even produce that many clusters.

This leads to empty clusters, due to another code bug which let empty clusters be produced in the above case, and then prg construction fails.

Comments:

Patch:
diff --git a/make_prg/from_msa/cluster_sequences.py b/make_prg/from_msa/cluster_sequences.py
index bc3cc81..aa67ce6 100644
--- a/make_prg/from_msa/cluster_sequences.py
+++ b/make_prg/from_msa/cluster_sequences.py
@@ -121,6 +121,10 @@ def extract_clusters(
             f"Mismatch between number of sequences/ID lists and number of cluster assignments"
         )
     num_clusters = max(cluster_assignment) + 1
+    if set(cluster_assignment) != set(range(num_clusters)):
+        raise ValueError(
+            "Inconsistent cluster numbering (likely reason: more input sequences that clustered data points)"
+        )
     result = [list() for _ in range(num_clusters)]
     for cluster_num, clustered_elem in zip(cluster_assignment, value_pool):
         result[cluster_num].extend(clustered_elem)
@@ -141,7 +145,9 @@ def merge_clusters(clusters: List[ClusteredIDs], first_id: str) -> ClusteredIDs:
 
 
 def kmeans_cluster_seqs_in_interval(
-    interval: List[int], alignment: MSA, kmer_size: int,
+    interval: List[int],
+    alignment: MSA,
+    kmer_size: int,
 ) -> ClusteredIDs:
     """Divide sequences in interval into subgroups of similar sequences."""
     interval_alignment = alignment[:, interval[0] : interval[1] + 1]
@@ -181,7 +187,15 @@ def kmeans_cluster_seqs_in_interval(
                 break
             start = time.time()
             kmeans = KMeans(n_clusters=num_clusters, random_state=2).fit(count_matrix)
+            prev_cluster_assignment = cluster_assignment
             cluster_assignment = list(kmeans.predict(count_matrix))
+            num_fitted_clusters = len(set(cluster_assignment))
+            # Below holds when alignments are different, but kmer counts are identical
+            # (due to repeats), making kmeans unable to fit requested number of clusters
+            if num_fitted_clusters < num_clusters:
+                cluster_assignment = prev_cluster_assignment
+                num_clusters -= 1
+                break
             seqclustering = extract_clusters(seq_to_gapped_seqs, cluster_assignment)
 
     if num_clusters == 1 or num_clusters == num_sequences:
diff --git a/make_prg/from_msa/prg_builder.py b/make_prg/from_msa/prg_builder.py
index 8185a46..f79490e 100644
--- a/make_prg/from_msa/prg_builder.py
+++ b/make_prg/from_msa/prg_builder.py
@@ -74,7 +74,7 @@ class PrgBuilder(object):
 
     @classmethod
     def get_consensus(cls, alignment: MSA):
-        """ Produces a 'consensus string' from an MSA: at each position of the
+        """Produces a 'consensus string' from an MSA: at each position of the
         MSA, the string has a base if all aligned sequences agree, and a "*" if not.
         IUPAC ambiguous bases result in non-consensus and are later expanded in the prg.
         N results in consensus at that position unless they are all N."""
@@ -157,10 +157,6 @@ class PrgBuilder(object):
                         recur = False
                     elif interval.start not in self.subAlignedSeqs:
                         self.subAlignedSeqs[interval.start] = []
-                        logging.debug(
-                            "subAlignedSeqs now has keys: %s",
-                            list(self.subAlignedSeqs.keys()),
-                        )
                     else:
                         logging.debug(
                             "subAlignedSeqs already had key %d in keys: %s. This shouldn't happen.",
@@ -210,35 +206,14 @@ class PrgBuilder(object):
     def max_nesting_level_reached(self):
         max_nesting = []
         if self.subAlignedSeqs == {}:
-            logging.debug(
-                "self.subAlignedSeqs == {} at nesting level %d for interval %s",
-                self.nesting_level,
-                self.interval,
-            )
             max_nesting.append(self.nesting_level)
         else:
-            logging.debug(
-                "self.subAlignedSeqs.keys(): %s", list(self.subAlignedSeqs.keys())
-            )
-            logging.debug(
-                "self.subAlignedSeqs[self.subAlignedSeqs.keys()[0]]: %s",
-                self.subAlignedSeqs[list(self.subAlignedSeqs.keys())[0]],
-            )
             for interval_start in list(self.subAlignedSeqs.keys()):
                 logging.debug("interval start: %d", interval_start)
                 for subaseq in self.subAlignedSeqs[interval_start]:
-                    logging.debug(
-                        "type of subAlignedSeqs object in list: %s", type(subaseq)
-                    )
                     recur = subaseq.max_nesting_level_reached
-                    logging.debug(
-                        "recur max level nesting returned: %d, which has type %s",
-                        recur,
-                        type(recur),
-                    )
                     max_nesting.append(recur)
         m = max(max_nesting)
-        logging.debug("found the max of %s is %d", max_nesting, m)
         return m
 
     @property


Test patch:
diff --git a/tests/from_msa/test_cluster_sequences.py b/tests/from_msa/test_cluster_sequences.py
index 4809a54..9d29ade 100644
--- a/tests/from_msa/test_cluster_sequences.py
+++ b/tests/from_msa/test_cluster_sequences.py
@@ -7,7 +7,7 @@ from numpy import array, array_equal
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 
-from make_prg.seq_utils import standard_bases
+from make_prg.seq_utils import standard_bases, ungap
 from make_prg.from_msa.cluster_sequences import (
     count_distinct_kmers,
     count_kmer_occurrences,
@@ -223,6 +223,20 @@ class TestExtractClusters(TestCase):
         expected = [["s1", "s2", "s3"], ["s4"]]
         self.assertEqual(actual, expected)
 
+    def test_GivenInconsistentClusterNumbers_RaisesError(self):
+        """
+        `gapped_clusters` can occur when number of distinct data points < requested number of
+        clusters in kmeans, leading to production of empty clusters and then
+        messing up prg construction.
+        `large_clusters` should never occur, cluster assigment should be consecutive integers from 0.
+        """
+        gapped_clusters = [0, 2, 2]
+        with self.assertRaises(ValueError):
+            extract_clusters(self.seqdict_ids, gapped_clusters)
+        large_clusters = [5, 6, 6]
+        with self.assertRaises(ValueError):
+            extract_clusters(self.seqdict_ids, large_clusters)
+
 
 class TestClustering_Trivial(TestCase):
     def test_one_seq_returns_single_id(self):
@@ -256,7 +270,10 @@ class TestClustering_SmallSequences(TestCase):
     )
     def test_ambiguous_sequences_in_short_interval_separate_clusters(self):
         alignment = MSA(
-            [SeqRecord(Seq("ARAT"), id="s1"), SeqRecord(Seq("WAAT"), id="s2"),]
+            [
+                SeqRecord(Seq("ARAT"), id="s1"),
+                SeqRecord(Seq("WAAT"), id="s2"),
+            ]
         )
         result = kmeans_cluster_seqs_in_interval([0, 3], alignment, 5)
         self.assertEqual([["s1"], ["s2"]], result)
@@ -390,6 +407,32 @@ class TestClustering_RunKmeans(TestCase):
         result = kmeans_cluster_seqs_in_interval([0, 4], alignment, 4)
         self.assertEqual(len(result), MAX_CLUSTERS)
 
+    def test_GivenSequencesWithSameKmerCounts_ClusteringInterrupted(self):
+        """
+        Sequences below are not 'one-ref-like', yet kmer counts are identical.
+        This is because the sequences contain repeats and gaps, making them
+        not identical from the point of view of edit distance.
+        Number of clusters will try to be increased, but kmeans will only find one,
+        as there is a single data point in kmer space.
+        This test checks the code deals with this by aborting further clustering.
+        """
+        sequences = [
+            "TTTTTTTGGGGGGGAAAAAAATTTTTTT-------AAAAAAATTTTTTTAAAAAAA-------",
+            "-------TTTTTTTAAAAAAATTTTTTTGGGGGGGAAAAAAATTTTTTT-------AAAAAAA",
+            "TTTTTTTAAAAAAATTTTTTTAAAAAAATTTTTTT-------GGGGGGG-------AAAAAAA",
+        ]
+        ungapped_sequences = list(map(ungap, sequences))
+        distinct_kmers = count_distinct_kmers(ungapped_sequences, kmer_size=7)
+        count_matrix = count_kmer_occurrences(ungapped_sequences, distinct_kmers)
+        distinct_count_patterns = set(map(str, count_matrix))
+        assert len(distinct_count_patterns) == 1
+        assert not sequences_are_one_reference_like(sequences)
+
+        expected_clustering = [["s0"], ["s1"], ["s2"]]
+        alignment = make_alignment(sequences)
+        result = kmeans_cluster_seqs_in_interval([0, len(sequences[0])], alignment, 7)
+        self.assertEqual(result, expected_clustering)
+
 
 class TestMergeClusters(TestCase):
     clusters = [["s1", "s2"], ["s3"]]