python 中函数的复杂广播

complex broadcasting to a function in python

这是基于我之前的一个问题,但更复杂,额外的并发症让我不知所措。

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

Broadcasting function calls in np arrays

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

类似于为 链接的更简单问题创建的 vectorized/Broadcasted 版本,我试图摆脱 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 时,没有传递标量,而是出现了数组。这告诉我我需要进一步填充数组(可能增加 2 个维度?)?我正在努力正确填充数组。我想我需要让数组有更多的维度,然后再压缩它们,但我不清楚该怎么做。我的无数次尝试都是错误的,并且总是以将数组发送到 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]]

解决办法是利用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 

注:

M1M2 在试运行中比原始问题更大。