python - 仅使用数字 `matrix[:,mask]` 和 `matrix[mask,:]` 将 `0` 转置为 `1` ？

Question

我只想编写一个函数来计算 matrix 的每一行或每一列的真实平均值(对行或列中的数字求平均值时不计算零元素) 。我尝试使用 axis 来控制是按行计算还是按列计算参数为1或0 ，分别。

这是按列计算的函数

def true_mean(matrix, axis):
  countnonzero = (matrix!=0).sum(axis) 
  mask = countnonzero!=0
  output_mat = np.zeros(matrix.T.shape[axis])
  output_mat[mask] = matrix[:,mask].sum(axis)/countnonzero[mask] # line4
  return output_mat

测试功能

eachPSM = np.ones([5,4])
eachPSM[0] = 0
eachPSM[2,2:4] = 5
print each PSM
> [[ 0.  0.  0.  0.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  5.  5.]
 [ 1.  1.  1.  1.]
 [ 1.  1.  1.  1.]]
ans = true_mean(eachPSM,0)
print ans
> [ 1.  1.  2.  2.]

但是，如果我想按行计算( axis = 1 )，则只需 line4必须更改为

output_mat[mask] = matrix[mask,:].sum(axis)/countnonzero[mask]

有没有办法翻转matrix[:,mask]至matrix[mask,:]仅使用号码 0和1 ？因此我只能使用一个函数来计算行和列的真实平均值。

Answer 1

您可以使用 [] 运算符将元组作为输入参数的事实:

indexer = [slice(None), slice(None)]
indexer[axis] = mask

print(x[indexer])

slice(None) 相当于 :，因此我们构造一个采用完整矩阵 [:, :] 的元组并替换使用掩码输入所需的轴。

完整示例:

import numpy as np

x = np.arange(9).reshape(3, 3)
mask = np.array([True, False, True])

for axis in [0, 1]:
    indexer = [slice(None)] * x.ndim
    indexer[axis] = mask    
    print(x[indexer])

打印

[[0 1 2]
 [6 7 8]]

和

[[0 2]
 [3 5]
 [6 8]]