I am looking for advice to see if the following code performance could be further improved. This is and example using a 4x3 numpy 2d array:

import numpy as np

x = np.arange(12).reshape((4,3))
n, m = x.shape
y = np.zeros((n, m))

for j in range(m):
    x_j = x[:, :j+1]
    y[:,j] = np.linalg.norm(x_j, axis=1)

print x
print y

Which is printing

[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
[[  0.           1.           2.23606798]
 [  3.           5.           7.07106781]
 [  6.           9.21954446  12.20655562]
 [  9.          13.45362405  17.3781472 ]]

As you can see the code is computing the norms of the vectors considering increasing number of columns, so that y[i,j] represent the norm of the vector x[i,:j+1]. I couldn't find if this operation has a name and if it is possible to vectorize further the process and get rid of the for loop.

I only found in this post that using np.sqrt(np.einsum('ij,ij->i', x_j, x_j)) is a bit faster than using np.linalg.norm(x_j, axis=1).


Indeed, there is a better way. Exponentiation aside, we can see that this operation is equivalent to multiplication by an upper triangular matrix of 1. The former is about 100x faster. Here the code (Run it online !):


from time import time as time
import numpy as np
n = 1000
m = 500
size = n * m

triu = np.triu(np.ones((m, m)))
x = np.arange(size).reshape((n, m))
y = np.zeros((n, m))

# Your implementation
tic = time()
for j in range(m):
    x_j = x[:, :j + 1]
    y[:, j] = np.linalg.norm(x_j, axis=1)
tac = time()
print('Operation took {} ms'.format((tac - tic) * 1e3))

# Optimized implementation
tic = time()
y1 = np.sqrt(np.dot(x**2, triu))
tac = time()
print('Operation took {} ms'.format((tac - tic) * 1e3))

# Optimized implementation using cumsum method
tic = time()
y2= np.sqrt(np.cumsum(np.square(x), axis=1))
tac = time()
print('Operation took {} ms'.format((tac - tic) * 1e3))


Operation took 1690.1559829711914 ms
Operation took 18.942832946777344 ms
Operation took 6.124973297119141 ms
  • \$\begingroup\$ I think that you made an error copying the code. As you can see here the code runs without warnings and the expression np.sum(y==y1) return size value. \$\endgroup\$ – Amir Dib Oct 24 '17 at 23:28
  • \$\begingroup\$ @Reti43 An easy way to check for equivalence (disregarding rounding) is to use assert np.allclose(y, y1). For me this passes for n = 1000, m = 500, n = 10000, m = 500 and n = 1000, m = 5000 (this last one took more than a minute for OP's code). Python 3.6.3, numpy 1.13.3, no warnings for me. \$\endgroup\$ – Graipher Oct 25 '17 at 8:32
  • \$\begingroup\$ The different results observed are due to 32-bit int overflow. On 32-bit Python arange defaults to dtype=int32 and computing x**2 leads to overflow.. \$\endgroup\$ – Janne Karila Oct 25 '17 at 8:44
  • \$\begingroup\$ @JanneKarila Great catch! Addressing that works. I'm on the 64bit version though, but regardless. \$\endgroup\$ – Reti43 Oct 25 '17 at 9:24
  • \$\begingroup\$ Thanks for your answer. For the record, I found that np.sqrt(np.cumsum(np.square(x), axis=1)) also do the job with similar performance. \$\endgroup\$ – Delforge Oct 25 '17 at 11:53

This is faster:

  • \$\begingroup\$ Welcome to code review. Good answers contain at least one observation about the code, and don't necessarily contain any code at all. This answer is rather border line and might be down voted or removed. On code review we are here to help improve code, not write code. \$\endgroup\$ – pacmaninbw Apr 19 '20 at 16:17
  • \$\begingroup\$ Can you expand on your answer - perhaps explain what code this would replace? And what value would b1 be assigned? \$\endgroup\$ – Sᴀᴍ Onᴇᴌᴀ Apr 21 '20 at 15:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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