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

Code before:

    def sum_I(a, b):
        return (b * (b + 1) - a * (a - 1)) / 2

    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] + (m,))
    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[..., :n]
    M = np.sum((rx - (N + 1.0) / 2) ** 2, axis=-1)
    mnM = n * (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
    n = x.shape[axis]
    m = y.shape[axis]
    N = m + n
    return N < 3


@xp_capabilities(np_only=True)
@_axis_nan_policy_factory(SignificanceResult, n_samples=2, too_small=_mood_too_small)

Review comment (suggestion):
This puts them in the order that corresponds with `phi` - along with zeros for the terms we don't care about.

Code after:
        def func(x):
            return xp.mean(x, axis=-1, keepdims=True)

    else:  # center == 'trimmed'

        def func(x):
            # keepdims=True doesn't currently work for lazy arrays
            return _stats_py.trim_mean(x, proportiontocut, axis=-1)[..., xp.newaxis]

    ni = [sample.shape[-1] for sample in samples]
    N = sum(ni)

    # Implementation follows [3] pg 355 F-K.
    Xibar = [func(sample) for sample in samples]
    Xij_Xibar = [xp.abs(sample - Xibar_) for sample, Xibar_ in zip(samples, Xibar)]
    Xij_Xibar = xp.concat(Xij_Xibar, axis=-1)
    ranks = _stats_py._rankdata(Xij_Xibar, method='average', xp=xp)
    ranks = xp.astype(ranks, dtype)
    a_Ni = special.ndtri(ranks / (2*(N + 1.0)) + 0.5)

    # [3] Equation 2.1
    splits = list(itertools.accumulate(ni, initial=0))
    Ai = [a_Ni[..., i:j] for i, j in zip(splits[:-1], splits[1:])]
    Aibar = [xp.mean(Ai_, axis=-1) for Ai_ in Ai]
    abar = xp.mean(a_Ni, axis=-1)
    V2 = xp.var(a_Ni, axis=-1, correction=1)
    statistic = sum(ni_ * (Aibar_ - abar)**2 for ni_, Aibar_ in zip(ni, Aibar)) / V2

    chi2 = _SimpleChi2(xp.asarray(k-1, dtype=dtype))
    pval = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=xp)
    return FlignerResult(statistic, pval)


def _mood_statistic_with_ties(x, y, t, m, n, N):
    # First equation of "Mood's Squared Rank Test", Mielke pg 313
    E_0_T = m * (N * N - 1) / 12

    # m, n, N, t, and S are defined in the second paragraph of Mielke pg 312
    # The only difference is that our `t` has zeros interspersed with the relevant
    # numbers to keep the array rectangular, but these terms add nothing to the sum.
    np = array_namespace(x, y, t)  # needed to use `cumulative_sum`
    S = np.cumulative_sum(t, include_initial=True, axis=-1)
    S_i, S_i_m1 = S[..., 1:], S[..., :-1]
    # Second equation of "Mood's Squared Rank Test", Mielke pg 313
    varM = (m * n * (N + 1.0) * (N**2 - 4) / 180
            - m * n / (180 * N * (N - 1))
            * np.sum(t * (t ** 2 - 1) * (t ** 2 - 4 + (15 * (N - S_i - S_i_m1) ** 2)),
                     axis=-1))

    # There is a formula for Phi (`phi` in code) in terms of t, S, and Psi(I) at the
    # bottom of Mielke pg 312. Psi(I) = [I - (N+1)/2]^2 is defined (with a mistake in