Repository: numpy/numpy
File: numpy/lib/tests/test_function_base.py (Python)

Code before:
        assert_equal(2, np.quantile(arr[0:5], q, method=m))
        assert_equal(3, np.quantile(arr[0:6], q, method=m))
        assert_equal(4, np.quantile(arr[0:7], q, method=m))
        assert_equal(4, np.quantile(arr[0:8], q, method=m))
        assert_equal(4, np.quantile(arr[0:9], q, method=m))
        assert_equal(5, np.quantile(arr, q, method=m))

    def test_inf_err(self):

        m = "inverted_cdf"
        q = 0.5
        arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        wgt = np.ones(10)

        for i in range(len(arr)):
            wgt[i] = np.inf
            with pytest.raises(ValueError) as ex:
                a = np.quantile(arr, q, weights=wgt, method=m)
            assert "Weights must be non-infinite" in str(ex)
            wgt[i] = 1
        
        for i in range(len(arr)):
            wgt[i] = np.inf
            with pytest.raises(ValueError) as ex:
                a = np.quantile(arr, q, weights=wgt, method=m)
            assert "Weights must be non-infinite" in str(ex)
    
    def test_nan_err(self):

        m = "inverted_cdf"
        q = 0.5
        arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        wgt = np.ones(10)

        for i in range(len(arr)):
            wgt[i] = np.nan
            with pytest.raises(ValueError) as ex:
                a = np.quantile(arr, q, weights=wgt, method=m)
            assert "At least one weight is nan" in str(ex)
            wgt[i] = 1
        
        for i in range(len(arr)):
            wgt[i] = np.nan
            with pytest.raises(ValueError) as ex:
                a = np.quantile(arr, q, weights=wgt, method=m)
            assert "At least one weight is nan" in str(ex)
    
    def test_all_zeroes_err(self):

        m = "inverted_cdf"
        q = 0.5

Review comment (suggestion):
could parametrize the test to check one weights array with a few NaNs and another with just 1, but probably don't need these two exhaustive loops

Code after:
                assert_allclose(q, np.quantile(y, alpha, method=method))
        elif method == "lower":
            assert_allclose(q, np.quantile(y, alpha, method="higher"))
        elif method == "higher":
            assert_allclose(q, np.quantile(y, alpha, method="lower"))
        else:
            # "averaged_inverted_cdf", "hazen", "weibull", "linear",
            # "median_unbiased", "normal_unbiased", "midpoint"
            assert_allclose(q, np.quantile(y, alpha, method=method))

    @pytest.mark.parametrize("method", methods_supporting_weights)
    @pytest.mark.parametrize("alpha", [0.2, 0.5, 0.9])
    def test_quantile_constant_weights(self, method, alpha):
        rng = np.random.default_rng(4321)
        # We choose n and alpha such that we have cases for
        #  - n * alpha is an integer
        #  - n * alpha is a float that gets rounded down
        #  - n * alpha is a float that gest rounded up
        n = 102  # n * alpha = 20.4, 51. , 91.8
        y = rng.random(n)
        q = np.quantile(y, alpha, method=method)

        w = np.ones_like(y)
        qw = np.quantile(y, alpha, method=method, weights=w)
        assert_allclose(qw, q)

        w = 8.125 * np.ones_like(y)
        qw = np.quantile(y, alpha, method=method, weights=w)
        assert_allclose(qw, q)

    @pytest.mark.parametrize("method", methods_supporting_weights)
    @pytest.mark.parametrize("alpha", [0, 0.2, 0.5, 0.9, 1])
    def test_quantile_with_integer_weights(self, method, alpha):
        # Integer weights can be interpreted as repeated observations.
        rng = np.random.default_rng(4321)
        # We choose n and alpha such that we have cases for
        #  - n * alpha is an integer
        #  - n * alpha is a float that gets rounded down
        #  - n * alpha is a float that gest rounded up
        n = 102  # n * alpha = 20.4, 51. , 91.8
        y = rng.random(n)
        w = rng.integers(low=0, high=10, size=n, dtype=np.int32)

        qw = np.quantile(y, alpha, method=method, weights=w)
        q = np.quantile(np.repeat(y, w), alpha, method=method)
        assert_allclose(qw, q)

    @pytest.mark.parametrize("method", methods_supporting_weights)
    def test_quantile_with_weights_and_axis(self, method):
        rng = np.random.default_rng(4321)