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).


2 Answers 2


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
    Commented Oct 24, 2017 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
    Commented Oct 25, 2017 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\$ Commented Oct 25, 2017 at 8:44
  • \$\begingroup\$ @JanneKarila Great catch! Addressing that works. I'm on the 64bit version though, but regardless. \$\endgroup\$
    – Reti43
    Commented Oct 25, 2017 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
    Commented Oct 25, 2017 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
    Commented Apr 19, 2020 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\$ Commented Apr 21, 2020 at 15:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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