首页 > 解决方案 > 复杂的广播到python中的函数

问题描述

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

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

广播自定义功能

在 np 数组中广播函数调用

在python中广播对aa 3D数组的函数调用

将 python 函数广播到 np 数组上

这是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

类似于为开始时链接的更简单问题创建的矢量化/广播版本,我试图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]]

标签: pythonpython-3.xnumpyvectorizationarray-broadcasting

解决方案


解决方案是利用np.frompyfunc. 它可以用作 的直接替代品np.vectorize,但不像 hello 那样慢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在试运行中比原始问题更大。


推荐阅读