numpy binned mean、余分な軸を節約

Question

numpyで次の問題が発生しているようです。

I形状の配列X有する：私はbinvalsの値（及びいくつかのbins）によれば、npart寸法に沿って、このアレイ上のビニング統計を計算する必要があるが、私は持っているので、そこにすべての他の寸法を維持X.shape = (nexp, ntime, ndim, npart) をbinned統計量を使用して元の配列Xの偏りを除去します。ビニング値は、形状がbinvals.shape = (nexp, ntime, npart)です。

私が何をしようとしているかを説明する完全な、最小限の例。 （ので、この実装は永遠にかかります）実際に、私は大規模な配列にし、ビンのいくつかのhunderdsで働いている、ということに注意してください：この結果を見ると

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) 

それをより明確にできますか？

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]]]) 

ベクター化されたソリューションはありますか？私はscipy.stats.binned_statisticを使用することを考えましたが、私はこの目的のためにどのように使用するのか理解できないようです。ありがとう！

Answer 1

私は、主に@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] 

これは私の元のコードと同じ結果を示します。 私がここに示した小さな配列では、この解決策は元のコードの約3倍高速です。私はそれがより大きい配列でより速いと予想します。

更新：

確かにそれは（正式なテストをしなかった）より速く、より大きなアレイ上だが、これにもかかわらず、それだけでパフォーマンスの面で許容できるレベルに達して...上の任意の更なる提案を余分なベクトル化は大歓迎です。

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)