Alternative rapide pour numpy.median.reduceat

En ce qui concerne cette réponse , existe-t-il un moyen rapide de calculer les médianes sur un tableau qui a des groupes avec un nombre inégal d'éléments?

Par exemple:

data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67, ... ]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3,    ... ]

Et puis je veux calculer la différence entre le nombre et la médiane par groupe (par exemple, la médiane du groupe 0est 1.025donc le premier résultat est 1.00 - 1.025 = -0.025). Donc, pour le tableau ci-dessus, les résultats apparaissent comme:

result = [-0.025, 0.025, 0.05, -0.05, -0.19, 0.29, 0.00, 0.10, -0.10, ...]

Puisqu'il np.median.reduceatn'existe pas (encore), existe-t-il un autre moyen rapide d'y parvenir? Mon tableau contiendra des millions de lignes, la vitesse est donc cruciale!

Les indices peuvent être supposés contigus et ordonnés (il est facile de les transformer s'ils ne le sont pas).

Exemples de données pour les comparaisons de performances:

import numpy as np

np.random.seed(0)
rows = 10000
cols = 500
ngroup = 100

# Create random data and groups (unique per column)
data = np.random.rand(rows,cols)
groups = np.random.randint(ngroup, size=(rows,cols)) + 10*np.tile(np.arange(cols),(rows,1))

# Flatten
data = data.ravel()
groups = groups.ravel()

# Sort by group
idx_sort = groups.argsort()
data = data[idx_sort]
groups = groups[idx_sort]

Parfois, vous devez écrire du code numpy non idiomatique si vous voulez vraiment accélérer votre calcul, ce que vous ne pouvez pas faire avec numpy natif.

numbacompile votre code python en bas niveau C. Étant donné que beaucoup de numpy lui-même est généralement aussi rapide que C, cela finit surtout par être utile si votre problème ne se prête pas à la vectorisation native avec numpy. Ceci est un exemple (où j'ai supposé que les indices sont contigus et triés, ce qui se reflète également dans les données d'exemple):

import numpy as np
import numba

# use the inflated example of roganjosh https://stackoverflow.com/a/58788534
data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3] 

data = np.array(data * 500) # using arrays is important for numba!
index = np.sort(np.random.randint(0, 30, 4500))               

# jit-decorate; original is available as .py_func attribute
@numba.njit('f8[:](f8[:], i8[:])') # explicit signature implies ahead-of-time compile
def diffmedian_jit(data, index): 
    res = np.empty_like(data) 
    i_start = 0 
    for i in range(1, index.size): 
        if index[i] == index[i_start]: 
            continue 

        # here: i is the first _next_ index 
        inds = slice(i_start, i)  # i_start:i slice 
        res[inds] = data[inds] - np.median(data[inds]) 

        i_start = i 

    # also fix last label 
    res[i_start:] = data[i_start:] - np.median(data[i_start:])

    return res

Et voici quelques synchronisations utilisant la %timeitmagie d'IPython :

>>> %timeit diffmedian_jit.py_func(data, index)  # non-jitted function
... %timeit diffmedian_jit(data, index)  # jitted function
...
4.27 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
65.2 µs ± 1.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

En utilisant les données d'exemple mises à jour dans la question, ces nombres (c'est-à-dire le temps d'exécution de la fonction python par rapport au temps d'exécution de la fonction accélérée JIT) sont

>>> %timeit diffmedian_jit.py_func(data, groups) 
... %timeit diffmedian_jit(data, groups)
2.45 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
93.6 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Cela équivaut à une accélération de 65x dans le petit cas et une accélération de 26x dans le plus grand cas (par rapport au code en boucle lente, bien sûr) en utilisant le code accéléré. Un autre avantage est que (contrairement à la vectorisation typique avec numpy natif), nous n'avions pas besoin de mémoire supplémentaire pour atteindre cette vitesse, il s'agit de code de bas niveau optimisé et compilé qui finit par être exécuté.

La fonction ci-dessus suppose que les tableaux numpy int sont int64par défaut, ce qui n'est pas réellement le cas sous Windows. Une alternative consiste donc à supprimer la signature de l'appel à numba.njit, déclenchant une compilation juste à temps appropriée. Mais cela signifie que la fonction sera compilée lors de la première exécution, ce qui peut interférer avec les résultats de synchronisation (nous pouvons soit exécuter la fonction une fois manuellement, en utilisant des types de données représentatifs, soit simplement accepter que la première exécution de synchronisation soit beaucoup plus lente, ce qui devrait Etre ignoré). C'est exactement ce que j'ai essayé d'empêcher en spécifiant une signature, ce qui déclenche une compilation anticipée.

Quoi qu'il en soit, dans le cas JIT correctement, le décorateur dont nous avons besoin est juste

@numba.njit
def diffmedian_jit(...):

