Merge pull request #2126 from Ericgig/operkettrace

Add operation to trace oper-ket without reshaping

doc/changes/2126.feature

@@ -0,0 +1 @@
+Create the `trace_oper_ket` operation

qutip/core/data/trace.pxd

@@ -4,3 +4,6 @@ from qutip.core.data cimport CSR, Dense
 
 cpdef double complex trace_csr(CSR matrix) nogil except *
 cpdef double complex trace_dense(Dense matrix) nogil except *
+
+cpdef double complex trace_oper_ket_csr(CSR matrix) nogil except *
+cpdef double complex trace_oper_ket_dense(Dense matrix) nogil except *

qutip/core/data/trace.pyx

@@ -2,11 +2,14 @@
 #cython: boundscheck=False, wraparound=False, initializedcheck=False
 
 cimport cython
+from libc.math cimport sqrt
 
 from qutip.core.data cimport Data, CSR, Dense
+from qutip.core.data cimport base
 
 __all__ = [
     'trace', 'trace_csr', 'trace_dense',
+    'trace_oper_ket', 'trace_oper_ket_csr', 'trace_oper_ket_dense',
 ]
 
 
@@ -17,6 +20,13 @@ cdef void _check_shape(Data matrix) nogil except *:
         ]))
 
 
+cdef void _check_shape_oper_ket(int N, Data matrix) nogil except *:
+    if matrix.shape[0] != N * N or matrix.shape[1] != 1:
+        raise ValueError("".join([
+            "matrix ", str(matrix.shape), " is not a stacked square matrix."
+        ]))
+
+
 cpdef double complex trace_csr(CSR matrix) nogil except *:
     _check_shape(matrix)
     cdef size_t row, ptr
@@ -39,6 +49,28 @@ cpdef double complex trace_dense(Dense matrix) nogil except *:
     return trace
 
 
+cpdef double complex trace_oper_ket_csr(CSR matrix) nogil except *:
+    cdef size_t N = <size_t>sqrt(matrix.shape[0])
+    _check_shape_oper_ket(N, matrix)
+    cdef size_t row
+    cdef double complex trace = 0
+    cdef size_t stride = N + 1
+    for row in range(N):
+        if matrix.row_index[row * stride] != matrix.row_index[row * stride + 1]:
+            trace += matrix.data[matrix.row_index[row * stride]]
+    return trace
+
+cpdef double complex trace_oper_ket_dense(Dense matrix) nogil except *:
+    cdef size_t N = <size_t>sqrt(matrix.shape[0])
+    _check_shape_oper_ket(N, matrix)
+    cdef double complex trace = 0
+    cdef size_t ptr = 0
+    cdef size_t stride = N + 1
+    for ptr in range(N):
+        trace += matrix.data[ptr * stride]
+    return trace
+
+
 from .dispatch import Dispatcher as _Dispatcher
 import inspect as _inspect
 
@@ -58,4 +90,20 @@ trace.add_specialisations([
     (Dense, trace_dense),
 ], _defer=True)
 
+trace_oper_ket = _Dispatcher(
+    _inspect.Signature([
+        _inspect.Parameter('matrix', _inspect.Parameter.POSITIONAL_OR_KEYWORD),
+    ]),
+    name='trace_oper_ket',
+    module=__name__,
+    inputs=('matrix',),
+    out=False,
+)
+trace_oper_ket.__doc__ =\
+    """Compute the trace (sum of digaonal elements) of a stacked square matrix ."""
+trace_oper_ket.add_specialisations([
+    (CSR, trace_oper_ket_csr),
+    (Dense, trace_oper_ket_dense),
+], _defer=True)
+
 del _inspect, _Dispatcher

qutip/solver/mcsolve.py

@@ -225,12 +225,12 @@ def reset(self, hard=False):
 
     def _prob_func(self, state):
         if self.issuper:
-            return _data.norm.trace(unstack_columns(state))
+            return _data.trace_oper_ket(state).real
         return _data.norm.l2(state)**2
 
     def _norm_func(self, state):
         if self.issuper:
-            return _data.norm.trace(unstack_columns(state))
+            return _data.trace_oper_ket(state).real
         return _data.norm.l2(state)
 
     def _find_collapse_time(self, norm_old, norm, t_prev, t_final):
@@ -339,14 +339,8 @@ class MCSolver(MultiTrajSolver):
         (see :class:`qutip.QobjEvo`'s documentation). They must be operators
         even if ``H`` is a superoperator.
 
-    options : SolverOptions, [optional]
+    options : dict, [optional]
         Options for the evolution.
-
-    seed : int, SeedSequence, list, [optional]
-        Seed for the random number generator. It can be a single seed used to
-        spawn seeds for each trajectory or a list of seed, one for each
-        trajectory. Seeds are saved in the result and can be reused with::
-            seeds=prev_result.seeds
     """
     name = "mcsolve"
     resultclass = McResult

qutip/tests/core/data/test_mathematics.py

@@ -804,6 +804,26 @@ def op_numpy(self, matrix):
     ]
 
 
+class TestTrace_oper_ket(UnaryOpMixin):
+    def op_numpy(self, matrix):
+        N = int(matrix.shape[0] ** 0.5)
+        return np.sum(np.diag(matrix.reshape((N, N))))
+
+    shapes = [
+        (pytest.param((100, 1), id="oper-ket"),),
+    ]
+    bad_shapes = [
+        (pytest.param((1, 100), id="bra"),),
+        (pytest.param((99, 1), id="ket"),),
+        (pytest.param((99, 99), id="ket"),),
+        (pytest.param((2, 99), id="nonsquare"),),
+    ]
+    specialisations = [
+        pytest.param(data.trace_oper_ket_csr, CSR, complex),
+        pytest.param(data.trace_oper_ket_dense, Dense, complex),
+    ]
+
+
 class TestPow(UnaryOpMixin):
     def op_numpy(self, matrix, n):
         return np.linalg.matrix_power(matrix, n)