自动屏蔽给定操作的numpy数组

Question

例如，我有两个numpy数组

import numpy as np

A = np.arange(5*3*3*2).reshape(5, 3, 3, 2)
B = np.arange(3*3).reshape(3, 3)

如果我想在共享轴上添加A和B，我会这样做

C = A + B[None, :, :, None]
# C has shape (5, 3, 3, 2) which is what I want

我想编写一个写函数来概括这种求和，但不是入门方法。看起来像

def mask(M, Mshape, out_shape):
    # not sure what to put here
    pass

def add_tensors(A, B, Ashape, Bshape, out_shape):
    # Here I mask A, B so that it has shape out_shape
    A = mask(A, Aaxis, out_shape)
    B = mask(B, Baxis, out_shape)
    return A + B

有什么建议吗？是否可以将其设为ufunc？

Answer 1

In [447]: A = np.arange(5*3*3*2).reshape(5, 3, 3, 2) 
     ...: B = np.arange(3*3).reshape(3, 3)                                                             

这些都是等效的：

In [448]: A + B[None,:, :, None];                                                                      
In [449]: A + B[:, :, None];         # initial None is automatic                                                                           

从列表构建索引元组：

In [454]: tup = [slice(None)]*3; tup[-1] = None; tup = tuple(tup)                                      
In [455]: tup                                                                                          
Out[455]: (slice(None, None, None), slice(None, None, None), None)
In [456]: A + B[tup];                                           

或等效形状：

In [457]: sh = B.shape + (1,)                                                                          
In [458]: sh                                                                                           
Out[458]: (3, 3, 1)
In [459]: A + B.reshape(sh); 

[expand_dims还使用参数化的reshape：

In [462]: np.expand_dims(B,2).shape                                                                    
Out[462]: (3, 3, 1)
In [463]: A+np.expand_dims(B,2);