Dia data layer

.gitignore

@@ -29,6 +29,7 @@ rhs*.pyx
 
 qutip/cy/*.c
 *.cpp
+!qutip/core/data/src/*.cpp
 *.dat
 qutip/core/*.h
 

doc/changes/2196.feature

@@ -0,0 +1 @@
+Add Dia data layer

qutip/core/_brtools.pyx

@@ -201,7 +201,7 @@ cdef class _EigenBasisTransform:
     def __init__(self, QobjEvo oper, bint sparse=False):
         if oper.dims[0] != oper.dims[1]:
             raise ValueError
-        if type(oper(0).data) is _data.CSR and not sparse:
+        if type(oper(0).data) in (_data.CSR, _data.Dia) and not sparse:
             oper = oper.to(Dense)
         self.oper = oper
         self.isconstant = oper.isconstant

qutip/core/data/__init__.pxd

@@ -5,6 +5,7 @@ from qutip.core.data cimport dense, csr
 from qutip.core.data.base cimport Data, idxint
 from qutip.core.data.dense cimport Dense
 from qutip.core.data.csr cimport CSR
+from qutip.core.data.dia cimport Dia
 
 from qutip.core.data.add cimport *
 from qutip.core.data.adjoint cimport *

qutip/core/data/__init__.py

@@ -3,6 +3,7 @@
 from . import dense, csr
 from .dense import Dense
 from .csr import CSR
+from .dia import Dia
 from .base import Data
 
 from .add import *
@@ -37,16 +38,22 @@
 to.add_conversions([
     (Dense, CSR, dense.from_csr, 1),
     (CSR, Dense, csr.from_dense, 1.4),
+    (Dia, Dense, dia.from_dense, 1.4),
+    (Dense, Dia, dense.from_dia, 1.2),
+    (Dia, CSR, dia.from_csr, 1),
+    (CSR, Dia, csr.from_dia, 1),
 ])
 to.register_aliases(['csr', 'CSR'], CSR)
 to.register_aliases(['Dense', 'dense'], Dense)
+to.register_aliases(['DIA', 'Dia', 'dia', 'diag'], Dia)
 
 
 from . import _creator_utils
 import numpy as np
 create.add_creators([
     (_creator_utils.is_data, _creator_utils.data_copy, 100),
     (_creator_utils.isspmatrix_csr, CSR, 80),
+    (_creator_utils.isspmatrix_dia, Dia, 80),
     (_creator_utils.is_nparray, Dense, 80),
     (_creator_utils.issparse, CSR, 20),
     (_creator_utils.true, Dense, -np.inf),

qutip/core/data/_creator_utils.py

@@ -2,13 +2,20 @@
 Define functions to use as Data creator for `create` in `convert.py`.
 """
 
-from scipy.sparse import isspmatrix_csr, issparse
+from scipy.sparse import isspmatrix_csr, issparse, isspmatrix_dia
 import numpy as np
 from .csr import CSR
 from .base import Data
 from .dense import Dense
 
-__all__ = ['is_data', 'is_nparray', 'data_copy', 'isspmatrix_csr', 'issparse']
+__all__ = [
+    'data_copy',
+    'is_data',
+    'is_nparray',
+    'isspmatrix_csr',
+    'isspmatrix_dia',
+    'issparse'
+]
 
 
 def is_data(arg):

qutip/core/data/add.pxd

@@ -2,11 +2,14 @@
 
 from qutip.core.data.csr cimport CSR
 from qutip.core.data.dense cimport Dense
+from qutip.core.data.dia cimport Dia
 
 cdef void add_dense_eq_order_inplace(Dense left, Dense right, double complex scale)
 cpdef CSR add_csr(CSR left, CSR right, double complex scale=*)
 cpdef Dense add_dense(Dense left, Dense right, double complex scale=*)
+cpdef Dia add_dia(Dia left, Dia right, double complex scale=*)
 cpdef Dense iadd_dense(Dense left, Dense right, double complex scale=*)
 
 cpdef CSR sub_csr(CSR left, CSR right)
 cpdef Dense sub_dense(Dense left, Dense right)
+cpdef Dia sub_dia(Dia left, Dia right)

qutip/core/data/add.pyx

@@ -9,10 +9,12 @@ from qutip.settings import settings
 
 from qutip.core.data.base cimport idxint, Data
 from qutip.core.data.dense cimport Dense
+from qutip.core.data.dia cimport Dia
+from qutip.core.data.tidyup cimport tidyup_dia
 from qutip.core.data.csr cimport (
     CSR, Accumulator, acc_alloc, acc_free, acc_scatter, acc_gather, acc_reset,
 )
-from qutip.core.data cimport csr, dense
+from qutip.core.data cimport csr, dense, dia
 
 cnp.import_array()
 
@@ -23,8 +25,8 @@ cdef extern from *:
 
 
 __all__ = [
-    'add', 'add_csr', 'add_dense', 'iadd_dense',
-    'sub', 'sub_csr', 'sub_dense',
+    'add', 'add_csr', 'add_dense', 'iadd_dense', 'add_dia',
+    'sub', 'sub_csr', 'sub_dense', 'sub_dia',
 ]
 
 
@@ -217,6 +219,79 @@ cpdef Dense iadd_dense(Dense left, Dense right, double complex scale=1):
     return left
 
 
+cpdef Dia add_dia(Dia left, Dia right, double complex scale=1):
+    _check_shape(left, right)
+    cdef idxint diag_left=0, diag_right=0, out_diag=0, i
+    cdef double complex *ptr_out,
+    cdef double complex *ptr_left
+    cdef double complex *ptr_right
+    cdef bint sorted=True
+    cdef Dia out = dia.empty(left.shape[0], left.shape[1], left.num_diag + right.num_diag)
+    cdef int length, size=left.shape[1]
+
+    ptr_out = out.data
+    ptr_left = left.data
+    ptr_right = right.data
+
+    with nogil:
+        while diag_left < left.num_diag and diag_right < right.num_diag:
+            if left.offsets[diag_left] == right.offsets[diag_right]:
+                out.offsets[out_diag] = left.offsets[diag_left]
+                blas.zcopy(&size, ptr_left, &_ONE, ptr_out, &_ONE)
+                blas.zaxpy(&size, &scale, ptr_right, &_ONE, ptr_out, &_ONE)
+                ptr_left += size
+                diag_left += 1
+                ptr_right += size
+                diag_right += 1
+            elif left.offsets[diag_left] <= right.offsets[diag_right]:
+                out.offsets[out_diag] = left.offsets[diag_left]
+                blas.zcopy(&size, ptr_left, &_ONE, ptr_out, &_ONE)
+                ptr_left += size
+                diag_left += 1
+            else:
+                out.offsets[out_diag] = right.offsets[diag_right]
+                blas.zcopy(&size, ptr_right, &_ONE, ptr_out, &_ONE)
+                if scale != 1:
+                    blas.zscal(&size, &scale, ptr_out, &_ONE)
+                ptr_right += size
+                diag_right += 1
+            if out_diag != 0 and out.offsets[out_diag-1] >= out.offsets[out_diag]:
+                sorted=False
+            ptr_out += size
+            out_diag += 1
+
+        if diag_left < left.num_diag:
+            for i in range(left.num_diag - diag_left):
+                out.offsets[out_diag] = left.offsets[diag_left + i]
+                if out_diag != 0 and out.offsets[out_diag-1] >= out.offsets[out_diag]:
+                    sorted=False
+                out_diag += 1
+
+            length = size * (left.num_diag - diag_left)
+            blas.zcopy(&length, ptr_left, &_ONE, ptr_out, &_ONE)
+
+
+        if diag_right < right.num_diag:
+            for i in range(right.num_diag - diag_right):
+                out.offsets[out_diag] = right.offsets[diag_right + i]
+                if out_diag != 0 and out.offsets[out_diag-1] >= out.offsets[out_diag]:
+                    sorted=False
+                out_diag += 1
+
+            length = size * (right.num_diag - diag_right)
+            blas.zcopy(&length, ptr_right, &_ONE, ptr_out, &_ONE)
+            if scale != 1:
+                blas.zscal(&length, &scale, ptr_out, &_ONE)
+
+        out.num_diag = out_diag
+
+    if not sorted:
+        dia.clean_dia(out, True)
+    if settings.core['auto_tidyup']:
+        tidyup_dia(out, settings.core['auto_tidyup_atol'], True)
+    return out
+
+
 cpdef CSR sub_csr(CSR left, CSR right):
     return add_csr(left, right, -1)
 
@@ -225,6 +300,10 @@ cpdef Dense sub_dense(Dense left, Dense right):
     return add_dense(left, right, -1)
 
 
+cpdef Dia sub_dia(Dia left, Dia right):
+    return add_dia(left, right, -1)
+
+
 from .dispatch import Dispatcher as _Dispatcher
 import inspect as _inspect
 
@@ -250,6 +329,7 @@ add.__doc__ =\
 add.add_specialisations([
     (Dense, Dense, Dense, add_dense),
     (CSR, CSR, CSR, add_csr),
+    (Dia, Dia, Dia, add_dia),
 ], _defer=True)
 
 sub = _Dispatcher(
@@ -271,6 +351,7 @@ sub.__doc__ =\
 sub.add_specialisations([
     (Dense, Dense, Dense, sub_dense),
     (CSR, CSR, CSR, sub_csr),
+    (Dia, Dia, Dia, sub_dia),
 ], _defer=True)
 
 del _inspect, _Dispatcher

qutip/core/data/adjoint.pxd

@@ -2,6 +2,7 @@
 
 from qutip.core.data.csr cimport CSR
 from qutip.core.data.dense cimport Dense
+from qutip.core.data.dia cimport Dia
 
 cpdef CSR adjoint_csr(CSR matrix)
 cpdef CSR transpose_csr(CSR matrix)
@@ -10,3 +11,7 @@ cpdef CSR conj_csr(CSR matrix)
 cpdef Dense adjoint_dense(Dense matrix)
 cpdef Dense transpose_dense(Dense matrix)
 cpdef Dense conj_dense(Dense matrix)
+
+cpdef Dia adjoint_dia(Dia matrix)
+cpdef Dia transpose_dia(Dia matrix)
+cpdef Dia conj_dia(Dia matrix)

qutip/core/data/adjoint.pyx

@@ -8,16 +8,17 @@ cimport cython
 from qutip.core.data.base cimport idxint
 from qutip.core.data.csr cimport CSR
 from qutip.core.data.dense cimport Dense
-from qutip.core.data cimport csr, dense
+from qutip.core.data.dia cimport Dia
+from qutip.core.data cimport csr, dense, dia
 
 # Import std::conj as `_conj` to avoid clashing with our 'conj' dispatcher.
 cdef extern from "<complex>" namespace "std" nogil:
     double complex _conj "conj"(double complex x)
 
 __all__ = [
-    'adjoint', 'adjoint_csr', 'adjoint_dense',
-    'conj', 'conj_csr', 'conj_dense',
-    'transpose', 'transpose_csr', 'transpose_dense',
+    'adjoint', 'adjoint_csr', 'adjoint_dense', 'adjoint_dia',
+    'conj', 'conj_csr', 'conj_dense', 'conj_dia',
+    'transpose', 'transpose_csr', 'transpose_dense', 'transpose_dia',
 ]
 
 
@@ -109,6 +110,52 @@ cpdef Dense conj_dense(Dense matrix):
     return out
 
 
+cpdef Dia adjoint_dia(Dia matrix):
+    cdef Dia out = dia.empty(matrix.shape[1], matrix.shape[0], matrix.num_diag)
+    cdef size_t i, new_i,
+    cdef idxint new_offset, j
+    with nogil:
+        out.num_diag = matrix.num_diag
+        for i in range(matrix.num_diag):
+            new_i = matrix.num_diag - i - 1
+            new_offset = out.offsets[new_i] = -matrix.offsets[i]
+            for j in range(out.shape[1]):
+                if (j < new_offset) or (j - new_offset >= matrix.shape[1]):
+                    out.data[new_i * out.shape[1] + j] = 0.
+                else:
+                    out.data[new_i * out.shape[1] + j] = _conj(matrix.data[i * matrix.shape[1] + j - new_offset])
+    return out
+
+
+cpdef Dia transpose_dia(Dia matrix):
+    cdef Dia out = dia.empty(matrix.shape[1], matrix.shape[0], matrix.num_diag)
+    cdef size_t i, new_i,
+    cdef idxint new_offset, j
+    with nogil:
+        out.num_diag = matrix.num_diag
+        for i in range(matrix.num_diag):
+            new_i = matrix.num_diag - i - 1
+            new_offset = out.offsets[new_i] = -matrix.offsets[i]
+            for j in range(out.shape[1]):
+                if (j < new_offset) or (j - new_offset >= matrix.shape[1]):
+                    out.data[new_i * out.shape[1] + j] = 0.
+                else:
+                    out.data[new_i * out.shape[1] + j] = matrix.data[i * matrix.shape[1] + j - new_offset]
+    return out
+
+
+cpdef Dia conj_dia(Dia matrix):
+    cdef Dia out = dia.empty_like(matrix)
+    cdef size_t i, j
+    with nogil:
+        out.num_diag = matrix.num_diag
+        for i in range(matrix.num_diag):
+            out.offsets[i] = matrix.offsets[i]
+            for j in range(matrix.shape[1]):
+                out.data[i * matrix.shape[1] + j] = _conj(matrix.data[i * matrix.shape[1] + j])
+    return out
+
+
 from .dispatch import Dispatcher as _Dispatcher
 import inspect as _inspect
 
@@ -125,6 +172,7 @@ adjoint.__doc__ = """Hermitian adjoint (matrix conjugate transpose)."""
 adjoint.add_specialisations([
     (Dense, Dense, adjoint_dense),
     (CSR, CSR, adjoint_csr),
+    (Dia, Dia, adjoint_dia),
 ], _defer=True)
 
 transpose = _Dispatcher(
@@ -140,6 +188,7 @@ transpose.__doc__ = """Transpose of a matrix."""
 transpose.add_specialisations([
     (Dense, Dense, transpose_dense),
     (CSR, CSR, transpose_csr),
+    (Dia, Dia, transpose_dia),
 ], _defer=True)
 
 conj = _Dispatcher(
@@ -155,6 +204,7 @@ conj.__doc__ = """Element-wise conjugation of a matrix."""
 conj.add_specialisations([
     (Dense, Dense, conj_dense),
     (CSR, CSR, conj_csr),
+    (Dia, Dia, conj_dia),
 ], _defer=True)
 
 del _inspect, _Dispatcher

qutip/core/data/constant.py

@@ -3,8 +3,9 @@
 # (e.g. `create`) should not be here, but should be defined within the
 # higher-level components of QuTiP instead.
 
-from . import csr, dense
+from . import csr, dense, dia
 from .csr import CSR
+from .dia import Dia
 from .dense import Dense
 from .base import Data
 from .dispatch import Dispatcher as _Dispatcher
@@ -40,6 +41,7 @@
     """
 zeros.add_specialisations([
     (CSR, csr.zeros),
+    (Dia, dia.zeros),
     (Dense, dense.zeros),
 ], _defer=True)
 
@@ -70,6 +72,7 @@
     """
 identity.add_specialisations([
     (CSR, csr.identity),
+    (Dia, dia.identity),
     (Dense, dense.identity),
 ], _defer=True)
 

qutip/core/data/csr.pxd

@@ -12,6 +12,7 @@ cimport numpy as cnp
 
 from qutip.core.data cimport base
 from qutip.core.data.dense cimport Dense
+from qutip.core.data.dia cimport Dia
 
 cdef class CSR(base.Data):
     cdef double complex *data
@@ -161,6 +162,7 @@ cpdef CSR from_dense(Dense matrix)
 cdef CSR from_coo_pointers(base.idxint *rows, base.idxint *cols, double complex *data,
                            base.idxint n_rows, base.idxint n_cols, base.idxint nnz,
                            double tol=*)
+cpdef CSR from_dia(Dia matrix)
 
 cpdef CSR _from_csr_blocks(base.idxint[:] block_rows, base.idxint[:] block_cols, CSR[:] block_ops,
                           base.idxint n_blocks, base.idxint block_size)