algos_rank_helper.pxi.in

"""
Template for each `dtype` helper function for rank

WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
"""

#----------------------------------------------------------------------
# rank_1d, rank_2d
#----------------------------------------------------------------------

{{py:

# dtype ctype pos_nan_value neg_nan_value
dtypes = [('object', 'object', 'Infinity()', 'NegInfinity()'),
          ('float64', 'float64_t', 'np.inf', '-np.inf'),
          ('uint64', 'uint64_t', '', ''),
          ('int64', 'int64_t', 'np.iinfo(np.int64).max',
           'np.iinfo(np.int64).min')]

}}

{{for dtype, ctype, pos_nan_value, neg_nan_value in dtypes}}


@cython.wraparound(False)
@cython.boundscheck(False)
{{if dtype == 'object'}}


def rank_1d_{{dtype}}(object in_arr, bint retry=1, ties_method='average',
                      ascending=True, na_option='keep', pct=False):
{{else}}


def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
                      na_option='keep', pct=False):
{{endif}}
    """
    Fast NaN-friendly version of scipy.stats.rankdata
    """

    cdef:
        Py_ssize_t i, j, n, dups = 0, total_tie_count = 0

        {{if dtype == 'object'}}
        ndarray sorted_data, values
        {{else}}
        ndarray[{{ctype}}] sorted_data, values
        {{endif}}

        ndarray[float64_t] ranks
        ndarray[int64_t] argsorted

        {{if dtype == 'uint64'}}
        {{ctype}} val
        {{else}}
        {{ctype}} val, nan_value
        {{endif}}

        float64_t sum_ranks = 0
        int tiebreak = 0
        bint keep_na = 0
        float count = 0.0
    tiebreak = tiebreakers[ties_method]

    {{if dtype == 'float64'}}
    values = np.asarray(in_arr).copy()
    {{elif dtype == 'object'}}
    values = np.array(in_arr, copy=True)

    if values.dtype != np.object_:
        values = values.astype('O')
    {{else}}
    values = np.asarray(in_arr)
    {{endif}}

    keep_na = na_option == 'keep'

    {{if dtype != 'uint64'}}
    if ascending ^ (na_option == 'top'):
        nan_value = {{pos_nan_value}}
    else:
        nan_value = {{neg_nan_value}}

    {{if dtype == 'object'}}
    mask = lib.isnaobj(values)
    {{elif dtype == 'float64'}}
    mask = np.isnan(values)
    {{elif dtype == 'int64'}}
    mask = values == iNaT
    {{endif}}

    np.putmask(values, mask, nan_value)
    {{endif}}

    n = len(values)
    ranks = np.empty(n, dtype='f8')

    {{if dtype == 'object'}}
    try:
        _as = values.argsort()
    except TypeError:
        if not retry:
            raise

        valid_locs = (~mask).nonzero()[0]
        ranks.put(valid_locs, rank_1d_object(values.take(valid_locs), 0,
                                             ties_method=ties_method,
                                             ascending=ascending))
        np.putmask(ranks, mask, np.nan)
        return ranks
    {{else}}
    if tiebreak == TIEBREAK_FIRST:
        # need to use a stable sort here
        _as = values.argsort(kind='mergesort')
        if not ascending:
            tiebreak = TIEBREAK_FIRST_DESCENDING
    else:
        _as = values.argsort()
    {{endif}}

    if not ascending:
        _as = _as[::-1]

    sorted_data = values.take(_as)
    argsorted = _as.astype('i8')

    {{if dtype == 'object'}}
    for i in range(n):
        sum_ranks += i + 1
        dups += 1

        val = util.get_value_at(sorted_data, i)

        if (val is nan_value) and keep_na:
            ranks[argsorted[i]] = nan
            continue

        count += 1.0

        if (i == n - 1 or
            are_diff(util.get_value_at(sorted_data, i + 1), val)):
            if tiebreak == TIEBREAK_AVERAGE:
                for j in range(i - dups + 1, i + 1):
                    ranks[argsorted[j]] = sum_ranks / dups
            elif tiebreak == TIEBREAK_MIN:
                for j in range(i - dups + 1, i + 1):
                    ranks[argsorted[j]] = i - dups + 2
            elif tiebreak == TIEBREAK_MAX:
                for j in range(i - dups + 1, i + 1):
                    ranks[argsorted[j]] = i + 1
            elif tiebreak == TIEBREAK_FIRST:
                raise ValueError('first not supported for non-numeric data')
            elif tiebreak == TIEBREAK_FIRST_DESCENDING:
                for j in range(i - dups + 1, i + 1):
                    ranks[argsorted[j]] = 2 * i - j - dups + 2
            elif tiebreak == TIEBREAK_DENSE:
                total_tie_count += 1
                for j in range(i - dups + 1, i + 1):
                    ranks[argsorted[j]] = total_tie_count
            sum_ranks = dups = 0
    {{else}}
    with nogil:
        for i in range(n):
            sum_ranks += i + 1
            dups += 1

            val = sorted_data[i]

            {{if dtype != 'uint64'}}
            if (val == nan_value) and keep_na:
                ranks[argsorted[i]] = nan
                continue
            {{endif}}

            count += 1.0

            if i == n - 1 or sorted_data[i + 1] != val:
                if tiebreak == TIEBREAK_AVERAGE:
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = sum_ranks / dups
                elif tiebreak == TIEBREAK_MIN:
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = i - dups + 2
                elif tiebreak == TIEBREAK_MAX:
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = i + 1
                elif tiebreak == TIEBREAK_FIRST:
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = j + 1
                elif tiebreak == TIEBREAK_FIRST_DESCENDING:
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = 2 * i - j - dups + 2
                elif tiebreak == TIEBREAK_DENSE:
                    total_tie_count += 1
                    for j in range(i - dups + 1, i + 1):
                        ranks[argsorted[j]] = total_tie_count
                sum_ranks = dups = 0
    {{endif}}
    if pct:
        return ranks / count
    else:
        return ranks


