Repository: scipy/scipy
Pull request: ENH: spatial: add `axis` and `keepdims` optional argument for `distance.jensenshannon`.
<!-- 
Thanks for contributing a pull request! Please ensure that
your PR satisfies the checklist before submitting:
http://scipy.github.io/devdocs/dev/contributor/development_workflow.html#checklist-before-submitting-a-pr

Also, please name and describe your PR as you would write a
commit message:
https://docs.scipy.org/doc/numpy/dev/development_workflow.html#writing-the-commit-message

Note that we are a team of volunteers; we appreciate your
patience during the review process.

Again, thanks for contributing!
-->

#### Reference issue
<!--Example: Closes gh-WXYZ.-->
fix #11325

#### What does this implement/fix?
<!--Please explain your changes.-->
I added `axis` optional argument for `spatial.distance.jensenshannon`.
I added tests for it and update docstrings.


Problem statement:
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.

Hints from issue 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

Commit message: ENH: spatial: fix doctests.

Patch to review:
diff --git a/scipy/spatial/distance.py b/scipy/spatial/distance.py
index 1c89741d965c..97fbed15eadf 100644
--- a/scipy/spatial/distance.py
+++ b/scipy/spatial/distance.py
@@ -1211,7 +1211,7 @@ def canberra(u, v, w=None):
     return d
 
 
-def jensenshannon(p, q, base=None):
+def jensenshannon(p, q, base=None, axis=0):
     """
     Compute the Jensen-Shannon distance (metric) between
     two 1-D probability arrays. This is the square root
@@ -1239,6 +1239,9 @@ 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.
 
     Returns
     -------
@@ -1256,16 +1259,20 @@ def jensenshannon(p, q, base=None):
     0.46450140402245893
     >>> distance.jensenshannon([1.0, 0.0, 0.0], [1.0, 0.0, 0.0])
     0.0
+    >>> distance.jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=0)
+    array([0.0, 0.0])
+    >>> distance.jensenshannon([[1.0, 2.0]], [[0.5, 1.5]], axis=1)
+    array([0.0649045])
 
     """
     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)
+    js = np.sum(left, axis=axis) + np.sum(right, axis=axis)
     if base is not None:
         js /= np.log(base)
     return np.sqrt(js / 2.0)
diff --git a/scipy/spatial/tests/test_distance.py b/scipy/spatial/tests/test_distance.py
index 3d88083ce55f..f50aa447e09c 100644
--- a/scipy/spatial/tests/test_distance.py
+++ b/scipy/spatial/tests/test_distance.py
@@ -2200,3 +2200,28 @@ 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])
+
+    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])


Code review comments:
diff_hunk: @@ -1211,7 +1211,7 @@ def canberra(u, v, w=None):
     return d
 
 
-def jensenshannon(p, q, base=None):
+def jensenshannon(p, q, base=None, axis=0):
     """
     Compute the Jensen-Shannon distance (metric) between
     two 1-D probability arrays. This is the square root
path: scipy/spatial/distance.py
text: @user1:
Need to be updated probably.

@author:
Thanks. I updated.

diff_hunk: @@ -1211,7 +1211,7 @@ def canberra(u, v, w=None):
     return d
 
 
-def jensenshannon(p, q, base=None):
+def jensenshannon(p, q, base=None, axis=0):
path: scipy/spatial/distance.py
text: @user1:
```suggestion
def jensenshannon(p, q, base=None, *, axis=0, keepdim=False):
```

@author:
I'm not sure what the `*` is for?

@user2:
This is to say that in order to use the parameters after it, you must call them explicitly `axis=0`, you cannot pass just `0`. 

@user2:
I am not sure about the need for `keepdim`. @user1 ?

If we add it, we will need to add `.. versionadded:: 1.7.0` here too.

@author:
Thanks @user2 I didn't know that.

> If we add it, we will need to add .. versionadded:: 1.7.0 here too.

I added it.

diff_hunk: @@ -1239,6 +1239,9 @@ 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.
 
     Returns
     -------
path: scipy/spatial/distance.py
text: @user1:
The return type is not just `double` anymore.