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

Code before:
    Examples
    --------
    >>> from scipy import stats
    >>> 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 = 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)

Review comment (refactor):
~~Implementation below avoids inserting a `0` at the begining of `cumsum` because `insert` is not available in the standard and `concatenate` would require creating an array of zeros of the right shape. I suppose the `include_initial` argument of `cumulative_sum` would have worked - but `ssbn` needs to be initialized to an array of the right shape anyway, so might as well refactor as below.~~ Not thinking straight - this is 1d. I will refactor this back.

Code after:

    Examples
    --------
    >>> from scipy import stats
    >>> 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.