ENH: Enforce Qobj data is fast_csr (#609)

* fast_csr_safety

* changes

* fix overlap

* updates

* remove unnecessary tests

* fix operator test name

qutip/bloch_redfield.py

@@ -41,9 +41,10 @@
 from qutip.expect import expect
 from qutip.solver import Options, _solver_safety_check
 from qutip.cy.spmatfuncs import cy_ode_rhs
-from qutip.cy.sparse_utils import dense2D_to_fastcsr_fmode
+from qutip.cy.spconvert import dense2D_to_fastcsr_fmode
 from qutip.solver import Result
 from qutip.superoperator import liouvillian
+from qutip.cy.spconvert import arr_coo2fast
 
 
 # -----------------------------------------------------------------------------
@@ -355,8 +356,9 @@ def bloch_redfield_tensor(H, a_ops, spectra_cb, c_ops=[], use_secular=True):
                 cols.append(J)
                 data.append(elem)
 
-    R = sp.coo_matrix((np.array(data),(np.array(rows),np.array(cols))),
-            shape=(N**2,N**2),dtype=complex).tocsr()
+    R = arr_coo2fast(np.array(data, dtype=complex),
+                    np.array(rows, dtype=np.int32),
+                    np.array(cols, dtype=np.int32), N**2, N**2)
 
     L.data = L.data + R
 

qutip/cy/setup.py

@@ -9,7 +9,7 @@
 dir_path = os.path.dirname(os.path.realpath(__file__))
 
 exts = ['spmatfuncs', 'stochastic', 'sparse_utils', 'graph_utils', 'interpolate',
-        'spmath', 'heom', 'math']
+        'spmath', 'heom', 'math', 'spconvert']
 
 _compiler_flags = ['-w', '-ffast-math', '-O3', '-march=native', '-funroll-loops']
 

qutip/cy/sparse_struct.pxi

@@ -148,7 +148,8 @@ cdef inline int int_max(int a, int b) nogil:
 
 @cython.boundscheck(False)
 @cython.wraparound(False)        
-cdef void init_CSR(CSR_Matrix * mat, int nnz, int nrows, int max_length = 0, int init_zeros = 1):
+cdef void init_CSR(CSR_Matrix * mat, int nnz, int nrows, int ncols = 0,
+                int max_length = 0, int init_zeros = 1):
     """
     Initialize CSR_Matrix struct. Matrix is assumed to be square with
     shape nrows x nrows.  Manually set mat.ncols otherwise
@@ -162,6 +163,8 @@ cdef void init_CSR(CSR_Matrix * mat, int nnz, int nrows, int max_length = 0, int
     nrows : int
         Number of rows in matrix. Also gives length 
         of indptr array (nrows+1).
+    ncols : int (default = 0)
+        Number of cols in matrix. Default is ncols = nrows.
     max_length : int (default = 0)
         Maximum length of data and indices arrays.  Used for resizing.
         Default value of zero indicates no resizing.
@@ -184,7 +187,10 @@ cdef void init_CSR(CSR_Matrix * mat, int nnz, int nrows, int max_length = 0, int
         mat.indptr = <int *>PyDataMem_NEW((nrows+1) * sizeof(int))
     mat.nnz = nnz
     mat.nrows = nrows
-    mat.ncols = nrows
+    if ncols == 0:
+        mat.ncols = nrows
+    else:
+        mat.ncols = ncols
     mat.is_set = 1
     mat.max_length = max_length
     mat.numpy_lock = 0
@@ -201,8 +207,7 @@ cdef void copy_CSR(CSR_Matrix * out, CSR_Matrix * mat):
         raise_error_CSR(-3)
     elif out.is_set:
         raise_error_CSR(-2)
-    init_CSR(out, mat.nnz, mat.nrows, mat.max_length)
-    out.ncols = mat.ncols
+    init_CSR(out, mat.nnz, mat.nrows, mat.nrows, mat.max_length)
     # We cannot use memcpy here since there are issues with
     # doing so on Win with the GCC compiler
     for kk in range(mat.nnz):
@@ -214,7 +219,8 @@ cdef void copy_CSR(CSR_Matrix * out, CSR_Matrix * mat):
 
 @cython.boundscheck(False)
 @cython.wraparound(False)        
-cdef void init_COO(COO_Matrix * mat, int nnz, int nrows, int max_length = 0, int init_zeros = 0):
+cdef void init_COO(COO_Matrix * mat, int nnz, int nrows, int ncols = 0, 
+                int max_length = 0, int init_zeros = 1):
     """
     Initialize COO_Matrix struct. Matrix is assumed to be square with
     shape nrows x nrows.  Manually set mat.ncols otherwise
