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

Code before:

@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:

        * '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


Review comment (refactor):
It is not SO bad to compare the implementation of the code above to the math if you assume take `Si` to represent $S_i$ and `S_i_m1` to represent $S_{i-1}$ at face value. The tricky thing is those zeros in `t` and corresponding duplicate elements in `S_i` and `S_i_m1`. But it all works out:
- The unique terms we care about are calculated because `t`, `S_i`, and `S_i_m1` are all aligned as we expect.
- The terms we don't care about turn into NaNs because at these locations, `t == 0` and `S_i == S_i_m1`.
- We replace the NaNs with zeros so they don't bite us later.

Code after:

    # 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

    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)