为给定操作自动屏蔽 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);

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

Automatically masking a numpy array for a given operation

python

numpy

numpy-ufunc

array-broadcasting