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

Code before:
    >>> y = [2, 4, 6, 8, 10]
    >>> stats.kruskal(x, y)
    KruskalResult(statistic=0.2727272727272734, pvalue=0.6015081344405895)

    >>> x = [1, 1, 1]
    >>> y = [2, 2, 2]
    >>> z = [2, 2]
    >>> stats.kruskal(x, y, z)
    KruskalResult(statistic=7.0, pvalue=0.0301973834223185)

    """
    num_groups = len(samples)
    if num_groups < 2:
        raise ValueError("Need at least two groups in stats.kruskal()")

    n = np.asarray([sample.shape[-1] for sample in samples])
    totaln = np.sum(n)

    if np.any(n) < 1:  # Only needed for `test_axis_nan_policy`
        raise ValueError("Inputs must not be empty.")

    alldata = np.concatenate(samples, axis=-1)
    ranked, t = _rankdata(alldata, method='average', return_ties=True)

    # ties = tiecorrect(ranked)
    ties = 1 - (t**3 - t).sum(axis=-1) / (totaln**3 - totaln)

    # Compute sum^2/n for each group and sum
    j = np.cumsum(n, axis=-1)
    ssbn = np.sum(ranked[..., :j[0]], axis=-1)**2 / n[0]
    for i in range(1, num_groups):
        ssbn += np.sum(ranked[..., j[i-1]:j[i]], axis=-1)**2 / n[i]

    h = 12.0 / (totaln * (totaln + 1)) * ssbn - 3 * (totaln + 1)
    df = num_groups - 1
    h /= ties

    chi2 = _SimpleChi2(df)
    pvalue = _get_pvalue(h, chi2, alternative='greater', symmetric=False, xp=np)
    return KruskalResult(h, pvalue)


FriedmanchisquareResult = namedtuple('FriedmanchisquareResult',
                                     ('statistic', 'pvalue'))


@xp_capabilities(np_only=True)
@_axis_nan_policy_factory(FriedmanchisquareResult, n_samples=None, paired=True)
def friedmanchisquare(*samples):
    """Compute the Friedman test for repeated samples.


Review comment (suggestion):
Compare against
https://github.com/scipy/scipy/blob/de80faf9d3480b9dbb9b888568b64499e0e70c19/scipy/stats/_stats_py.py#L8097-L8102
Note that if `size < 2`, the `_axis_nan_policy` decorator is going to produce NaNs for us, so the special case is not needed.

Code after:
    >>> x = [1, 3, 5, 7, 9]
    >>> y = [2, 4, 6, 8, 10]
    >>> stats.kruskal(x, y)
    KruskalResult(statistic=0.2727272727272734, pvalue=0.6015081344405895)

    >>> x = [1, 1, 1]
    >>> y = [2, 2, 2]
    >>> z = [2, 2]
    >>> stats.kruskal(x, y, z)
    KruskalResult(statistic=7.0, pvalue=0.0301973834223185)

    """
    num_groups = len(samples)
    if num_groups < 2:
        raise ValueError("Need at least two groups in stats.kruskal()")

    n = [sample.shape[-1] for sample in samples]
    totaln = sum(n)
    if np.any(n) < 1:  # Only needed for `test_axis_nan_policy`
        raise ValueError("Inputs must not be empty.")

    alldata = np.concatenate(samples, axis=-1)
    ranked, t = _rankdata(alldata, method='average', return_ties=True)
    ties = 1 - (t**3 - t).sum(axis=-1) / (totaln**3 - totaln)  # tiecorrect(ranked)

    # Compute sum^2/n for each group and sum
    j = list(itertools.accumulate(n, initial=0))
    ssbn = sum(np.sum(ranked[..., j[i]:j[i + 1]], axis=-1)**2 / n[i]
               for i in range(num_groups))

    h = 12.0 / (totaln * (totaln + 1)) * ssbn - 3 * (totaln + 1)
    df = num_groups - 1
    h /= ties

    chi2 = _SimpleChi2(df)
    pvalue = _get_pvalue(h, chi2, alternative='greater', symmetric=False, xp=np)
    return KruskalResult(h, pvalue)


FriedmanchisquareResult = namedtuple('FriedmanchisquareResult',
                                     ('statistic', 'pvalue'))


@xp_capabilities(np_only=True)
@_axis_nan_policy_factory(FriedmanchisquareResult, n_samples=None, paired=True)
def friedmanchisquare(*samples):
    """Compute the Friedman test for repeated samples.

    The Friedman test tests the null hypothesis that repeated samples of
    the same individuals have the same distribution.  It is often used
    to test for consistency among samples obtained in different ways.