Notez que les timings ci-dessus que j'ai montrés pour la fonction compilée jit ne s'appliquent qu'une fois la fonction compilée. Cela se produit soit lors de la définition (avec une compilation désirée, lorsqu'une signature explicite est passée à numba.njit), soit pendant le premier appel de fonction (avec une compilation différée, lorsqu'aucune signature n'est transmise à numba.njit). Si la fonction ne doit être exécutée qu'une seule fois, le temps de compilation doit également être pris en compte pour la vitesse de cette méthode. Cela ne vaut généralement la peine de compiler des fonctions que si le temps total de compilation + exécution est inférieur au temps d'exécution non compilé (ce qui est en fait vrai dans le cas ci-dessus, où la fonction python native est très lente). Cela se produit principalement lorsque vous appelez souvent votre fonction compilée.

Comme le note max9111 dans un commentaire, une des caractéristiques importantes de numbaest le cachemot - clé to jit. Passer cache=Trueà numba.jitstockera la fonction compilée sur le disque, de sorte que lors de la prochaine exécution du module python donné, la fonction sera chargée à partir de là plutôt que recompilée, ce qui peut encore vous épargner l'exécution à long terme.

Une approche consisterait à utiliser Pandas ici uniquement pour faire usage de groupby. J'ai un peu gonflé les tailles d'entrée pour donner une meilleure compréhension des timings (car il y a des frais généraux dans la création du DF).

import numpy as np
import pandas as pd

data =  [1.00, 1.05, 1.30, 1.20, 1.06, 1.54, 1.33, 1.87, 1.67]
index = [0,    0,    1,    1,    1,    1,    2,    3,    3]

data = data * 500
index = np.sort(np.random.randint(0, 30, 4500))

def df_approach(data, index):
    df = pd.DataFrame({'data': data, 'label': index})
    df['median'] = df.groupby('label')['data'].transform('median')
    df['result'] = df['data'] - df['median']

Donne ce qui suit timeit :

%timeit df_approach(data, index)
5.38 ms ± 50.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Pour la même taille d'échantillon, j'obtiens le considère que l'approche dict d'Aryerez est:

%timeit dict_approach(data, index)
8.12 ms ± 3.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Cependant, si nous augmentons les entrées d'un autre facteur de 10, les synchronisations deviennent:

%timeit df_approach(data, index)
7.72 ms ± 85 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit dict_approach(data, index)
30.2 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Cependant, au détriment d'une certaine réactivité, la réponse de Divakar utilisant le numpy pur se présente comme suit:

%timeit bin_median_subtract(data, index)
573 µs ± 7.48 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

À la lumière du nouvel ensemble de données (qui aurait vraiment dû être défini au début):

%timeit df_approach(data, groups)
472 ms ± 2.52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit bin_median_subtract(data, groups) #https://stackoverflow.com/a/58788623/4799172
3.02 s ± 31.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit dict_approach(data, groups) #https://stackoverflow.com/a/58788199/4799172
<I gave up after 1 minute>

# jitted (using @numba.njit('f8[:](f8[:], i4[:]') on Windows) from  https://stackoverflow.com/a/58788635/4799172
%timeit diffmedian_jit(data, groups)
132 ms ± 3.12 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Peut-être que vous l'avez déjà fait, mais sinon, voyez si c'est assez rapide:

median_dict = {i: np.median(data[index == i]) for i in np.unique(index)}
def myFunc(my_dict, a): 
    return my_dict[a]
vect_func = np.vectorize(myFunc)
median_diff = data - vect_func(median_dict, index)
median_diff

Production:

array([-0.025,  0.025,  0.05 , -0.05 , -0.19 ,  0.29 ,  0.   ,  0.1  ,
   -0.1  ])

Voici une approche basée sur NumPy pour obtenir la médiane groupée pour les valeurs de bacs / index positives -

def bin_median(a, i):
    sidx = np.lexsort((a,i))

    a = a[sidx]
    i = i[sidx]

    c = np.bincount(i)
    c = c[c!=0]

    s1 = c//2

    e = c.cumsum()
    s1[1:] += e[:-1]

    firstval = a[s1-1]
    secondval = a[s1]
    out = np.where(c%2,secondval,(firstval+secondval)/2.0)
    return out

Pour résoudre notre cas spécifique de soustraits -

def bin_median_subtract(a, i):
    sidx = np.lexsort((a,i))

    c = np.bincount(i)

    valid_mask = c!=0
    c = c[valid_mask]    

    e = c.cumsum()
    s1 = c//2
    s1[1:] += e[:-1]
    ssidx = sidx.argsort()
    starts = c%2+s1-1
    ends = s1

    starts_orgindx = sidx[np.searchsorted(sidx,starts,sorter=ssidx)]
    ends_orgindx  = sidx[np.searchsorted(sidx,ends,sorter=ssidx)]
    val = (a[starts_orgindx] + a[ends_orgindx])/2.
    out = a-np.repeat(val,c)
    return out