Repository: scipy/scipy
Pull request: MAINT: stats: fix truncnorm stats with array shapes
#### Reference issue
Closes gh-13580. 
Closes gh-13642
(Fixes the `mean`/`var` parts; the rest is a duplicate of gh-12192 to be fixed by gh-12197)

#### What does this implement/fix?
`truncnorm._stats` was not written to handle array shape parameters. This vectorizes it and enhances the `check_moments` test to catch this sort of problem in other distributions. 

(The strengthened test exposes a similar bug in `kappa3` and `kappa4`. That is now documented in gh-13582. In the meantime, the enhanced `check_loc_scale` test is `xfail`ed for those distributions.)

`test_continuous_basic` was `xfail`ing `truncnorm`, but no explanation was given. It passes all CI checks, so I have removed the `xfail` mark.


Problem statement:
truncnorm gives incorrect means and variances 
`truncnorm.mean` and `truncnorm.var` appear to give incorrect outputs when `loc` or `scale` are array-like and contain more than one number: 
```
>>> import scipy.stats as stats
>>> stats.truncnorm.mean(-1, 2, loc=0)
0.229637179091329
>>> stats.truncnorm.mean(-1, 2, loc=[0])
array([0.22963718])
>>> stats.truncnorm.mean(-1, 2, loc=[0, 0]) # expect to get 0.2296 twice
array([0., 0.])
>>> stats.truncnorm.var(-1, 1, loc=0)
0.29112509477279314
>>> stats.truncnorm.var(-1, 1, loc=[0])
array([0.29112509])
>>> stats.truncnorm.var(-1, 1, loc=[0, 0]) # expect to get 0.2911 twice
array([1., 1.])
```
`truncnorm.moment` doesn't seem to support array-likes at all, either:
```
>>> stats.truncnorm.moment(1, -1, 1, loc=[0])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\ProgramData\Miniconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py", line 1266, in moment
    fac = float(scale) / float(loc)
TypeError: float() argument must be a string or a number, not 'list'
```

This is on Scipy version 1.6.0. This problem was even worse on version 1.4.1, where `truncnorm.var()` could produce negative numbers:
```
>>> stats.truncnorm.var(-1, 1, loc=0)
0.29112509477279314
>>> stats.truncnorm.var(-1, 1, loc=[0, 0])
array([-0.41774981, -0.41774981])
```

I don't think I'm using these functions incorrectly, as the docstrings for `mean()`, `var()`, and `moment()` all say that `loc` and `scale` can be array-like.

stats.truncnorm variance works incorrectly when input is an array.
When trying to compute variance of `stats.truncnorm` it gives incorrect results when shape parameters `a` and `b` are arrays.

It works correctly when `a` and `b` are floats.

I am working with `scipy-1.6.1` and Python 3.9.1.

#### Reproducing code example:

```python
import numpy as np
from scipy import stats

a = np.arange(10, dtype=float)
b = a + 1

# This will print correct mean
print('mean:', stats.truncnorm.stats(a, b, moments='m'))

# This will print incorrect variance
print('var:', stats.truncnorm.stats(a, b, moments='v'))

# This will print correct variance
print('var:', [stats.truncnorm.stats(aa, bb, moments='v') for (aa, bb) in zip(a, b)])
```

**Output of this program:**
```
mean: array([0.45986223, 1.38316905, 2.31582133, 3.26045429, 4.21683078,
       5.18314709, 6.1572109 , 7.13706716, 8.12118899, 9.10845629])
var: array([-8.60292538, -8.60292538, -8.60292538, -8.60292538, -8.60292538,
       -8.60292538, -8.60292538, -8.60292538, -8.60292538, -8.60292538])
var: [array(0.07965182), array(0.07274289), array(0.06152078), array(0.04927779), array(0.03829029), array(0.02945243), array(0.02274838), array(0.01779288), array(0.01414854), array(0.01144881)]
```

Hints from issue discussion:
gh-12192 fixes `moment`'s issues with input arrays. That's a problem for all stats distributions. (Potential reviewers: that PR is waiting for review. I was reviewing, but ended up writing a fair amount of code, so I think it needs an independent review.) 

