-
Notifications
You must be signed in to change notification settings - Fork 627
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add coo format #2314
base: master
Are you sure you want to change the base?
Add coo format #2314
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
Add alternative COO format |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
for element in self.elements: | ||
part = (<_BaseElement> element) | ||
out = part.matmul_data_t(t, state, out) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -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) | ||||||
Comment on lines
+120
to
+123
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These are defined at the file level and not needed here. |
||||||
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 | ||||||
Comment on lines
+167
to
+176
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There should be a clean-up step: if Also it would be better to keep the entry sorted. If the index were sorted, this collision checking would be a lot smoother. |
||||||
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): | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. When only one data type is used, the function name is {operation}_{dtype}. |
||||||
_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): | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This can be inline in the function below. |
||||||
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 | ||||||
Comment on lines
+437
to
+442
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The scale is not applied. |
||||||
|
||||||
cpdef Dense iadd_dense(Dense left, Data right, double complex scale=1): | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. For function with multiple type:
Suggested change
|
||||||
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) | ||||||
|
||||||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if
self.elements
is empty, callingself.elements[0]
is a bug.