Repository: scipy/scipy
File: scipy/stats/_stats_py.py (Python)

Code before:
            valid_index = low_index >= 0
            low_index[~valid_index] = 0
            low = np.where(valid_index, np.take_along_axis(x, low_index), np.nan)
            high_index = bd.isf(p).astype(int)
            valid_index = high_index < n
            high_index[~valid_index] = 0
            high = np.where(valid_index, np.take_along_axis(x, high_index), np.nan)

        args = self._axis, self._axis_none, self._keepdims, self._ndim, self._nan_out
        low = _quantile_test_postprocess(low, *args)
        high = _quantile_test_postprocess(high, *args)
        return ConfidenceInterval(low, high)


def quantile_test_iv(x, q, p, alternative, axis, keepdims):

    x = np.atleast_1d(x)
    message = '`x` must be an array of numbers.'
    if not np.issubdtype(x.dtype, np.number):
        raise ValueError(message)

    q = np.array(q)[()]
    message = "`q` must be a scalar or array of numbers."
    if not np.issubdtype(q.dtype, np.number):
        raise ValueError(message)

    p = np.array(p)[()]
    message = "`p` must be a scalar or array of floats."
    if not np.issubdtype(p.dtype, np.inexact):
        raise ValueError(message)

    axis_none = axis is None
    ndim = max(x.ndim, p.ndim)
    if axis_none:
        x = np.ravel(x)
        q = np.ravel(q)
        p = np.ravel(p)
        axis = 0
    elif np.iterable(axis) or int(axis) != axis:
        message = "`axis` must be an integer or None."
        raise ValueError(message)
    elif (axis >= ndim) or (axis < -ndim):
        message = "`axis` is not compatible with the shapes of the inputs."
        raise ValueError(message)
    axis = int(axis)

    alternatives = {'two-sided', 'less', 'greater'}
    message = f"`alternative` must be one of {alternatives}"
    if alternative not in alternatives:
        raise ValueError(message)


Review comment (suggestion):
This takes care of reducing away dimensions if `_keepdims=False`, adding dimensions back if `axis=None` and `keepdims=True` was passed, and NaNing-out values based on the nan-pattern of the input (or if the working axis. has zero length).

Code after:
            high_index[~valid_index] = 0
            x_high = np.take_along_axis(x, high_index, axis=-1)
            high = np.where(valid_index, x_high, np.nan)
        elif alternative == 'greater':
            p = 1 - confidence_level
            low_index = (bd.ppf(p)).astype(int) - 1
            valid_index = low_index >= 0
            low_index[~valid_index] = 0
            x_low = np.take_along_axis(x, low_index, axis=-1)
            low = np.where(valid_index, x_low, np.nan)
            high = np.full(shape, np.inf)
        elif alternative == 'two-sided':
            p = (1 - confidence_level) / 2
            low_index = (bd.ppf(p)).astype(int) - 1
            valid_index = low_index >= 0
            low_index[~valid_index] = 0
            x_low = np.take_along_axis(x, low_index, axis=-1)
            low = np.where(valid_index, x_low, np.nan)
            high_index = bd.isf(p).astype(int)
            valid_index = high_index < n
            high_index[~valid_index] = 0
            x_high = np.take_along_axis(x, high_index, axis=-1)
            high = np.where(valid_index, x_high, np.nan)

        args = self._axis, self._axis_none, self._keepdims, self._ndim, self._nan_out
        low = _quantile_test_postprocess(low, *args)
        high = _quantile_test_postprocess(high, *args)
        return ConfidenceInterval(low, high)


def quantile_test_iv(x, q, p, alternative, axis, keepdims):

    x = np.atleast_1d(x)
    message = '`x` must be an array of numbers.'
    if not np.issubdtype(x.dtype, np.number):
        raise ValueError(message)

    q = np.array(q)[()]
    message = "`q` must be a scalar or array of numbers."
    if not np.issubdtype(q.dtype, np.number):
        raise ValueError(message)

    p = np.array(p)[()]
    message = "`p` must be a scalar or array of floats."
    if not np.issubdtype(p.dtype, np.inexact):
        raise ValueError(message)

    axis_none = axis is None
    ndim = max(x.ndim, p.ndim)
    if axis_none:
        x = np.ravel(x)