介绍
我有一个函数func
这是可矢量化的,我使用它对其进行矢量化np.frompyfunc
。而不是使用嵌套for
循环,我只想调用它一次,因此我需要用以下内容填充输入np.newaxis
's.
我的目标是摆脱两个嵌套for
循环并使用numpy.array
而是广播功能。
这是 MWEfor
循环(我想摆脱 for 循环,而是填充变量c_0
, c_1
, rn_1
, rn_2
, and factor
打电话时myfunc
.
感兴趣问题的 MWE
for i, b_0 in enumerate(b_00):
for j, b_1 in enumerate(b_11):
factor = b_0 + b_1
rn = (b_0 * coord + b_1 * coord2) / factor
rn_1 = coord - rn
rn_2 = coord2 - rn
results = np.add( results,np.prod(myfunc(c_0[:,None,:], c_1[None,:,:], rn_1, rn_2, factor), axis=2) )
上面的显式 for 循环是正确的,并且包含在块代码中以保证质量。
我目前的努力
factor = b_00[:, None] + b_11[None, :]
rn = np.add( (b_00[:,None] * coord[None,:])[:, None, :], (b_11[:,None] * coord2[None,:])[None, :, :] ) / factor[:,:,None]
rn_1 = coord - rn
rn_2 = coord2 - rn
results2 = np.prod(myfunc(c_0[:,None,:, None, None], c_1[None,:,:, None, None], rn_1[:,:,:],rn_2[:,:, :], factor[None,:, :, None]), axis=2)
results2 = np.squeeze(results2)
results2 = np.sum(results2, axis=(2,3))
整体代码块
import numpy as np
myfunc = np.frompyfunc(func,5,1)
################################################################
# Prep arrays needed for MWE
################################################################
results = np.zeros((5,3))
coord = np.array([1,1,2])
coord2 = np.array([3,3,3])
c_0 = np.array([[0,0,2],[0,2,0],[2,0,0],[1,1,0], [1,0,1]])
c_1 = np.array([[0,0,2],[0,2,0],[2,0,0]])
b_00 = np.array([2.2, 1.1, 4.4]) # np.array([2.2, 3.3, 40.4])
b_11 = np.array([1.2, 3.3]) # np.array([1.2, 5.3])
################################################################
# This is only for comparison. `results` is the correct answer
################################################################
for i, b_0 in enumerate(b_00):
for j, b_1 in enumerate(b_11):
factor = b_0 + b_1
rn = (b_0 * coord + b_1 * coord2) / factor
rn_1 = coord - rn
rn_2 = coord2 - rn
results = np.add( results,np.prod(myfunc(c_0[:,None,:], c_1[None,:,:], rn_1, rn_2, factor), axis=2) )
################################################################
# Prep for broadcasting (My attempt)
################################################################
factor = b_00[:, None] + b_11[None, :]
rn = np.add( (b_00[:,None] * coord[None,:])[:, None, :], (b_11[:,None] * coord2[None,:])[None, :, :] ) / factor[:,:,None]
rn_1 = coord - rn
rn_2 = coord2 - rn
# The following all get the same *wrong* answer
# results2 = np.prod(myfunc(c_0[:,None,:,None, None], c_1[None,:,:, None, None], rn_1[:, None, None],rn_2[:,None, None], factor[None,:, :]), axis=2)
# results2 = np.prod(myfunc(c_0[:,None,:, None, None, None], c_1[None,:,:, None, None, None], rn_1[None, None,:,:,:, None],rn_2[None, None,:,:, :, None], factor[None,:, :, None, None]), axis=2)
# results2 = np.prod(myfunc(c_0[:,None,:, None, None], c_1[None,:,:, None, None], rn_1[None, None,:,:,:],rn_2[None, None,:,:, :], factor[None,:, :, None]), axis=2)
# this should be the only line needing work!
results2 = np.prod(myfunc(c_0[:,None,:, None, None], c_1[None,:,:, None, None], rn_1[:,:,:],rn_2[:,:, :], factor[None,:, :, None]), axis=2)
results2 = np.squeeze(results2)
results2 = np.sum(results2, axis=(2,3))
assert np.allclose(results, results2)
# Vectorized function to be sent broadcasted arrays
def func(r0, r1, x, y, at):
val = 0.0
for i in range(r0+1):
for j in range(r1+1):
val += x + i*j + at * y
return val
Problems
- 通过上面的代码,我得到了结果数组的正确形状(
results2
是我在广播方面的尝试,并且results
是给出正确答案的缓慢 for 循环),但它有错误的值。
- As
@hpaulj
指出,如果我改变尺寸b_00
如果长度为 4(或任何更大的长度),我的解决方案甚至无法获得正确的形状。
Update
请确保它适用于当前的b_00 = np.array([2.2, 1.1, 4.4])
以及更一般的b_00 = np.array([2.2, 1.1, 4.4, 5.1, 6.2])
。我想要一个广播解决方案,但会接受一个比实际解决方案更快的解决方案for loops
.