import numpy as np
import scipy as sp
from scipy import sparse
from numba import jit

# matvec implementation is based on
# http://stackoverflow.com/questions/18595981/improving-performance-of-multiplication-of-scipy-sparse-matrices

@jit(nopython=True, nogil=True)
def matvec2dense(m_ptr, m_ind, m_val, v_nnz, v_val, out):
    l = len(v_nnz)
    for j in xrange(l):
        col_start = v_nnz[j]
        col_end = col_start + 1
        ind_start = m_ptr[col_start]
        ind_end = m_ptr[col_end]
        if ind_start != ind_end:
            out[m_ind[ind_start:ind_end]] += m_val[ind_start:ind_end] * v_val[j]


@jit(nopython=True, nogil=True)
def matvec2sparse(m_ptr, m_ind, m_val, v_nnz, v_val, sizes, indices, data):
    l = len(sizes) - 1
    for j in xrange(l):
        col_start = v_nnz[j]
        col_end = col_start + 1
        ind_start = m_ptr[col_start]
        ind_end = m_ptr[col_end]
        data_start = sizes[j]
        data_end = sizes[j+1]
        if ind_start != ind_end:
            indices[data_start:data_end] = m_ind[ind_start:ind_end]
            data[data_start:data_end] = m_val[ind_start:ind_end] * v_val[j]


def csc_matvec(mat_csc, vec, dense_output=True, dtype=None):
    v_nnz = vec.indices
    v_val = vec.data

    m_val = mat_csc.data
    m_ind = mat_csc.indices
    m_ptr = mat_csc.indptr

    res_dtype = dtype or np.result_type(mat_csc.dtype, vec.dtype)
    if dense_output:
        res = np.zeros((mat_csc.shape[0],), dtype=res_dtype)
        matvec2dense(m_ptr, m_ind, m_val, v_nnz, v_val, res)
    else:
        sizes = m_ptr.take(v_nnz+1) - m_ptr.take(v_nnz)
        sizes = np.concatenate(([0], np.cumsum(sizes)))
        n = sizes[-1]
        data = np.empty((n,), dtype=res_dtype)
        indices = np.empty((n,), dtype=np.intp)
        indptr = np.array([0, n], dtype=np.intp)
        matvec2sparse(m_ptr, m_ind, m_val, v_nnz, v_val, sizes, indices, data)
        res = sp.sparse.csr_matrix((data, indices, indptr), shape=(1, mat_csc.shape[0]), dtype=res_dtype)
        res.sum_duplicates() # expensive operation
    return res

@@ -2068,8 +2068,620 @@
 urn res%0A
+%0Ajit(nopython=True)%0Adef _blockify(ind, ptr, major_dim):%0A    # convenient function to compute only diagonal%0A    # elements of the product of 2 matrices;%0A    # indices must be intp in order to avoid overflow%0A    # major_dim is shape%5B0%5D for csc format and shape%5B1%5D for csr format%0A    n = len(ptr) - 1%0A    for i in xrange(1, n): #first row/col is unchanged%0A        lind = ptr%5Bi%5D%0A        rind = ptr%5Bi+1%5D%0A        for j in xrange(lind, rind):%0A            shift_ind = i * major_dim%0A            ind%5Bj%5D += shift_ind%0A%0A%0Adef inverse_permutation(p):%0A    s = np.empty(p.size, p.dtype)%0A    s%5Bp%5D = np.arange(p.size)%0A    return s%0A

add 2 new functions for manipulation with sparse matrices