仅使用数字'0'和'1'将'matrix[：，mask]`转换为'matrix[mask，：]`吗？

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]

有没有办法通过只使用数字0和1将matrix[:,mask]转换为matrix[mask,:]？所以我只能有一个函数来计算行和列的真实平均值

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