利用广播提高此减法效率

Question

我有一个形状为 (N, T, d) 的数组 x。我有两个函数 f 和 g，它们都接受一个形状为 (some_dimension, d) 的数组，并返回一个形状为 (some_dimension, ) 的数组。

我想对整个 x 计算 f。这很简单：只需使用 f(x.reshape(-1, d))。

接下来，我想只对第二个维度的第一个切片计算 g，也就是说，使用 g(x[:, 0, :])，然后将这个结果从对所有维度的 f 计算中减去。下面的代码展示了这个过程。

最小示例 - 低效的方法

import numpy as np

# Reproducibility
seed = 1234
rng = np.random.default_rng(seed=seed)

# Generate x
N = 100
T = 10
d = 2
x = rng.normal(loc=0.0, scale=1.0, size=(N, T, d))

# In practice the functions are not this simple
def f(x):
    return x[:, 0] + x[:, 1]

def g(x):
    return x[:, 0]**2 - x[:, 1]**2

# Compute f on all the (flattened) array
fx = f(x.reshape(-1, d)).reshape(N, T)

# Compute g only on the first slice of second dimension. Here are two ways of doing so
gx = np.tile(g(x[:, 0])[:, None], reps=(1, T))
gx = np.repeat(g(x[:, 0]), axis=0, repeats=T).reshape(N, T)

# Finally compute what I really want to compute
diff = fx - gx

有没有更高效的方法？我觉得使用广播（broadcasting）应该可以，但我想不出来。

Answer 1

我们把例子简化一下，这样可以更好地查看(5,4)的数组：

In [138]: 
     ...: # Generate x
     ...: N = 5
     ...: T = 4
     ...: d = 2
     ...: x = np.arange(40).reshape(N,T,d) #(rng.normal(loc=0.0, scale=1.0, size=(N, T, d))
     ...: 
     ...: # In practice the functions are not this simple
     ...: def f(x):
     ...:     return x[:, 0] + x[:, 1]
     ...: 
     ...: def g(x):
     ...:     return x[:, 0]**2 - x[:, 1]**2
     ...: 
     ...: # Compute f on all the (flattened) array
     ...: fx = f(x.reshape(-1, d)).reshape(N, T)
     ...: 
     ...: # Compute g only on the first slice of second dimension. Here are two ways of doing so
     ...: gx1 = np.tile(g(x[:, 0])[:, None], reps=(1, T))
     ...: gx2 = np.repeat(g(x[:, 0]), axis=0, repeats=T).reshape(N, T)

In [139]: fx.shape,gx1.shape,gx2.shape
Out[139]: ((5, 4), (5, 4), (5, 4))

这里的fx中的所有元素都不一样，所以不能再进行“广播”了。

In [140]: fx
Out[140]: 
array([[ 1,  5,  9, 13],
       [17, 21, 25, 29],
       [33, 37, 41, 45],
       [49, 53, 57, 61],
       [65, 69, 73, 77]])

你使用的tile和repeat其实是做同样的事情。tile是用repeat实现的，所以并没有增加什么：

In [141]: gx1
Out[141]: 
array([[ -1,  -1,  -1,  -1],
       [-17, -17, -17, -17],
       [-33, -33, -33, -33],
       [-49, -49, -49, -49],
       [-65, -65, -65, -65]])

In [142]: gx2
Out[142]: 
array([[ -1,  -1,  -1,  -1],
       [-17, -17, -17, -17],
       [-33, -33, -33, -33],
       [-49, -49, -49, -49],
       [-65, -65, -65, -65]])

gx只是把5个g()的值重复了4次。

In [143]: g(x[:, 0])
Out[143]: array([ -1, -17, -33, -49, -65])

In [144]: fx-gx1
Out[144]: 
array([[  2,   6,  10,  14],
       [ 34,  38,  42,  46],
       [ 66,  70,  74,  78],
       [ 98, 102, 106, 110],
       [130, 134, 138, 142]])

所以gx可以用一个(5,1)的数组来替代，这样它就可以和(5,4)的fx进行广播了：

In [145]: fx-g(x[:,0])[:,None]
Out[145]: 
array([[  2,   6,  10,  14],
       [ 34,  38,  42,  46],
       [ 66,  70,  74,  78],
       [ 98, 102, 106, 110],
       [130, 134, 138, 142]])

我还没有搞清楚我之前提到的T和d维度之间的区别。

这个回答可能有点啰嗦，但它展示了我如何想象和发现一个broadcasting的解决办法。