How to make vectorized computation instead of 'for' loops for all grid?

user_7805308 0
Viewed 174

I have a double for loop among all the grid, and I want to make it work faster. r, vec1, vec2, theta are the vectors of the same length N. c is a constant.

import numpy as np

N = 30
x_coord, y_coord = 300, 300
m1 = np.zeros((x_coord, y_coord))

vec1, vec2 = np.ones(N), np.ones(N)
theta = np.ones(N)

for x in np.arange(x_coord):
    for y in np.arange(y_coord):
        m1[x,y] = np.sum(np.cos(2.*np.pi*(r*(vec1*x + vec2*y))+theta)) * c

The time for two loops was:

1.03 s ± 8.96 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Also I tried to use np.meshgrid:

def f1(x, y):
    sum1 = vec1*x + vec2*y
    mltpl1 = r * sum1
    sum2 = 2.*np.pi * mltpl1 + theta
    sum3 = np.sum(np.cos(sum2))
    mltpl2 = sum3 * c
    return mltpl2

msh1, msh2 = np.meshgrid(range(x_coord), range(y_coord))
pairs = np.vstack((np.ravel(msh1), np.ravel(msh2))).T

m1 = np.reshape(list(map(lambda x: f1(x[0], x[1]), pairs)), (m1.shape[0], m1.shape[1])).T

Trying meshgrid time was more:

1.25 s ± 48.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

So I need a solution how to do it on a vectors and matrices level. Are there any ideas? Thank you in advance.

vectorization numpy python
2 Answers

We can use a trigonometric trick here -

cos(A + B) = cos A cos B − sin A sin B

This lets us leverage matrix-multiplication for a solution that would look something like this -

# Get x and y as 1D arrays
x = np.arange(x_coord)
y = np.arange(y_coord)

# Get the common constant for scaling vec1 and vec2 parts
k1 = 2.*np.pi*r

# Use outer multiplications for the two vectors against x,y and also scale them
p1 = k1*vec1*x[:,None] + theta
p2 = (k1*vec2)[:,None]*y

# Finally use trigonometry+matrix-multiplication for sum reductions
out = c*(np.cos(p1).dot(np.cos(p2)) - np.sin(p1).dot(np.sin(p2)))

Timings -

# Setup
In [151]: np.random.seed(0)
     ...: c = 2.3
     ...: N = 30
     ...: x_coord, y_coord = 300, 300
     ...: vec1 = np.random.rand(N)
     ...: vec2 = np.random.rand(N)
     ...: r = np.random.rand(N)
     ...: theta = np.ones(N)

# Original solution
In [152]: %%timeit
     ...: m1 = np.zeros((x_coord, y_coord))
     ...: for x in np.arange(x_coord):
     ...:     for y in np.arange(y_coord):
     ...:         m1[x,y] = np.sum(np.cos(2.*np.pi*(r*(vec1*x + vec2*y))+theta)) * c
1 loop, best of 3: 960 ms per loop

# Proposed solution
In [153]: %%timeit
     ...: x = np.arange(x_coord)
     ...: y = np.arange(y_coord)
     ...: k1 = 2.*np.pi*r
     ...: p1 = k1*vec1*x[:,None] + theta
     ...: p2 = (k1*vec2)[:,None]*y
     ...: out = c*(np.cos(p1).dot(np.cos(p2)) - np.sin(p1).dot(np.sin(p2)))
100 loops, best of 3: 2.54 ms per loop

375x+ speedup!

user_32938810

Are r, theta, and c constants? If you post a minimum working example it will make it easier for others to help you.

In any case, you can take advantage of broadcasting to do something like:

X = np.outer(x,vec1)
Y = np.outer(y,vec2)

Z = np.reshape(X[:,np.newaxis] + Y[np.newaxis],(x_coord*y_coord,N))

M = np.sum(np.cos(2.*np.pi*r*Z+theta),axis=1)*c
m = np.reshape(M,(x_coord,y_coord))

When I tried this with r, theta, and c as constants it gave the same result. I think it would also work fine if they are vectors applied pointwise.

How it works

The key observation is that almost all the operations are done point-wise and are the same across all x,y pairs. Therefore, we only need to figure out how to vectorize the addition of all pairs of vec1*x and vec2*y.

We first make lists of X=vec1*x and Y=vec2*y using the outer product.

We now add all pairs of rows of X and Y using broadcasting.

We then reshape it to a list of all the pairs, apply the rest of the function pointwise and sum along the correct axis.

Finally, we reshape it from a ndarray of length x_coord*y_coord to a 2d array of shape (x_coord,y_coord).

This is not very memory efficient since we construct all elements of all the sums at once, but unless you're working with wildly big data it should be okay. Even if you are working with enough data to not fit into memory, it probably makes sense to block this method as opposed to using a python loop.

user_39012160
Related