current community

your communities

more stack exchange communities

Stack Exchange
tour help
Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.
up vote 1 down vote favorite

How can I vectorize this code snippet and eliminate this double loop? Even with "only" 1001 elements in a this takes almost 30s.

a = np.linspace(-.5, .5, 1001)
S = np.zeros((a.size, a.size))

# east, north and tp are np.arrays of equal length
for l, sy in enumerate(a):
    for k, sx in enumerate(a):
            S[l,k] = np.sum((east*sx + north*sy - tp)**2)
share|improve this question

1 Answer 1

up vote 0 down vote accepted

This should do it:

S = np.sum((east*a[:, None, None] + north*a[:, None] - tp)**2, axis=-1)

But if the east north and tp arrays are large it may actually perform worse, as it avoids looping but creates an intermediate array of a.size**2 * east.size items.

If you expand your expression, you could also do:

e2 = np.dot(east, east)
n2 = np.dot(north, north)
t2 = np.dot(tp, tp)
en = np.dot(east, north)
et = np.dot(east, tp)
nt = np.dot(north, tp)

a2 = a*a

S = (e2*a2[:, None, None] + b2*a2[:, None] + t2 +
     2*en*a[:, None, None]*a[:, None] - 2*et*a[:, None, None] -
     2*nt*a[:, None])

which should give the same result without the ginormous intermediate array.

share|improve this answer

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.