Repository: scipy/scipy
File: scipy/stats/tests/test_stats.py (Python)

Code before:
        y = stats.moment(xp.zeros((1, 0)), axis=0)
        xp_assert_equal(y, xp.empty((0,)))

        y = stats.moment(xp.asarray([[]]), order=[0, 1], axis=0)
        xp_assert_equal(y, xp.empty((2, 0)))

    def test_nan_policy(self):
        x = np.arange(10.)
        x[9] = np.nan
        assert_equal(stats.moment(x, 2), np.nan)
        assert_almost_equal(stats.moment(x, nan_policy='omit'), 0.0)
        assert_raises(ValueError, stats.moment, x, nan_policy='raise')
        assert_raises(ValueError, stats.moment, x, nan_policy='foobar')

    @pytest.mark.parametrize('dtype', ['float32', 'float64', 'complex128'])
    @pytest.mark.parametrize('expect, order', [(0, 1), (1, 0)])
    def test_constant_moments(self, dtype, expect, order, xp):
        if dtype=='complex128' and is_torch(xp):
            pytest.skip()
        dtype = getattr(xp, dtype)
        rng = np.random.default_rng(3824693518)
        x = xp.asarray(rng.random(5), dtype=dtype)
        y = stats.moment(x, order=order)
        xp_assert_equal(y, xp.asarray(expect, dtype=dtype))

        message = "Precision loss occurred in moment calculation"
        with eager_warns(RuntimeWarning, match=message, xp=xp):
            y = stats.moment(xp.broadcast_to(x, (6, 5)), axis=0, order=order)
            xp_assert_equal(y, xp.full((5,), expect, dtype=dtype))

        with eager_warns(RuntimeWarning, match=message, xp=xp):
            y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=2,
                             order=order)
            xp_assert_equal(y, xp.full((1, 2, 4, 5), expect, dtype=dtype))

        y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=None,
                         order=order)
        xp_assert_equal(y, xp.full((), expect, dtype=dtype))


    def test_moment_propagate_nan(self, xp):
        # Check that the shape of the result is the same for inputs
        # with and without nans, cf gh-5817
        a = xp.reshape(xp.arange(8.), (2, -1))
        a = xpx.at(a)[1, 0].set(xp.nan)

        mm = stats.moment(xp.asarray(a), 2, axis=1)
        xp_assert_close(mm, xp.asarray([1.25, xp.nan]), atol=1e-15)

    def test_moment_empty_order(self, xp):
        # tests moment with empty `order` list

Review comment (performance):
We're taking a *central* moment along slices in which all elements are identical. The mean is (theoretically) equal to these elements, but numerically it can be slightly different. The subtraction results in catastrophic cancellation.

For full disclosure: 
- even if the mean calculation is exact and the result of subtracting the mean is exactly zero for all elements, we still produce the warning. We didn't bother to special-case this in `_demean` because if all elements of a slice are exactly the same, there may still be something worth calling attention to.
- catastrophic cancellation doesn't affect the precision of the final result for orders zero and one because they are identically `1.0` and `0.0`. We emit the washing anyways because it's hard to selectively suppress it. The work-around would be to only de-mean slices where order is >1, since those results will be replaced anyway. But using `apply_where` is out because this is not an elementwise operation, so it would be more involved. There are options I can think of that would be feasible, but I don't think any are worth it. Why are users using `moment` to take the zeroth moment of an array of all 6s anyway? I don't think we should sacrifice maintainability, performance in the general case, or developer brain cycles/cells to avoid a warning in non-use cases.

Code after:
        y = stats.moment(xp.zeros((1, 0)), axis=0)
        xp_assert_equal(y, xp.empty((0,)))

        y = stats.moment(xp.asarray([[]]), order=[0, 1], axis=0)
        xp_assert_equal(y, xp.empty((2, 0)))

    def test_nan_policy(self):
        x = np.arange(10.)
        x[9] = np.nan
        assert_equal(stats.moment(x, 2), np.nan)
        assert_almost_equal(stats.moment(x, nan_policy='omit'), 0.0)
        assert_raises(ValueError, stats.moment, x, nan_policy='raise')
        assert_raises(ValueError, stats.moment, x, nan_policy='foobar')

    @pytest.mark.parametrize('dtype', ['float32', 'float64', 'complex128'])
    @pytest.mark.parametrize('expect, order', [(0, 1), (1, 0)])
    def test_constant_moments(self, dtype, expect, order, xp):
        if dtype=='complex128' and is_torch(xp):
            pytest.skip()
        dtype = getattr(xp, dtype)
        rng = np.random.default_rng(3824693518)
        x = xp.asarray(rng.random(5), dtype=dtype)
        y = stats.moment(x, order=order)
        xp_assert_equal(y, xp.asarray(expect, dtype=dtype))

        message = "Precision loss occurred in moment calculation"
        with eager_warns(RuntimeWarning, match=message, xp=xp):
            y = stats.moment(xp.broadcast_to(x, (6, 5)), axis=0, order=order)
            xp_assert_equal(y, xp.full((5,), expect, dtype=dtype))

        with eager_warns(RuntimeWarning, match=message, xp=xp):
            y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=2,
                             order=order)
            xp_assert_equal(y, xp.full((1, 2, 4, 5), expect, dtype=dtype))

        y = stats.moment(xp.broadcast_to(x, (1, 2, 3, 4, 5)), axis=None,
                         order=order)
        xp_assert_equal(y, xp.full((), expect, dtype=dtype))

    def test_moment_propagate_nan(self, xp):
        # Check that the shape of the result is the same for inputs
        # with and without nans, cf gh-5817
        a = xp.reshape(xp.arange(8.), (2, -1))
        a = xpx.at(a)[1, 0].set(xp.nan)

        mm = stats.moment(xp.asarray(a), 2, axis=1)
        xp_assert_close(mm, xp.asarray([1.25, xp.nan]), atol=1e-15)

    def test_moment_empty_order(self, xp):
        # tests moment with empty `order` list
        with pytest.raises(ValueError, match=r"`order` must be a scalar or a"