When using broadcasting, rather than broadcasting scalars to match the arrays, the vectorized function is instead, for some reason, shrinking the arrays to scalars.
Below is a MWE. It contains a double for loop. I am having trouble writing faster code that does not use the for loops, but instead, uses broadcasting/vectorized numpy.
import numpy as np
def OneD(x, y, z):
ret = np.exp(x)**(y+1) / (z+1)
return ret
def ThreeD(a,b,c):
value = OneD(a[0],b[0], c)
value *= OneD(a[1],b[1], c)
value *= OneD(a[2],b[2], c)
return value
M_1 = M_2 = [[0,0,0],[0,0,1], [1,1,1], [1,0,2]]
scales0 = scales1 = [1.1, 2.2, 3.3, 4.4]
cc0 = cc1 = 1.77
results = np.zeros((4,4))
for s0, n0, in enumerate(M_1):
for s1, n1, in enumerate(M_2):
v = ThreeD(n0, n1, s1)
v *= cc0 * cc1 * scales0[s0] * scales1[s1]
results[s0, s1] += v
While I want to remove both for loops, to keep it simple I am trying first to get rid of the inner loop. Feel free to answer with both removed however.
Here is how I changed the loop
rr = [0,1,2,3]
myfun = np.vectorize(ThreeD)
for s0, n0, in enumerate(M_1):
#for s1, n1, in enumerate(M_2):
v = myfun(n0, M_2, rr)
v *= cc0 * cc1 * scales0[s0] * scales1[rr]
results[s0, rr] += v
Traceback (most recent call last):
File "main.py", line 36, in <module>
v = myfun(n0, M_2, rr)
File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 1573, in __call__
return self._vectorize_call(func=func, args=vargs)
File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 1633, in _vectorize_call
ufunc, otypes = self._get_ufunc_and_otypes(func=func, args=args)
File "/usr/lib/python3/dist-packages/numpy/lib/function_base.py", line 1597, in _get_ufunc_and_otypes
outputs = func(*inputs)
File "main.py", line 18, in ThreeD
value = OneD(a[0],b[0], c)
IndexError: invalid index to scalar variable.
Do I also need to vectorize the OneD function? I was hoping by vectorizing the ThreeD function, it would do the proper bookkeeping.
In your loops, n0 and n1 are elements of the nested M_ lists, each 3 elements.
In [78]: ThreeD(np.arange(3),np.arange(3),3)
Out[78]: 46.577468547527005
OneD works with arrays, so can get the full n lists/arrays:
In [79]: OneD(np.arange(3), np.arange(3),3)
Out[79]: array([ 0.25 , 1.84726402, 100.85719837])
In [80]: np.prod(_)
Out[80]: 46.577468547527005
and the product matches ThreeD.
Looking just at the ThreeD part of your double loop:
In [81]: for s0, n0, in enumerate(M_1):
...: for s1, n1, in enumerate(M_2):
...: print(n0,n1,s1, ThreeD(n0, n1, s1))
...:
[0, 0, 0] [0, 0, 0] 0 1.0
[0, 0, 0] [0, 0, 1] 1 0.125
[0, 0, 0] [1, 1, 1] 2 0.037037037037037035
[0, 0, 0] [1, 0, 2] 3 0.015625
[0, 0, 1] [0, 0, 0] 0 2.718281828459045
...
[1, 0, 2] [1, 0, 2] 3 46.577468547527005
Making arrays from your lists:
In [82]: M1 = np.array(M_1); M2 = np.array(M_2)
In [83]: M1.shape
Out[83]: (4, 3)
I replicate those ThreeD results with this broadcasted call:
In [87]: np.prod(OneD(M1[:,None,:], M2[None,:,:], np.arange(4)[None,:,None]), axis=2)
Out[87]:
array([[1.00000000e+00, 1.25000000e-01, 3.70370370e-02, 1.56250000e-02],
[2.71828183e+00, 9.23632012e-01, 2.73668744e-01, 3.13836514e-01],
[2.00855369e+01, 6.82476875e+00, 1.49418072e+01, 6.30357490e+00],
[2.00855369e+01, 1.85516449e+01, 1.49418072e+01, 4.65774685e+01]])
I am passing (4,1,3), (1,4,3) and (1,4,1) arrays to OneD. The result is (4,4,3), which I then multiply on the last axis to make a (4,4).
The rest of the calculation is:
In [88]: (cc0*cc1*np.array(scales0)[:,None]*np.array(scales1)[None,:])
Out[88]:
array([[ 3.790809, 7.581618, 11.372427, 15.163236],
[ 7.581618, 15.163236, 22.744854, 30.326472],
[11.372427, 22.744854, 34.117281, 45.489708],
[15.163236, 30.326472, 45.489708, 60.652944]])
In [89]: _87*_88 # multiplying these two 4x4 arrays
Out[89]:
array([[3.79080900e+00, 9.47702250e-01, 4.21201000e-01, 2.36925563e-01],
[2.06089744e+01, 1.40052502e+01, 6.22455564e+00, 9.51755427e+00],
[2.28421302e+02, 1.55228369e+02, 5.09773834e+02, 2.86747781e+02],
[3.04561737e+02, 5.62605939e+02, 6.79698445e+02, 2.82506059e+03]])
which matches `results:
In [90]: results
Out[90]:
array([[3.79080900e+00, 9.47702250e-01, 4.21201000e-01, 2.36925563e-01],
[2.06089744e+01, 1.40052502e+01, 6.22455564e+00, 9.51755427e+00],
[2.28421302e+02, 1.55228369e+02, 5.09773834e+02, 2.86747781e+02],
[3.04561737e+02, 5.62605939e+02, 6.79698445e+02, 2.82506059e+03]])