Search code examples
pythonscipysparse-matrix

How to best row-scale a SciPy sparse matrix?


I've got a SciPy sparse matrix A, let's say in CSR format, and a vector v of matching length.

What's the best way of row-scaling A with v, i.e., performing diag(v) * A?


Solution

  • The easy way is to let scipy handle the gory details, and simply do:

    scipy.sparse.spdiags(v, 0, len(v), len(v)) * A
    

    EDIT If (and only if) your matrix is stored in CSC format, you can do the operation in place as follows:

    A_csc.data = A_csc.data * v[A_csc.indices]
    

    I've done some timings, at it wildly depends on the sparsity of the matrix and its size, feel free to play with the following code:

    from __future__ import division
    import numpy as np
    import scipy.sparse as sps
    import timeit
    
    A_csr = None
    A_csc = None
    v = None
    
    def time_row_scaling(n, dens) :
        global A_csr, A_csc, v
        v = np.random.rand(n)
        A_csr = sps.rand(n, n, density=dens, format='csr')
        A_csc = A_csr.tocsc()
        def row_scale(A_csc, v) :
            A_csc.data = A_csc.data * v[A_csc.indices]
        row_scaled_1 = sps.spdiags(v, 0, n , n) * A_csr
        row_scaled_2 = sps.spdiags(v, 0, n , n) * A_csc
        row_scale(A_csc, v)
        if n < 1000 :
            np.testing.assert_almost_equal(row_scaled_1.toarray(),
                                           row_scaled_2.toarray())
            np.testing.assert_almost_equal(row_scaled_1.toarray(),
                                           A_csc.toarray())
        A_csc = A_csr.tocsc()
        t1 = timeit.timeit('sps.spdiags(v, 0, len(v) , len(v)) * A_csr',
                           'from __main__ import sps, v, A_csr',
                           number=1)
        t2 = timeit.timeit('sps.spdiags(v, 0, len(v), len(v)) * A_csc',
                           'from __main__ import sps, v, A_csc',
                           number=1)
        t3 = timeit.timeit('A_csc.data = A_csc.data * v[A_csc.indices]',
                           'from __main__ import A_csc, v',
                           number=1)
        print t1, t2, t3
    
    >>> time_row_scaling(1000, 0.01)
    0.00100659830939 0.00102425072673 0.000231944553347
    >>> time_row_scaling(1000, 0.1)
    0.0017328105481 0.00311339379218 0.00239826562947
    >>> time_row_scaling(10000, 0.01)
    0.0162369397769 0.0359325217874 0.0216837368279
    >>> time_row_scaling(10000, 0.1)
    0.167978350747 0.492032396702 0.209231639536
    

    Summary seems to be, if it is CSR, or really big, go with the simple first method. If it is a smallish, very sparse matrix, then the in place method will be faster, although all times are small then.