MAINT: stats.wilcoxon: fix failure with multidimensional x with NaN and slice length > 50

scipy/stats/_wilcoxon.py

@@ -102,6 +102,7 @@ def _wilcoxon_iv(x, y, zero_method, correction, alternative, method, axis):
                    "an instance of `stats.PermutationMethod`.")
         if method not in methods:
             raise ValueError(message)
+    output_z = True if method == 'approx' else False
 
     # logic unchanged here for backward compatibility
     n_zero = np.sum(d == 0, axis=-1)
@@ -127,7 +128,7 @@ def _wilcoxon_iv(x, y, zero_method, correction, alternative, method, axis):
     if 0 < d.shape[-1] < 10 and method == "approx":
         warnings.warn("Sample size too small for normal approximation.", stacklevel=2)
 
-    return d, zero_method, correction, alternative, method, axis
+    return d, zero_method, correction, alternative, method, axis, output_z
 
 
 def _wilcoxon_statistic(d, zero_method='wilcox'):
@@ -196,7 +197,7 @@ def _wilcoxon_nd(x, y=None, zero_method='wilcox', correction=True,
                  alternative='two-sided', method='auto', axis=0):
 
     temp = _wilcoxon_iv(x, y, zero_method, correction, alternative, method, axis)
-    d, zero_method, correction, alternative, method, axis = temp
+    d, zero_method, correction, alternative, method, axis, output_z = temp
 
     if d.size == 0:
         NaN = _get_nan(d)
@@ -232,6 +233,6 @@ def _wilcoxon_nd(x, y=None, zero_method='wilcox', correction=True,
     z = -np.abs(z) if (alternative == 'two-sided' and method == 'approx') else z
 
     res = _morestats.WilcoxonResult(statistic=statistic, pvalue=p[()])
-    if method == 'approx':
+    if output_z:
         res.zstatistic = z[()]
     return res

scipy/stats/tests/test_morestats.py

@@ -1658,6 +1658,22 @@ def test_permutation_method(self, size):
         assert_equal(np.round(res.pvalue, 2), res.pvalue)  # n_resamples used
         assert_equal(res.pvalue, ref.pvalue)  # random_state used
 
+    def test_method_auto_nan_propagate_ND_length_gt_50_gh20591(self):
+        # When method!='approx', nan_policy='propagate', and a slice of
+        # a >1 dimensional array input contained NaN, the result object of
+        # `wilcoxon` could (under yet other conditions) return `zstatistic`
+        # for some slices but not others. This resulted in an error because
+        # `apply_along_axis` would have to create a ragged array.
+        # Check that this is resolved.
+        rng = np.random.default_rng(235889269872456)
+        A = rng.normal(size=(51, 2))  # length along slice > exact threshold
+        A[5, 1] = np.nan
+        res = stats.wilcoxon(A)
+        ref = stats.wilcoxon(A, method='approx')
+        assert_allclose(res, ref)
+        assert hasattr(ref, 'zstatistic')
+        assert not hasattr(res, 'zstatistic')
+
 
 class TestKstat:
     def test_moments_normal_distribution(self):