Search code examples
pythonperformancevectorization

how to efficiently vectorize a function and retrieve values from it


Suppose I have a very computationally expensive function f(x). I want to compute some values of it, and then just access them instead of evaluating the function every time with new x values.

See the following simple example to illustrate what I mean:

import numpy as np

x = np.linspace(-3, 3, 6001)
fx = x**2
x = np.round(x, 3)

#I want to evaluate the function for the following w:
w = np.random.rand(10000)

#Rounding is necessary, so that the w match the x.
w = np.round(w, 3)

fx_w = []
for i in range(w.size):
    fx_w.append(fx[x==w[i]])

fx_w = np.asarray(fx_w)

So, I'd like to have f(w) computed from the values already generated for x. Of course, a for loop is out of the question, so my question is: how can I implement this somewhat efficiently?


Solution

  • You can use searchsorted to find the corresponding indices of your prepared function array. This will be an approximation. Rounding is not necessary.

    import numpy as np
    np.random.seed(42)
    
    x = np.linspace(-3, 3, 6001)
    
    fx = x ** 2
    
    w = np.random.rand(10000)
    result = fx[np.searchsorted(x, w)]
    
    print('aprox. F(x):', result)
    print('real F(x):', w ** 2)
    

    Output

    aprox. F(x): [0.140625   0.904401   0.535824   ... 0.896809   0.158404   0.047524  ]
    real F(x):   [0.1402803  0.90385769 0.53581513 ... 0.89625588 0.1579967  0.04714996]
    

    Your function has to be much more computationally intensive to justify this approach

    %timeit fx[np.searchsorted(x, w)] #1000 loops, best of 5: 992 µs per loop
    %timeit w ** 2                    #100000 loops, best of 5: 3.81 µs per loop