在不使用for或while的情况下重写numpy中的循环函数

Question

我正试图使用numpy库重现以下函数，我想在不使用关键字for或while的情况下生成一个等效的定义。我猜你需要使用广播、更新和从numpy重塑。但我对numpy和不使用"；对于"；或"；而"；对我来说是一件令人费解的事，尤其是在尝试使用嵌套循环时。

def _bcast(x):
x1, x2 = x
y = np.empty(x1.shape)
for i in range(x1.shape[0]):
for j in range(x1.shape[1]):
for k in range(x1.shape[2]):
y[i,j,k] = (x1[i,j,k]+4)*(4*x2[j,k] - 4)
return y

---

def _bcast_ax(x):
x1, x2 = x
y = np.empty((x1.shape[0], x2.shape[0], x2.shape[1]))
for i in range(x1.shape[0]):
for j in range(x2.shape[0]):
for k in range(x2.shape[1]):
y[i,j,k] = (4+x1[i,k])*(4*x2[j,k]-4)
return y

---

def bcast(x):

return (x1+4) * (4*x2 -4) 

---

def bcast_ax(x):

return (x**2)*(x[1]*2)*(x[2]**4) 

我尝试为这两个函数执行以下操作，但它们不起作用。

为了澄清，我需要这个测试通过bcast和bcast，以产生相同的结果。与_bcast_ax和bcast_ax 相同

def test_bcast(self):
def _bcast(x):
x1, x2 = x
y = np.empty(x1.shape)
for i in range(x1.shape[0]):
for j in range(x1.shape[1]):
for k in range(x1.shape[2]):
y[i,j,k] = (x1[i,j,k]+4)*(4*x2[j,k] - 4)
return y
X = [(np.random.randn(3,4,5), np.random.randn(4,5)) for _ in range(3)]
self._test_fun(ac.bcast, _bcast, X)

Answer 1

专注于

y[i,j,k] = (x1[i,j,k]+4)*(4*x2[j,k] - 4)

这意味着y和x1具有相同的形状。CCD_ 3具有相同的最后2个维度。我们可以重塑x2，使其具有新的领先维度x2[None,...]

y = (x1+4)*(4*x2[None,...] - 4)

但根据广播规则，新的领先维度是自动

y = (x1+4)*(4*x2-4)

应该起作用。

关键是要理解broadcasting。

测试

In [169]: x1, x2 = np.arange(24).reshape(2,3,4), np.arange(12).reshape(3,4)
In [170]:             y = np.empty(x1.shape)
...:             for i in range(x1.shape[0]):
...:                 for j in range(x1.shape[1]):
...:                     for k in range(x1.shape[2]):
...:                         y[i,j,k] = (x1[i,j,k]+4)*(4*x2[j,k] - 4)
...: 
In [171]: y
Out[171]: 
array([[[ -16.,    0.,   24.,   56.],
[  96.,  144.,  200.,  264.],
[ 336.,  416.,  504.,  600.]],
[[ -64.,    0.,   72.,  152.],
[ 240.,  336.,  440.,  552.],
[ 672.,  800.,  936., 1080.]]])
In [172]: (x1+4)*(4*x2-4)
Out[172]: 
array([[[ -16,    0,   24,   56],
[  96,  144,  200,  264],
[ 336,  416,  504,  600]],
[[ -64,    0,   72,  152],
[ 240,  336,  440,  552],
[ 672,  800,  936, 1080]]])