Update a Numpy 1D array B in each loop to solve the matrix expression A*x = B

user_2616060 0
Viewed 126

I need to solve x from sparse matrix expression A*x = B in loops, where A is a Scipy CSC sparse-matrix and B is Numpy 1D array. Both A and B are large about 500K rows. Basically, I need to update B in each loop. So the speed to update B is critical. Right now, my way is to define csc_matrix in each loop, and then convert it to 1D Numpy array as below which is really expensive in terms of time:

B = csc_matrix((data,(row, col)),shape=(500000, 1), dtype='complex128').toarray()[:,0];

Please note:

  • row has lots of the repeated index, such as [0,1,2,0,2,2,3,3....],
  • col is [0,0, 0,.......0];

Is there fast way to update B in each loop?

scipy sparse-matrix numpy python performance
1 Answers

Assuming col contains only zeros, data/row/col are Numpy arrays and you want B stored as a Numpy array. You can use Numba to generate B efficiently. Here is how:

import numba

# Works in-place to avoid any slow allocation in the critical loop.
# Note that the type of row may be different.
@nb.njit(void(nb.complex128[:], nb.complex128[:], nb.int64[:]))
def updateVector(B, data, row):
    B.fill(0.)
    for i in range(len(row)):
        B[row[i]] += data[i]

updateVector update the value of B in-place. This assume B has been allocated at the correct size before (using for example B = np.empty(500000, dtype=np.complex128)).

On my machine this is 14 times faster with the following configuration:

row = np.random.randint(0, 500000, size=100000)
col = np.zeros(100000, dtype=np.int64)
data = np.random.rand(100000) + np.random.rand(100000) * 1j
user_129395570
Related