computing column norms of sparse matrix with numba

norms_X_col_numba

import time
import numpy as np
from numba import njit
from numpy.linalg import norm
from scipy.sparse.linalg import norm as snorm
from libsvmdata import fetch_libsvm

X, y = fetch_libsvm("finance")

@njit
def norm_v1(data, indices, indptr):
    n_features = len(indptr) - 1
    norms = np.zeros(n_features)
    for j in range(n_features):
        norms[j] = norm(data[indptr[j]:indptr[j + 1]]) ** 2
    return norms


@njit
def norm_v2(data, indices, indptr):
    n_features = len(indptr) - 1
    norms = np.zeros(n_features)
    for j in range(n_features):
        tmp = 0
        for idx in range(indptr[j], indptr[j + 1]):
            tmp += data[idx] ** 2
        norms[j] = tmp
    return norms


v2 = norm_v2(X.data, X.indices, X.indptr)
scip = snorm(X, axis=0) ** 2


np.testing.assert_allclose(scip, v2)

np.testing.assert_allclose(
    norm_v1(X.data, X.indices, X.indptr),
    v2,
)

for vers in (norm_v1, norm_v2):
    # compile
    vers(X.data, X.indices, X.indptr)
    t0 = time.time()
    vers(X.data, X.indices, X.indptr)
    t1 = time.time()
    print(f'{1000 * (t1 - t0):.3f} ms for {vers.__name__}')