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

Code before:

    >>> sample1 = stats.norm.rvs(size=105, loc=0.5, random_state=rng)
    >>> stats.ks_2samp(sample1, sample2, alternative='less')
    KstestResult(statistic=0.4055137844611529,
                 pvalue=3.5474563068855554e-08,
                 statistic_location=-0.13249370614972575,
                 statistic_sign=-1)

    and indeed, with p-value smaller than our threshold, we reject the null
    hypothesis in favor of the alternative.

    """
    # because of the _axis_nan_policy decorator, we can assume the arrays
    # are broadcasted and `axis=-1`
    mode = method

    if mode not in ['auto', 'exact', 'asymp']:
        raise ValueError(f'Invalid value for mode: {mode}')
    alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get(
        alternative.lower()[0], alternative)
    if alternative not in ['two-sided', 'less', 'greater']:
        raise ValueError(f'Invalid value for alternative: {alternative}')
    MAX_AUTO_N = 10000  # 'auto' will attempt to be exact if n1,n2 <= MAX_AUTO_N

    data1 = np.sort(data1, axis=-1, stable=False)
    data2 = np.sort(data2, axis=-1, stable=False)
    n1 = data1.shape[-1]
    n2 = data2.shape[-1]
    if min(n1, n2) == 0:
        raise ValueError('Data passed to ks_2samp must not be empty')
    n = n1 + n2

    data_all = np.concatenate([data1, data2], axis=-1)
    ranks, data_all = _rankdata(data_all, method='min')  # assumes axis=-1

    cdf_counts1 = np.diff(ranks[..., :n1], prepend=1, append=n + 1, axis=-1)
    cdf_counts2 = np.diff(ranks[..., -n2:], prepend=1, append=n + 1, axis=-1)
    cdf1_vals, cdf1_counts = np.broadcast_arrays(np.linspace(0, 1, n1 + 1), cdf_counts1)
    cdf2_vals, cdf2_counts = np.broadcast_arrays(np.linspace(0, 1, n2 + 1), cdf_counts2)
    cdf1 = np.repeat(np.ravel(cdf1_vals), np.ravel(cdf1_counts), axis=-1)
    cdf2 = np.repeat(np.ravel(cdf2_vals), np.ravel(cdf2_counts), axis=-1)
    cddiffs = np.reshape(cdf1 - cdf2, ranks.shape[:-1] + (-1,))

    # Identify the location of the statistic
    argminS = np.argmin(cddiffs, axis=-1, keepdims=True)
    argmaxS = np.argmax(cddiffs, axis=-1, keepdims=True)
    loc_minS = np.squeeze(np.take_along_axis(data_all, argminS, axis=-1), axis=-1)
    loc_maxS = np.squeeze(np.take_along_axis(data_all, argmaxS, axis=-1), axis=-1)

    # Ensure sign of minS is not negative.
    minS = -np.squeeze(np.take_along_axis(cddiffs, argminS, axis=-1), axis=-1)

Review comment (suggestion):
```suggestion
    data1 = np.sort(data1, axis=-1)
    data2 = np.sort(data2, axis=-1)
```
Oops, guess I can't do this until array API kicks in.

Code after:
    Suppose, however, that the first sample were drawn from
    a normal distribution shifted toward greater values. In this case,
    the cumulative density function (CDF) of the underlying distribution tends
    to be *less* than the CDF underlying the second sample. Therefore, we would
    expect the null hypothesis to be rejected with ``alternative='less'``:

    >>> sample1 = stats.norm.rvs(size=105, loc=0.5, random_state=rng)
    >>> stats.ks_2samp(sample1, sample2, alternative='less')
    KstestResult(statistic=0.4055137844611529,
                 pvalue=3.5474563068855554e-08,
                 statistic_location=-0.13249370614972575,
                 statistic_sign=-1)

    and indeed, with p-value smaller than our threshold, we reject the null
    hypothesis in favor of the alternative.

    """
    # because of the _axis_nan_policy decorator, we can assume the arrays
    # are broadcasted and `axis=-1`
    mode = method

    if mode not in ['auto', 'exact', 'asymp']:
        raise ValueError(f'Invalid value for mode: {mode}')
    alternative = {'t': 'two-sided', 'g': 'greater', 'l': 'less'}.get(
        alternative.lower()[0], alternative)
    if alternative not in ['two-sided', 'less', 'greater']:
        raise ValueError(f'Invalid value for alternative: {alternative}')
    MAX_AUTO_N = 10000  # 'auto' will attempt to be exact if n1,n2 <= MAX_AUTO_N

    data1 = np.sort(data1, axis=-1)
    data2 = np.sort(data2, axis=-1)
    n1 = data1.shape[-1]
    n2 = data2.shape[-1]
    if min(n1, n2) == 0:
        raise ValueError('Data passed to ks_2samp must not be empty')
    n = n1 + n2

    data_all = np.concatenate([data1, data2], axis=-1)
    batch_shape = data_all.shape[:-1]

    # Previously, this part of the code was just:
    # cdf1 = _xp_searchsorted(data1, data_all, side='right') / n1
    # cdf2 = _xp_searchsorted(data2, data_all, side='right') / n2
    # cddiffs = cdf1 - cdf2
    # but the switch from `np.searchsorted` to `_xp_searchsorted` would come with a
    # pretty steep performance hit. When `np.searchsorted` is vectorized, we can
    # probably use that. In the meantime, the following algorithm is typically faster.

    # We want the ECDF of each sample evaluated at *all* the points in the pooled
    # sample. The values each ECDF can assume are given by:
    cdf1_vals = np.broadcast_to(np.linspace(0, 1, n1 + 1), batch_shape + (n1 + 1,))