1#!python
2#cython: language_level=3
3# This file is part of QuTiP: Quantum Toolbox in Python.
4#
5#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
6#    All rights reserved.
7#
8#    Redistribution and use in source and binary forms, with or without
9#    modification, are permitted provided that the following conditions are
10#    met:
11#
12#    1. Redistributions of source code must retain the above copyright notice,
13#       this list of conditions and the following disclaimer.
14#
15#    2. Redistributions in binary form must reproduce the above copyright
16#       notice, this list of conditions and the following disclaimer in the
17#       documentation and/or other materials provided with the distribution.
18#
19#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
20#       of its contributors may be used to endorse or promote products derived
21#       from this software without specific prior written permission.
22#
23#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
24#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
25#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
26#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
27#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
28#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
29#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
30#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
31#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
32#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
33#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
34###############################################################################
35import numpy as np
36from qiskit.providers.aer.pulse.qutip_lite.fastsparse import fast_csr_matrix
37cimport numpy as cnp
38from libc.math cimport abs, fabs, sqrt
39from libcpp cimport bool
40cimport cython
41cnp.import_array()
42from libc.string cimport memset
43
44cdef extern from "numpy/arrayobject.h" nogil:
45    void PyArray_ENABLEFLAGS(cnp.ndarray arr, int flags)
46    void PyDataMem_FREE(void * ptr)
47    void PyDataMem_RENEW(void * ptr, size_t size)
48    void PyDataMem_NEW(size_t size)
49
50
51cdef extern from "<complex>" namespace "std" nogil:
52    double abs(double complex x)
53    double real(double complex x)
54    double imag(double complex x)
55
56cdef extern from "<complex>" namespace "std" nogil:
57    double cabs "abs" (double complex x)
58
59cdef inline int int_max(int x, int y):
60    return x ^ ((x ^ y) & -(x < y))
61
62include "parameters.pxi"
63
64@cython.boundscheck(False)
65@cython.wraparound(False)
66def _sparse_bandwidth(
67        int[::1] idx,
68        int[::1] ptr,
69        int nrows):
70    """
71    Calculates the max (mb), lower(lb), and upper(ub) bandwidths of a
72    csr_matrix.
73    """
74    cdef int ldist
75    cdef int lb = -nrows
76    cdef int ub = -nrows
77    cdef int mb = 0
78    cdef size_t ii, jj
79
80    for ii in range(nrows):
81        for jj in range(ptr[ii], ptr[ii + 1]):
82            ldist = ii - idx[jj]
83            lb = int_max(lb, ldist)
84            ub = int_max(ub, -ldist)
85            mb = int_max(mb, ub + lb + 1)
86
87    return mb, lb, ub
88
89
90@cython.boundscheck(False)
91@cython.wraparound(False)
92def _sparse_profile(int[::1] idx,
93        int[::1] ptr,
94        int nrows):
95    cdef int ii, jj, temp, ldist=0
96    cdef LTYPE_t pro = 0
97    for ii in range(nrows):
98        temp = 0
99        for jj in range(ptr[ii], ptr[ii + 1]):
100            ldist = idx[jj] - ii
101            temp = int_max(temp, ldist)
102        pro += temp
103    return pro
104
105
106@cython.boundscheck(False)
107@cython.wraparound(False)
108def _sparse_permute(
109        cnp.ndarray[cython.numeric, ndim=1] data,
110        int[::1] idx,
111        int[::1] ptr,
112        int nrows,
113        int ncols,
114        cnp.ndarray[ITYPE_t, ndim=1] rperm,
115        cnp.ndarray[ITYPE_t, ndim=1] cperm,
116        int flag):
117    """
118    Permutes the rows and columns of a sparse CSR or CSC matrix according to
119    the permutation arrays rperm and cperm, respectively.
120    Here, the permutation arrays specify the new order of the rows and columns.
121    i.e. [0,1,2,3,4] -> [3,0,4,1,2].
122    """
123    cdef int ii, jj, kk, k0, nnz
124    cdef cnp.ndarray[cython.numeric] new_data = np.zeros_like(data)
125    cdef cnp.ndarray[ITYPE_t] new_idx = np.zeros_like(idx)
126    cdef cnp.ndarray[ITYPE_t] new_ptr = np.zeros_like(ptr)
127    cdef cnp.ndarray[ITYPE_t] perm_r
128    cdef cnp.ndarray[ITYPE_t] perm_c
129    cdef cnp.ndarray[ITYPE_t] inds
130
131    if flag == 0:  # CSR matrix
132        if rperm.shape[0] != 0:
133            inds = np.argsort(rperm).astype(ITYPE)
134            perm_r = np.arange(rperm.shape[0], dtype=ITYPE)[inds]
135
136            for jj in range(nrows):
137                ii = perm_r[jj]
138                new_ptr[ii + 1] = ptr[jj + 1] - ptr[jj]
139
140            for jj in range(nrows):
141                new_ptr[jj + 1] = new_ptr[jj+1] + new_ptr[jj]
142
143            for jj in range(nrows):
144                k0 = new_ptr[perm_r[jj]]
145                for kk in range(ptr[jj], ptr[jj + 1]):
146                    new_idx[k0] = idx[kk]
147                    new_data[k0] = data[kk]
148                    k0 = k0 + 1
149
150        if cperm.shape[0] != 0:
151            inds = np.argsort(cperm).astype(ITYPE)
152            perm_c = np.arange(cperm.shape[0], dtype=ITYPE)[inds]
153            nnz = new_ptr[new_ptr.shape[0] - 1]
154            for jj in range(nnz):
155                new_idx[jj] = perm_c[new_idx[jj]]
156
157    elif flag == 1:  # CSC matrix
158        if cperm.shape[0] != 0:
159            inds = np.argsort(cperm).astype(ITYPE)
160            perm_c = np.arange(cperm.shape[0], dtype=ITYPE)[inds]
161
162            for jj in range(ncols):
163                ii = perm_c[jj]
164                new_ptr[ii + 1] = ptr[jj + 1] - ptr[jj]
165
166            for jj in range(ncols):
167                new_ptr[jj + 1] = new_ptr[jj + 1] + new_ptr[jj]
168
169            for jj in range(ncols):
170                k0 = new_ptr[perm_c[jj]]
171                for kk in range(ptr[jj], ptr[jj + 1]):
172                    new_idx[k0] = idx[kk]
173                    new_data[k0] = data[kk]
174                    k0 = k0 + 1
175
176        if rperm.shape[0] != 0:
177            inds = np.argsort(rperm).astype(ITYPE)
178            perm_r = np.arange(rperm.shape[0], dtype=ITYPE)[inds]
179            nnz = new_ptr[new_ptr.shape[0] - 1]
180            for jj in range(nnz):
181                new_idx[jj] = perm_r[new_idx[jj]]
182
183    return new_data, new_idx, new_ptr
184
185
186@cython.boundscheck(False)
187@cython.wraparound(False)
188def _sparse_reverse_permute(
189        cnp.ndarray[cython.numeric, ndim=1] data,
190        int[::1] idx,
191        int[::1] ptr,
192        int nrows,
193        int ncols,
194        cnp.ndarray[ITYPE_t, ndim=1] rperm,
195        cnp.ndarray[ITYPE_t, ndim=1] cperm,
196        int flag):
197    """
198    Reverse permutes the rows and columns of a sparse CSR or CSC matrix
199    according to the original permutation arrays rperm and cperm, respectively.
200    """
201    cdef int ii, jj, kk, k0, nnz
202    cdef cnp.ndarray[cython.numeric, ndim=1] new_data = np.zeros_like(data)
203    cdef cnp.ndarray[ITYPE_t, ndim=1] new_idx = np.zeros_like(idx)
204    cdef cnp.ndarray[ITYPE_t, ndim=1] new_ptr = np.zeros_like(ptr)
205
206    if flag == 0:  # CSR matrix
207        if rperm.shape[0] != 0:
208            for jj in range(nrows):
209                ii = rperm[jj]
210                new_ptr[ii + 1] = ptr[jj + 1] - ptr[jj]
211
212            for jj in range(nrows):
213                new_ptr[jj + 1] = new_ptr[jj + 1] + new_ptr[jj]
214
215            for jj in range(nrows):
216                k0 = new_ptr[rperm[jj]]
217                for kk in range(ptr[jj], ptr[jj + 1]):
218                    new_idx[k0] = idx[kk]
219                    new_data[k0] = data[kk]
220                    k0 = k0 + 1
221
222        if cperm.shape[0] > 0:
223            nnz = new_ptr[new_ptr.shape[0] - 1]
224            for jj in range(nnz):
225                new_idx[jj] = cperm[new_idx[jj]]
226
227    if flag == 1:  # CSC matrix
228        if cperm.shape[0] != 0:
229            for jj in range(ncols):
230                ii = cperm[jj]
231                new_ptr[ii + 1] = ptr[jj + 1] - ptr[jj]
232
233            for jj in range(ncols):
234                new_ptr[jj + 1] = new_ptr[jj + 1] + new_ptr[jj]
235
236            for jj in range(ncols):
237                k0 = new_ptr[cperm[jj]]
238                for kk in range(ptr[jj], ptr[jj + 1]):
239                    new_idx[k0] = idx[kk]
240                    new_data[k0] = data[kk]
241                    k0 = k0 + 1
242
243        if cperm.shape[0] != 0:
244            nnz = new_ptr[new_ptr.shape[0] - 1]
245            for jj in range(nnz):
246                new_idx[jj] = rperm[new_idx[jj]]
247
248    return new_data, new_idx, new_ptr
249
250
251@cython.boundscheck(False)
252@cython.wraparound(False)
253def _isdiag(int[::1] idx,
254        int[::1] ptr,
255        int nrows):
256
257    cdef int row, num_elems
258    for row in range(nrows):
259        num_elems = ptr[row+1] - ptr[row]
260        if num_elems > 1:
261            return 0
262        elif num_elems == 1:
263            if idx[ptr[row]] != row:
264                return 0
265    return 1
266
267
268@cython.boundscheck(False)
269@cython.wraparound(False)
270cpdef cnp.ndarray[complex, ndim=1, mode='c'] _csr_get_diag(complex[::1] data,
271    int[::1] idx, int[::1] ptr, int k=0):
272
273    cdef size_t row, jj
274    cdef int num_rows = ptr.shape[0]-1
275    cdef int abs_k = abs(k)
276    cdef int start, stop
277    cdef cnp.ndarray[complex, ndim=1, mode='c'] out = np.zeros(num_rows-abs_k, dtype=complex)
278
279    if k >= 0:
280        start = 0
281        stop = num_rows-abs_k
282    else: #k < 0
283        start = abs_k
284        stop = num_rows
285
286    for row in range(start, stop):
287        for jj in range(ptr[row], ptr[row+1]):
288            if idx[jj]-k == row:
289                out[row-start] = data[jj]
290                break
291    return out
292
293
294@cython.boundscheck(False)
295@cython.wraparound(False)
296@cython.cdivision(True)
297def unit_row_norm(complex[::1] data, int[::1] ptr, int nrows):
298    cdef size_t row, ii
299    cdef double total
300    for row in range(nrows):
301        total = 0
302        for ii in range(ptr[row], ptr[row+1]):
303            total += real(data[ii]) * real(data[ii]) + imag(data[ii]) * imag(data[ii])
304        total = sqrt(total)
305        for ii in range(ptr[row], ptr[row+1]):
306            data[ii] /= total
307
308
309
310@cython.boundscheck(False)
311@cython.wraparound(False)
312cpdef double zcsr_one_norm(complex[::1] data, int[::1] ind, int[::1] ptr,
313                     int nrows, int ncols):
314
315    cdef int k
316    cdef size_t ii, jj
317    cdef double * col_sum = <double *>PyDataMem_NEW(ncols * sizeof(double))
318    cdef double max_col = 0
319
320    memset(&col_sum[0],0,ncols * sizeof(double))
321    for ii in range(nrows):
322        for jj in range(ptr[ii], ptr[ii+1]):
323            k = ind[jj]
324            col_sum[k] += cabs(data[jj])
325    for ii in range(ncols):
326        if col_sum[ii] > max_col:
327            max_col = col_sum[ii]
328    PyDataMem_FREE(col_sum)
329    return max_col
330
331
332@cython.boundscheck(False)
333@cython.wraparound(False)
334cpdef double zcsr_inf_norm(complex[::1] data, int[::1] ind, int[::1] ptr,
335                     int nrows, int ncols):
336
337    cdef int k
338    cdef size_t ii, jj
339    cdef double * row_sum = <double *>PyDataMem_NEW(nrows * sizeof(double))
340    cdef double max_row = 0
341
342    memset(&row_sum[0],0,nrows * sizeof(double))
343    for ii in range(nrows):
344        for jj in range(ptr[ii], ptr[ii+1]):
345            row_sum[ii] += cabs(data[jj])
346    for ii in range(nrows):
347        if row_sum[ii] > max_row:
348            max_row = row_sum[ii]
349    PyDataMem_FREE(row_sum)
350    return max_row
351
352
353@cython.boundscheck(False)
354@cython.wraparound(False)
355cpdef bool cy_tidyup(complex[::1] data, double atol, unsigned int nnz):
356    """
357    Performs an in-place tidyup of CSR matrix data
358    """
359    cdef size_t kk
360    cdef double re, im
361    cdef bool re_flag, im_flag, out_flag = 0
362    for kk in range(nnz):
363        re_flag = 0
364        im_flag = 0
365        re = real(data[kk])
366        im = imag(data[kk])
367        if fabs(re) < atol:
368            re = 0
369            re_flag = 1
370        if fabs(im) < atol:
371            im = 0
372            im_flag = 1
373
374        if re_flag or im_flag:
375            data[kk] = re + 1j*im
376
377        if re_flag and im_flag:
378            out_flag = 1
379    return out_flag
380