Merge pull request #465 from jamesmartini/develop_entropy_API

Change entropy_per_pos API, issue #459
cogent3 · Jan 9, 2020 · 79f9da9 · 79f9da9
2 parents 4ddbfd7 + 7963ce0
commit 79f9da9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/src/cogent3/core/alignment.py b/src/cogent3/core/alignment.py
@@ -2337,9 +2337,12 @@ def probs_per_pos(
         )
         return counts.to_freq_array()
 
-    def entropy_per_pos(self, motif_length=1):
+    def entropy_per_pos(self, motif_length=1, include_ambiguity=False, allow_gap=False, alert=False):
         """returns shannon entropy per position"""
-        probs = self.probs_per_pos(motif_length=motif_length)
+        probs = self.probs_per_pos(motif_length=motif_length,
+                                   include_ambiguity=include_ambiguity,
+                                   allow_gap=allow_gap,
+                                   alert=alert)
         return probs.entropy()
 
     def probs_per_seq(

diff --git a/tests/test_core/test_alignment.py b/tests/test_core/test_alignment.py
@@ -2978,6 +2978,23 @@ def test_entropy_per_pos(self):
         f = a.entropy_per_pos()
         e = array([0, 0, 1, 1])
         self.assertEqual(f, e)
+        f = a.entropy_per_pos(motif_length=2)
+        e = array([0, 1])
+        self.assertEqual(f, e)
+        seqs = []
+        for s in ["-GAT", "ACCT", "GAGT"]:
+            seqs.append(make_seq(s, moltype="dna"))
+        a = ArrayAlignment(seqs)
+        f = a.entropy_per_pos(allow_gap=True)
+        e = array([1.584962500721156, 1.584962500721156, 1.584962500721156, 0])
+        self.assertEqual(f, e)
+        seqs = []
+        for s in ["-RAT", "ACCT", "GTGT"]:
+            seqs.append(make_seq(s, moltype="dna"))
+        a = ArrayAlignment(seqs)
+        f = a.entropy_per_pos(include_ambiguity=True)
+        e = array([1.584962500721156, 1.584962500721156, 1.584962500721156, 0])
+        self.assertEqual(f, e)
 
     def test_coevolution_segments(self):
         """specifying coordinate segments produces matrix with just those"""