Add COO format

doc/changes/2314.feature

@@ -0,0 +1 @@
+Add alternative COO format

qutip/core/cy/qobjevo.pxd

@@ -22,3 +22,4 @@ cdef class QobjEvo:
     cdef double complex _expect_dense(QobjEvo self, double t, Dense state) except *
 
     cpdef Data matmul_data(QobjEvo self, object t, Data state, Data out=*)
+

qutip/core/cy/qobjevo.pyx

@@ -10,11 +10,11 @@ import qutip
 from .. import Qobj
 from .. import data as _data
 from ..dimensions import Dimensions
-from ..coefficient import coefficient, CompilationOptions
-from ._element import *
+from ..coefficient import coefficient
 from qutip.settings import settings
 
-from qutip.core.cy._element cimport _BaseElement
+from qutip.core.cy._element cimport (_BaseElement, _ConstantElement, _EvoElement,
+        _FuncElement)
 from qutip.core.data cimport Dense, Data, dense
 from qutip.core.data.expect cimport *
 from qutip.core.data.reshape cimport (column_stack_dense, column_unstack_dense)
@@ -1045,11 +1045,12 @@ cdef class QobjEvo:
                     # `state` was reshaped inplace, restore it's original shape
                     column_unstack_dense(state, nrow, inplace=state.fortran)
         else:
+            expect_func = expect_data_dense_ket if state.shape[1] == 1 else expect_data_dense_dm
             for element in self.elements:
                 part = (<_BaseElement> element)
                 coeff = part.coeff(t)
                 part_data = part.data(t)
-                out += coeff * expect_data_dense(part_data, state)
+                out += coeff * expect_func(part_data, state)
         return out
 
     def matmul(self, t, state):
@@ -1083,14 +1084,17 @@ cdef class QobjEvo:
 
     cpdef Data matmul_data(QobjEvo self, object t, Data state, Data out=None):
         """Compute ``out += self(t) @ state``"""
-        cdef _BaseElement part
         t = self._prepare(t, state)
+        if len(self.elements) == 0:
+            return self.elements[0].matmul_data_t(t, state, out=out)
+
         if out is None and type(state) is Dense:
             out = dense.zeros(self.shape[0], state.shape[1],
                               (<Dense> state).fortran)
         elif out is None:
             out = _data.zeros[type(state)](self.shape[0], state.shape[1])
 
+        cdef _BaseElement part
         for element in self.elements:
             part = (<_BaseElement> element)
             out = part.matmul_data_t(t, state, out)

qutip/core/data/__init__.pxd

@@ -1,10 +1,11 @@
 #cython: language_level=3
 
 # Package-level relative imports in Cython (0.29.17) are temperamental.
-from qutip.core.data cimport dense, csr
+from qutip.core.data cimport dense, csr, coo
 from qutip.core.data.base cimport Data, idxint
-from qutip.core.data.dense cimport Dense
+from qutip.core.data.coo cimport COO
 from qutip.core.data.csr cimport CSR
+from qutip.core.data.dense cimport Dense
 from qutip.core.data.dia cimport Dia
 
 from qutip.core.data.add cimport *

qutip/core/data/__init__.py

@@ -1,7 +1,8 @@
 # First-class type imports
 
-from . import dense, csr
+from . import dense, csr, coo
 from .dense import Dense
+from .coo import COO
 from .csr import CSR
 from .dia import Dia
 from .base import Data
@@ -36,13 +37,18 @@
 
 from .convert import to, create
 to.add_conversions([
+    (Dense, COO, dense.from_coo, 1),
+    (COO, Dense, coo.from_dense, 1.6),
     (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),
+    (COO, CSR, coo.from_csr, 1),
+    (CSR, COO, csr.from_coo, 1),
     (Dia, CSR, dia.from_csr, 1),
     (CSR, Dia, csr.from_dia, 1),
 ])
+to.register_aliases(['coo', 'COO'], COO)
 to.register_aliases(['csr', 'CSR'], CSR)
 to.register_aliases(['Dense', 'dense'], Dense)
 to.register_aliases(['DIA', 'Dia', 'dia', 'diag'], Dia)
