numpy performance of norm calculation using variable dimension

Question

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

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

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

Output

Operation took 1690.1559829711914 ms
Operation took 18.942832946777344 ms
Operation took 6.124973297119141 ms