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

Code before:
        
        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
        arr = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        wgt = np.zeros(10)
        with pytest.raises(ValueError) as ex:
            a = np.quantile(arr, q, weights=wgt, method=m)

        assert "At least one weight must be non-zero" in str(ex)
            
            
class TestLerp:
    @hypothesis.given(t0=st.floats(allow_nan=False, allow_infinity=False,
                                   min_value=0, max_value=1),
                      t1=st.floats(allow_nan=False, allow_infinity=False,
                                   min_value=0, max_value=1),
                      a = st.floats(allow_nan=False, allow_infinity=False,
                                    min_value=-1e300, max_value=1e300),
                      b = st.floats(allow_nan=False, allow_infinity=False,
                                    min_value=-1e300, max_value=1e300))
    def test_linear_interpolation_formula_monotonic(self, t0, t1, a, b):
        l0 = nfb._lerp(a, b, t0)
        l1 = nfb._lerp(a, b, t1)

Review comment (suggestion):
Could also probably just parametrize over `np.nan` and `np.inf` + expected message above for concision (and same comment about not needing the exhaustive loops

Code after:
        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)

        # 1d weight and single alpha
        y = rng.random((2, 10, 3))
        w = np.abs(rng.random(10))
        alpha = 0.5
        q = np.quantile(y, alpha, weights=w, method=method, axis=1)
        q_res = np.zeros(shape=(2, 3))
        for i in range(2):
            for j in range(3):
                q_res[i, j] = np.quantile(
                    y[i, :, j], alpha, method=method, weights=w
                )
        assert_allclose(q, q_res)

        # 1d weight and 1d alpha
        alpha = [0, 0.2, 0.4, 0.6, 0.8, 1]  # shape (6,)
        q = np.quantile(y, alpha, weights=w, method=method, axis=1)
        q_res = np.zeros(shape=(6, 2, 3))
        for i in range(2):
            for j in range(3):
                q_res[:, i, j] = np.quantile(