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

Code before:
    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)
def mood(x, y, axis=0, alternative="two-sided"):
    """Perform Mood's test for equal scale parameters.

    Mood's two-sample test for scale parameters is a non-parametric
    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:


Review comment (suggestion):
This formula can be implemented just as it appears in the paper:
<img width="688" height="126" alt="image" src="https://github.com/user-attachments/assets/0989618e-7b90-400a-bafc-e3125ff9b945" />

because terms where `t` is `0.` aren't included in the sum.

Code after:
    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
    # the location of the ^2) at the beginning of "Mood's Squared Rank Test" (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=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