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

Code before:
    test for the null hypothesis that two samples are drawn from the
    same distribution with the same scale parameter.

    Parameters
    ----------
    x, y : array_like
        Arrays of sample data. There must be at least three observations
        total.
    axis : int, optional
        The axis along which the samples are tested.  `x` and `y` can be of
        different length along `axis`.
        If `axis` is None, `x` and `y` are flattened and the test is done on
        all values in the flattened arrays.
    alternative : {'two-sided', 'less', 'greater'}, optional
        Defines the alternative hypothesis. Default is 'two-sided'.
        The following options are available:

        * 'two-sided': the scales of the distributions underlying `x` and `y`
          are different.
        * 'less': the scale of the distribution underlying `x` is less than
          the scale of the distribution underlying `y`.
        * 'greater': the scale of the distribution underlying `x` is greater
          than the scale of the distribution underlying `y`.

        .. versionadded:: 1.7.0

    Returns
    -------
    res : SignificanceResult
        An object containing attributes:

        statistic : scalar or ndarray
            The z-score for the hypothesis test.  For 1-D inputs a scalar is
            returned.
        pvalue : scalar ndarray
            The p-value for the hypothesis test.

    See Also
    --------
    fligner : A non-parametric test for the equality of k variances
    ansari : A non-parametric test for the equality of 2 variances
    bartlett : A parametric test for equality of k variances in normal samples
    levene : A parametric test for equality of k variances

    Notes
    -----
    The data are assumed to be drawn from probability distributions ``f(x)``
    and ``f(x/s) / s`` respectively, for some probability density function f.
    The null hypothesis is that ``s == 1``.

    For multi-dimensional arrays, if the inputs are of shapes

Review comment (performance):
This finds the locations where the relevant `a_i` need to go. It might be easiest to think through how `rankdata` would have found `t` to begin with: we concatenated the samples, and then rankdata produces `t` that has the multiplicity of each element in a sorted order (see https://github.com/scipy/scipy/pull/23837#discussion_r2491120035). So if `x` starts in a sorted order, and then concatenate `y` (important: _append_ the elements of `y`), and then do a _stable_ argsort, then the indices tell us where the elements of `x` (and therefore the elemnts of `a` below) ended up.

It's probably possible to refactor the code (including code internal to `_rankdata`) so that we don't have to re-calculate these indices. That would be more computationally efficient since we avoid sorting the same data twice, but I don't really want to pick apart `_rankdata` or have it return any other information at this time.

Code after:
    def sum_I2(a, b=None):
        return (a * (a + 1) * (2 * a + 1) / 6 if b is None
                else sum_I2(b) - sum_I2(a) + a**2)

    def sum_I(a, b=None):
        return (a * (a + 1) / 2 if b is None
                else sum_I(b) - sum_I(a) + a)

    def sum_1(a, b):
        return (b - a) + 1

    with np.errstate(invalid='ignore', divide='ignore'):
        sum_I2 = sum_I2(S_i_m1 + 1, S_i)
        sum_I = sum_I(S_i_m1 + 1, S_i)
        sum_1 = sum_1(S_i_m1 + 1, S_i)
        c = (N + 1) / 2
        phi = (sum_I2 - 2*c*sum_I + sum_1*c**2) / t

    phi[t == 0] = 0  # where t = 0 we get NaNs; eliminate them

    # Mielke pg 312 defines `a` as the count of elements in sample `x` for each of the
    # unique values in the combined sample. The tricky thing is getting these to line
    # up with the locations of nonzero elements in `t`/`phi`.
    x = np.sort(x, axis=-1)
    xy = np.concatenate((x, y), axis=-1)
    i = np.argsort(xy, stable=True, axis=-1)
    _, a = _stats_py._rankdata(x, method='average', return_ties=True)
    zeros = np.zeros(a.shape[:-1] + (n,))
    a = np.concatenate((a, zeros), axis=-1)
    a = np.take_along_axis(a, i, axis=-1)

    # Mielke pg 312 defines test statistic `T` as the inner product `a` and `phi`
    T = np.vecdot(a, phi, axis=-1)

    return (T - E_0_T) / np.sqrt(varM)


def _mood_statistic_no_ties(r, m, n, N):
    rx = r[..., :m]
    M = np.sum((rx - (N + 1.0) / 2) ** 2, axis=-1)
    mnM = m * (N * N - 1.0) / 12
    varM = m * n * (N + 1.0) * (N + 2) * (N - 2) / 180
    return (M - mnM) / sqrt(varM)


def _mood_too_small(samples, kwargs, axis=-1):
    x, y = samples
    m = x.shape[axis]
    n = y.shape[axis]
    N = m + n
    return N < 3