import numpy as np
import scipy.sparse as sparse
from ctypes import POINTER, c_double, c_int, c_int64, c_float, c_int
from pynao.m_libnao import libsparsetools
[docs]def count_nnz_spmat_denmat(csr, ncolB):
nrow, ncol = csr.shape
if not sparse.isspmatrix_csr(csr):
raise Exception("Matrix must be in csr format")
nnz = libsparsetools.count_nnz_spmat_denmat(
c_int(nrow), c_int(ncolB),
csr.indptr.ctypes.data_as(POINTER(c_int)))
return nnz
[docs]def spmat_denmat(csr, B):
"""
sparse matrix dense matrix multiplication directly giving a sparse matrix
in CSR format.
"""
if not sparse.isspmatrix_csr(csr):
raise Exception("Matrix must be in csr format")
nrow, ncol = csr.shape
nnz = count_nnz_spmat_denmat(csr, B.shape[1])
indptr_new = np.zeros((csr.shape[0]+1), dtype=np.int32)
indices_new = np.zeros((nnz), dtype=np.int32)
data_new = np.zeros((nnz), dtype=csr.data.dtype)
if csr.dtype == np.float32:
libsparsetools.scsr_spmat_denmat(
c_int(nrow), c_int(ncol), c_int(csr.nnz),
csr.indptr.ctypes.data_as(POINTER(c_int)),
csr.indices.ctypes.data_as(POINTER(c_int)),
csr.data.ctypes.data_as(POINTER(c_float)),
c_int(B.shape[0]), c_int(B.shape[1]),
B.ctypes.data_as(POINTER(c_float)),
indptr_new.ctypes.data_as(POINTER(c_int)),
indices_new.ctypes.data_as(POINTER(c_int)),
data_new.ctypes.data_as(POINTER(c_float)))
elif csr.dtype == np.float64:
libsparsetools.dcsr_spmat_denmat(
c_int(nrow), c_int(ncol), c_int(csr.nnz),
csr.indptr.ctypes.data_as(POINTER(c_int)),
csr.indices.ctypes.data_as(POINTER(c_int)),
csr.data.ctypes.data_as(POINTER(c_double)),
c_int(B.shape[0]), c_int(B.shape[1]),
B.ctypes.data_as(POINTER(c_double)),
indptr_new.ctypes.data_as(POINTER(c_int)),
indices_new.ctypes.data_as(POINTER(c_int)),
data_new.ctypes.data_as(POINTER(c_double)))
else:
raise ValueError("Not implemented")
ret = sparse.csr_matrix((data_new, indices_new, indptr_new), shape=(nrow, B.shape[1]))
return ret
"""
Wrapper to sparse matrix operations from scipy implemented with openmp
"""
[docs]def csr_matvec(csr, x, y=None):
nrow, ncol = csr.shape
nnz = csr.data.shape[0]
# for small sizes scipy is way faster
if nnz < 50000:
return csr.dot(x)
if not sparse.isspmatrix_csr(csr):
raise Exception("Matrix must be in csr format")
if x.size != ncol:
print(x.size, ncol)
raise ValueError("wrong dimension!")
xx = np.require(x, requirements=["A", "O"], dtype=csr.dtype)
if y is None:
y = np.zeros((nrow), dtype=csr.dtype)
if csr.dtype == np.float32:
libsparsetools.scsr_matvec(c_int(nrow), c_int(ncol), c_int(nnz),
csr.indptr.ctypes.data_as(POINTER(c_int)),
csr.indices.ctypes.data_as(POINTER(c_int)),
csr.data.ctypes.data_as(POINTER(c_float)),
xx.ctypes.data_as(POINTER(c_float)),
y.ctypes.data_as(POINTER(c_float)))
elif csr.dtype == np.float64:
libsparsetools.dcsr_matvec(c_int(nrow), c_int(ncol), c_int(nnz),
csr.indptr.ctypes.data_as(POINTER(c_int)),
csr.indices.ctypes.data_as(POINTER(c_int)),
csr.data.ctypes.data_as(POINTER(c_double)),
xx.ctypes.data_as(POINTER(c_double)),
y.ctypes.data_as(POINTER(c_double)))
else:
raise ValueError("Not implemented")
return y
[docs]def csc_matvec(csc, x):
"""
Matrix vector multiplication
using csc format
"""
if not sparse.isspmatrix_csc(csc):
raise Exception("Matrix must be in csc format")
nrow, ncol = csc.shape
nnz = csc.data.shape[0]
if x.size != ncol:
print(x.size, ncol)
raise ValueError("wrong dimension!")
xx = np.require(x, requirements="C")
if csc.dtype == np.float32:
y = np.zeros((nrow), dtype=np.float32)
libsparsetools.scsc_matvec(c_int(nrow), c_int(ncol), c_int(nnz),
csc.indptr.ctypes.data_as(POINTER(c_int)),
csc.indices.ctypes.data_as(POINTER(c_int)),
csc.data.ctypes.data_as(POINTER(c_float)),
xx.ctypes.data_as(POINTER(c_float)),
y.ctypes.data_as(POINTER(c_float)))
elif csc.dtype == np.float64:
y = np.zeros((nrow), dtype=np.float64)
libsparsetools.dcsc_matvec(c_int(nrow), c_int(ncol), c_int(nnz),
csc.indptr.ctypes.data_as(POINTER(c_int)),
csc.indices.ctypes.data_as(POINTER(c_int)),
csc.data.ctypes.data_as(POINTER(c_double)),
xx.ctypes.data_as(POINTER(c_double)),
y.ctypes.data_as(POINTER(c_double)))
else:
raise ValueError("Not implemented")
return y
[docs]def csc_matvecs(csc, B, transB = False, order="C"):
"""
Matrix matrix multiplication
using csc format
"""
if not sparse.isspmatrix_csc(csc):
raise Exception("Matrix must be in csc format")
if transB:
# Here need to be careful, since using the transpose of B
# will change from row major to col major and vice-versa
mat = np.require(B.T, dtype=B.dtype, requirements=["A", "O", order])
else:
mat = np.require(B, dtype=B.dtype, requirements=["A", "O", order])
nrow, ncol = csc.shape
nvecs = mat.shape[1]
if csc.dtype == np.float32:
C = np.zeros((nrow, nvecs), dtype=np.float32, order=order)
libsparsetools.scsc_matvecs(c_int(nrow), c_int(ncol), c_int(nvecs),
csc.indptr.ctypes.data_as(POINTER(c_int)),
csc.indices.ctypes.data_as(POINTER(c_int)),
csc.data.ctypes.data_as(POINTER(c_float)),
mat.ctypes.data_as(POINTER(c_float)),
C.ctypes.data_as(POINTER(c_float)))
elif csc.dtype == np.float64:
C = np.zeros((nrow, nvecs), dtype=np.float64, order=order)
libsparsetools.dcsc_matvecs(c_int(nrow), c_int(ncol), c_int(nvecs),
csc.indptr.ctypes.data_as(POINTER(c_int)),
csc.indices.ctypes.data_as(POINTER(c_int)),
csc.data.ctypes.data_as(POINTER(c_double)),
mat.ctypes.data_as(POINTER(c_double)),
C.ctypes.data_as(POINTER(c_double)))
else:
raise ValueError("Not implemented")
return C