Arrays numpy装箱平均值，保留额外轴

似乎我被

numpy

的以下问题困住了

我有一个带有形状的数组

X

：

X.shape=（nexp、ntime、ndim、npart）

我需要根据

binvals

（以及一些

bin

）中的值，沿着

npart

维度计算此数组上的装箱统计信息，但保留所有其他维度，因为我必须使用装箱统计信息来消除原始数组

X

中的一些偏差。装箱值具有shape

binvals.shape=（nexp、ntime、npart）

一个完整的，最小的例子，来解释我想做什么。请注意，实际上，我使用的是大型阵列和数百个存储箱（因此此实现需要花费很长时间）：

看看这一结果可能会更清楚吗

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\u statistic

，但我似乎无法理解如何将其用于此目的。谢谢

好的，我想我明白了，主要是基于@jdehesa的回答

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)

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]

它给出了与我的原始代码相同的结果。 在我在这里的示例中使用的小型阵列上，此解决方案的速度大约是原始代码的三倍。我希望在更大的阵列上速度更快

更新：

---

事实上，它在更大的阵列上速度更快（没有进行任何正式测试），但尽管如此，它在性能方面还是达到了可接受的水平。。。如果您对额外的矢量导航有任何进一步的建议，我们将非常欢迎。

您能提供一些虚拟输入吗？您的意思是什么？任何事情都可以做到：

X=np.random.randn（120）.重塑（3,4,2,5）

，

binvals=np.random.randn（24）.重塑（3,4,2）

和

bins=np.linspace（binvals.min（），binvals.max（），10）

我得到的是

索引器：布尔索引与发布代码上的样本数据不匹配。

，那是因为我把宾瓦尔的形状弄错了。它应该是：

X=np.random.randn（120）.重塑（3,4,2,5）

，

binvals=np.random.randn（60）.重塑（3,4,5）

和

bins=np.linspace（binvals.min（），binvals.max（），10）

。如果确定一些输入和预期输出，解决这些问题就会容易得多。嗯，我必须多考虑一下，但在我看来，这并不是我想要的东西。@user6760680我添加了一个没有循环的替代解决方案（应该更快），以牺牲更多内存为代价。我花了一段时间才明白什么不能说服我，但重点是你正在使用，虽然我必须计算一个数组上的统计信息，但要存储一个不同的数组。@user6760680好吧，我明白你的意思，我误解了这个问题，我会解决它。顺便说一句，答案在最后使用了，这显然是限制在32个不同的选择（所以你最多只能使用32个存储箱）。。。如果您需要更多，您将需要为最后一步选择不同的路径…我也更新了我的答案。我的代码没有给出相同的结果，但是。。。当我运行它时，它会产生接近于零的值（我想这就是点），而原始代码会产生高达+/-6的值（这很奇怪，因为

X

值位于

[0,1]

）。。。我不知道有什么区别！以防对您有用…@jdehesa X值来自标准正态分布，因此它们不限于[0,1]。我检查了我的代码，它做了我需要的，即使可能没有我希望的那么快。无论如何，非常感谢您的建议，它非常有用，至少可以显著提高性能！

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]