numpy performance of norm calculation using variable dimension

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 !):

Source

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

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

Output

Operation took 1690.1559829711914 ms
Operation took 18.942832946777344 ms
Operation took 6.124973297119141 ms
• 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. – Amir Dib Oct 24 '17 at 23:28
• @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. – Graipher Oct 25 '17 at 8:32
• 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.. – Janne Karila Oct 25 '17 at 8:44
• @JanneKarila Great catch! Addressing that works. I'm on the 64bit version though, but regardless. – Reti43 Oct 25 '17 at 9:24
• 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. – Delforge Oct 25 '17 at 11:53