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

Code before:
        x = np.atleast_2d(x)
        counts = rng.integers(10, size=x.shape[-1])
        x_rep = []
        weights = []
        for x_ in x:
            counts_ = rng.permuted(counts)
            x_rep.append(np.repeat(x_, counts_))
            weights.append(counts_)
        x_rep, weights = np.stack(x_rep), np.stack(weights)
        if ndim < 2:
            x_rep, weights = np.squeeze(x_rep, axis=0), np.squeeze(weights, axis=0)
        x_rep, weights = np.swapaxes(x_rep, -1, axis), np.swapaxes(weights, -1, axis)
        weights = np.asarray(weights, dtype=x.dtype)
        return weights, x_rep

    @pytest.mark.parametrize('method',
         ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
          'hazen', 'interpolated_inverted_cdf', 'linear',
          'median_unbiased', 'normal_unbiased', 'weibull',
          '_lower', '_higher', '_midpoint', '_nearest'])
    @pytest.mark.parametrize('shape_x, shape_p, axis',
         [(10, None, -1), (10, 10, -1), (10, (2, 3), -1), ((10, 2), None, 0)])
    @pytest.mark.parametrize('weights', [False, True])
    def test_against_numpy(self, method, shape_x, shape_p, axis, weights, xp):
        if weights and method.startswith('_'):
            pytest.skip('`weights=True` not supported by private (legacy) methods.')
        if weights and is_numpy(xp) and xp.__version__ < "2.0":
            pytest.skip('`weights` not supported by NumPy < 2.0.')
        dtype = xp_default_dtype(xp)
        rng = np.random.default_rng(23458924568734956)
        x = rng.random(size=shape_x)
        p = rng.random(size=shape_p)

        if weights:
            weights, x_rep = self._get_weights_x_rep(x, axis, rng)
        else:
            weights, x_rep = None, x

        ref = np.quantile(x_rep, p, axis=axis,
                          method=method[1:] if method.startswith('_') else method)

        x, p = xp.asarray(x, dtype=dtype), xp.asarray(p, dtype=dtype)
        weights = weights if weights is None else xp.asarray(weights, dtype=dtype)
        res = stats.quantile(x, p, method=method, weights=weights, axis=axis)

        xp_assert_close(res, xp.asarray(ref, dtype=dtype))

    @skip_xp_backends(cpu_only=True, reason="PyTorch doesn't have `betainc`.",
                      exceptions=['cupy', 'jax.numpy'])
    @pytest.mark.parametrize('axis', [0, 1])
    @pytest.mark.parametrize('keepdims', [False, True])

Review comment (performance):
The purpose of this function is to generate integral weights and use them to produce an `x_rep` array that is `x` with elements repeated `weights` times. The somewhat tricky thing is to ensure that the weights have the same sum along `axis` so we don't have a ragged `x_rep` array. I thought about using `random_table`... maybe that or another function would help here? It felt very clumsy, but I didn't really try to optimize this.

Code after:
            message = "The input contains nan values"
            with pytest.raises(ValueError, match=message):
                stats.quantile(xp.asarray([xp.nan, 1, 2]), p, nan_policy='raise')

        message = "`method` must be one of..."
        with pytest.raises(ValueError, match=message):
            stats.quantile(x, p, method='a duck')

        message = "`method='harrell-davis'` does not support `weights`."
        with pytest.raises(ValueError, match=message):
            stats.quantile(x, p, weights=x, method='harrell-davis')

        message = "`method='round_nearest'` does not support `weights`."
        with pytest.raises(ValueError, match=message):
            stats.quantile(x, p, weights=x, method='round_nearest')

        message = "If specified, `keepdims` must be True or False."
        with pytest.raises(ValueError, match=message):
            stats.quantile(x, p, keepdims=42)

        message = "`keepdims` may be False only if the length of `p` along `axis` is 1."
        with pytest.raises(ValueError, match=message):
            stats.quantile(x, xp.asarray([0.5, 0.6]), keepdims=False)


    def _get_weights_x_rep(self, x, axis, rng):
        x = np.swapaxes(x, axis, -1)
        ndim = x.ndim
        x = np.atleast_2d(x)
        counts = rng.integers(10, size=x.shape[-1], dtype=np.int32)
        x_rep = []
        weights = []
        for x_ in x:
            counts_ = rng.permuted(counts)
            x_rep.append(np.repeat(x_, counts_))
            weights.append(counts_)
        x_rep, weights = np.stack(x_rep), np.stack(weights)
        if ndim < 2:
            x_rep, weights = np.squeeze(x_rep, axis=0), np.squeeze(weights, axis=0)
        x_rep, weights = np.swapaxes(x_rep, -1, axis), np.swapaxes(weights, -1, axis)
        weights = np.asarray(weights, dtype=x.dtype)
        return weights, x_rep


    @skip_xp_backends(cpu_only=True, reason="PyTorch doesn't have `betainc`.",
                      exceptions=['cupy'])
    @pytest.mark.parametrize('method',
         ['inverted_cdf', 'averaged_inverted_cdf', 'closest_observation',
          'hazen', 'interpolated_inverted_cdf', 'linear',
          'median_unbiased', 'normal_unbiased', 'weibull',
          'harrell-davis', 'round_nearest', 'round_outward', 'round_inward',