def rank_2d_{{dtype}}(object in_arr, axis=0, ties_method='average',
                      ascending=True, na_option='keep', pct=False):
    """
    Fast NaN-friendly version of scipy.stats.rankdata
    """

    cdef:
        Py_ssize_t i, j, z, k, n, dups = 0, total_tie_count = 0

        {{if dtype == 'object'}}
        Py_ssize_t infs
        {{endif}}

        ndarray[float64_t, ndim=2] ranks
        {{if dtype == 'int64' or dtype == 'uint64'}}
        ndarray[{{ctype}}, ndim=2, cast=True] values
        {{else}}
        ndarray[{{ctype}}, ndim=2] values
        {{endif}}

        ndarray[int64_t, ndim=2] argsorted

        {{if dtype == 'uint64'}}
        {{ctype}} val
        {{else}}
        {{ctype}} val, nan_value
        {{endif}}

        float64_t sum_ranks = 0
        int tiebreak = 0
        bint keep_na = 0
        float count = 0.0

    tiebreak = tiebreakers[ties_method]

    keep_na = na_option == 'keep'

    in_arr = np.asarray(in_arr)

    if axis == 0:
        values = in_arr.T.copy()
    else:
        values = in_arr.copy()

    {{if dtype == 'object'}}
    if values.dtype != np.object_:
        values = values.astype('O')
    {{endif}}

    {{if dtype != 'uint64'}}
    if ascending ^ (na_option == 'top'):
        nan_value = {{pos_nan_value}}
    else:
        nan_value = {{neg_nan_value}}

    {{if dtype == 'object'}}
    mask = lib.isnaobj2d(values)
    {{elif dtype == 'float64'}}
    mask = np.isnan(values)
    {{elif dtype == 'int64'}}
    mask = values == iNaT
    {{endif}}

    np.putmask(values, mask, nan_value)
    {{endif}}

    n, k = (<object> values).shape
    ranks = np.empty((n, k), dtype='f8')

    {{if dtype == 'object'}}
    try:
        _as = values.argsort(1)
    except TypeError:
        values = in_arr
        for i in range(len(values)):
            ranks[i] = rank_1d_object(in_arr[i], ties_method=ties_method,
                                      ascending=ascending, pct=pct)
        if axis == 0:
            return ranks.T
        else:
            return ranks
    {{else}}
    if tiebreak == TIEBREAK_FIRST:
        # need to use a stable sort here
        _as = values.argsort(axis=1, kind='mergesort')
        if not ascending:
            tiebreak = TIEBREAK_FIRST_DESCENDING
    else:
        _as = values.argsort(1)
    {{endif}}

    if not ascending:
        _as = _as[:, ::-1]

    values = _take_2d_{{dtype}}(values, _as)
    argsorted = _as.astype('i8')

    for i in range(n):
        {{if dtype == 'object'}}
        dups = sum_ranks = infs = 0
        {{else}}
        dups = sum_ranks = 0
        {{endif}}

        total_tie_count = 0
        count = 0.0
        for j in range(k):
            {{if dtype != 'object'}}
            sum_ranks += j + 1
            dups += 1
            {{endif}}

            val = values[i, j]

            {{if dtype != 'uint64'}}
            {{if dtype == 'object'}}
            if (val is nan_value) and keep_na:
            {{else}}
            if (val == nan_value) and keep_na:
            {{endif}}
                ranks[i, argsorted[i, j]] = nan

                {{if dtype == 'object'}}
                infs += 1
                {{endif}}

                continue
            {{endif}}

            count += 1.0

            {{if dtype == 'object'}}
            sum_ranks += (j - infs) + 1
            dups += 1
            {{endif}}

            {{if dtype == 'object'}}
            if j == k - 1 or are_diff(values[i, j + 1], val):
            {{else}}
            if j == k - 1 or values[i, j + 1] != val:
            {{endif}}
                if tiebreak == TIEBREAK_AVERAGE:
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = sum_ranks / dups
                elif tiebreak == TIEBREAK_MIN:
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = j - dups + 2
                elif tiebreak == TIEBREAK_MAX:
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = j + 1
                elif tiebreak == TIEBREAK_FIRST:
                    {{if dtype == 'object'}}
                    raise ValueError('first not supported '
                                     'for non-numeric data')
                    {{else}}
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = z + 1
                    {{endif}}
                elif tiebreak == TIEBREAK_FIRST_DESCENDING:
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = 2 * j - z - dups + 2
                elif tiebreak == TIEBREAK_DENSE:
                    total_tie_count += 1
                    for z in range(j - dups + 1, j + 1):
                        ranks[i, argsorted[i, z]] = total_tie_count
                sum_ranks = dups = 0
        if pct:
            ranks[i, :] /= count
    if axis == 0:
        return ranks.T
    else:
        return ranks

{{endfor}}