Repository: scipy/scipy
Pull request: BUG: Correctly use `axis` in `scipy.stats.tmean`
Closes #12106

Fixes `scipy.stats.tmean` so that it correctly uses the `axis` keyword.

(_Please check that this addresses gh-9770, too_ - @mdhaber)

Problem statement:
`axis` option for `stats.tmean` do not appear to be working as expected
<!-- 
Thank you for taking the time to file a bug report. 
Please fill in the fields below, deleting the sections that 
don't apply to your issue. You can view the final output
by clicking the preview button above.
Note: This is a comment, and won't appear in the output.
-->

My issue is about (see title)

#### Reproducing code example:
<!-- 
If you place your code between the triple backticks below, 
it will be rendered as a code block. 
-->

```
import scipy as sp
import numpy as np

r = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])

print(sp.stats.tmean(r, axis=1))
print(sp.stats.tstd(r, axis=1))
```
Output:
```
2.0
[0. 0. 0.]
```

Expectation:
```
[1.0, 2.0, 3.0]
[0., 0., 0.]
```

#### Scipy/Numpy/Python version information:
```
1.4.1 1.18.1 sys.version_info(major=3, minor=7, micro=6, releaselevel='final', serial=0)
```

Hints from issue discussion:
@jtstanley, thanks for the report. That's a bug. `tmean` is ignoring the `axis` argument.
Hello!

There seem to be two issues with `tmean`. 

1. First, as @jtstanley pointed out, it ignores `axis` if no `limits` are provided. The culprit here is this line, where `axis` has been forced to `None`: 
```
if limits is None:
        return np.mean(a, None)
```

2. If `limits` is not `None`, then this happens:
```
    am = _mask_to_limits(a.ravel(), limits, inclusive)
```
The array is forcibly `ravel`-ed, and consequently, the `axis` keyword loses any meaning. 

Both fixes seem simple -- change `np.mean(a, None)` to `np.mean(a, axis)` for (1), and removing the call to `ravel`, filling missing entries with `np.nan`, and then using  `np.nanmean(am, axis)` for (2) (this is exactly what's done in `tvar`). 

I'm happy to make both changes an open a PR if no one else is working on this. 
Thanks @siddhantwahal, please go for it!

Commit message: MAINT: no need to cast in asarray

Patch to review:
diff --git a/THANKS.txt b/THANKS.txt
index 49bbd6d4db53..a7fb9cecafb8 100644
--- a/THANKS.txt
+++ b/THANKS.txt
@@ -236,7 +236,9 @@ Sam Wallan for scipy.linalg.lapack enhancements
 Richard Weiss for a bug fix in scipy.optimize._differentialevolution.py.
 Luigi F. Cruz for adding time/frequency domain option to scipy.signal.resample.
 Wesley Alves for improvements to scipy.stats.jarque_bera and scipy.stats.shapiro
+<<<<<<< HEAD
 Mark Borgerding for contributing linalg.convolution_matrix.
+Siddhant Wahal for a bug fix in scipy.stats.tmean.
 
 Institutions
 ------------
diff --git a/scipy/stats/stats.py b/scipy/stats/stats.py
index e844c3f57f50..5558d7037bbb 100644
--- a/scipy/stats/stats.py
+++ b/scipy/stats/stats.py
@@ -540,7 +540,7 @@ def tmean(a, limits=None, inclusive=(True, True), axis=None):
 
     Returns
     -------
-    tmean : float
+    tmean : ndarray
         Trimmed mean.
 
     See Also
@@ -559,10 +559,10 @@ def tmean(a, limits=None, inclusive=(True, True), axis=None):
     """
     a = asarray(a)
     if limits is None:
-        return np.mean(a, None)
-
-    am = _mask_to_limits(a.ravel(), limits, inclusive)
-    return am.mean(axis=axis)
+        return np.mean(a, axis)
+    am = _mask_to_limits(a, limits, inclusive)
+    mean = np.ma.filled(am.mean(axis=axis), fill_value=np.nan)
+    return mean if mean.ndim > 0 else mean.item()
 
 
 def tvar(a, limits=None, inclusive=(True, True), axis=0, ddof=1):
diff --git a/scipy/stats/tests/test_stats.py b/scipy/stats/tests/test_stats.py
index 2bc14085ec3d..d8357111804a 100644
--- a/scipy/stats/tests/test_stats.py
+++ b/scipy/stats/tests/test_stats.py
@@ -62,6 +62,45 @@ def test_tmean(self):
         y2 = stats.tmean(X, limits=None)
         assert_approx_equal(y1, y2, significant=self.dprec)
 
+        x_2d = arange(63, dtype=float64).reshape(9, 7)
+        y = stats.tmean(x_2d, axis=None)
+        assert_approx_equal(y, x_2d.mean(), significant=self.dprec)
+
+        y = stats.tmean(x_2d, axis=0)
+        assert_array_almost_equal(y, x_2d.mean(axis=0), decimal=8)
+
+        y = stats.tmean(x_2d, axis=1)
+        assert_array_almost_equal(y, x_2d.mean(axis=1), decimal=8)
+
+        y = stats.tmean(x_2d, limits=(2, 61), axis=None)
+        assert_approx_equal(y, 31.5, significant=self.dprec)
+
+        y = stats.tmean(x_2d, limits=(2, 21), axis=0)
+        y_true = [14, 11.5, 9, 10, 11, 12, 13]
+        assert_array_almost_equal(y, y_true, decimal=8)
+
+        y = stats.tmean(x_2d, limits=(2, 21), inclusive=(True, False), axis=0)
+        y_true = [10.5, 11.5, 9, 10, 11, 12, 13]
+        assert_array_almost_equal(y, y_true, decimal=8)
+
+        x_2d_with_nan = np.array(x_2d)
+        x_2d_with_nan[-1, -3:] = np.nan
+        y = stats.tmean(x_2d_with_nan, limits=(1, 13), axis=0)
+        y_true = [7, 4.5, 5.5, 6.5, np.nan, np.nan, np.nan]
+        assert_array_almost_equal(y, y_true, decimal=8)
+
+        with suppress_warnings() as sup:
+            sup.record(RuntimeWarning, "Mean of empty slice")
+
+            y = stats.tmean(x_2d, limits=(2, 21), axis=1)
+            y_true = [4, 10, 17, 21, np.nan, np.nan, np.nan, np.nan, np.nan]
+            assert_array_almost_equal(y, y_true, decimal=8)
+
+            y = stats.tmean(x_2d, limits=(2, 21),
+                            inclusive=(False, True), axis=1)
+            y_true = [4.5, 10, 17, 21, np.nan, np.nan, np.nan, np.nan, np.nan]
+            assert_array_almost_equal(y, y_true, decimal=8)
+
     def test_tvar(self):
         y = stats.tvar(X, limits=(2, 8), inclusive=(True, True))
         assert_approx_equal(y, 4.6666666666666661, significant=self.dprec)


Code review comments:
diff_hunk: @@ -236,7 +236,9 @@ Sam Wallan for scipy.linalg.lapack enhancements
 Richard Weiss for a bug fix in scipy.optimize._differentialevolution.py.
 Luigi F. Cruz for adding time/frequency domain option to scipy.signal.resample.
 Wesley Alves for improvements to scipy.stats.jarque_bera and scipy.stats.shapiro
+<<<<<<< HEAD
path: THANKS.txt
text: @user1:
rebase error

@author:
Aargh, don't know how I missed that, sorry.