@@ -52,6 +58,7 @@
 import numpy as np
 create.add_creators([
     (_creator_utils.is_data, _creator_utils.data_copy, 100),
+    (_creator_utils.isspmatrix_coo, COO, 80),
     (_creator_utils.isspmatrix_csr, CSR, 80),
     (_creator_utils.isspmatrix_dia, Dia, 80),
     (_creator_utils.is_nparray, Dense, 80),

qutip/core/data/_creator_utils.py

@@ -2,19 +2,25 @@
 Define functions to use as Data creator for `create` in `convert.py`.
 """
 
-from scipy.sparse import isspmatrix_csr, issparse, isspmatrix_dia
+from scipy.sparse import (
+    isspmatrix_coo,
+    isspmatrix_csr,
+    issparse,
+    isspmatrix_dia,
+)
 import numpy as np
 from .csr import CSR
 from .base import Data
 from .dense import Dense
 
 __all__ = [
-    'data_copy',
-    'is_data',
-    'is_nparray',
-    'isspmatrix_csr',
-    'isspmatrix_dia',
-    'issparse'
+    "data_copy",
+    "is_data",
+    "is_nparray",
+    "isspmatrix_coo",
+    "isspmatrix_csr",
+    "isspmatrix_dia",
+    "issparse",
 ]
 
 
@@ -32,7 +38,7 @@ def true(arg):
 
 def data_copy(arg, shape, copy=True):
     if shape is not None and shape != arg.shape:
-        raise ValueError("".join([
-            "shapes do not match: ", str(shape), " and ", str(arg.shape),
-        ]))
+        raise ValueError(
+            f"shapes do not match: {str(shape)} and {str(arg.shape)}"
+        )
     return arg.copy() if copy else arg

qutip/core/data/add.pxd

@@ -1,15 +1,21 @@
 #cython: language_level=3
-
+from qutip.core.data.coo cimport COO
 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 COO add_coo(COO left, COO right, double complex scale=*)
+cpdef Dense iadd_dense_coo(Dense left, COO 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 Dense add_dense(Dense left, Dense right, double complex scale=*)
+cpdef Dense iadd_dense_dense(Dense left, Dense right, double complex scale=*)
+
+cdef void add_dense_eq_order_inplace(Dense left, Dense right, double complex scale)
+
+cpdef COO sub_coo(COO left, COO right)
 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

@@ -2,24 +2,26 @@
 #cython: boundscheck=False, wraparound=False, initializedcheck=False
 
 cimport cython
-import numpy as np
 cimport numpy as cnp
 from scipy.linalg cimport cython_blas as blas
 from qutip.settings import settings
+from libc.string cimport memcpy
 
 from qutip.core.data.base cimport idxint, Data
 from qutip.core.data.dense cimport Dense
+from qutip.core.data.coo cimport COO
 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, dia
+from qutip.core.data cimport coo, csr, dia
 
 cnp.import_array()
 
 __all__ = [
-    'add', 'add_csr', 'add_dense', 'iadd_dense', 'add_dia',
+    'add', 'add_coo', 'add_csr', 'add_dense', 'add_dia',
+    'iadd_dense', 'iadd_dense_coo', 'iadd_dense_dense', 
     'sub', 'sub_csr', 'sub_dense', 'sub_dia',
 ]
 
@@ -115,6 +117,106 @@ cdef idxint _add_csr_scale(Accumulator *acc, CSR a, CSR b, CSR c,
         c.row_index[row + 1] = nnz
     return nnz
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.profile(False)
+@cython.linetrace(False)   
+cpdef COO add_coo(COO left, COO right, double complex scale=1):
+    """
+    Matrix addition of `left` and `right` for COO inputs and output.  If given,
+    `right` is multiplied by `scale`, so the full operation is
+        ``out := left + scale*right``
+    The two matrices must be of exactly the same shape.
+
+    Parameters
+    ----------
+    left : COO
+        Matrix to be added.
+    right : COO
+        Matrix to be added.  If `scale` is given, this matrix will be
+        multiplied by `scale` before addition.
+    scale : optional double complex (1)
+        The scalar value to multiply `right` by before addition.
+
+    Returns
+    -------
+    out : COO
+        The result `left + scale*right`.
+    """
+    _check_shape(left, right)
+    cdef idxint left_nnz = coo.nnz(left)
+    cdef idxint right_nnz = coo.nnz(right)
+    cdef idxint nnz = left_nnz + right_nnz
+    cdef idxint i
+    cdef COO out
+
+    # Fast paths for zero matrices.
+    if right_nnz == 0 or scale == 0:
+        return left.copy()
+    if left_nnz == 0:
+        out = right.copy()
+        # Fast path if the multiplication is a no-op.
+        if scale != 1:
+            for i in range(right_nnz):
+                out.data[i] *= scale
+        return out
+
+    # Main path.
+    out = coo.empty(left.shape[0], left.shape[1], nnz)
+
+    for ptr in range(left_nnz):
+        out.data[ptr] = left.data[ptr]
+        out.col_index[ptr] = left.col_index[ptr]
+        out.row_index[ptr] = left.row_index[ptr]
+
+    for ptr in range(right_nnz):
+        out.data[ptr + left_nnz] = scale * right.data[ptr]
+        out.col_index[ptr + left_nnz] = right.col_index[ptr]
+        out.row_index[ptr + left_nnz] = right.row_index[ptr]
+    out._nnz = nnz
+    return out
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.profile(False)
+@cython.linetrace(False)   
+cpdef Dense iadd_dense_coo(Dense left, COO right, double complex scale=1):
+    """
+    Matrix addition of `left` and `right` for COO inputs and output.  If given,
+    `right` is multiplied by `scale`, so the full operation is
+        ``out := left + scale*right``
+    The two matrices must be of exactly the same shape.
+
+    Parameters
+    ----------
+    left : COO
+        Matrix to be added.
+    right : COO
+        Matrix to be added.  If `scale` is given, this matrix will be
+        multiplied by `scale` before addition.
+    scale : optional double complex (1)
+        The scalar value to multiply `right` by before addition.
+
+    Returns
+    -------
+    out : COO
+        The result `left + scale*right`.
+    """
+    _check_shape(left, right)
+
+    cdef idxint ptr, row, col, row_stride, col_stride
+    row_stride = 1 if left.fortran else left.shape[1]
+    col_stride = left.shape[0] if left.fortran else 1
+    # Fast paths for zero matrices.
+    if coo.nnz(right) == 0 or scale == 0:
+        return left
+
+    for ptr in range(coo.nnz(right)):
+        row = right.row_index[ptr]
+        col = right.col_index[ptr]
+        left.data[row_stride * row + col_stride * col] = right.data[ptr]
+    return left
 
 cpdef CSR add_csr(CSR left, CSR right, double complex scale=1):
     """
@@ -198,7 +300,7 @@ cpdef Dense add_dense(Dense left, Dense right, double complex scale=1):
     return out
 
 
-cpdef Dense iadd_dense(Dense left, Dense right, double complex scale=1):
+cpdef Dense iadd_dense_dense(Dense left, Dense right, double complex scale=1):
     _check_shape(left, right)
     cdef int size = left.shape[0] * left.shape[1]
     cdef int dim1, dim2
@@ -286,6 +388,8 @@ cpdef Dia add_dia(Dia left, Dia right, double complex scale=1):
         tidyup_dia(out, settings.core['auto_tidyup_atol'], True)
     return out
 
+cpdef COO sub_coo(COO left, COO right):
+    return add_coo(left, right, -1)
 
 cpdef CSR sub_csr(CSR left, CSR right):
     return add_csr(left, right, -1)
@@ -322,11 +426,29 @@ add.__doc__ =\
     scalar.
     """
 add.add_specialisations([
-    (Dense, Dense, Dense, add_dense),
+    (COO, COO, COO, add_coo),
     (CSR, CSR, CSR, add_csr),
+    (Dense, Dense, Dense, add_dense),
     (Dia, Dia, Dia, add_dia),
 ], _defer=True)
 
+from .convert import to as _to
+
+cpdef Dense iadd_dense_data(Dense left, Data right, double complex scale=1):
+    cdef Data out = add(left, right)
+    cdef idxint size = left.shape[0]* left.shape[1]
+    cdef Dense out_dense = _to(Dense, out)
+    memcpy(left.data, out_dense.data, size*sizeof(double complex))
+    return left
+
+cpdef Dense iadd_dense(Dense left, Data right, double complex scale=1):
+    if type(right) == Dense:
+        return iadd_dense_dense(left, right, scale)
+    elif type(right) == COO:
+        return iadd_dense_coo(left, right, scale)
+    else:
+        return iadd_dense_data(left, right, scale)
+
 sub = _Dispatcher(
     _inspect.Signature([
         _inspect.Parameter('left', _inspect.Parameter.POSITIONAL_ONLY),
@@ -344,8 +466,9 @@ sub.__doc__ =\
     where `left` and `right` are matrices.
     """
 sub.add_specialisations([
-    (Dense, Dense, Dense, sub_dense),
+    (COO, COO, COO, sub_coo),
     (CSR, CSR, CSR, sub_csr),
+    (Dense, Dense, Dense, sub_dense),
     (Dia, Dia, Dia, sub_dia),
 ], _defer=True)
 

qutip/core/data/adjoint.pxd

@@ -3,6 +3,11 @@
 from qutip.core.data.csr cimport CSR
 from qutip.core.data.dense cimport Dense
 from qutip.core.data.dia cimport Dia
+from qutip.core.data.coo cimport COO
+
+cpdef COO adjoint_coo(COO matrix)
+cpdef COO transpose_coo(COO matrix)
+cpdef COO conj_coo(COO matrix)
 
 cpdef CSR adjoint_csr(CSR matrix)
 cpdef CSR transpose_csr(CSR matrix)