I'll take a look into `truncnorm` mean and variance. I hope that part of the issue is distribution specific.
While smoking out this issue, I'm going to collect the other bugs I find here to avoid slowing things down. I'll open separate issues if need be. 
```python3
mean, var, skew, kurt = stats.kappa4.stats(0.1, 0, moments='mvsk', loc=np.array([1, 2]))  # ValueError: The truth value of
mean, var, skew, kurt = stats.kappa3.stats(1, moments='mvsk', loc=np.array([1, 2]))  # ValueError: The truth value of
```
OK, I was concerned that there might be something wrong with the way `rv_continuous` was dealing with `loc`, but no, it is that `truncnorm.stats` is not dealing with array shape parameters correctly. 
```python3
stats.truncnorm.stats([-1], [2])  # right: (array([0.22963718]), array([0.51976254]))
stats.truncnorm.stats([-1, -1], [2, 2])  # wrong: (array([0., 0.]), array([1., 1.]))
```
The shapes are arrays in the original issue because `rv_continuous` broadcasts them with `loc`, etc.
One problem is line 7917:
```
        pA, pB = self._pdf(np.array([a, b]), a, b)
```
fix for mean is:
```
        pA, pB = self.pdf(np.array([a, b]), a, b)
```
But that didn't fix `var` so I'll just use `np.vectorize`. I'll open a PR.
Hi @mdhaber Can I try your fix and open a PR for this issue? Please let me know. Thanks.
I opened gh-13581 to fix this in February. It just hasn't been reviewed by another maintainer. 
Ah! I missed it. No problems. Thanks for highlighting your work on this issue.
Thanks. I think this is a duplicate of gh-13580, which would be fixed by gh-13581. Reviews welcome.
Thank you for your answer! Really looking forward for this to be merged.
@hocop You can help make that happen! Reviews from users are valuable, too. Please take a look at gh-13581. If you can check out  that branch and confirm that it fixes the issue for you, that would be great. If you can only look at the code and confirm whether it looks right, that would be valuable too. Please add your comments/approval there. Thanks!
@mdhaber thank you for your work! I checked it out and it does fix my issue. Left my comment in [gh-13581](https://github.com/scipy/scipy/pull/13581) approving your pull request.
Thanks!

Commit message: Update scipy/stats/tests/test_continuous_basic.py

Co-authored-by: Samuel Wallan <44255917+swallan@users.noreply.github.com>

Patch to review:
diff --git a/scipy/stats/_continuous_distns.py b/scipy/stats/_continuous_distns.py
index 025917174722..b9a87476c8ee 100644
--- a/scipy/stats/_continuous_distns.py
+++ b/scipy/stats/_continuous_distns.py
@@ -8222,31 +8222,37 @@ def n_th_moment(n, a, b):
                           np.nan)
 
     def _stats(self, a, b, moments='mv'):
-        pA, pB = self._pdf(np.array([a, b]), a, b)
-        m1 = pA - pB
-        mu = m1
-        # use _lazywhere to avoid nan (See detailed comment in _munp)
-        probs = [pA, -pB]
-        vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y,
-                          fillvalue=0)
-        m2 = 1 + np.sum(vals)
-        vals = _lazywhere(probs, [probs, [a-mu, b-mu]], lambda x, y: x*y,
-                          fillvalue=0)
-        # mu2 = m2 - mu**2, but not as numerically stable as:
-        # mu2 = (a-mu)*pA - (b-mu)*pB + 1
-        mu2 = 1 + np.sum(vals)
-        vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y**2,
-                          fillvalue=0)
-        m3 = 2*m1 + np.sum(vals)
-        vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y**3,
-                          fillvalue=0)
-        m4 = 3*m2 + np.sum(vals)
-
-        mu3 = m3 + m1 * (-3*m2 + 2*m1**2)
-        g1 = mu3 / np.power(mu2, 1.5)
-        mu4 = m4 + m1*(-4*m3 + 3*m1*(2*m2 - m1**2))
-        g2 = mu4 / mu2**2 - 3
-        return mu, mu2, g1, g2
+        pA, pB = self.pdf(np.array([a, b]), a, b)
+
+        def _truncnorm_stats_scalar(a, b, pA, pB, moments):
+            m1 = pA - pB
+            mu = m1
+            # use _lazywhere to avoid nan (See detailed comment in _munp)
+            probs = [pA, -pB]
+            vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y,
+                              fillvalue=0)
+            m2 = 1 + np.sum(vals)
+            vals = _lazywhere(probs, [probs, [a-mu, b-mu]], lambda x, y: x*y,
+                              fillvalue=0)
+            # mu2 = m2 - mu**2, but not as numerically stable as:
+            # mu2 = (a-mu)*pA - (b-mu)*pB + 1
+            mu2 = 1 + np.sum(vals)
+            vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y**2,
+                              fillvalue=0)
+            m3 = 2*m1 + np.sum(vals)
+            vals = _lazywhere(probs, [probs, [a, b]], lambda x, y: x*y**3,
+                              fillvalue=0)
+            m4 = 3*m2 + np.sum(vals)
+
+            mu3 = m3 + m1 * (-3*m2 + 2*m1**2)
+            g1 = mu3 / np.power(mu2, 1.5)
+            mu4 = m4 + m1*(-4*m3 + 3*m1*(2*m2 - m1**2))
+            g2 = mu4 / mu2**2 - 3
+            return mu, mu2, g1, g2
+
+        _truncnorm_stats = np.vectorize(_truncnorm_stats_scalar,
+                                        excluded=('moments',))
+        return _truncnorm_stats(a, b, pA, pB, moments)
 
     def _rvs(self, a, b, size=None, random_state=None):
         # if a and b are scalar, use _rvs_scalar, otherwise need to create
