scoreatpercentiles vectorized

try_scoreatpercentiles.py

# -*- coding: utf-8 -*-
"""scoreatpercentiles vectorized

Created on Mon Dec 24 09:29:31 2012

Author: Josef Perktold
"""

import numpy as np

def percentiles(x, per, interpolation_method='fraction', issorted=False, 
                axis=0):
    '''score_at_percentiles vectorized
    '''
    TINY = 1e-15
    if not issorted:
        xsorted = np.sort(x, axis=axis)
    else:
        xsorted = x
    scalar = np.isscalar(per)
    per = np.atleast_1d(per)
    
    per_invalid = (per < 0) | (per > 100)
    if per_invalid.any():
        per = np.clip(per, 0, 100)

    n = (xsorted.shape[axis] - 1)
    idx = per / 100. * n

    if interpolation_method == 'lower':
        idx = np.floor(idx + TINY)
    elif interpolation_method == 'higher':
        idx = np.ceil(idx - TINY)
    elif interpolation_method == 'fraction':
        pass # keep idx as fraction and interpolate
    else:
        raise ValueError("interpolation_method can only be 'fraction', " \
                         "'lower' or 'higher'")

    return_shape = list(xsorted.shape)
    return_shape[axis] = len(per)
    
    sumval = np.empty(return_shape)
    sumval.fill(np.nan)
    
    i = idx.astype(int)

    indexer1 = [slice(None)] * xsorted.ndim
    indexer1[axis] = np.atleast_1d(i)
    indexer2 = [slice(None)] * xsorted.ndim
    indexer2[axis] = np.atleast_1d(np.minimum(i+1, n-1))
    
    weights1 = np.array((i+1 - idx), float)
    weights2 = np.array((idx - i), float)

    mask = (i == idx)
    weights1[mask] = 1
    weights2[mask] = 0
    
    wshape = [1] * xsorted.ndim
    wshape[axis] = len(per)
    weights1.shape = wshape
    weights2.shape = wshape
    
    
    result = xsorted[indexer1] * weights1 + xsorted[indexer2] * weights2
    result /= (weights1 + weights2)
    
    if per_invalid.any():
        sl = indexer1 = [slice(None)] * result.ndim
        sl[axis] = per_invalid
        result[sl] = np.nan
    
    if result.size == 1 and scalar:
        return result.item()

    return result


x = np.arange(11)
np.random.shuffle(x)
per = [0, 10, 21, 25, 29.9, 50, 75, 90, 100]
interpolation_method = 'fraction'
axis = 0
print percentiles(x, per, interpolation_method, axis=axis)
print percentiles(x, 50, interpolation_method, axis=axis)
print percentiles(x, 10*np.sqrt(3)**2, interpolation_method, axis=axis)

print percentiles(np.repeat([x],3,0), per, interpolation_method, axis=1)
print percentiles(np.repeat([x],3,0).T, per, interpolation_method, axis=0)
print percentiles(np.repeat([x],3,0).T, per, 'lower', axis=0)
print percentiles(np.repeat([x],3,0).T, per, 'higher', axis=0)
print percentiles(np.repeat([x],3,0).T, [-5, 50, 120], 'higher', axis=0)