Summing up scipy rows and filling it in the diagonal - scipy

Question

I have a sparse matrix of large dimension and I want to take sum of elements of each row, i, and fill diagonal with the results. How do I do this?
My numpy approach:
A = np.sum(A, axis=1)
D = np.diag(A)
How do I approach this with scipy?
Edit: I am using scipy.sparse and initializing A as a csr_matrix. I have read the docs but perhaps I am misread or missing something. With sparse.diags(D), I just got a column array, i.e it didn't give me a diagonal matrix when I tried to turn the sparse matrix into np.array.

Answer 1

Make a sparse matrix:
In [149]: from scipy import sparse
In [150]: M = sparse.csr_matrix(np.arange(6).reshape(2,3))
In [151]: M
Out[151]:
<2x3 sparse matrix of type '<class 'numpy.int64'>'
with 5 stored elements in Compressed Sparse Row format>
In [152]: M.A
Out[152]:
array([[0, 1, 2],
[3, 4, 5]])
Sum on an axis:
In [153]: M.sum(axis=0)
Out[153]: matrix([[3, 5, 7]])
Note this is a dense np.matrix. Since sum increases density, sparse returns this rather than a sparse matrix.
Convert that matrix to 1d ndarray:
In [154]: M.sum(axis=0).A1
Out[154]: array([3, 5, 7])
Use diags to make a sparse matrix:
In [155]: M1=sparse.diags(M.sum(axis=0).A1)
In [156]: M1
Out[156]:
<3x3 sparse matrix of type '<class 'numpy.float64'>'
with 3 stored elements (1 diagonals) in DIAgonal format>
In [157]: M1.A
Out[157]:
array([[3., 0., 0.],
[0., 5., 0.],
[0., 0., 7.]])
Read sparse.diags to see its required arguments. sparse.dia_matrix can also be used (but also read its docs).
Math on the Dia format matrix is likely to produce a csr format matrix.
Other diagonal creation commands:
M1=sparse.diags(M.sum(axis=0),[0], shape=(3,3))
M1=sparse.dia_matrix((M.sum(axis=0), 0),shape=(3,3))