Scaling and incrementing non-zero elements of a NumPy matrix

Question

I have a NumPy matrix C and want create a copy of it cPrime, which has some operation of the original matrix to all non-zero values of it. In the code below, for each non-zero element of C, I multiply by 30 and then add 1:

import numpy as np

size = 6
C = np.zeros((size,size), dtype=int)

C[4][0] = 2
C[4][1] = 5
C[4][2] = 3
C[0][3] = 1
C[1][3] = 1
C[2][3] = 1
C[3][5] = 3

cPrime = np.zeros((size, size),dtype=int)
for i in range(size):
    for j in range(size):
        if C[i][j] != 0:
            cPrime[i][j] = C[i][j]*30 + 1

This code works, but it feels inefficient. I'm about 99% sure that there's a really efficient way to achieve this goal, maybe using a masked array, but I haven't been able to figure it out.

Answer 1

This is a typical use case for numpy.where:

cPrime = np.where(C, C * 30 + 1, C)

This is about twice as fast as (30 * C + 1) * (C != 0) and generalizes more easily to other conditions.

Answer 2

cPrime = 30 * C + 1, which uses broadcasting, comes close. Unfortunately, it adds 1 indiscriminately to every single element, even the elements that were originally zero.

C != 0 gives you the indexes of all the locations you want to operate on:

>>> C
array([[0, 0, 0, 1, 0, 0],
       [0, 0, 0, 1, 0, 0],
       [0, 0, 0, 1, 0, 0],
       [0, 0, 0, 0, 0, 3],
       [2, 5, 3, 0, 0, 0],
       [0, 0, 0, 0, 0, 0]])
>>> C != 0
array([[False, False, False,  True, False, False],
       [False, False, False,  True, False, False],
       [False, False, False,  True, False, False],
       [False, False, False, False, False,  True],
       [ True,  True,  True, False, False, False],
       [False, False, False, False, False, False]], dtype=bool)

Combining the two concepts, you can write either

cPrime = 30 * C + 1 * (C != 0)

or

cPrime = (30 * C + 1) * (C != 0)

Answer 3

For really large sparse matrices, convert numpy dense to scipy.sparse. These store only the non-zeros (well, 2 ints + 1 double, 24 bytes):

import scipy.sparse

S = scipy.sparse.csr_matrix( C )  # dense to sparse
print( "S: %s  %d non-0" % (S.shape, S.nnz) )
S *= 30
S.data += 1  # increment only the non-0
# Dense = S.toarray()  # sparse to dense

(Be careful, there are corner cases where sparse and numpy dense behave differently; SO has almost 1000 questions/tagged/scipy+sparse-matrix .)