Repository: scipy/scipy
Pull request: ENH: stats: Add nan_policy to zmap.
The change moves the zscore code to zmap, with the obvious
change at the end to compute the z-scores for `scores`.
zscore is now just a one-liner that calls zmap.

Closes gh-12403.


Problem statement:
Add nan_policy to stats.zmap
`stats.zscore` has a `nan_policy`. `stats.zmap`, being an almost identical function, should probably have the same option.

As an alternative, it might be better to refactor `zscore` in terms of `zmap`, as I think `zscore(x) == zmap(x, x)`.

Hints from issue discussion:
@mdhaber, I see you added this to the list of stats issues to watch.  After working on `zscore` recently, I started looking at `zmap`.  As @chrism2671 pointed out, we can refactor the code so that all the work currently in `zscore` can be done in `zmap` instead, and `zscore` will just call `zmap`.  I'll have a PR for this pretty quickly.

Commit message: ENH: stats: Add nan_policy to zmap.

The change moves the zscore code to zmap, with the obvious
change at the end to compute the z-scores for `scores`.
zscore is now just a one-liner that calls zmap.

Closes gh-12403.

Patch to review:
diff --git a/scipy/stats/stats.py b/scipy/stats/stats.py
index 79e0930bad11..6f8ea4f4a73f 100644
--- a/scipy/stats/stats.py
+++ b/scipy/stats/stats.py
@@ -2539,39 +2539,10 @@ def zscore(a, axis=0, ddof=0, nan_policy='propagate'):
     array([[-1.13490897, -0.37830299,         nan, -0.08718406,  1.60039602],
            [-0.91611681, -0.89090508,  1.4983032 ,  0.88731639, -0.5785977 ]])
     """
-    a = np.asanyarray(a)
-
-    if a.size == 0:
-        return np.empty(a.shape)
-
-    contains_nan, nan_policy = _contains_nan(a, nan_policy)
-
-    if contains_nan and nan_policy == 'omit':
-        if axis is None:
-            mn = _quiet_nanmean(a.ravel())
-            std = _quiet_nanstd(a.ravel())
-            isconst = _isconst(a.ravel())
-        else:
-            mn = np.apply_along_axis(_quiet_nanmean, axis, a)
-            std = np.apply_along_axis(_quiet_nanstd, axis, a)
-            isconst = np.apply_along_axis(_isconst, axis, a)
-    else:
-        mn = a.mean(axis=axis, keepdims=True)
-        std = a.std(axis=axis, ddof=ddof, keepdims=True)
-        if axis is None:
-            isconst = (a.item(0) == a).all()
-        else:
-            isconst = (_first(a, axis) == a).all(axis=axis, keepdims=True)
-
-    # Set std deviations that are 0 to 1 to avoid division by 0.
-    std[isconst] = 1.0
-    z = (a - mn) / std
-    # Set the outputs associated with a constant input to nan.
-    z[np.broadcast_to(isconst, z.shape)] = np.nan
-    return z
+    return zmap(a, a, axis=axis, ddof=ddof, nan_policy=nan_policy)
 
 
-def zmap(scores, compare, axis=0, ddof=0):
+def zmap(scores, compare, axis=0, ddof=0, nan_policy='propagate'):
     """
     Calculate the relative z-scores.
 
@@ -2593,6 +2564,13 @@ def zmap(scores, compare, axis=0, ddof=0):
     ddof : int, optional
         Degrees of freedom correction in the calculation of the
         standard deviation. Default is 0.
+    nan_policy : {'propagate', 'raise', 'omit'}, optional
+        Defines how to handle the occurrence of nans in `compare`.
+        'propagate' returns nan, 'raise' raises an exception, 'omit'
+        performs the calculations ignoring nan values. Default is
+        'propagate'. Note that when the value is 'omit', nans in `scores`
+        also propagate to the output, but they do not affect the z-scores
+        computed for the non-nan values.
 
     Returns
     -------
