在numpy中从联合pmf计算条件概率，速度太慢。有什么建议？

Question

我有一个联合概率质量函数数组，比如说它的形状是 (1,2,3,4,5,6)，我想计算一个概率表，条件是某些维度的值（导出条件概率表），这样可以帮助我做决策。

目前我写的代码是这样的（输入是一个字典 "vdict"，格式是 {'variable_1': value_1, 'variable_2': value_2 ... }）

for i in vdict:
   dim = self.invardict.index(i) # The index of the dimension that our Variable resides in
   val = self.valdict[i][vdict[i]] # The value we want it to be
   d = d.swapaxes(0, dim)
   **d = array([d[val]])**
   d = d.swapaxes(0, dim)

...

所以，我现在的做法是：

1. 我把变量转换成对应的条件概率表中的维度。
2. 我把第零个轴和之前找到的轴交换位置。
3. 我用想要的值替换整个零轴。

然后，我把维度放回到原来的位置。

现在，问题是，为了完成第二步，我需要 (a.) 计算一个子数组，(b.) 把它放进一个列表，然后再转换成数组，这样我才能得到新的数组。

问题是，粗体字部分意味着我在创建新的对象，而不是直接使用旧对象的引用。如果维度 d 非常大（我正好遇到这种情况），而且使用 d 的方法被调用很多次（这也是我遇到的情况），那么整个过程就会变得非常慢。

所以，有没有人想出什么办法可以替代这段代码，并且运行得更快？也许有什么方法可以让我在原地计算条件概率。

注意：我必须保持原始的轴顺序（或者至少要知道在移除某个轴时如何更新变量到维度的字典）。我希望不需要使用自定义的数据类型。

Answer 1

好的，我自己找到了解决办法，经过一番尝试，发现了numpy的数组操作可以直接在原地进行。

我把循环中的最后三行改成了：

    d = conditionalize(d, dim, val)

其中，conditionalize的定义是：

    def conditionalize(arr, dim, val):
        arr = arr.swapaxes(dim, 0)
        shape = arr.shape[1:]       # shape of the sub-array when we omit the desired dimension.
        count = array(shape).prod() # count of elements omitted the desired dimension.
        arr = arr.reshape(array(arr.shape).prod()) # flatten the array in-place.
        arr = arr[val*count:(val+1)*count] # take the needed elements
        arr = arr.reshape((1,)+shape) # the desired sub-array shape.
        arr = arr. swapaxes(0, dim)   # fix dimensions

        return arr

这样一来，我程序的执行时间从15分钟减少到了6秒，效果真是太明显了。

希望这能帮助到遇到同样问题的人。