1# This file is part of QuTiP: Quantum Toolbox in Python.
2#
3#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
4#    All rights reserved.
5#
6#    Redistribution and use in source and binary forms, with or without
7#    modification, are permitted provided that the following conditions are
8#    met:
9#
10#    1. Redistributions of source code must retain the above copyright notice,
11#       this list of conditions and the following disclaimer.
12#
13#    2. Redistributions in binary form must reproduce the above copyright
14#       notice, this list of conditions and the following disclaimer in the
15#       documentation and/or other materials provided with the distribution.
16#
17#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
18#       of its contributors may be used to endorse or promote products derived
19#       from this software without specific prior written permission.
20#
21#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
22#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
23#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
24#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
25#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
26#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
27#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
28#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
29#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
30#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
31#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
32###############################################################################
33
34from functools import partial
35
36import numpy as np
37from numpy.testing import assert_, run_module_suite, assert_allclose
38import pytest
39
40# disable the MC progress bar
41import os
42
43from qutip import *
44from qutip.random_objects import rand_ket
45
46os.environ['QUTIP_GRAPHICS'] = "NO"
47
48
49class TestJCModelEvolution:
50    """
51    A test class for the QuTiP functions for the evolution of JC model
52    """
53
54    def qubit_integrate(self, tlist, psi0, epsilon, delta, g1, g2):
55
56        H = epsilon / 2.0 * sigmaz() + delta / 2.0 * sigmax()
57
58        c_op_list = []
59
60        rate = g1
61        if rate > 0.0:
62            c_op_list.append(np.sqrt(rate) * sigmam())
63
64        rate = g2
65        if rate > 0.0:
66            c_op_list.append(np.sqrt(rate) * sigmaz())
67
68        output = mesolve(
69            H, psi0, tlist, c_op_list, [sigmax(), sigmay(), sigmaz()])
70        expt_list = output.expect[0], output.expect[1], output.expect[2]
71
72        return expt_list[0], expt_list[1], expt_list[2]
73
74    def jc_steadystate(self, N, wc, wa, g, kappa, gamma,
75                       pump, psi0, use_rwa, tlist):
76
77        # Hamiltonian
78        a = tensor(destroy(N), identity(2))
79        sm = tensor(identity(N), destroy(2))
80
81        if use_rwa:
82            # use the rotating wave approxiation
83            H = wc * a.dag(
84            ) * a + wa * sm.dag() * sm + g * (a.dag() * sm + a * sm.dag())
85        else:
86            H = wc * a.dag() * a + wa * sm.dag() * sm + g * (
87                a.dag() + a) * (sm + sm.dag())
88
89        # collapse operators
90        c_op_list = []
91
92        n_th_a = 0.0  # zero temperature
93
94        rate = kappa * (1 + n_th_a)
95        c_op_list.append(np.sqrt(rate) * a)
96
97        rate = kappa * n_th_a
98        if rate > 0.0:
99            c_op_list.append(np.sqrt(rate) * a.dag())
100
101        rate = gamma
102        if rate > 0.0:
103            c_op_list.append(np.sqrt(rate) * sm)
104
105        rate = pump
106        if rate > 0.0:
107            c_op_list.append(np.sqrt(rate) * sm.dag())
108
109        # find the steady state
110        rho_ss = steadystate(H, c_op_list)
111
112        return expect(a.dag() * a, rho_ss), expect(sm.dag() * sm, rho_ss)
113
114    def jc_integrate(self, N, wc, wa, g, kappa, gamma,
115                     pump, psi0, use_rwa, tlist):
116
117        # Hamiltonian
118        a = tensor(destroy(N), identity(2))
119        sm = tensor(identity(N), destroy(2))
120
121        if use_rwa:
122            # use the rotating wave approxiation
123            H = wc * a.dag() * a + wa * sm.dag() * sm + g * (
124                a.dag() * sm + a * sm.dag())
125        else:
126            H = wc * a.dag() * a + wa * sm.dag() * sm + g * (
127                a.dag() + a) * (sm + sm.dag())
128
129        # collapse operators
130        c_op_list = []
131
132        n_th_a = 0.0  # zero temperature
133
134        rate = kappa * (1 + n_th_a)
135        c_op_list.append(np.sqrt(rate) * a)
136
137        rate = kappa * n_th_a
138        if rate > 0.0:
139            c_op_list.append(np.sqrt(rate) * a.dag())
140
141        rate = gamma
142        if rate > 0.0:
143            c_op_list.append(np.sqrt(rate) * sm)
144
145        rate = pump
146        if rate > 0.0:
147            c_op_list.append(np.sqrt(rate) * sm.dag())
148
149        # evolve and calculate expectation values
150        output = mesolve(
151            H, psi0, tlist, c_op_list, [a.dag() * a, sm.dag() * sm])
152        expt_list = output.expect[0], output.expect[1]
153        return expt_list[0], expt_list[1]
154
155    def testQubitDynamics1(self):
156        "mesolve: qubit with dissipation"
157
158        epsilon = 0.0 * 2 * np.pi   # cavity frequency
159        delta = 1.0 * 2 * np.pi   # atom frequency
160        g2 = 0.1
161        g1 = 0.0
162        psi0 = basis(2, 0)        # initial state
163        tlist = np.linspace(0, 5, 200)
164
165        sx, sy, sz = self.qubit_integrate(tlist, psi0, epsilon, delta, g1, g2)
166
167        sx_analytic = np.zeros(np.shape(tlist))
168        sy_analytic = -np.sin(2 * np.pi * tlist) * np.exp(-tlist * g2)
169        sz_analytic = np.cos(2 * np.pi * tlist) * np.exp(-tlist * g2)
170
171        assert_(max(abs(sx - sx_analytic)) < 0.05)
172        assert_(max(abs(sy - sy_analytic)) < 0.05)
173        assert_(max(abs(sz - sz_analytic)) < 0.05)
174
175    def testQubitDynamics2(self):
176        "mesolve: qubit without dissipation"
177
178        epsilon = 0.0 * 2 * np.pi   # cavity frequency
179        delta = 1.0 * 2 * np.pi   # atom frequency
180        g2 = 0.0
181        g1 = 0.0
182        psi0 = basis(2, 0)        # initial state
183        tlist = np.linspace(0, 5, 200)
184
185        sx, sy, sz = self.qubit_integrate(tlist, psi0, epsilon, delta, g1, g2)
186
187        sx_analytic = np.zeros(np.shape(tlist))
188        sy_analytic = -np.sin(2 * np.pi * tlist) * np.exp(-tlist * g2)
189        sz_analytic = np.cos(2 * np.pi * tlist) * np.exp(-tlist * g2)
190
191        assert_(max(abs(sx - sx_analytic)) < 0.05)
192        assert_(max(abs(sy - sy_analytic)) < 0.05)
193        assert_(max(abs(sz - sz_analytic)) < 0.05)
194
195    def testCase1(self):
196        "mesolve: cavity-qubit interaction, no dissipation"
197
198        use_rwa = True
199        N = 4           # number of cavity fock states
200        wc = 2 * np.pi * 1.0   # cavity frequency
201        wa = 2 * np.pi * 1.0   # atom frequency
202        g = 2 * np.pi * 0.01  # coupling strength
203        kappa = 0.0     # cavity dissipation rate
204        gamma = 0.0     # atom dissipation rate
205        pump = 0.0     # atom pump rate
206
207        # start with an excited atom and maximum number of photons
208        n = N - 2
209        psi0 = tensor(basis(N, n), basis(2, 1))
210        tlist = np.linspace(0, 1000, 2000)
211
212        nc, na = self.jc_integrate(
213            N, wc, wa, g, kappa, gamma, pump, psi0, use_rwa, tlist)
214
215        nc_ex = (n + 0.5 * (1 - np.cos(2 * g * np.sqrt(n + 1) * tlist)))
216        na_ex = 0.5 * (1 + np.cos(2 * g * np.sqrt(n + 1) * tlist))
217
218        assert_(max(abs(nc - nc_ex)) < 0.005, True)
219        assert_(max(abs(na - na_ex)) < 0.005, True)
220
221    def testCase2(self):
222        "mesolve: cavity-qubit without interaction, decay"
223
224        use_rwa = True
225        N = 4           # number of cavity fock states
226        wc = 2 * np.pi * 1.0   # cavity frequency
227        wa = 2 * np.pi * 1.0   # atom frequency
228        g = 2 * np.pi * 0.0   # coupling strength
229        kappa = 0.005   # cavity dissipation rate
230        gamma = 0.01    # atom dissipation rate
231        pump = 0.0     # atom pump rate
232
233        # start with an excited atom and maximum number of photons
234        n = N - 2
235        psi0 = tensor(basis(N, n), basis(2, 1))
236        tlist = np.linspace(0, 1000, 2000)
237
238        nc, na = self.jc_integrate(
239            N, wc, wa, g, kappa, gamma, pump, psi0, use_rwa, tlist)
240
241        nc_ex = (n + 0.5 * (1 - np.cos(2 * g * np.sqrt(n + 1) * tlist))) * \
242            np.exp(-kappa * tlist)
243        na_ex = 0.5 * (1 + np.cos(2 * g * np.sqrt(n + 1) * tlist)) * \
244            np.exp(-gamma * tlist)
245
246        assert_(max(abs(nc - nc_ex)) < 0.005, True)
247        assert_(max(abs(na - na_ex)) < 0.005, True)
248
249    def testCase3(self):
250        "mesolve: cavity-qubit with interaction, decay"
251
252        use_rwa = True
253        N = 4           # number of cavity fock states
254        wc = 2 * np.pi * 1.0   # cavity frequency
255        wa = 2 * np.pi * 1.0   # atom frequency
256        g = 2 * np.pi * 0.1   # coupling strength
257        kappa = 0.05    # cavity dissipation rate
258        gamma = 0.001   # atom dissipation rate
259        pump = 0.25    # atom pump rate
260
261        # start with an excited atom and maximum number of photons
262        n = N - 2
263        psi0 = tensor(basis(N, n), basis(2, 1))
264        tlist = np.linspace(0, 200, 500)
265
266        nc, na = self.jc_integrate(
267            N, wc, wa, g, kappa, gamma, pump, psi0, use_rwa, tlist)
268
269        # we don't have any analytics for this parameters, so
270        # compare with the steady state
271        nc_ss, na_ss = self.jc_steadystate(
272            N, wc, wa, g, kappa, gamma, pump, psi0, use_rwa, tlist)
273
274        nc_ss = nc_ss * np.ones(np.shape(nc))
275        na_ss = na_ss * np.ones(np.shape(na))
276
277        assert_(abs(nc[-1] - nc_ss[-1]) < 0.005, True)
278        assert_(abs(na[-1] - na_ss[-1]) < 0.005, True)
279
280# percent error for failure
281me_error = 1e-8
282
283
284class TestMESolverConstDecay:
285    """
286    A test class for the time-dependent ode check function.
287    """
288
289    def testMEDecay(self):
290        "mesolve: simple constant decay"
291
292        N = 10  # number of basis states to consider
293        a = destroy(N)
294        H = a.dag() * a
295        psi0 = basis(N, 9)  # initial state
296        kappa = 0.2  # coupling to oscillator
297        c_op_list = [np.sqrt(kappa) * a]
298        tlist = np.linspace(0, 10, 100)
299        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
300        expt = medata.expect[0]
301        actual_answer = 9.0 * np.exp(-kappa * tlist)
302        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
303        assert_(avg_diff < me_error)
304
305    def testMEDecaySingleCollapse(self):
306        "mesolve: simple constant decay"
307
308        N = 10  # number of basis states to consider
309        a = destroy(N)
310        H = a.dag() * a
311        psi0 = basis(N, 9)  # initial state
312        kappa = 0.2  # coupling to oscillator
313        c_op = np.sqrt(kappa) * a
314        tlist = np.linspace(0, 10, 100)
315        medata = mesolve(H, psi0, tlist, [c_op], [a.dag() * a])
316        expt = medata.expect[0]
317        actual_answer = 9.0 * np.exp(-kappa * tlist)
318        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
319        assert_(avg_diff < me_error)
320
321
322    def testMEDecayAsFuncList(self):
323        "mesolve: constant decay as function list"
324
325        N = 10  # number of basis states to consider
326        a = destroy(N)
327        H = a.dag() * a
328        psi0 = basis(N, 9)  # initial state
329        kappa = 0.2  # coupling to oscillator
330
331        def sqrt_kappa(t, args):
332            return np.sqrt(kappa)
333        c_op_list = [[a, sqrt_kappa]]
334        tlist = np.linspace(0, 10, 100)
335        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
336        expt = medata.expect[0]
337        actual_answer = 9.0 * np.exp(-kappa * tlist)
338        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
339        assert_(avg_diff < me_error)
340
341    def testMEDecayAsStrList(self):
342        "mesolve: constant decay as string list"
343
344        N = 10  # number of basis states to consider
345        a = destroy(N)
346        H = a.dag() * a
347        psi0 = basis(N, 9)  # initial state
348        kappa = 0.2  # coupling to oscillator
349        c_op_list = [[a, 'sqrt(k)']]
350        args = {'k': kappa}
351        tlist = np.linspace(0, 10, 100)
352        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a], args=args)
353        expt = medata.expect[0]
354        actual_answer = 9.0 * np.exp(-kappa * tlist)
355        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
356        assert_(avg_diff < me_error)
357
358
359# average error for failure
360me_error = 1e-6
361
362
363class TestMESolveTDDecay:
364    """
365    A test class for the time-dependent odes.  Comparing to analytic answer
366
367    N(t)=9 * exp[ -kappa*( 1-exp(-t) ) ]
368
369    """
370
371    def testMETDDecayAsFuncList(self):
372        "mesolve: time-dependence as function list"
373
374        N = 10  # number of basis states to consider
375        a = destroy(N)
376        H = a.dag() * a
377        psi0 = basis(N, 9)  # initial state
378        kappa = 0.2  # coupling to oscillator
379
380        def sqrt_kappa(t, args):
381            return np.sqrt(kappa * np.exp(-t))
382        c_op_list = [[a, sqrt_kappa]]
383        tlist = np.linspace(0, 10, 100)
384        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
385        expt = medata.expect[0]
386        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
387        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
388        assert_(avg_diff < me_error)
389
390    def testMETDDecayAsPartFuncList(self):
391        "mesolve: time-dependence as partial function list"
392
393        me_error = 1e-5
394        N = 10
395        a = destroy(N)
396        H = num(N)
397        psi0 = basis(N, 9)
398        tlist = np.linspace(0, 10, 100)
399        c_ops = [[[a, partial(lambda t, args, k:
400                              np.sqrt(k * np.exp(-t)), k=kappa)]]
401                 for kappa in [0.05, 0.1, 0.2]]
402
403        for idx, kappa in enumerate([0.05, 0.1, 0.2]):
404            medata = mesolve(H, psi0, tlist, c_ops[idx], [H])
405            ref = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
406            avg_diff = np.mean(abs(ref - medata.expect[0]) / ref)
407            assert_(avg_diff < me_error)
408
409    def testMETDDecayAsStrList(self):
410        "mesolve: time-dependence as string list"
411
412        N = 10  # number of basis states to consider
413        a = destroy(N)
414        H = a.dag() * a
415        psi0 = basis(N, 9)  # initial state
416        kappa = 0.2  # coupling to oscillator
417        c_op_list = [[a, 'sqrt(k*exp(-t))']]
418        args = {'k': kappa}
419        tlist = np.linspace(0, 10, 100)
420        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a], args=args)
421        expt = medata.expect[0]
422        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
423        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
424        assert_(avg_diff < me_error)
425
426    def testMETDDecayAsArray(self):
427        "mesolve: time-dependence as array"
428
429        N = 10
430        a = destroy(N)
431        H = a.dag() * a
432        psi0 = basis(N, 9)
433        kappa = 0.2
434        tlist = np.linspace(0, 10, 1000)
435        c_op_list = [[a, np.sqrt(kappa * np.exp(-tlist))]]
436        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
437        expt = medata.expect[0]
438        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
439        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
440        assert_(avg_diff < 100 * me_error)
441
442    def testMETDDecayAsFunc(self):
443        "mesolve: time-dependent Liouvillian as single function"
444
445        N = 10  # number of basis states to consider
446        a = destroy(N)
447        H = a.dag() * a
448        rho0 = ket2dm(basis(N, 9))  # initial state
449        kappa = 0.2  # coupling to oscillator
450
451        def Liouvillian_func(t, args):
452            c = np.sqrt(kappa * np.exp(-t))*a
453            return liouvillian(H, [c])
454
455        tlist = np.linspace(0, 10, 100)
456        args = {'kappa': kappa}
457        out1 = mesolve(Liouvillian_func, rho0, tlist, [], [], args=args)
458        expt = expect(a.dag()*a, out1.states)
459        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
460        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
461        assert_(avg_diff < me_error)
462
463
464# average error for failure
465# me_error = 1e-6
466
467class TestMESolveSuperInit:
468    """
469    A test class comparing mesolve run with an identity super-operator and
470    a density matrix as initial conditions, respectively.
471    """
472
473    def fidelitycheck(self, out1, out2, rho0vec):
474        fid = np.zeros(len(out1.states))
475        for i, E in enumerate(out2.states):
476            rhot = vector_to_operator(E*rho0vec)
477            fid[i] = fidelity(out1.states[i], rhot)
478        return fid
479
480    def jc_integrate(self, N, wc, wa, g, kappa, gamma,
481                     pump, psi0, use_rwa, tlist):
482        # Hamiltonian
483        a = tensor(destroy(N), identity(2))
484        sm = tensor(identity(N), destroy(2))
485
486        # Identity super-operator
487        E0 = sprepost(tensor(qeye(N), qeye(2)), tensor(qeye(N), qeye(2)))
488
489        if use_rwa:
490            # use the rotating wave approxiation
491            H = wc * a.dag() * a + wa * sm.dag() * sm + g * (
492                a.dag() * sm + a * sm.dag())
493        else:
494            H = wc * a.dag() * a + wa * sm.dag() * sm + g * (
495                a.dag() + a) * (sm + sm.dag())
496
497        # collapse operators
498        c_op_list = []
499
500        n_th_a = 0.0  # zero temperature
501
502        rate = kappa * (1 + n_th_a)
503        c_op_list.append(np.sqrt(rate) * a)
504
505        rate = kappa * n_th_a
506        if rate > 0.0:
507            c_op_list.append(np.sqrt(rate) * a.dag())
508
509        rate = gamma
510        if rate > 0.0:
511            c_op_list.append(np.sqrt(rate) * sm)
512
513        rate = pump
514        if rate > 0.0:
515            c_op_list.append(np.sqrt(rate) * sm.dag())
516
517        # evolve and calculate expectation values
518        output1 = mesolve(H, psi0, tlist, c_op_list, [])
519        output2 = mesolve(H, E0, tlist, c_op_list, [])
520        return output1, output2
521
522    def testSuperJC(self):
523        "mesolve: super vs. density matrix as initial condition"
524        me_error = 1e-6
525
526        use_rwa = True
527        N = 4           # number of cavity fock states
528        wc = 2 * np.pi * 1.0   # cavity frequency
529        wa = 2 * np.pi * 1.0   # atom frequency
530        g = 2 * np.pi * 0.1   # coupling strength
531        kappa = 0.05    # cavity dissipation rate
532        gamma = 0.001   # atom dissipation rate
533        pump = 0.25    # atom pump rate
534
535        # start with an excited atom and maximum number of photons
536        n = N - 2
537        psi0 = tensor(basis(N, n), basis(2, 1))
538        rho0vec = operator_to_vector(psi0*psi0.dag())
539        tlist = np.linspace(0, 100, 50)
540
541        out1, out2 = self.jc_integrate(
542            N, wc, wa, g, kappa, gamma, pump, psi0, use_rwa, tlist)
543
544        fid = self.fidelitycheck(out1, out2, rho0vec)
545        assert_(max(abs(1.0-fid)) < me_error, True)
546
547    def testMETDDecayAsFuncList(self):
548        "mesolve: time-dependence as function list with super as init cond"
549        me_error = 1e-6
550
551        N = 10  # number of basis states to consider
552        a = destroy(N)
553        H = a.dag() * a
554        psi0 = basis(N, 9)  # initial state
555        rho0vec = operator_to_vector(psi0*psi0.dag())
556        E0 = sprepost(qeye(N), qeye(N))
557        kappa = 0.2  # coupling to oscillator
558
559        def sqrt_kappa(t, args):
560            return np.sqrt(kappa * np.exp(-t))
561        c_op_list = [[a, sqrt_kappa]]
562        tlist = np.linspace(0, 10, 100)
563        out1 = mesolve(H, psi0, tlist, c_op_list, [])
564        out2 = mesolve(H, E0, tlist, c_op_list, [])
565
566        fid = self.fidelitycheck(out1, out2, rho0vec)
567        assert_(max(abs(1.0-fid)) < me_error, True)
568
569    def testMETDDecayAsPartFuncList(self):
570        "mesolve: time-dep. as partial function list with super as init cond"
571        me_error = 1e-5
572
573        N = 10
574        a = destroy(N)
575        H = num(N)
576        psi0 = basis(N, 9)
577        rho0vec = operator_to_vector(psi0*psi0.dag())
578        E0 = sprepost(qeye(N), qeye(N))
579        tlist = np.linspace(0, 10, 100)
580        c_ops = [[[a, partial(lambda t, args, k:
581                              np.sqrt(k * np.exp(-t)), k=kappa)]]
582                 for kappa in [0.05, 0.1, 0.2]]
583
584        for idx, kappa in enumerate([0.05, 0.1, 0.2]):
585            out1 = mesolve(H, psi0, tlist, c_ops[idx], [])
586            out2 = mesolve(H, E0, tlist, c_ops[idx], [])
587            fid = self.fidelitycheck(out1, out2, rho0vec)
588            assert_(max(abs(1.0-fid)) < me_error, True)
589
590    def testMETDDecayAsStrList(self):
591        "mesolve: time-dependence as string list with super as init cond"
592        me_error = 1e-6
593
594        N = 10  # number of basis states to consider
595        a = destroy(N)
596        H = a.dag() * a
597        psi0 = basis(N, 9)  # initial state
598        rho0vec = operator_to_vector(psi0*psi0.dag())
599        E0 = sprepost(qeye(N), qeye(N))
600        kappa = 0.2  # coupling to oscillator
601        c_op_list = [[a, 'sqrt(k*exp(-t))']]
602        args = {'k': kappa}
603        tlist = np.linspace(0, 10, 100)
604        out1 = mesolve(H, psi0, tlist, c_op_list, [], args=args)
605        out2 = mesolve(H, E0, tlist, c_op_list, [], args=args)
606        fid = self.fidelitycheck(out1, out2, rho0vec)
607        assert_(max(abs(1.0-fid)) < me_error, True)
608
609    def testMETDDecayAsArray(self):
610        "mesolve: time-dependence as array with super as init cond"
611        me_error = 1e-5
612
613        N = 10
614        a = destroy(N)
615        H = a.dag() * a
616        psi0 = basis(N, 9)
617        rho0vec = operator_to_vector(psi0*psi0.dag())
618        E0 = sprepost(qeye(N), qeye(N))
619        kappa = 0.2
620        tlist = np.linspace(0, 10, 1000)
621        c_op_list = [[a, np.sqrt(kappa * np.exp(-tlist))]]
622        out1 = mesolve(H, psi0, tlist, c_op_list, [])
623        out2 = mesolve(H, E0, tlist, c_op_list, [])
624        fid = self.fidelitycheck(out1, out2, rho0vec)
625        assert_(max(abs(1.0-fid)) < me_error, True)
626
627    def testMETDDecayAsFunc(self):
628        "mesolve: time-dependence as function with super as init cond"
629
630        N = 10  # number of basis states to consider
631        a = destroy(N)
632        H = a.dag() * a
633        rho0 = ket2dm(basis(N, 9))  # initial state
634        rho0vec = operator_to_vector(rho0)
635        E0 = sprepost(qeye(N), qeye(N))
636        kappa = 0.2  # coupling to oscillator
637
638        def Liouvillian_func(t, args):
639            c = np.sqrt(kappa * np.exp(-t))*a
640            data = liouvillian(H, [c])
641            return data
642
643        tlist = np.linspace(0, 10, 100)
644        args = {'kappa': kappa}
645        out1 = mesolve(Liouvillian_func, rho0, tlist, [], [], args=args)
646        out2 = mesolve(Liouvillian_func, E0, tlist, [], [], args=args)
647
648        fid = self.fidelitycheck(out1, out2, rho0vec)
649        assert_(max(abs(1.0-fid)) < me_error, True)
650
651    def test_me_interp1(self):
652        "mesolve: interp time-dependent collapse operator #1"
653
654        N = 10  # number of basis states to consider
655        kappa = 0.2  # coupling to oscillator
656        tlist = np.linspace(0, 10, 100)
657        a = destroy(N)
658        H = a.dag() * a
659        psi0 = basis(N, 9)  # initial state
660        S = Cubic_Spline(tlist[0],tlist[-1], np.sqrt(kappa*np.exp(-tlist)))
661        c_op_list = [[a, S]]
662        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
663        expt = medata.expect[0]
664        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
665        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
666        assert_(avg_diff < 1e-5)
667
668    def test_me_interp2(self):
669         "mesolve: interp time-dependent collapse operator #2"
670
671         N = 10  # number of basis states to consider
672         kappa = 0.2  # coupling to oscillator
673         tlist = np.linspace(0, 10, 100)
674         C = Cubic_Spline(tlist[0], tlist[-1], np.ones_like(tlist))
675         S = Cubic_Spline(tlist[0],tlist[-1], np.sqrt(kappa*np.exp(-tlist)))
676         a = destroy(N)
677         H = [[a.dag() * a, C]]
678         psi0 = basis(N, 9)  # initial state
679         c_op_list = [[a, S]]
680         medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
681         expt = medata.expect[0]
682         actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
683         avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
684         assert_(avg_diff < 1e-5)
685
686    def test_me_interp3(self):
687        "mesolve: interp time-dependent collapse operator #3"
688
689        N = 10  # number of basis states to consider
690        kappa = 0.2  # coupling to oscillator
691        tlist = np.linspace(0, 10, 100)
692        C = Cubic_Spline(tlist[0], tlist[-1], np.ones_like(tlist))
693        S = Cubic_Spline(tlist[0],tlist[-1], np.sqrt(kappa*np.exp(-tlist)))
694        a = destroy(N)
695        H = [a.dag() * a, [a.dag() * a, C]]
696        psi0 = basis(N, 9)  # initial state
697        c_op_list = [[a, S]]
698        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
699        expt = medata.expect[0]
700        actual_answer = 9.0 * np.exp(-kappa * (1.0 - np.exp(-tlist)))
701        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
702        assert_(avg_diff < 1e-5)
703
704    def test_me_interp4(self):
705        "mesolve: interp time-dependent collapse operator #4"
706
707        N = 10  # number of basis states to consider
708        kappa = 0.2  # coupling to oscillator
709        tlist = np.linspace(0, 10, 100)
710        C = Cubic_Spline(tlist[0], tlist[-1], np.ones_like(tlist))
711        S = Cubic_Spline(tlist[0],tlist[-1], np.sqrt(kappa*np.exp(-tlist)))
712        a = destroy(N)
713        H = [a.dag() * a, [a.dag() * a, C]]
714        psi0 = basis(N, 9)  # initial state
715        c_op_list = [[a, S],[a, S]]
716        medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
717        expt = medata.expect[0]
718        actual_answer = 9.0 * np.exp(-2*kappa * (1.0 - np.exp(-tlist)))
719        avg_diff = np.mean(abs(actual_answer - expt) / actual_answer)
720        assert_(avg_diff < 1e-5)
721
722
723class TestMESolverMisc:
724    """
725    A test class for the misc mesolve features.
726    """
727    def testMEFinalState(self):
728        "mesolve: final_state has correct dims"
729
730        N = 5
731        psi0 = tensor(basis(N+1,0), basis(N+1,0), basis(N+1,N))
732        a = tensor(destroy(N+1), qeye(N+1), qeye(N+1))
733        b = tensor(qeye(N+1), destroy(N+1), qeye(N+1))
734        c = tensor(qeye(N+1), qeye(N+1), destroy(N+1))
735        H = a*b*c.dag() * c.dag() - a.dag()*b.dag()*c * c
736
737        times = np.linspace(0.0, 2.0, 100)
738        opts = Options(store_states=False, store_final_state=True)
739        rho0 = ket2dm(psi0)
740        result = mesolve(H, rho0, times, [], [a.dag()*a, b.dag()*b, c.dag()*c],
741                         options=opts)
742        assert_(rho0.dims == result.final_state.dims)
743
744
745    def testSEFinalState(self):
746        "sesolve: final_state has correct dims"
747
748        N = 5
749        psi0 = tensor(basis(N+1,0), basis(N+1,0), basis(N+1,N))
750        a = tensor(destroy(N+1), qeye(N+1), qeye(N+1))
751        b = tensor(qeye(N+1), destroy(N+1), qeye(N+1))
752        c = tensor(qeye(N+1), qeye(N+1), destroy(N+1))
753        H = a*b*c.dag() * c.dag() - a.dag()*b.dag()*c * c
754
755        times = np.linspace(0.0, 2.0, 100)
756        opts = Options(store_states=False, store_final_state=True)
757        result = mesolve(H, psi0, times, [], [a.dag()*a, b.dag()*b, c.dag()*c],
758                         options=opts)
759        assert_(psi0.dims == result.final_state.dims)
760
761    def test_num_collapse_set(self):
762        H = sigmaz()
763        psi = (basis(2, 0) + basis(2, 1)).unit()
764        ts = [0, 1]
765        for c_ops in (
766            sigmax(),
767            [sigmax()],
768            [sigmay(), sigmax()],
769        ):
770            res = qutip.mesolve(H, psi, ts, c_ops=c_ops)
771            if not isinstance(c_ops, list):
772                c_ops = [c_ops]
773            assert res.num_collapse == len(c_ops)
774
775    # All Hamiltonians should have dimensions [3, 2], so "bad" states can have
776    # [2, 3] instead - the aim is to test mismatch in most cases.
777    @pytest.mark.parametrize(['state'], [
778        pytest.param(basis([2, 3], [0, 0]), id='ket_bad_tensor'),
779        pytest.param(basis([2, 3], [0, 0]).proj(), id='dm_bad_tensor'),
780        pytest.param(basis([2, 3], [0, 0]).dag(), id='bra_bad_tensor'),
781        pytest.param(to_super(basis([2, 3], [0, 0]).proj()),
782                     id='super_bad_tensor'),
783        pytest.param(basis([3, 2], [0, 0]).dag(), id='bra_good_tensor'),
784        pytest.param(operator_to_vector(basis([3, 2], [0, 0]).proj()),
785                     id='operket_good_tensor'),
786        pytest.param(operator_to_vector(basis([3, 2], [0, 0]).proj()).dag(),
787                     id='operbra_good_tensor'),
788        pytest.param(tensor(basis(2, 0), qeye(3)), id='nonsquare_operator'),
789    ])
790    @pytest.mark.parametrize(['operator'], [
791        pytest.param(qeye([3, 2]), id='constant_hamiltonian'),
792        pytest.param(liouvillian(qeye([3, 2]), []), id='constant_liouvillian'),
793        pytest.param([[qeye([3, 2]), lambda t, args: 1]], id='py_scalar'),
794        pytest.param(lambda t, args: qeye([3, 2]), id='py_hamiltonian'),
795        pytest.param(lambda t, args: liouvillian(qeye([3, 2]), []),
796                     id='py_liouvillian'),
797    ])
798    def test_incorrect_state_caught(self, state, operator):
799        """
800        Test that mesolve will correctly catch an input state that is not a
801        correctly shaped state, relative to the Hamiltonian or Liouvillian.
802
803        Regression test for gh-1456.
804        """
805        times = [0, 1e-5]
806        c_op = qeye([3, 2])
807        with pytest.raises(ValueError):
808            mesolve(operator, state, times, c_ops=[c_op])
809
810
811class TestMESolveStepFuncCoeff:
812    """
813    A Test class for using time-dependent array coefficients
814    as step functions instead of doing interpolation
815    """
816    def python_coeff(self, t, args):
817        if t < np.pi/2:
818            return 1.
819        else:
820            return 0.
821
822    def test_py_coeff(self):
823        """
824        Test for Python function as coefficient as step function coeff
825        """
826        rho0 = rand_ket(2)
827        tlist = np.array([0, np.pi/2])
828        qu = QobjEvo([[sigmax(), self.python_coeff]],
829                     tlist=tlist, args={"_step_func_coeff": 1})
830        result = mesolve(qu, rho0=rho0, tlist=tlist)
831        assert(qu.type == "func")
832        assert_allclose(
833            fidelity(result.states[-1], sigmax()*rho0), 1, rtol=1.e-7)
834
835    def test_array_cte_coeff(self):
836        """
837        Test for Array coefficient with uniform tlist as step function coeff
838        """
839        rho0 = rand_ket(2)
840        tlist = np.array([0., np.pi/2, np.pi], dtype=float)
841        npcoeff = np.array([0.25, 0.75, 0.75])
842        qu = QobjEvo([[sigmax(), npcoeff]],
843                     tlist=tlist, args={"_step_func_coeff": 1})
844        result = mesolve(qu, rho0=rho0, tlist=tlist)
845        assert(qu.type == "array")
846        assert_allclose(
847            fidelity(result.states[-1], sigmax()*rho0), 1, rtol=1.e-7)
848
849    def test_array_t_coeff(self):
850        """
851        Test for Array with non-uniform tlist as step function coeff
852        """
853        rho0 = rand_ket(2)
854        tlist = np.array([0., np.pi/2, np.pi*3/2], dtype=float)
855        npcoeff = np.array([0.5, 0.25, 0.25])
856        qu = QobjEvo([[sigmax(), npcoeff]],
857                     tlist=tlist, args={"_step_func_coeff": 1})
858        result = mesolve(qu, rho0=rho0, tlist=tlist)
859        assert(qu.type == "array")
860        assert_allclose(
861            fidelity(result.states[-1], sigmax()*rho0), 1, rtol=1.e-7)
862
863    def test_array_str_coeff(self):
864        """
865        Test for Array and string as step function coeff.
866        qobjevo_codegen is used and uniform tlist
867        """
868        rho0 = rand_ket(2)
869        tlist = np.array([0., np.pi/2, np.pi], dtype=float)
870        npcoeff1 = np.array([0.25, 0.75, 0.75], dtype=complex)
871        npcoeff2 = np.array([0.5, 1.5, 1.5], dtype=float)
872        strcoeff = "1."
873        qu = QobjEvo(
874            [[sigmax(), npcoeff1], [sigmax(), strcoeff], [sigmax(), npcoeff2]],
875            tlist=tlist, args={"_step_func_coeff": 1})
876        result = mesolve(qu, rho0=rho0, tlist=tlist)
877        assert_allclose(
878            fidelity(result.states[-1], sigmax()*rho0), 1, rtol=1.e-7)
879
880    def test_array_str_py_coeff(self):
881        """
882        Test for Array, string and Python function as step function coeff.
883        qobjevo_codegen is used and non non-uniform tlist
884        """
885        rho0 = rand_ket(2)
886        tlist = np.array([0., np.pi/4, np.pi/2, np.pi], dtype=float)
887        npcoeff1 = np.array([0.4, 1.6, 1.0, 1.0], dtype=complex)
888        npcoeff2 = np.array([0.4, 1.6, 1.0, 1.0], dtype=float)
889        strcoeff = "1."
890        qu = QobjEvo(
891            [[sigmax(), npcoeff1], [sigmax(), npcoeff2],
892             [sigmax(), self.python_coeff], [sigmax(), strcoeff]],
893            tlist=tlist, args={"_step_func_coeff": 1})
894        result = mesolve(qu, rho0=rho0, tlist=tlist)
895        assert(qu.type == "mixed_callable")
896        assert_allclose(
897            fidelity(result.states[-1], sigmax()*rho0), 1, rtol=1.e-7)
898
899    def test_dynamic_args(self):
900        "sesolve: state feedback"
901        tol = 1e-3
902        def f(t, args):
903            return np.sqrt(args["state_vec"][3])
904
905        H = [qeye(2), [destroy(2)+create(2), f]]
906        res = mesolve(H, basis(2,1), tlist=np.linspace(0,10,11),
907                      c_ops=[qeye(2)],
908                      e_ops=[num(2)], args={"state_vec":basis(2,1)})
909        assert_(max(abs(res.expect[0][5:])) < tol,
910            msg="evolution with feedback not proceding as expected")
911
912        def f(t, args):
913            return np.sqrt(args["expect_op_0"])
914
915        H = [qeye(2), [destroy(2)+create(2), f]]
916        res = mesolve(H, basis(2,1), tlist=np.linspace(0,10,11),
917                      c_ops=[qeye(2)],
918                      e_ops=[num(2)], args={"expect_op_0":num(2)})
919        assert_(max(abs(res.expect[0][5:])) < tol,
920            msg="evolution with feedback not proceding as expected")
921
922
923def test_non_hermitian_dm():
924    """Test that mesolve works correctly for density matrices that are
925    not Hermitian.
926    See Issue #1460
927    """
928    N = 2
929    a = destroy(N)
930    x = (a + a.dag())/np.sqrt(2)
931    H = a.dag() * a
932
933    # Create non-Hermitian initial state.
934    rho0 = x*fock_dm(N, 0)
935
936    tlist = np.linspace(0, 0.1, 2)
937
938    options = Options()
939    options.store_final_state = True
940    options.store_states = True
941
942    result = mesolve(H, rho0, tlist, e_ops=[x], options=options)
943
944    msg = ('Mesolve is not working properly with a non Hermitian density' +
945       ' matrix as input. Check computation of '
946      )
947
948    imag_part = np.abs(np.imag(result.expect[0][-1]))
949    # Since we used an initial state that is not Hermitian, the expectation of
950    # x must be imaginary for t>0.
951    assert_(imag_part > 0,
952            msg + "expectation values. They should be imaginary")
953
954    # Check that the output state is not hermitian since the input was not
955    # Hermitian either.
956    assert_(not result.final_state.isherm,
957            msg + " final density  matrix. It should not be hermitian")
958    assert_(not result.states[-1].isherm,
959            msg + " states. They should not be hermitian.")
960
961    # Check that when suing a callable we get imaginary expectation values.
962    def callable_x(t, rho):
963        "Dummy callable_x expectation operator."
964        return expect(rho, x)
965    result = mesolve(H, rho0, tlist, e_ops=callable_x)
966
967    imag_part = np.abs(np.imag(result.expect[-1]))
968    assert_(imag_part > 0,
969            msg + "expectation values when using callable operator." +
970            "They should be imaginary.")
971
972
973def test_tlist_h_with_constant_c_ops():
974    """
975    Test that it's possible to mix a time-dependent Hamiltonian given as a
976    QobjEvo with interpolated coefficients with time-independent collapse
977    operators, if the solver times are not equal to the interpolation times of
978    the Hamiltonian.
979
980    See gh-1560.
981    """
982    state = basis(2, 0)
983    all_times = np.linspace(0, 1, 11)
984    few_times = np.linspace(0, 1, 3)
985    dependence = np.cos(2*np.pi * all_times)
986    hamiltonian = QobjEvo([[sigmax(), dependence]], tlist=all_times)
987    collapse = qeye(2)
988    result = mesolve(hamiltonian, state, few_times, c_ops=[collapse])
989    assert result.num_collapse == 1
990    assert len(result.states) == len(few_times)
991
992
993def test_tlist_h_with_other_tlist_c_ops_raises():
994    state = basis(2, 0)
995    all_times = np.linspace(0, 1, 11)
996    few_times = np.linspace(0, 1, 3)
997    dependence = np.cos(2*np.pi * all_times)
998    hamiltonian = QobjEvo([[sigmax(), dependence]], tlist=all_times)
999    collapse = [qeye(2), np.cos(2*np.pi * few_times)]
1000    with pytest.raises(ValueError) as exc:
1001        mesolve(hamiltonian, state, few_times, c_ops=[collapse])
1002    assert str(exc.value) == "Time lists are not compatible"
1003
1004
1005if __name__ == "__main__":
1006    run_module_suite()
1007