@@ -2614,10 +2592,36 @@ def zmap(scores, compare, axis=0, ddof=0):
     array([-1.06066017,  0.        ,  0.35355339,  0.70710678])
 
     """
-    scores, compare = map(np.asanyarray, [scores, compare])
-    mns = compare.mean(axis=axis, keepdims=True)
-    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
-    return (scores - mns) / sstd
+    a = np.asanyarray(compare)
+
+    if a.size == 0:
+        return np.empty(a.shape)
+
+    contains_nan, nan_policy = _contains_nan(a, nan_policy)
+
+    if contains_nan and nan_policy == 'omit':
+        if axis is None:
+            mn = _quiet_nanmean(a.ravel())
+            std = _quiet_nanstd(a.ravel())
+            isconst = _isconst(a.ravel())
+        else:
+            mn = np.apply_along_axis(_quiet_nanmean, axis, a)
+            std = np.apply_along_axis(_quiet_nanstd, axis, a)
+            isconst = np.apply_along_axis(_isconst, axis, a)
+    else:
+        mn = a.mean(axis=axis, keepdims=True)
+        std = a.std(axis=axis, ddof=ddof, keepdims=True)
+        if axis is None:
+            isconst = (a.item(0) == a).all()
+        else:
+            isconst = (_first(a, axis) == a).all(axis=axis, keepdims=True)
+
+    # Set std deviations that are 0 to 1 to avoid division by 0.
+    std[isconst] = 1.0
+    z = (scores - mn) / std
+    # Set the outputs associated with a constant input to nan.
+    z[np.broadcast_to(isconst, z.shape)] = np.nan
+    return z
 
 
 def gstd(a, axis=0, ddof=1):
diff --git a/scipy/stats/tests/test_stats.py b/scipy/stats/tests/test_stats.py
index 5bdd7d37ccd3..057de77fc48a 100644
--- a/scipy/stats/tests/test_stats.py
+++ b/scipy/stats/tests/test_stats.py
@@ -1886,7 +1886,7 @@ def test_obj_arrays_ndim(self):
         assert np.all(m.count == 2) and m.count.shape == (1, 1)
 
 
-class TestVariability(object):
+class TestSEM:
 
     testcase = [1, 2, 3, 4]
     scalar_testcase = 4.
@@ -1915,12 +1915,22 @@ def test_sem(self):
         assert_raises(ValueError, stats.sem, x, nan_policy='raise')
         assert_raises(ValueError, stats.sem, x, nan_policy='foobar')
 
-    def test_zmap(self):
-        # not in R, so tested by using:
-        #     (testcase[i] - mean(testcase, axis=0)) / sqrt(var(testcase) * 3/4)
-        y = stats.zmap(self.testcase,self.testcase)
-        desired = ([-1.3416407864999, -0.44721359549996, 0.44721359549996, 1.3416407864999])
-        assert_array_almost_equal(desired,y,decimal=12)
+
+class TestZmapZscore:
+
+    # Expected values were computed independently "by hand"
+    # using the formula for the z-score.
+    @pytest.mark.parametrize(
+        'x, y, expected',
+        [([1, 2, 3, 4], [1, 2, 3, 4],
+          [-1.3416407864999, -0.44721359549996, 0.44721359549996,
+           1.3416407864999]),
+         ([1, 2, 3], [0, 1, 2, 3, 4],
+          [-0.7071067811865476, 0, 0.7071067811865476]),]
+    )
+    def test_zmap(self, x, y, expected):
+        z = stats.zmap(x, y)
+        assert_allclose(z, expected, rtol=1e-12)
 
     def test_zmap_axis(self):
         # Test use of 'axis' keyword in zmap.
@@ -1957,12 +1967,40 @@ def test_zmap_ddof(self):
         assert_array_almost_equal(z[0], z0_expected)
         assert_array_almost_equal(z[1], z1_expected)
 
+    def test_zmap_nan_policy_omit(self):
+        # nans in `scores` are propagated, regardless of `nan_policy`.
+        # `nan_policy` only affects how nans in `compare` are handled.
+        scores = np.array([-3, -1, 2, np.nan])
+        compare = np.array([-8, -3, 2, 7, 12, np.nan])
+        z = stats.zmap(scores, compare, nan_policy='omit')
+        assert_allclose(z, stats.zmap(scores, compare[~np.isnan(compare)]))
+
+    def test_zmap_nan_policy_omit_with_axis(self):
+        scores = np.arange(-5.0, 9.0).reshape(2, -1)
+        compare = np.linspace(-8, 6, 24).reshape(2, -1)
+        compare[0, 4] = np.nan
+        compare[0, 6] = np.nan
+        compare[1, 1] = np.nan
+        z = stats.zmap(scores, compare, nan_policy='omit', axis=1)
+        expected = np.array([stats.zmap(scores[0],
+                                        compare[0][~np.isnan(compare[0])]),
+                             stats.zmap(scores[1],
+                                        compare[1][~np.isnan(compare[1])])])
+        assert_allclose(z, expected, rtol=1e-14)
+
+    def test_zmap_nan_policy_raise(self):
+        scores = np.array([1, 2, 3])
+        compare = np.array([-8, -3, 2, 7, 12, np.nan])
+        with pytest.raises(ValueError, match='input contains nan'):
+            stats.zmap(scores, compare, nan_policy='raise')
+
     def test_zscore(self):
         # not in R, so tested by using:
         #    (testcase[i] - mean(testcase, axis=0)) / sqrt(var(testcase) * 3/4)
-        y = stats.zscore(self.testcase)
-        desired = ([-1.3416407864999, -0.44721359549996, 0.44721359549996, 1.3416407864999])
-        assert_array_almost_equal(desired,y,decimal=12)
+        y = stats.zscore([1, 2, 3, 4])
+        desired = ([-1.3416407864999, -0.44721359549996, 0.44721359549996,
+                    1.3416407864999])
+        assert_array_almost_equal(desired, y, decimal=12)
 
     def test_zscore_axis(self):
         # Test use of 'axis' keyword in zscore.


Code review comments:
diff_hunk: @@ -2614,10 +2592,36 @@ def zmap(scores, compare, axis=0, ddof=0):
     array([-1.06066017,  0.        ,  0.35355339,  0.70710678])
 
     """
