Repository: cogent3/cogent3
Pull request #1007 (2022-01-20 03:25:32)

ENH: generalise Jensen-Shannon calculations to > 2 distributions
[NEW] algorithm within function made more general, no change to API


Patch:
diff --git a/src/cogent3/maths/measure.py b/src/cogent3/maths/measure.py
index 1ac26e379..920c2826b 100644
--- a/src/cogent3/maths/measure.py
+++ b/src/cogent3/maths/measure.py
@@ -87,43 +87,39 @@ def paralinear_continuous_time(P, pi, Q, validate=False):
     return pl
 
 
-def jsd(freqs1, freqs2, validate=False):
-    """calculate Jensen–Shannon divergence between two probability distributions
+def jsd(*vectors, validate=False):
+    """calculate Jensen–Shannon divergence between two or more probability distributions
 
     Parameters
     ----------
-    freqs1 : one dimensional array
-        row vector frequencies, sum to 1
-    freqs2 : one dimensional array
-        row vector frequencies, sum to 1
+    *vectors
+        >= 2 frequency vectors, each vector must sum to 1
     validate : bool
-
+        check the consistency of the provided vectors, namely they're
+        all 1D, have the same number of elements and sum to 1
     """
     # Convert input arrays into numpy arrays
-    freqs1 = array(freqs1)
-    freqs2 = array(freqs2)
-
+    vectors = [array(v) for v in vectors]
+    num_vectors = len(vectors)
     if validate:
-        assert_equal(
-            freqs1.shape, freqs2.shape, err_msg="freqs1/freqs2 mismatched shape"
-        )
-        assert freqs1.ndim == 1, "freqs1 has incorrect dimension"
-        assert freqs2.ndim == 1, "freqs2 has incorrect dimension"
+        assert len({f.shape for f in vectors}) == 1, "mismatched shape"
+        dims = {v.ndim for v in vectors}
+        assert dims == {1}, "incorrect dimension"
         try:
-            validate_freqs_array(freqs1)
-            validate_freqs_array(freqs2)
+            for v in vectors:
+                validate_freqs_array(v)
         except ValueError as err:
             raise AssertionError("freqs not valid") from err
 
-    H_mn = fsum(safe_p_log_p(freqs1 / 2 + freqs2 / 2))
-    mn_H = fsum([fsum(i) for i in map(safe_p_log_p, [freqs1, freqs2])]) / 2
+    H_mn = fsum(safe_p_log_p(array(vectors).mean(axis=0)))
+    mn_H = fsum([fsum(i) for i in map(safe_p_log_p, vectors)]) / num_vectors
     jsd_ = H_mn - mn_H
-    if jsd_ < 0 and isclose(jsd_, 0, atol=1e-10):
+    if jsd_ < 0:
+        if not isclose(jsd_, 0, atol=1e-10):
+            raise ArithmeticError(
+                f"{jsd_} is negative and below defined precision threshold"
+            )
         jsd_ = 0
-    elif jsd_ < 0:
-        raise ArithmeticError(
-            f"{jsd_} is negative and below defined precision threshold"
-        )
 
     return jsd_
 


Test patch:
diff --git a/tests/test_maths/test_measure.py b/tests/test_maths/test_measure.py
index 16933e6ea..c5d60355e 100644
--- a/tests/test_maths/test_measure.py
+++ b/tests/test_maths/test_measure.py
@@ -1,6 +1,6 @@
 from unittest import TestCase, main
 
-from numpy import diag_indices, dot, finfo, float64
+from numpy import arange, array, diag_indices, dot, finfo, float64
 from numpy.random import random
 from numpy.testing import assert_allclose
 
@@ -238,6 +238,18 @@ class TestJensenShannon(TestCase):
         result = jsd(pi_0, pi_1)
         self.assertTrue(result >= 0)
 
+    def test_general_jsd(self):
+        """check correctness of JSD for > 2 distributions"""
+        freqs = (0.1, 0.2, 0.3, 0.4), (0.4, 0.3, 0.2, 0.1), (0.1, 0.4, 0.2, 0.3)
+        got = jsd(*freqs, validate=True)
+        # expected value from the R-package philentropy gJSD implementation
+        assert_allclose(got, 0.1374318, atol=1e-7)
+
+        # with invalid freqs
+        freqs = (0.1, 0.2, 0.3, 0.4), (0.4, 0.3, 0.1, 0.2), (0.1, 0.4, 0.4, 0.3)
+        with self.assertRaises(AssertionError):
+            jsd(*freqs, validate=True)
+
     def test_jsm(self):
         """evaluate jsm between identical, and non-identical distributions"""
         case1 = [