0

Problem:

When using broadcasting, rather than broadcasting scalars to match the arrays, the vectorized function is instead, for some reason, shrinking the arrays to scalars.

MWE:

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.

Failed Attempt:

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

Error Message:

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.

Charlie Crown
  • 1,071
  • 2
  • 11
  • 28
  • I am starting to wonder if using np.vectorize will make this faster, even once it does work... – Charlie Crown Jun 14 '20 at 01:37
  • What kinds of speedup does `np.vectorize` promise? – hpaulj Jun 14 '20 at 01:45
  • None, that is the problem. So I am at a complete loss on how to speed this up. – Charlie Crown Jun 14 '20 at 01:45
  • `np.vectorize` does broadcast the inputs against each other, passing scalar values to the function. Your `OneD` is already fully 'vectorized'. It will work `x,y,z` of any dimension, provided that 'broadcast'. Without examining it carefully, I suspect `ThreeD` can be something like `OneD(a, b, c).sum(axis=0)`. – hpaulj Jun 14 '20 at 01:59
  • I have played around with vectorizing the `ThreeD`, but it actually slowed it down abit. I had `np.prod(OneD)` where I had used `np.vectorize(OneD)` prior to calling `OneD`. That is okay though. My main concern is getting rid of the double `for loops` entirely if I can. This requires broadcasting and my brain is not...broadcasting. – Charlie Crown Jun 14 '20 at 02:25

1 Answers1

1

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]])
hpaulj
  • 221,503
  • 14
  • 230
  • 353