Vectorize python function

Question

I have an old function named old_func that takes as input two positional arguments, x and y. The input of the function was written like this using a tuple as the input:

def old_func(position):

    x, y = position 
    return x**2 + y**2

I now want a quick and easy way of calling the function over a grid of values:

xx = numpy.linspace(0, 1, 100)
yy = numpy.linspace(0, 1, 100)
X, Y = numpy.meshgrid(xx, yy)

array_positions = (X,Y)
old_fun(array_positions)

The intent is that each operation on x in the function is done on all of the X and the same for y. I tried vectorizing the function using numpy.vectorize but that does not work. I prefer not to change the function to accept NumPy arrays because that will take too long.

EelkeSpaak · Accepted Answer · 2015-07-21 15:46:41Z

5

Your own code should (and does) work fine:

def old_fun(position):

    x, y = position 
    z = x**2 + y**2
    return z

xx = numpy.linspace(0,1,100)
yy = numpy.linspace(0,1,100)
X,Y = numpy.meshgrid(xx, yy)

array_positions = (X,Y)
Z = old_fun(array_positions)

Z.shape is now (100, 100).

In general, numpy arrays will work with any standard operator, like + and **. Note that although your old_fun takes a tuple as input, and this tuple needs to consist of two values, the type of the two values can be anything, as long as this type supports math operators. Both standard Python scalars and numpy arrays support these, so the code works fine.

A note about @JuanManuel's answer: although it works perfectly fine as well, apply_along_axis is specifically not meant for vectorization in the sense that people usually want to use it, i.e. getting good performance out of numpy. apply_along_axis will result in a slow Python loop, not a fast C-loop like properly vectorized code will. Your own code uses proper vectorization, so use that instead.

answered Jul 21, 2015 at 15:46

EelkeSpaak

2,8652 gold badges21 silver badges38 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Dipole Over a year ago

You are right. Hmm I should have tested my example to make sure it actually gave the original error that my actual function did. I will try to edit my question.

Gyan Veda · Accepted Answer · 2015-07-21 15:33:55Z

1

The code below does the trick and saves you the trouble of creating array_positions.

First, use ravel to flatten X and Y into NumPy arrays of shape (10000,).

X_flattened = X.ravel()
Y_flattened = Y.ravel()

Then, use apply_along_axis to implement the custom function iteratively down the length of these flattened arrays.

float_array = np.apply_along_axis(old_func, 0, (X_flattened, Y_flattened))

Finally, reshape the output array into the desired shape of (100, 100).

np.reshape(float_array, (100, 100))

edited Jul 21, 2015 at 15:33

answered Jul 21, 2015 at 14:55

Gyan Veda

6,62912 gold badges45 silver badges68 bronze badges

2 Comments

Dipole Over a year ago

Thanks I think you are right, however since my old_func function is technically a function of one single tuple I would need np.apply_along_axis(old_func, 0, (X_flattened, Y_flattened)) right?

Gyan Veda Over a year ago

Yes, that's correct! I modified the answer accordingly.

Collectives™ on Stack Overflow

Vectorize python function

2 Answers 2

1 Comment

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related