diff --git a/scipy/stats/tests/test_continuous_basic.py b/scipy/stats/tests/test_continuous_basic.py
index 41acde727402..592c9d7fe4fe 100644
--- a/scipy/stats/tests/test_continuous_basic.py
+++ b/scipy/stats/tests/test_continuous_basic.py
@@ -109,7 +109,7 @@
                    'ncf', 'nct', 'ncx2', 'norminvgauss', 'pearson3', 'rdist',
                    'reciprocal', 'rice', 'skewnorm', 't', 'tukeylambda',
                    'vonmises', 'vonmises_line', 'rv_histogram_instance',
-                   'studentized_range'])
+                   'truncnorm', 'studentized_range'])
 
 _h = np.histogram([1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5, 6,
                    6, 6, 6, 7, 7, 7, 8, 8, 9], bins=8)
@@ -133,9 +133,6 @@ def cases_test_cont_basic():
 def test_cont_basic(distname, arg, sn, n_fit_samples):
     # this test skips slow distributions
 
-    if distname == 'truncnorm':
-        pytest.xfail(reason=distname)
-
     try:
         distfn = getattr(stats, distname)
     except TypeError:
@@ -245,6 +242,7 @@ def test_levy_stable_random_state_property():
 def cases_test_moments():
     fail_normalization = set(['vonmises'])
     fail_higher = set(['vonmises', 'ncf'])
+    fail_loc_scale = set(['kappa3', 'kappa4'])  # see gh-13582
 
     for distname, arg in distcont[:] + [(histogram_test_instance, tuple())]:
         if distname == 'levy_stable':
@@ -254,25 +252,28 @@ def cases_test_moments():
             msg = ("studentized_range is far too slow for this test and it is "
                    "redundant with test_distributions::TestStudentizedRange::"
                    "test_moment_against_mp")
-            yield pytest.param(distname, arg, True, True, True,
+            yield pytest.param(distname, arg, True, True, True, True,
                                marks=pytest.mark.xslow(reason=msg))
             continue
         cond1 = distname not in fail_normalization
         cond2 = distname not in fail_higher
+        cond3 = distname not in fail_loc_scale
 
-        yield distname, arg, cond1, cond2, False
+        yield distname, arg, cond1, cond2, cond3, False
 
-        if not cond1 or not cond2:
+        if not cond1 or not cond2 or not cond3:
             # Run the distributions that have issues twice, once skipping the
             # not_ok parts, once with the not_ok parts but marked as knownfail
-            yield pytest.param(distname, arg, True, True, True,
+            yield pytest.param(distname, arg, True, True, True, True,
                                marks=pytest.mark.xfail)
 
 
 @pytest.mark.slow
-@pytest.mark.parametrize('distname,arg,normalization_ok,higher_ok,is_xfailing',
+@pytest.mark.parametrize('distname,arg,normalization_ok,higher_ok,'
+                         'loc_scale_ok,is_xfailing',
                          cases_test_moments())
-def test_moments(distname, arg, normalization_ok, higher_ok, is_xfailing):
+def test_moments(distname, arg, normalization_ok, higher_ok, loc_scale_ok,
+                 is_xfailing):
     try:
         distfn = getattr(stats, distname)
     except TypeError:
@@ -296,7 +297,9 @@ def test_moments(distname, arg, normalization_ok, higher_ok, is_xfailing):
             check_var_expect(distfn, arg, m, v, distname)
             check_kurt_expect(distfn, arg, m, v, k, distname)
 
-        check_loc_scale(distfn, arg, m, v, distname)
+        if loc_scale_ok:
+            check_loc_scale(distfn, arg, m, v, distname)
+
         check_moment(distfn, arg, m, v, distname)
 
 
@@ -633,7 +636,9 @@ def check_vecentropy(distfn, args):
 
 
 def check_loc_scale(distfn, arg, m, v, msg):
-    loc, scale = 10.0, 10.0
+    # Make `loc` and `scale` arrays to catch bugs like gh-13580 where
+    # `loc` and `scale` arrays improperly broadcast with shapes. 
+    loc, scale = np.array([10.0, 20.0]), np.array([10.0, 20.0])
     mt, vt = distfn.stats(loc=loc, scale=scale, *arg)
     npt.assert_allclose(m*scale + loc, mt)
     npt.assert_allclose(v*scale*scale, vt)


Code review comments:
diff_hunk: @@ -633,7 +636,9 @@ def check_vecentropy(distfn, args):
 
 
 def check_loc_scale(distfn, arg, m, v, msg):
-    loc, scale = 10.0, 10.0
+    # Make `loc` and `scale` arrays to catch bugs like gh-13580 where
+    # `loc` and `scale` arrays improperly broadcast with shapes. 
path: scipy/stats/tests/test_continuous_basic.py
text: @user1:
```suggestion
    # `loc` and `scale` arrays improperly broadcast with shapes.
```

This should fix the lint failure.