Merge pull request #475 from jamesmartini/develop_counts474

Fixed Alignment.counts_per_pos to handle gaps, and added appropriate test, fixes #474
cogent3 · Jan 8, 2020 · bcf2567 · bcf2567
2 parents 855160b + 6dfe128
commit bcf2567
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
@@ -2936,11 +2936,11 @@ def counts_per_pos(
         for i in range(0, len(self) - motif_length + 1, motif_length):
             counts = CategoryCounter([s[i : i + motif_length] for s in data])
             if all_motifs is not None:
-                allow_gap.update(list(counts))
+                all_motifs.update(list(counts))
             result.append(counts)
 
         if all_motifs:
-            alpha.extend(list(sorted(set(alpha) ^ all_motifs)))
+            alpha += tuple(sorted(set(alpha) ^ all_motifs))
 
         for i, counts in enumerate(result):
             result[i] = counts.tolist(alpha)

diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
@@ -2248,9 +2248,21 @@ def test_counts_per_pos(self):
             ]
         )
 
+        exp_gap = array(
+            [
+                [1, 1, 0, 1, 0],
+                [0, 2, 0, 0, 1],
+                [0, 0, 3, 0, 0],
+                [0, 2, 0, 1, 0],
+                [0, 1, 2, 0, 0],
+                [0, 2, 0, 1, 0]
+            ]
+        )
+
         s1 = DNA.make_seq("TCAGAG", name="s1")
         s2 = DNA.make_seq("CCACAC", name="s2")
         s3 = DNA.make_seq("AGATAT", name="s3")
+        s4 = DNA.make_seq("G-ACCC", name="s4")
         aln = self.Class([s1, s2, s3], moltype=DNA)
         obs = aln.counts_per_pos()
         self.assertEqual(obs.array, exp)
@@ -2259,6 +2271,9 @@ def test_counts_per_pos(self):
         self.assertEqual(obs[0, "TC"], 1)
         self.assertEqual(obs[1, "AC"], 1)
         self.assertEqual(obs[2, "AC"], 1)
+        aln = self.Class([s1, s2, s4], moltype=DNA)
+        obs = aln.counts_per_pos(allow_gap=True)
+        self.assertEqual(obs.array, exp_gap)
 
     def test_get_seq_entropy(self):
         """ArrayAlignment get_seq_entropy should get entropy of each seq"""