Repository: scipy/scipy

Issue:
Add axis parameter for scipy.spatial.distance.jensenshannon
My suggestion is to add an axis parameter for `scipy.spatial.distance.jensenshannon`, to be consistent with `scipy.stats.entropy`, instead of assuming an axis of `0` as it currently does.

I'm happy to open a PR when I get time, if this is something that would be useful.


Discussion:
`jensenshannon` docs say that only 1-D arrays are supported, which is the case for all metrics in `distance` - the only difference seems to be that the former doesn't raise.

adding an `axis` argument I'm not sure about. `stats.entropy` also doesn't have an `axis` keyword. 

do you have an actual problem? beyond input validation there doesn't seem much to do here
Does `stats.entropy` not have an `axis` parameter?

https://github.com/scipy/scipy/blob/977ee7ef558da28144722c4b07324395f3f226bb/scipy/stats/_distn_infrastructure.py#L2612-L2676

Of course I can transpose my arrays before and after running them through `jensenshannon`; I just wanted to be able to calculate multiple Jensen-Shannon divergences in the same way I could calculate multiple Kullback-Leibler divergences through `stats.entropy`.
Ah yes it does have an `axis` keyword - I checked in my installed SciPy (1.3.2), which didn't have it yet. So yes, please ignore my objection - PR welcome

Resolving patch:
diff --git a/scipy/spatial/distance.py b/scipy/spatial/distance.py
--- a/scipy/spatial/distance.py
+++ b/scipy/spatial/distance.py
@@ -1211,10 +1211,10 @@ def canberra(u, v, w=None):
     return d
 
 
-def jensenshannon(p, q, base=None):
+def jensenshannon(p, q, base=None, *, axis=0, keepdims=False):
     """
     Compute the Jensen-Shannon distance (metric) between
-    two 1-D probability arrays. This is the square root
+    two probability arrays. This is the square root
     of the Jensen-Shannon divergence.
 
     The Jensen-Shannon distance between two probability
@@ -1239,11 +1239,26 @@ def jensenshannon(p, q, base=None):
         the base of the logarithm used to compute the output
         if not given, then the routine uses the default base of
         scipy.stats.entropy.
+    axis : int, optional
+        Axis along which the Jensen-Shannon distances are computed. The default
+        is 0.
+
+        .. versionadded:: 1.7.0
+    keepdims : bool, optional
+        If this is set to `True`, the reduced axes are left in the
+        result as dimensions with size one. With this option,
+        the result will broadcast correctly against the input array.
+        Default is False.
+
+        .. versionadded:: 1.7.0
 
     Returns
     -------
-    js : double
-        The Jensen-Shannon distance between `p` and `q`
+    js : double or ndarray
+        The Jensen-Shannon distances between `p` and `q` along the `axis`.
+
+    Notes
+    -----
 
     .. versionadded:: 1.2.0
 
@@ -1256,16 +1271,28 @@ def jensenshannon(p, q, base=None):
     0.46450140402245893
     >>> distance.jensenshannon([1.0, 0.0, 0.0], [1.0, 0.0, 0.0])
     0.0
+    >>> a = np.array([[1, 2, 3, 4],
+    ...               [5, 6, 7, 8],
+    ...               [9, 10, 11, 12]])
+    >>> b = np.array([[13, 14, 15, 16],
+    ...               [17, 18, 19, 20],
+    ...               [21, 22, 23, 24]])
+    >>> distance.jensenshannon(a, b, axis=0)
+    array([0.1954288, 0.1447697, 0.1138377, 0.0927636])
+    >>> distance.jensenshannon(a, b, axis=1)
+    array([0.1402339, 0.0399106, 0.0201815])
 
     """
     p = np.asarray(p)
     q = np.asarray(q)
-    p = p / np.sum(p, axis=0)
-    q = q / np.sum(q, axis=0)
+    p = p / np.sum(p, axis=axis, keepdims=True)
+    q = q / np.sum(q, axis=axis, keepdims=True)
     m = (p + q) / 2.0
     left = rel_entr(p, m)
     right = rel_entr(q, m)
-    js = np.sum(left, axis=0) + np.sum(right, axis=0)
+    left_sum = np.sum(left, axis=axis, keepdims=keepdims)
+    right_sum = np.sum(right, axis=axis, keepdims=keepdims)
+    js = left_sum + right_sum
     if base is not None:
         js /= np.log(base)
     return np.sqrt(js / 2.0)


Test patch:
diff --git a/scipy/spatial/tests/test_distance.py b/scipy/spatial/tests/test_distance.py
--- a/scipy/spatial/tests/test_distance.py
+++ b/scipy/spatial/tests/test_distance.py
@@ -2200,3 +2200,32 @@ def test_yule_all_same():
 
     d = cdist(x[:1], x[:1], 'yule')
     assert_equal(d, [[0.0]])
+
+
+def test_jensenshannon():
+    assert_almost_equal(jensenshannon([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], 2.0),
+                        1.0)
+    assert_almost_equal(jensenshannon([1.0, 0.0], [0.5, 0.5]),
+                        0.46450140402245893)
+    assert_almost_equal(jensenshannon([1.0, 0.0, 0.0], [1.0, 0.0, 0.0]), 0.0)
+
+    assert_almost_equal(jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=0),
+                        [0.0, 0.0])
+    assert_almost_equal(jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=1),
+                        [0.0649045])
+    assert_almost_equal(jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=0,
+                                      keepdims=True), [[0.0, 0.0]])
+    assert_almost_equal(jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=1,
+                                      keepdims=True), [[0.0649045]])
+
+    a = np.array([[1, 2, 3, 4],
+                  [5, 6, 7, 8],
+                  [9, 10, 11, 12]])
+    b = np.array([[13, 14, 15, 16],
+                  [17, 18, 19, 20],
+                  [21, 22, 23, 24]])
+
+    assert_almost_equal(jensenshannon(a, b, axis=0),
+                        [0.1954288, 0.1447697, 0.1138377, 0.0927636])
+    assert_almost_equal(jensenshannon(a, b, axis=1),
+                        [0.1402339, 0.0399106, 0.0201815])