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

Code before:
    def test_weights_2d_axis0(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([[2, 5], [10, 5], [6, 5]])
        weights = np.array([[10, 1], [5, 1], [3, 1]])
        desired = np.array([3, 5])
        check_equal_hmean(a, desired, axis=0, weights=weights, rtol=1e-5, xp=xp)

    def test_weights_2d_axis1(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([[2, 10, 6], [7, 7, 7]])
        weights = np.array([[10, 5, 3], [1, 1, 1]])
        desired = np.array([3, 7])
        check_equal_hmean(a, desired, axis=1, weights=weights, rtol=1e-5, xp=xp)

    @skip_xp_invalid_arg
    def test_weights_masked_1d_array(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([2, 10, 6, 42])
        weights = np.ma.array([10, 5, 3, 42], mask=[0, 0, 0, 1])
        desired = 3
        xp = np.ma  # check_equal_hmean uses xp.asarray; this will preserve the mask
        check_equal_hmean(a, desired, weights=weights, rtol=1e-5,
                          dtype=np.float64, xp=xp)


@make_xp_test_case(stats.gmean)
class TestGMean:
    @pytest.mark.filterwarnings(
        "ignore:divide by zero encountered in log:RuntimeWarning:dask"
    )
    def test_0(self, xp):
        a = [1, 0, 2]
        desired = 0
        check_equal_gmean(a, desired, xp=xp)

    def test_1d(self, xp):
        #  Test a 1d case
        a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
        desired = 45.2872868812
        check_equal_gmean(a, desired, xp=xp)

        a = [1, 2, 3, 4]
        desired = power(1 * 2 * 3 * 4, 1. / 4.)
        check_equal_gmean(a, desired, rtol=1e-14, xp=xp)

        a = array([1, 2, 3, 4], float32)
        desired = power(1 * 2 * 3 * 4, 1. / 4.)
        check_equal_gmean(a, desired, dtype=xp.float32, xp=xp)

Review comment (suggestion):
Masked arrays still treat dtypes as they did in the past; this adjustment is just made to make testing easier.

Code after:
        a = np.asarray([2, 10, 6])
        weights = np.asarray([10, 5, 3])
        desired = 3
        check_equal_hmean(a, desired, weights=weights, rtol=1e-5, xp=xp)

    def test_weights_2d_axis0(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([[2, 5], [10, 5], [6, 5]])
        weights = np.array([[10, 1], [5, 1], [3, 1]])
        desired = np.array([3, 5])
        check_equal_hmean(a, desired, axis=0, weights=weights, rtol=1e-5, xp=xp)

    def test_weights_2d_axis1(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([[2, 10, 6], [7, 7, 7]])
        weights = np.array([[10, 5, 3], [1, 1, 1]])
        desired = np.array([3, 7])
        check_equal_hmean(a, desired, axis=1, weights=weights, rtol=1e-5, xp=xp)

    @skip_xp_invalid_arg
    def test_weights_masked_1d_array(self, xp):
        # Desired result from:
        # https://www.hackmath.net/en/math-problem/35871
        a = np.array([2., 10., 6., 42.])
        weights = np.ma.array([10., 5., 3., 42.], mask=[0, 0, 0, 1])
        desired = 3.
        xp = np.ma  # check_equal_hmean uses xp.asarray; this will preserve the mask
        check_equal_hmean(a, desired, weights=weights, rtol=1e-5,
                          dtype=np.float64, xp=xp)


@make_xp_test_case(stats.gmean)
class TestGMean:
    @pytest.mark.filterwarnings(
        "ignore:divide by zero encountered in log:RuntimeWarning:dask"
    )
    def test_0(self, xp):
        a = [1, 0, 2]
        desired = 0
        check_equal_gmean(a, desired, xp=xp)

    @pytest.mark.parametrize('dtype', [None, 'int64', 'float32', 'float64'])
    def test_1d(self, dtype, xp):
        dtype = dtype if dtype is None else getattr(xp, dtype)

        #  Test a 1d case
        a = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
        desired = 45.2872868812
        check_equal_gmean(a, desired, dtype=dtype, xp=xp)