-    scores, compare = map(np.asanyarray, [scores, compare])
-    mns = compare.mean(axis=axis, keepdims=True)
-    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
-    return (scores - mns) / sstd
+    a = np.asanyarray(compare)
+
+    if a.size == 0:
+        return np.empty(a.shape)
+
+    contains_nan, nan_policy = _contains_nan(a, nan_policy)
+
+    if contains_nan and nan_policy == 'omit':
+        if axis is None:
+            mn = _quiet_nanmean(a.ravel())
+            std = _quiet_nanstd(a.ravel())
path: scipy/stats/stats.py
text: @user1:
What about `ddof` here and below?

@author:
Yikes, that's a bug.  I'll fix it.

diff_hunk: @@ -2614,10 +2592,36 @@ def zmap(scores, compare, axis=0, ddof=0):
     array([-1.06066017,  0.        ,  0.35355339,  0.70710678])
 
     """
-    scores, compare = map(np.asanyarray, [scores, compare])
-    mns = compare.mean(axis=axis, keepdims=True)
-    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
-    return (scores - mns) / sstd
+    a = np.asanyarray(compare)
+
+    if a.size == 0:
+        return np.empty(a.shape)
+
+    contains_nan, nan_policy = _contains_nan(a, nan_policy)
+
+    if contains_nan and nan_policy == 'omit':
+        if axis is None:
+            mn = _quiet_nanmean(a.ravel())
+            std = _quiet_nanstd(a.ravel())
+            isconst = _isconst(a.ravel())
+        else:
+            mn = np.apply_along_axis(_quiet_nanmean, axis, a)
+            std = np.apply_along_axis(_quiet_nanstd, axis, a)
path: scipy/stats/stats.py
text: @user1:
Why don't we need to `keepdims` here? Or are they kept by default? I haven't thought about why we would need to keep them anyway, I just noticed that it's different below.

@author:
`_quiet_nanmean` and `_quiet_nanstd` will take care of `keepdims=True`.

We use `keepdims=True` when computing the mean and std. dev. so that broadcasting works correctly when we calculate `z = (scores - mn) / std`.  `mn` and `std` must maintain shapes that are compatible for broadcasting with `scores`.  If, for example, `scores` has shape `(3, 4, 5)` and we use `axis=1`, we want `mn` and `std` to have shape `(3, 1, 5)`.  If we didn't use `keepdims`, they would have shape `(3, 5)` and couldn't broadcast with `scores`.