numpy binned signifie, conserver des axes supplémentaires

Question

Il semble que je suis bloqué sur le problème suivant avec numpy.

je un tableau X avec la forme: X.shape = (nexp, ntime, ndim, npart) J'ai besoin pour calculer les statistiques regroupées par casiers sur ce tableau le long npart dimension, selon les valeurs binvals (et certains bins), mais en gardant toutes les autres dimensions, parce que j'ai pour utiliser la statistique binned pour supprimer certains biais dans le tableau d'origine X. Les valeurs Binning ont la forme binvals.shape = (nexp, ntime, npart).

Un exemple complet et minimal pour expliquer ce que j'essaie de faire. Notez que, en réalité, je travaille sur de grands tableaux et avec plusieurs hunderds de bacs (donc cette mise en œuvre prend une éternité):

import numpy as np 

np.random.seed(12345) 

X = np.random.randn(24).reshape(1,2,3,4) 
binvals = np.random.randn(8).reshape(1,2,4) 
bins = [-np.inf, 0, np.inf] 
nexp, ntime, ndim, npart = X.shape 

cleanX = np.zeros_like(X) 
for ne in range(nexp): 
    for nt in range(ntime): 
     indices = np.digitize(binvals[ne, nt, :], bins) 
     for nd in range(ndim): 
      for nb in range(1, len(bins)): 
       inds = indices==nb 
       cleanX[ne, nt, nd, inds] = X[ne, nt, nd, inds] - \ 
        np.mean(X[ne, nt, nd, inds], axis = -1) 

En regardant les résultats de cela peut le rendre plus clair?

In [8]: X 
Out[8]: 
array([[[[-0.20470766, 0.47894334, -0.51943872, -0.5557303 ], 
     [ 1.96578057, 1.39340583, 0.09290788, 0.28174615], 
     [ 0.76902257, 1.24643474, 1.00718936, -1.29622111]], 

     [[ 0.27499163, 0.22891288, 1.35291684, 0.88642934], 
     [-2.00163731, -0.37184254, 1.66902531, -0.43856974], 
     [-0.53974145, 0.47698501, 3.24894392, -1.02122752]]]]) 

In [10]: cleanX 
Out[10]: 
array([[[[ 0.  , 0.67768523, -0.32069682, -0.35698841], 
     [ 0.  , 0.80405255, -0.49644541, -0.30760713], 
     [ 0.  , 0.92730041, 0.68805503, -1.61535544]], 

     [[ 0.02303938, -0.02303938, 0.23324375, -0.23324375], 
     [-0.81489739, 0.81489739, 1.05379752, -1.05379752], 
     [-0.50836323, 0.50836323, 2.13508572, -2.13508572]]]]) 


In [12]: binvals 
Out[12]: 
array([[[ -5.77087303e-01, 1.24121276e-01, 3.02613562e-01, 
      5.23772068e-01], 
     [ 9.40277775e-04, 1.34380979e+00, -7.13543985e-01, 
      -8.31153539e-01]]]) 

Existe-t-il une solution vectorisée? J'ai pensé utiliser scipy.stats.binned_statistic, mais il semble que je ne sois pas capable de comprendre comment l'utiliser pour ce but. Merci!

Answer 1

Ok, je pense que je l'ai eu, principalement basé sur la réponse par @jdehesa.

clean2 = np.zeros_like(X) 
d = np.digitize(binvals, bins) 
for i in range(1, len(bins)): 
    m = d == i 
    minds = np.where(m) 
    sl = [*minds[:2], slice(None), minds[2]] 
    msum = m.sum(axis=-1) 
    clean2[sl] = (X - \ 
        (np.sum(X * m[...,np.newaxis,:], axis=-1)/
        msum[..., np.newaxis])[..., np.newaxis])[sl] 

Ce qui donne les mêmes résultats que mon code original. Sur les petits tableaux que j'ai dans l'exemple ici, cette solution est environ trois fois plus rapide que le code original. Je m'attends à ce que ce soit beaucoup plus rapide sur les tableaux plus grands.

Mise à jour:

En effet, il est plus rapide sur des réseaux plus importants (n'a pas fait de test formel), mais malgré cela, il atteint seulement le niveau d'acceptable en termes de performance ... toute suggestion plus loin des vectoriztaions supplémentaires seraient les bienvenus.

Answer 2

import numpy as np 

np.random.seed(100) 

nexp = 3 
ntime = 4 
ndim = 5 
npart = 100 
nbins = 4 

binvals = np.random.rand(nexp, ntime, npart) 
X = np.random.rand(nexp, ntime, ndim, npart) 
bins = np.linspace(0, 1, nbins + 1) 

d = np.digitize(binvals, bins)[:, :, np.newaxis, :] 
r = np.arange(1, len(bins)).reshape((-1, 1, 1, 1, 1)) 
m = d[np.newaxis, ...] == r 
counts = np.sum(m, axis=-1, keepdims=True).clip(min=1) 
means = np.sum(X[np.newaxis, ...] * m, axis=-1, keepdims=True)/counts 
cleanX = X - np.choose(d - 1, means)