@@ -227,6 +233,8 @@ cdef void init_COO(COO_Matrix * mat, int nnz, int nrows, int max_length = 0, int
         Number of nonzero elements.
     nrows : int
         Number of rows in matrix.
+    nrows : int (default = 0)
+        Number of cols in matrix. Default is ncols = nrows.
     max_length : int (default = 0)
         Maximum length of arrays.  Used for resizing.
         Default value of zero indicates no resizing.
@@ -254,7 +262,10 @@ cdef void init_COO(COO_Matrix * mat, int nnz, int nrows, int max_length = 0, int
         mat.cols = <int *>PyDataMem_NEW(nnz * sizeof(int))
     mat.nnz = nnz
     mat.nrows = nrows
-    mat.ncols = nrows
+    if ncols == 0:
+        mat.ncols = nrows
+    else:
+        mat.ncols = ncols
     mat.is_set = 1
     mat.max_length = max_length
     mat.numpy_lock = 0
@@ -384,22 +395,21 @@ cdef void COO_to_CSR(CSR_Matrix * out, COO_Matrix * mat):
     Conversion from COO to CSR. Not in place,
     but result is sorted correctly.
     """
-    cdef int i, j, iad, j0, nnz = mat.nnz, nrows = mat.nrows
+    cdef int i, j, iad, j0
     cdef double complex val
     cdef size_t kk
-    init_CSR(out, nnz, nrows, max_length=0, init_zeros=1)
-    out.ncols = mat.ncols
+    init_CSR(out, mat.nnz, mat.nrows, mat.ncols, max_length=0, init_zeros=1)
     # Determine row lengths
-    for kk in range(nnz):
+    for kk in range(mat.nnz):
         out.indptr[mat.rows[kk]] = out.indptr[mat.rows[kk]] + 1
     # Starting position of rows
     j = 0
-    for kk in range(nrows):
+    for kk in range(mat.nrows):
         j0 = out.indptr[kk]
         out.indptr[kk] = j
         j += j0
     #Do the data
-    for kk in range(nnz):
+    for kk in range(mat.nnz):
         i = mat.rows[kk]
         j = mat.cols[kk]
         val = mat.data[kk]
@@ -408,21 +418,21 @@ cdef void COO_to_CSR(CSR_Matrix * out, COO_Matrix * mat):
         out.indices[iad] = j
         out.indptr[i] = iad+1
     # Shift back
-    for kk in range(nrows,0,-1):
+    for kk in range(mat.nrows,0,-1):
         out.indptr[kk] = out.indptr[kk-1]
     out.indptr[0] = 0
 
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
 cdef void CSR_to_COO(COO_Matrix * out, CSR_Matrix * mat):
-    cdef int k1, k2, nrows=mat.nrows, nnz=mat.nnz
+    cdef int k1, k2
     cdef size_t jj, kk
-    init_COO(out, nnz, nrows)
-    for kk in range(nnz):
+    init_COO(out, mat.nnz, mat.nrows, mat.ncols)
+    for kk in range(mat.nnz):
         out.data[kk] = mat.data[kk]
         out.cols[kk] = mat.indices[kk]
-    for kk in range(nrows,0,-1):
+    for kk in range(mat.nrows,0,-1):
         k1 = mat.indptr[kk+1]
         k2 = mat.indptr[kk]
         for jj in range(k1,k2,-1):

qutip/cy/sparse_utils.pyx

@@ -288,49 +288,3 @@ def unit_row_norm(complex[::1] data, int[::1] ptr, int nrows):
         total = sqrt(total)
         for ii in range(ptr[row], ptr[row+1]):
             data[ii] /= total
-
-
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def dense2D_to_fastcsr_cmode(complex[:, ::1] mat, int nrows, int ncols):
-    cdef int nnz = 0
-    cdef size_t ii, jj
-    cdef np.ndarray[complex, ndim=1, mode='c'] data = np.zeros(nrows*ncols, dtype=complex)
-    cdef np.ndarray[int, ndim=1, mode='c'] ind = np.zeros(nrows*ncols, dtype=np.int32)
-    cdef np.ndarray[int, ndim=1, mode='c'] ptr = np.zeros(nrows+1, dtype=np.int32)
-
-    for ii in range(nrows):
-        for jj in range(ncols):
-            if mat[ii,jj] != 0:
-                ind[nnz] = jj
-                data[nnz] = mat[ii,jj]
-                nnz += 1
-        ptr[ii+1] = nnz
-
-    if nnz < (nrows*ncols):
-        return fast_csr_matrix((data[:nnz], ind[:nnz], ptr), shape=(nrows,ncols))
-    else:
-        return fast_csr_matrix((data, ind, ptr), shape=(nrows,ncols))
-
-
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def dense2D_to_fastcsr_fmode(complex[::1, :] mat, int nrows, int ncols):
-    cdef int nnz = 0
-    cdef size_t ii, jj
-    cdef np.ndarray[complex, ndim=1, mode='c'] data = np.zeros(nrows*ncols, dtype=complex)
-    cdef np.ndarray[int, ndim=1, mode='c'] ind = np.zeros(nrows*ncols, dtype=np.int32)
-    cdef np.ndarray[int, ndim=1, mode='c'] ptr = np.zeros(nrows+1, dtype=np.int32)
-
-    for ii in range(nrows):
-        for jj in range(ncols):
-            if mat[ii,jj] != 0:
-                ind[nnz] = jj
-                data[nnz] = mat[ii,jj]
-                nnz += 1
-        ptr[ii+1] = nnz
-
-    if nnz < (nrows*ncols):
-        return fast_csr_matrix((data[:nnz], ind[:nnz], ptr), shape=(nrows,ncols))
-    else:
-        return fast_csr_matrix((data, ind, ptr), shape=(nrows,ncols))

qutip/cy/spconvert.pyx

@@ -0,0 +1,161 @@
+# This file is part of QuTiP: Quantum Toolbox in Python.
+#
+#    Copyright (c) 2011 and later, The QuTiP Project.
+#    All rights reserved.
+#
+#    Redistribution and use in source and binary forms, with or without
+#    modification, are permitted provided that the following conditions are
+#    met:
+#
+#    1. Redistributions of source code must retain the above copyright notice,
+#       this list of conditions and the following disclaimer.
+#
+#    2. Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#
+#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
+#       of its contributors may be used to endorse or promote products derived
+#       from this software without specific prior written permission.
+#
+#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+###############################################################################
+import numpy as np
+from qutip.fastsparse import fast_csr_matrix
+cimport numpy as np
+cimport cython
+
+include "sparse_struct.pxi"
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def coo2fast(object A):
+    cdef int nnz = A.nnz
+    cdef int nrows = A.shape[0]
+    cdef int ncols = A.shape[1]
+    cdef complex[::1] data = A.data
+    cdef int[::1] rows = A.rows
+    cdef int[::1] cols = A.cols
+    cdef CSR_Matrix out
+    init_CSR(&out, nnz, nrows, ncols)
+
+    cdef int i, j, iad, j0
+    cdef double complex val
+    cdef size_t kk
+    # Determine row lengths
+    for kk in range(nnz):
+        out.indptr[rows[kk]] = out.indptr[rows[kk]] + 1
+    # Starting position of rows
+    j = 0
+    for kk in range(nrows):
+        j0 = out.indptr[kk]
+        out.indptr[kk] = j
+        j += j0
+    #Do the data
+    for kk in range(nnz):
+        i = rows[kk]
+        j = cols[kk]
+        val = data[kk]
+        iad = out.indptr[i]
+        out.data[iad] = val
+        out.indices[iad] = j
+        out.indptr[i] = iad+1
+    # Shift back
+    for kk in range(nrows,0,-1):
+        out.indptr[kk] = out.indptr[kk-1]
+    out.indptr[0] = 0
+
+    return CSR_to_scipy(&out)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def arr_coo2fast(complex[::1] data, int[::1] rows, int[::1] cols, int nrows, int ncols):
+    cdef int nnz = data.shape[0]
+    cdef CSR_Matrix out
+    init_CSR(&out, nnz, nrows, ncols)
+
+    cdef int i, j, iad, j0
+    cdef double complex val
+    cdef size_t kk
+    # Determine row lengths
+    for kk in range(nnz):
+        out.indptr[rows[kk]] = out.indptr[rows[kk]] + 1
+    # Starting position of rows
+    j = 0
+    for kk in range(nrows):
+        j0 = out.indptr[kk]
+        out.indptr[kk] = j
+        j += j0
+    #Do the data
+    for kk in range(nnz):
+        i = rows[kk]
+        j = cols[kk]
+        val = data[kk]
+        iad = out.indptr[i]
+        out.data[iad] = val
+        out.indices[iad] = j
+        out.indptr[i] = iad+1
+    # Shift back
+    for kk in range(nrows,0,-1):
+        out.indptr[kk] = out.indptr[kk-1]
+    out.indptr[0] = 0
+
+    return CSR_to_scipy(&out)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def dense2D_to_fastcsr_cmode(complex[:, ::1] mat, int nrows, int ncols):
+    cdef int nnz = 0
+    cdef size_t ii, jj
+    cdef np.ndarray[complex, ndim=1, mode='c'] data = np.zeros(nrows*ncols, dtype=complex)
+    cdef np.ndarray[int, ndim=1, mode='c'] ind = np.zeros(nrows*ncols, dtype=np.int32)
+    cdef np.ndarray[int, ndim=1, mode='c'] ptr = np.zeros(nrows+1, dtype=np.int32)
+
+    for ii in range(nrows):
+        for jj in range(ncols):
+            if mat[ii,jj] != 0:
+                ind[nnz] = jj
+                data[nnz] = mat[ii,jj]
+                nnz += 1
+        ptr[ii+1] = nnz
+
+    if nnz < (nrows*ncols):
+        return fast_csr_matrix((data[:nnz], ind[:nnz], ptr), shape=(nrows,ncols))
+    else:
+        return fast_csr_matrix((data, ind, ptr), shape=(nrows,ncols))
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def dense2D_to_fastcsr_fmode(complex[::1, :] mat, int nrows, int ncols):
+    cdef int nnz = 0
+    cdef size_t ii, jj
+    cdef np.ndarray[complex, ndim=1, mode='c'] data = np.zeros(nrows*ncols, dtype=complex)
+    cdef np.ndarray[int, ndim=1, mode='c'] ind = np.zeros(nrows*ncols, dtype=np.int32)
+    cdef np.ndarray[int, ndim=1, mode='c'] ptr = np.zeros(nrows+1, dtype=np.int32)
+
+    for ii in range(nrows):
+        for jj in range(ncols):
+            if mat[ii,jj] != 0:
+                ind[nnz] = jj
+                data[nnz] = mat[ii,jj]
+                nnz += 1
+        ptr[ii+1] = nnz
+
+    if nnz < (nrows*ncols):
+        return fast_csr_matrix((data[:nnz], ind[:nnz], ptr), shape=(nrows,ncols))
+    else:
+        return fast_csr_matrix((data, ind, ptr), shape=(nrows,ncols))
+