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

Code before:
    pval = _get_pvalue(statistic, chi2, alternative='greater', symmetric=False, xp=np)
    return FlignerResult(statistic, pval)


def _mood_statistic_with_ties(x, y, t, n, m, N):
    # First equation of "Mood's Squared Rank Test", Mielke pg 313
    E_0_T = n * (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 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 in terms of t, S, and Psi(I) = [I - (N+1)/2]^2 is
    # defined (with a mistake in the location of the ^2) at the beginning of
    # "Mood's Squared Rank Test", Mielke pg 313. To vectorize this calculation,
    # let c = (N + 1) / 2, so Psi(I) = I^2 - 2*c*I + c^2. We sum each of these three
    # parts of Psi separately using formulas for sums from a to b (inclusive) of
    # terms I^2, I, and 1 where I takes on successive integers.
    def sum_I2(a, b):
        return (b * (b + 1) * (2*b + 1) - a * (a - 1) * (2*a - 1)) / 6

    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)

Review comment (suggestion):
This comment is about one of the most significant changes required to vectorize this: use summation formulas rather than summing with NumPy.

Code after:
    """
    xp = array_namespace(*samples)

    if center not in ['mean', 'median', 'trimmed']:
        raise ValueError("center must be 'mean', 'median' or 'trimmed'.")

    k = len(samples)
    if k < 2:
        raise ValueError("Must provide at least two samples.")

    samples = xp_promote(*samples, force_floating=True, xp=xp)
    dtype = samples[0].dtype

    # Handle empty input
    for sample in samples:
        if sample.size == 0:
            NaN = _get_nan(*samples, xp=xp)
            return FlignerResult(NaN, NaN)

    if center == 'median':

        def func(x):
            return (xp.median(x, axis=-1, keepdims=True)
                    if (is_numpy(xp) or is_dask(xp))
                    else stats.quantile(x, 0.5, axis=-1, keepdims=True))

    elif center == 'mean':

        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:])]