复杂的广播到python中的函数

Question

这是基于我之前的问题Replacing for loops with function call inside with broadcasting/vectorized solution，但更复杂，额外的复杂性让我不知所措。

这里有一些类似的帖子，但是，在阅读了他们的问题的解决方案后，我无法解决我的问题：

Broadcasting custom function

Broadcasting function calls in np arrays

Broadcasting a function call to a a 3D array in python

Broadcasting a python function onto np arrays

这里是 MWE

def OneD(x, y, z, r_1, r_2):
    ret = 0.0
    for i in range(x+1):
        for j in range(y+1):
            m = i + j
            if m % 2 == 0:
                ret += np.exp(i+r_1)**(j+1) / (z+1+r_2)
    return ret 

def ThreeD(a,b,c, d, e):
    value = OneD(a[0],b[0], c, d[0], e[0])
    value *= OneD(a[1],b[1], c, d[1], e[1])
    value *= OneD(a[2],b[2], c, d[2], e[2])    
    return value

M1 = M2 = np.array(([[0,0,0],[0,0,1], [1,1,1], [1,0,2]]), dtype=int) 
scales0 = scales1 = np.array([1.1, 2.2, 3.3, 4.4])
cc0 = cc1 = cc2 =1.77    # for simplicity, all constants made the same
r1 = np.array([0.0,0.0,0.0])
r2 = r1 + 1.0   
results = np.zeros((4,4))

for s0, n0, in enumerate(M1):
    for s1, n1, in enumerate(M2):
        v = ThreeD(n0, n1, cc2, r1, r2)
        v *= cc0 * cc1 * scales0[s0] * scales1[s1]
        results[s0, s1] += v

类似于为start 链接的更简单问题创建的矢量化/广播版本，我试图摆脱for loops 并通过对函数OneD 执行矢量化调用来加快计算速度

v  = np.prod(OneD(M1[:,None,:], M2[None,:,:], np.arange(4)[None,:,None], r1, r2), axis=2)
results = v * cc0*cc1*np.array(scales0)[:,None]*np.array(scales1)[None,:] 

我遇到的问题是，现在，当使用OneD 时，没有传递标量，而是出现了数组。这告诉我我需要进一步填充数组（也许增加二维？）？我正在努力正确填充数组。我想我需要让数组有更多的维度，然后再挤压它们，但我不清楚如何做到这一点。我的无数次尝试都错了，并且总是不断将数组发送到OneD。对于此示例，我需要以 (4,4) 数组结尾。

[[  7.07469713e-02   1.41493943e-01   2.12240914e-01   5.65975771e-01]                                                                                                          
 [  1.41493943e-01   2.37400124e+00   3.56100187e+00   5.31397826e+00]                                                                                                          
 [  2.12240914e-01   3.56100187e+00   3.75915002e+02   6.68688926e+01]                                                                                                          
 [  2.37400124e+00   8.93002921e+00   1.12371774e+02   3.99028687e+03]]

Answer 1

解决方案是使用np.frompyfunc。可以直接替代np.vectorize，但不像np.vectorize那么慢

解决方案

myfunc3 = np.frompyfunc(OneD,5,1)

v  = np.prod(myfunc3(M1[:,None,:], M2[None,:,:], cc2, r1, r2), axis=2)
results = v * cc0*cc1*np.array(scales0)[:,None]*np.array(scales1)[None,:]

对于任何有兴趣的人，下面是一堆不同的试验方法，解决方案是最后一种方法

测试套件

import numpy as np
from timeit import default_timer as tm
############################################################################
#
#   Original code uses ThreeD and OneD in a double for loop
#     My goal is to first, get rid of ThreeD
#                   Second, get rid of for loops
#             Progress: checkmark to all. Final trial uses np.frompyfunc
#                       np.frompyfunc gets rid of double for loops.
#
###########################################################################

def OneD(x, y, z, r_1, r_2):
    ret = 0.0
    for i in range(x+1):
        for j in range(y+1):
            m = i + j
            if m % 2 == 0:
                ret += np.exp(i+r_1)**(j+1) / (z+1+r_2)
    return ret 

