1#cython: language_level=3
2# This file is part of QuTiP: Quantum Toolbox in Python.
3#
4#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
5#    All rights reserved.
6#
7#    Redistribution and use in source and binary forms, with or without
8#    modification, are permitted provided that the following conditions are
9#    met:
10#
11#    1. Redistributions of source code must retain the above copyright notice,
12#       this list of conditions and the following disclaimer.
13#
14#    2. Redistributions in binary form must reproduce the above copyright
15#       notice, this list of conditions and the following disclaimer in the
16#       documentation and/or other materials provided with the distribution.
17#
18#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
19#       of its contributors may be used to endorse or promote products derived
20#       from this software without specific prior written permission.
21#
22#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
23#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
24#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
25#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
26#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
27#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
28#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
29#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
30#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
31#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
32#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
33###############################################################################
34import numpy as np
35cimport numpy as np
36cimport cython
37cimport libc.math
38
39DTYPE = np.float64
40ctypedef np.float64_t DTYPE_t
41
42ITYPE = np.int32
43ctypedef np.int32_t ITYPE_t
44
45CTYPE = np.complex128
46ctypedef np.complex128_t CTYPE_t
47
48CTYPE = np.int64
49ctypedef np.int64_t LTYPE_t
50
51@cython.boundscheck(False)
52@cython.wraparound(False)
53cpdef CTYPE_t cy_overlap(object op1, object op2):
54
55    cdef Py_ssize_t row
56    cdef CTYPE_t tr = 0.0
57
58    op1 = op1.T.tocsr()
59
60    cdef int col1, row1_idx_start, row1_idx_end
61    cdef np.ndarray[CTYPE_t, ndim=1, mode="c"] data1 = op1.data.conj()
62    cdef np.ndarray[ITYPE_t, ndim=1, mode="c"] idx1 = op1.indices
63    cdef np.ndarray[ITYPE_t, ndim=1, mode="c"] ptr1 = op1.indptr
64
65    cdef int col2, row2_idx_start, row2_idx_end
66    cdef np.ndarray[CTYPE_t, ndim=1, mode="c"] data2 = op2.data
67    cdef np.ndarray[ITYPE_t, ndim=1, mode="c"] idx2 = op2.indices
68    cdef np.ndarray[ITYPE_t, ndim=1, mode="c"] ptr2 = op2.indptr
69
70    cdef int num_rows = ptr1.shape[0]-1
71
72    for row in range(num_rows):
73
74        row1_idx_start = ptr1[row]
75        row1_idx_end = ptr1[row + 1]
76        for row1_idx from row1_idx_start <= row1_idx < row1_idx_end:
77            col1 = idx1[row1_idx]
78
79            row2_idx_start = ptr2[col1]
80            row2_idx_end = ptr2[col1 + 1]
81            for row2_idx from row2_idx_start <= row2_idx < row2_idx_end:
82                col2 = idx2[row2_idx]
83
84                if col2 == row:
85                    tr += data1[row1_idx] * data2[row2_idx]
86
87    return tr / op1.shape[0]
88
89
90@cython.boundscheck(False)
91@cython.wraparound(False)
92cpdef cy_grape_inner(U, np.ndarray[DTYPE_t, ndim=3, mode="c"] u,
93                     int r, int J, int M, U_b_list, U_f_list, H_ops,
94                     float dt, float eps, float alpha, float beta,
95                     int phase_sensitive,
96                     int use_u_limits, float u_min, float u_max):
97
98    cdef int j, k
99
100    for m in range(M-1):
101        P = U_b_list[m] * U
102        for j in range(J):
103            Q = 1j * dt * H_ops[j] * U_f_list[m]
104
105            if phase_sensitive:
106                du = - cy_overlap(P, Q)
107            else:
108                du = - 2 * cy_overlap(P, Q) * cy_overlap(U_f_list[m], P)
109
110            if alpha > 0.0:
111                # penalty term for high power control signals u
112                du += -2 * alpha * u[r, j, m] * dt
113
114            if beta:
115                # penalty term for late control signals u
116                du += -2 * beta * m ** 2 * u[r, j, m] * dt
117
118            u[r + 1, j, m] = u[r, j, m] + eps * du.real
119
120            if use_u_limits:
121                if u[r + 1, j, m] < u_min:
122                    u[r + 1, j, m] = u_min
123                elif u[r + 1, j, m] > u_max:
124                    u[r + 1, j, m] = u_max
125
126    for j in range(J):
127        u[r + 1, j, M-1] = u[r + 1, j, M-2]
128