def ThreeD(a,b,c, d, e):
    value =  OneD(a[0],b[0], c, d[0], e[0])
    value *= OneD(a[1],b[1], c, d[1], e[1])
    value *= OneD(a[2],b[2], c, d[2], e[2])    
    return value

#######################################################
#
#   Variables used by all trials
#
#######################################################
M1 = M2 = np.array(([[0,0,0],[0,0,1], [0,1,0], [1,0,0], [2,0,0], [0,2,0],[0,0,2], [1,1,0],[1,0,1],[0,1,1]]), dtype=int) 
scales0 = scales1 = np.random.rand(len(M1))*10.0 # np.array([1.1, 2.2, 3.3, 4.4])
print(scales0)
cc0 = cc1 = cc2 =1.77    # for simplicity, all constants made the same
r1 = np.array([0.0,0.0,0.0])
r2 = r1 + 1.0   

#################################################################
#
#   Original double for loop method - base case
#  No dessert until it is slower than at least one other method
#
#################################################################
results = np.zeros((len(M1),len(M2)))
sb = tm()
for s0, n0, in enumerate(M1):
    for s1, n1, in enumerate(M2):
        v = ThreeD(n0, n1, cc2, r1, r2)
        v *= cc0 * cc1 * scales0[s0] * scales1[s1]
        results[s0, s1] += v
e1 = tm()   
print(results)
answer_longway = np.sum(results)
print("Value for long way is {} and time is {}".format(answer_longway, e1-sb) )

######################################################################
######################################################################
#
#                      np.vectorize (a slow way)
#
######################################################################
######################################################################
results = np.zeros((len(M1),len(M2)))
myfunc = np.vectorize(OneD)
v = 0
s2 = tm()
for s0, n0, in enumerate(M1):
    for s1, n1, in enumerate(M2):
        v = np.prod(myfunc(n0, n1, cc2, r1, r2))
        v *= cc0 * cc1 * scales0[s0] * scales1[s1]
        results[s0, s1] += v
e2 = tm()
answer_np_vectorize = np.sum(results)
print("Value for np.vectorize way is {} and time is {}".format(answer_np_vectorize, e2-s2))

######################################################################
######################################################################
#
#                      np.frompyfunc (a same as loop way)
#
######################################################################
######################################################################
results = np.zeros((len(M1),len(M2)))
v = 0
myfunc2 = np.frompyfunc(OneD,5,1)
s3 = tm()
for s0, n0, in enumerate(M1):
    for s1, n1, in enumerate(M2):
        v = np.prod(myfunc2(n0, n1, cc2, r1, r2))
        v *= cc0 * cc1 * scales0[s0] * scales1[s1]
        results[s0, s1] += v
e3 = tm()
answer_np_vectorize = np.sum(results)
print("Value for np.frompyfunc way is {} and time is {}".format(answer_np_vectorize, e3-s3))

######################################################################
######################################################################
#
#                      broadcasted (a fast way?)
#
######################################################################
######################################################################
results = np.zeros((len(M1),len(M2)))
myfunc3 = np.frompyfunc(OneD,5,1)

s4 = tm()
v  = np.prod(myfunc3(M1[:,None,:], M2[None,:,:], cc2, r1, r2), axis=2)
results = v * cc0*cc1*np.array(scales0)[:,None]*np.array(scales1)[None,:]
e4 = tm()

answer_broadcast = np.sum(results)
print("Value for broadcasted way is {} and time is {}".format(answer_broadcast, e4-s4))

样本输出是

Value for long way is 2069.471054878208 and time is 0.00458109400351532                                                       
Value for np.vectorize way is 2069.471054878208 and time is 0.01092407900316175                                               
Value for np.frompyfunc way is 2069.471054878208 and time is 0.00398122602999792434                                             
Value for broadcasted way is 2069.471054878208 and time is 0.0023705440033460036 

注意：

M1 和M2 在试运行中比原来的问题更大。