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"""
34Tests for Permutational Invariant Quantum solver (PIQS).
35"""
36import numpy as np
37from numpy.testing import (
38    assert_,
39    run_module_suite,
40    assert_raises,
41    assert_array_equal,
42    assert_array_almost_equal,
43    assert_almost_equal,
44    assert_equal,
45)
46from scipy.sparse import block_diag
47from qutip import Qobj, entropy_vn, sigmax, sigmaz
48from qutip.cy.piqs import (
49    get_blocks,
50    j_min,
51    j_vals,
52    m_vals,
53    _num_dicke_states,
54    _num_dicke_ladders,
55    get_index,
56    jmm1_dictionary,
57)
58from qutip.cy.piqs import Dicke as _Dicke
59from qutip.piqs import *
60
61import sys
62import unittest
63
64
65# Disable tests for python2 as qutip.piqs does not support python2.
66if sys.version_info[0] < 3:
67    raise unittest.SkipTest("qutip.piqs module is not tested for Python 2")
68
69
70class TestDicke:
71    """
72    Tests for `qutip.piqs.Dicke` class and auxiliary functions.
73    """
74
75    def test_num_dicke_states(self):
76        """
77        PIQS: Test the `num_dicke_state` function.
78        """
79        N_list = [1, 2, 3, 4, 5, 6, 9, 10, 20, 100, 123]
80        dicke_states = [num_dicke_states(i) for i in N_list]
81        assert_array_equal(
82            dicke_states, [2, 4, 6, 9, 12, 16, 30, 36, 121, 2601, 3906]
83        )
84        N = -1
85        assert_raises(ValueError, num_dicke_states, N)
86        N = 0.2
87        assert_raises(ValueError, num_dicke_states, N)
88
89    def test_num_tls(self):
90        """
91        PIQS: Test the `num_two_level` function.
92        """
93        N_dicke = [2, 4, 6, 9, 12, 16, 30, 36, 121, 2601, 3906]
94        N = [1, 2, 3, 4, 5, 6, 9, 10, 20, 100, 123]
95        calculated_N = [num_tls(i) for i in N_dicke]
96        assert_array_equal(calculated_N, N)
97
98    def test_num_dicke_ladders(self):
99        """
100        PIQS: Test the `_num_dicke_ladders` function.
101        """
102        ndl_true = [1, 2, 2, 3, 3, 4, 4, 5, 5]
103        ndl = [num_dicke_ladders(N) for N in range(1, 10)]
104        assert_array_equal(ndl, ndl_true)
105
106    def test_get_blocks(self):
107        """
108        PIQS: Test the function to get blocks.
109        """
110        N_list = [1, 2, 5, 7]
111        blocks = [
112            np.array([2]),
113            np.array([3, 4]),
114            np.array([6, 10, 12]),
115            np.array([8, 14, 18, 20]),
116        ]
117        calculated_blocks = [get_blocks(i) for i in N_list]
118        for (i, j) in zip(calculated_blocks, blocks):
119            assert_array_equal(i, j)
120
121    def test_j_vals(self):
122        """
123        PIQS: Test calculation of j values for given N.
124        """
125        N_list = [1, 2, 3, 4, 7]
126        j_vals_real = [
127            np.array([0.5]),
128            np.array([0.0, 1.0]),
129            np.array([0.5, 1.5]),
130            np.array([0.0, 1.0, 2.0]),
131            np.array([0.5, 1.5, 2.5, 3.5]),
132        ]
133        j_vals_calc = [j_vals(i) for i in N_list]
134
135        for (i, j) in zip(j_vals_calc, j_vals_real):
136            assert_array_equal(i, j)
137
138    def test_m_vals(self):
139        """
140        PIQS: Test calculation of m values for a particular j.
141        """
142        j_list = [0.5, 1, 1.5, 2, 2.5]
143        m_real = [
144            np.array([-0.5, 0.5]),
145            np.array([-1, 0, 1]),
146            np.array([-1.5, -0.5, 0.5, 1.5]),
147            np.array([-2, -1, 0, 1, 2]),
148            np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5]),
149        ]
150
151        m_calc = [m_vals(i) for i in j_list]
152        for (i, j) in zip(m_real, m_calc):
153            assert_array_equal(i, j)
154
155    def test_dicke_blocks(self):
156        """
157        PIQS: Test the `dicke_blocks` function.
158        """
159        # test 1
160        # mixed state with non-symmetrical block matrix elements
161        N = 3
162        true_matrix = (excited(N) + dicke(N, 0.5, 0.5)).unit()
163        test_blocks = dicke_blocks(true_matrix)
164        test_matrix = Qobj(block_diag(test_blocks))
165        assert_(test_matrix == true_matrix)
166        # test 2
167        # all elements in block-diagonal matrix
168        N = 4
169        true_matrix = Qobj(block_matrix(N)).unit()
170        test_blocks = dicke_blocks(true_matrix)
171        test_matrix = Qobj(block_diag(test_blocks))
172        assert_(test_matrix == true_matrix)
173
174    def test_dicke_blocks_full(self):
175        """
176        PIQS: Test the `dicke_blocks_full` function.
177        """
178        N = 3
179        test_blocks = dicke_blocks_full(excited(3))
180        test_matrix = Qobj(block_diag(test_blocks))
181        true_expanded = np.zeros((8, 8))
182        true_expanded[0, 0] = 1.0
183        assert_(test_matrix == Qobj(true_expanded))
184
185    def test_dicke_function_trace(self):
186        """
187        PIQS: Test the `dicke_function_trace` function.
188        """
189        ## test for N odd
190        N = 3
191        # test trace
192        rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
193        f = lambda x: x
194        test_val = dicke_function_trace(f, rho_mixed)
195        true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).tr()
196        true_val2 = (rho_mixed).tr()
197        # test with linear function (trace)
198        assert_almost_equal(test_val, true_val)
199        assert_almost_equal(test_val, true_val2)
200        # test with nonlinear function
201        f = lambda rho: rho ** 3
202        test_val3 = dicke_function_trace(f, rho_mixed)
203        true_val3 = (Qobj(block_diag(dicke_blocks_full(rho_mixed))) ** 3).tr()
204        assert_almost_equal(test_val3, true_val3)
205        ## test for N even
206        N = 4
207        # test trace
208        rho_mixed = (excited(N) + dicke(N, 0, 0)).unit()
209        f = lambda x: x
210        test_val = dicke_function_trace(f, rho_mixed)
211        true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).tr()
212        true_val2 = (rho_mixed).tr()
213        # test with linear function (trace)
214        assert_almost_equal(test_val, true_val)
215        assert_almost_equal(test_val, true_val2)
216        # test with nonlinear function
217        f = lambda rho: rho ** 3
218        test_val3 = dicke_function_trace(f, rho_mixed)
219        true_val3 = (Qobj(block_diag(dicke_blocks_full(rho_mixed))) ** 3).tr()
220        assert np.allclose(test_val3, true_val3)
221
222    def test_entropy_vn_dicke(self):
223        """
224        PIQS: Test the `entropy_vn_dicke` function.
225        """
226        N = 3
227        rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
228        test_val = entropy_vn_dicke(rho_mixed)
229        true_val = entropy_vn(Qobj(block_diag(dicke_blocks_full(rho_mixed))))
230        assert np.allclose(test_val, true_val)
231
232    def test_purity_dicke(self):
233        """
234        PIQS: Test the `purity_dicke` function.
235        """
236        N = 3
237        rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
238        test_val = purity_dicke(rho_mixed)
239        true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).purity()
240        assert np.allclose(test_val, true_val)
241
242    def test_get_index(self):
243        """
244        PIQS: Test the index fetching function for given j, m, m1 value.
245        """
246        N = 1
247        jmm1_list = [
248            (0.5, 0.5, 0.5),
249            (0.5, 0.5, -0.5),
250            (0.5, -0.5, 0.5),
251            (0.5, -0.5, -0.5),
252        ]
253        indices = [(0, 0), (0, 1), (1, 0), (1, 1)]
254
255        blocks = get_blocks(N)
256        calculated_indices = [
257            get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
258            for jmm1 in jmm1_list
259        ]
260        assert_array_almost_equal(calculated_indices, indices)
261        N = 2
262        blocks = get_blocks(N)
263        jmm1_list = [
264            (1, 1, 1),
265            (1, 1, 0),
266            (1, 1, -1),
267            (1, 0, 1),
268            (1, 0, 0),
269            (1, 0, -1),
270            (1, -1, 1),
271            (1, -1, 0),
272            (1, -1, -1),
273            (0, 0, 0),
274        ]
275        indices = [
276            (0, 0),
277            (0, 1),
278            (0, 2),
279            (1, 0),
280            (1, 1),
281            (1, 2),
282            (2, 0),
283            (2, 1),
284            (2, 2),
285            (3, 3),
286        ]
287        calculated_indices = [
288            get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
289            for jmm1 in jmm1_list
290        ]
291        assert_array_almost_equal(calculated_indices, indices)
292        N = 3
293        blocks = get_blocks(N)
294        jmm1_list = [
295            (1.5, 1.5, 1.5),
296            (1.5, 1.5, 0.5),
297            (1.5, 1.5, -0.5),
298            (1.5, 1.5, -1.5),
299            (1.5, 0.5, 0.5),
300            (1.5, -0.5, -0.5),
301            (1.5, -1.5, -1.5),
302            (1.5, -1.5, 1.5),
303            (0.5, 0.5, 0.5),
304            (0.5, 0.5, -0.5),
305            (0.5, -0.5, 0.5),
306            (0.5, -0.5, -0.5),
307        ]
308
309        indices = [
310            (0, 0),
311            (0, 1),
312            (0, 2),
313            (0, 3),
314            (1, 1),
315            (2, 2),
316            (3, 3),
317            (3, 0),
318            (4, 4),
319            (4, 5),
320            (5, 4),
321            (5, 5),
322        ]
323
324        calculated_indices = [
325            get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
326            for jmm1 in jmm1_list
327        ]
328        assert_array_almost_equal(calculated_indices, indices)
329
330    def test_jmm1_dictionary(self):
331        """
332        PIQS: Test the function to generate the mapping from jmm1 to ik matrix.
333        """
334        d1, d2, d3, d4 = jmm1_dictionary(1)
335
336        d1_correct = {
337            (0, 0): (0.5, 0.5, 0.5),
338            (0, 1): (0.5, 0.5, -0.5),
339            (1, 0): (0.5, -0.5, 0.5),
340            (1, 1): (0.5, -0.5, -0.5),
341        }
342
343        d2_correct = {
344            (0.5, -0.5, -0.5): (1, 1),
345            (0.5, -0.5, 0.5): (1, 0),
346            (0.5, 0.5, -0.5): (0, 1),
347            (0.5, 0.5, 0.5): (0, 0),
348        }
349
350        d3_correct = {
351            0: (0.5, 0.5, 0.5),
352            1: (0.5, 0.5, -0.5),
353            2: (0.5, -0.5, 0.5),
354            3: (0.5, -0.5, -0.5),
355        }
356
357        d4_correct = {
358            (0.5, -0.5, -0.5): 3,
359            (0.5, -0.5, 0.5): 2,
360            (0.5, 0.5, -0.5): 1,
361            (0.5, 0.5, 0.5): 0,
362        }
363
364        assert_equal(d1, d1_correct)
365        assert_equal(d2, d2_correct)
366        assert_equal(d3, d3_correct)
367        assert_equal(d4, d4_correct)
368
369        d1, d2, d3, d4 = jmm1_dictionary(2)
370
371        d1_correct = {
372            (3, 3): (0.0, -0.0, -0.0),
373            (2, 2): (1.0, -1.0, -1.0),
374            (2, 1): (1.0, -1.0, 0.0),
375            (2, 0): (1.0, -1.0, 1.0),
376            (1, 2): (1.0, 0.0, -1.0),
377            (1, 1): (1.0, 0.0, 0.0),
378            (1, 0): (1.0, 0.0, 1.0),
379            (0, 2): (1.0, 1.0, -1.0),
380            (0, 1): (1.0, 1.0, 0.0),
381            (0, 0): (1.0, 1.0, 1.0),
382        }
383
384        d2_correct = {
385            (0.0, -0.0, -0.0): (3, 3),
386            (1.0, -1.0, -1.0): (2, 2),
387            (1.0, -1.0, 0.0): (2, 1),
388            (1.0, -1.0, 1.0): (2, 0),
389            (1.0, 0.0, -1.0): (1, 2),
390            (1.0, 0.0, 0.0): (1, 1),
391            (1.0, 0.0, 1.0): (1, 0),
392            (1.0, 1.0, -1.0): (0, 2),
393            (1.0, 1.0, 0.0): (0, 1),
394            (1.0, 1.0, 1.0): (0, 0),
395        }
396
397        d3_correct = {
398            15: (0.0, -0.0, -0.0),
399            10: (1.0, -1.0, -1.0),
400            9: (1.0, -1.0, 0.0),
401            8: (1.0, -1.0, 1.0),
402            6: (1.0, 0.0, -1.0),
403            5: (1.0, 0.0, 0.0),
404            4: (1.0, 0.0, 1.0),
405            2: (1.0, 1.0, -1.0),
406            1: (1.0, 1.0, 0.0),
407            0: (1.0, 1.0, 1.0),
408        }
409
410        d4_correct = {
411            (0.0, -0.0, -0.0): 15,
412            (1.0, -1.0, -1.0): 10,
413            (1.0, -1.0, 0.0): 9,
414            (1.0, -1.0, 1.0): 8,
415            (1.0, 0.0, -1.0): 6,
416            (1.0, 0.0, 0.0): 5,
417            (1.0, 0.0, 1.0): 4,
418            (1.0, 1.0, -1.0): 2,
419            (1.0, 1.0, 0.0): 1,
420            (1.0, 1.0, 1.0): 0,
421        }
422
423        assert_equal(d1, d1_correct)
424        assert_equal(d2, d2_correct)
425        assert_equal(d3, d3_correct)
426        assert_equal(d4, d4_correct)
427
428    def test_lindbladian(self):
429        """
430        PIQS: Test the generation of the Lindbladian matrix.
431        """
432        N = 1
433        gCE = 0.5
434        gCD = 0.5
435        gCP = 0.5
436        gE = 0.1
437        gD = 0.1
438        gP = 0.1
439
440        system = Dicke(
441            N=N,
442            emission=gE,
443            pumping=gP,
444            dephasing=gD,
445            collective_emission=gCE,
446            collective_pumping=gCP,
447            collective_dephasing=gCD,
448        )
449
450        lindbladian = system.lindbladian()
451        Ldata = np.zeros((4, 4), dtype="complex")
452        Ldata[0] = [-0.6, 0, 0, 0.6]
453        Ldata[1] = [0, -0.9, 0, 0]
454        Ldata[2] = [0, 0, -0.9, 0]
455        Ldata[3] = [0.6, 0, 0, -0.6]
456
457        lindbladian_correct = Qobj(
458            Ldata, dims=[[[2], [2]], [[2], [2]]], shape=(4, 4)
459        )
460        assert_array_almost_equal(lindbladian.data.toarray(), Ldata)
461        N = 2
462        gCE = 0.5
463        gCD = 0.5
464        gCP = 0.5
465        gE = 0.1
466        gD = 0.1
467        gP = 0.1
468        system = Dicke(
469            N=N,
470            emission=gE,
471            pumping=gP,
472            dephasing=gD,
473            collective_emission=gCE,
474            collective_pumping=gCP,
475            collective_dephasing=gCD,
476        )
477
478        lindbladian = system.lindbladian()
479
480        Ldata = np.zeros((16, 16), dtype="complex")
481        Ldata[0][0], Ldata[0][5], Ldata[0][15] = -1.2, 1.1, 0.1
482        Ldata[1, 1], Ldata[1, 6] = -2, 1.1
483        Ldata[2, 2] = -2.2999999999999998
484        Ldata[4, 4], Ldata[4, 9] = -2, 1.1
485        Ldata[5, 0], Ldata[5, 5], Ldata[5, 10], Ldata[5, 15] = (
486            1.1,
487            -2.25,
488            1.1,
489            0.05,
490        )
491        Ldata[6, 1], Ldata[6, 6] = 1.1, -2
492        Ldata[8, 8] = -2.2999999999999998
493        Ldata[9, 4], Ldata[9, 9] = 1.1, -2
494        Ldata[10, 5], Ldata[10, 10], Ldata[10, 15] = 1.1, -1.2, 0.1
495        Ldata[15, 0], Ldata[15, 5], Ldata[15, 10], Ldata[15, 15] = (
496            0.1,
497            0.05,
498            0.1,
499            -0.25,
500        )
501        lindbladian_correct = Qobj(
502            Ldata, dims=[[[4], [4]], [[4], [4]]], shape=(16, 16)
503        )
504        assert_array_almost_equal(lindbladian.data.toarray(), Ldata)
505
506    def test_gamma(self):
507        """
508        PIQS: Test the calculation of various gamma values for diagonal system.
509
510        For N = 6 |j, m> would be :
511
512        | 3, 3>
513        | 3, 2> | 2, 2>
514        | 3, 1> | 2, 1> | 1, 1>
515        | 3, 0> | 2, 0> | 1, 0> |0, 0>
516        | 3,-1> | 2,-1> | 1,-1>
517        | 3,-2> | 2,-2>
518        | 3,-3>
519        """
520        N = 6
521        collective_emission = 1.0
522        emission = 1.0
523        dephasing = 1.0
524        pumping = 1.0
525        collective_pumping = 1.0
526        model = _Dicke(
527            N,
528            collective_emission=collective_emission,
529            emission=emission,
530            dephasing=dephasing,
531            pumping=pumping,
532            collective_pumping=collective_pumping,
533        )
534        tau_calculated = [
535            model.gamma3((3, 1, 1)),
536            model.gamma2((2, 1, 1)),
537            model.gamma4((1, 1, 1)),
538            model.gamma5((3, 0, 0)),
539            model.gamma1((2, 0, 0)),
540            model.gamma6((1, 0, 0)),
541            model.gamma7((3, -1, -1)),
542            model.gamma8((2, -1, -1)),
543            model.gamma9((1, -1, -1)),
544        ]
545        tau_real = [
546            2.0,
547            8.0,
548            0.333333,
549            1.5,
550            -19.5,
551            0.666667,
552            2.0,
553            8.0,
554            0.333333,
555        ]
556        assert_array_almost_equal(tau_calculated, tau_real)
557
558    def test_jspin(self):
559        """
560        PIQS: Test calculation of the j algebra relation for the total operators.
561
562        The jx, jy, jz, jp and jm for a given N in the (j, m, m1)
563        basis should follow the following algebra
564        [jx, jy] == 1j * jz, [jp, jm] == 2 * jz, jx^2 + jy^2 + jz^2 == j2^2.
565        """
566        N_list = [1, 2, 3, 4, 7]
567
568        for nn in N_list:
569            # tests 1
570            [jx, jy, jz] = jspin(nn)
571            jp, jm = jspin(nn, "+"), jspin(nn, "-")
572            test_jxjy = jx * jy - jy * jx
573            true_jxjy = 1j * jz
574            test_jpjm = jp * jm - jm * jp
575            true_jpjm = 2 * jz
576
577            assert_array_almost_equal(test_jxjy, true_jxjy)
578            assert_array_almost_equal(test_jpjm, true_jpjm)
579
580            # tests 2
581            [jx, jy, jz] = jspin(nn)
582            jp, jm = jspin(nn, "+"), jspin(nn, "-")
583            test_jxjy = jx * jy - jy * jx
584            true_jxjy = 1j * jz
585            test_jpjm = jp * jm - jm * jp
586            true_jpjm = 2 * jz
587
588            assert_array_almost_equal(test_jxjy, true_jxjy)
589            assert_array_almost_equal(test_jpjm, true_jpjm)
590
591            assert_array_almost_equal(jspin(nn, "x"), jx)
592            assert_array_almost_equal(jspin(nn, "y"), jy)
593            assert_array_almost_equal(jspin(nn, "z"), jz)
594            assert_array_almost_equal(jspin(nn, "+"), jp)
595            assert_array_almost_equal(jspin(nn, "-"), jm)
596            assert_raises(TypeError, jspin, nn, "q")
597
598    def test_j_min_(self):
599        """
600        PIQS: Test the `j_min` function.
601        """
602        even = [2, 4, 6, 8]
603        odd = [1, 3, 5, 7]
604
605        for i in even:
606            assert_(j_min(i) == 0)
607
608        for i in odd:
609            assert_(j_min(i) == 0.5)
610
611    def test_energy_degeneracy(self):
612        """
613        PIQS: Test the energy degeneracy (m) of Dicke state | j, m >.
614        """
615        true_en_deg = [1, 1, 1, 1, 1]
616        true_en_deg_even = [2, 6, 20]
617        true_en_deg_odd = [1, 1, 3, 3, 35, 35]
618        test_en_deg = []
619        test_en_deg_even = []
620        test_en_deg_odd = []
621
622        for nn in [1, 2, 3, 4, 7]:
623            test_en_deg.append(energy_degeneracy(nn, nn / 2))
624
625        for nn in [2, 4, 6]:
626            test_en_deg_even.append(energy_degeneracy(nn, 0))
627
628        for nn in [1, 3, 7]:
629            test_en_deg_odd.append(energy_degeneracy(nn, 1 / 2))
630            test_en_deg_odd.append(energy_degeneracy(nn, -1 / 2))
631
632        assert_array_equal(test_en_deg, true_en_deg)
633        assert_array_equal(test_en_deg_even, true_en_deg_even)
634        assert_array_equal(test_en_deg_odd, true_en_deg_odd)
635
636    def test_state_degeneracy(self):
637        """
638        PIQS: Test the calculation of the degeneracy of the Dicke state |j, m>,
639        state_degeneracy(N, j).
640        """
641        true_state_deg = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 14, 14, 42, 42]
642        state_deg = []
643        state_deg = []
644        for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
645            state_deg.append(state_degeneracy(nn, nn / 2))
646        for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
647            state_deg.append(state_degeneracy(nn, (nn / 2) % 1))
648        assert_array_equal(state_deg, true_state_deg)
649
650        # check error
651        assert_raises(ValueError, state_degeneracy, 2, -1)
652
653    def test_m_degeneracy(self):
654        """
655        PIQS: Test the degeneracy of TLS states with same m eigenvalue.
656        """
657        true_m_deg = [1, 2, 2, 3, 4, 5, 5, 6]
658        m_deg = []
659        for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
660            m_deg.append(m_degeneracy(nn, -(nn / 2) % 1))
661        assert_array_equal(m_deg, true_m_deg)
662
663        # check error
664        assert_raises(ValueError, m_degeneracy, 6, -6)
665
666    def test_ap(self):
667        """
668        PIQS: Test the calculation of the real coefficient A_{+}(j,m).
669
670        For given values of j, m. For a Dicke state,
671        J_{+} |j, m> = A_{+}(j,m) |j, m + 1>.
672        """
673        true_ap_list = [110, 108, 104, 98, 90, 54, 38, 20, 0]
674        ap_list = []
675        for m in [0, 1, 2, 3, 4, 7, 8, 9, 10]:
676            ap_list.append(ap(10, m) ** 2)
677
678        assert_almost_equal(ap_list, true_ap_list)
679
680    def test_am(self):
681        """
682        PIQS: Test the calculation of the real coefficient A_{-}(j,m).
683
684        For a Dicke state,  J_{-} |j, m> = A_{+}(j,m) |j, m - 1>.
685        """
686        true_am_list = [110, 110, 108, 104, 98, 68, 54, 38, 20]
687        am_list = []
688        for m in [0, 1, 2, 3, 4, 7, 8, 9, 10]:
689            am_list.append(am(10, m) ** 2)
690
691        assert_almost_equal(am_list, true_am_list)
692
693    def test_spin_algebra(self):
694        """
695        PIQS: Test the function that creates the SU2 algebra in uncoupled basis.
696        The list [sx, sy, sz, sp, sm] is checked for N = 2.
697        """
698        sx1 = [
699            [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
700            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
701            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
702            [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
703        ]
704
705        sx2 = [
706            [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
707            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
708            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
709            [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
710        ]
711
712        sy1 = [
713            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j, 0.0 + 0.0j],
714            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
715            [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
716            [0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j],
717        ]
718
719        sy2 = [
720            [0.0 + 0.0j, 0.0 - 0.5j, 0.0 + 0.0j, 0.0 + 0.0j],
721            [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
722            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
723            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.0j],
724        ]
725
726        sz1 = [
727            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
728            [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
729            [0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j, 0.0 + 0.0j],
730            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j],
731        ]
732
733        sz2 = [
734            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
735            [0.0 + 0.0j, -0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
736            [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
737            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j],
738        ]
739
740        sp1 = [
741            [0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
742            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
743            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
744            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
745        ]
746
747        sp2 = [
748            [0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
749            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
750            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
751            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
752        ]
753
754        sm1 = [
755            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
756            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
757            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
758            [0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
759        ]
760
761        sm2 = [
762            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
763            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
764            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
765            [0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
766        ]
767
768        assert_array_equal(spin_algebra(2, "x")[0].full(), sx1)
769        assert_array_equal(spin_algebra(2, "x")[1].full(), sx2)
770        assert_array_equal(spin_algebra(2, "y")[0].full(), sy1)
771        assert_array_equal(spin_algebra(2, "y")[1].full(), sy2)
772        assert_array_equal(spin_algebra(2, "z")[0].full(), sz1)
773        assert_array_equal(spin_algebra(2, "z")[1].full(), sz2)
774        assert_array_equal(spin_algebra(2, "+")[0].full(), sp1)
775        assert_array_equal(spin_algebra(2, "+")[1].full(), sp2)
776        assert_array_equal(spin_algebra(2, "-")[0].full(), sm1)
777        assert_array_equal(spin_algebra(2, "-")[1].full(), sm2)
778
779        # test error
780        assert_raises(TypeError, spin_algebra, 2, "q")
781
782    def test_collective_algebra(self):
783        """
784        PIQS: Test the generation of the collective algebra in uncoupled basis.
785
786        The list [jx, jy, jz] created in the 2^N Hilbert space is
787        checked for N = 2.
788        """
789
790        jx_n2 = [
791            [0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
792            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
793            [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
794            [0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
795        ]
796
797        jy_n2 = [
798            [0.0 + 0.0j, 0.0 - 0.5j, 0.0 - 0.5j, 0.0 + 0.0j],
799            [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
800            [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
801            [0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.5j, 0.0 + 0.0j],
802        ]
803
804        jz_n2 = [
805            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
806            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
807            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
808            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -1.0 + 0.0j],
809        ]
810
811        jp_n2 = [
812            [0.0 + 0.0j, 1.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
813            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
814            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
815            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
816        ]
817
818        jm_n2 = [
819            [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
820            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
821            [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
822            [0.0 + 0.0j, 1.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
823        ]
824
825        assert_array_equal(jspin(2, "x", basis="uncoupled").full(), jx_n2)
826        assert_array_equal(jspin(2, "y", basis="uncoupled").full(), jy_n2)
827        assert_array_equal(jspin(2, "z", basis="uncoupled").full(), jz_n2)
828        assert_array_equal(jspin(2, "+", basis="uncoupled").full(), jp_n2)
829        assert_array_equal(jspin(2, "-", basis="uncoupled").full(), jm_n2)
830
831        # error
832        assert_raises(TypeError, spin_algebra, 2, "q")
833
834    def test_block_matrix(self):
835        """
836        PIQS: Test the calculation of the block-diagonal matrix for given N.
837
838        If the matrix element |j,m><j,m'| is allowed it is 1, otherwise 0.
839        """
840        # N = 1 TLSs
841        block_1 = [[1.0, 1.0], [1.0, 1.0]]
842
843        # N = 2 TLSs
844        block_2 = [
845            [1.0, 1.0, 1.0, 0.0],
846            [1.0, 1.0, 1.0, 0.0],
847            [1.0, 1.0, 1.0, 0.0],
848            [0.0, 0.0, 0.0, 1.0],
849        ]
850
851        # N = 3 TLSs
852        block_3 = [
853            [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
854            [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
855            [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
856            [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
857            [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
858            [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
859        ]
860
861        assert_equal(Qobj(block_1), Qobj(block_matrix(1)))
862        assert_equal(Qobj(block_2), Qobj(block_matrix(2)))
863        assert_equal(Qobj(block_3), Qobj(block_matrix(3)))
864
865    def test_dicke_basis(self):
866        """
867        PIQS: Test if the Dicke basis (j, m, m') is constructed correctly.
868
869        We test the state with for N = 2,
870
871        0   0   0.3 0
872        0   0.5 0   0
873        0.3 0   0   0
874        0   0   0   0.5
875        """
876        N = 2
877        true_dicke_basis = np.zeros((4, 4))
878        true_dicke_basis[1, 1] = 0.5
879        true_dicke_basis[-1, -1] = 0.5
880        true_dicke_basis[0, 2] = 0.3
881        true_dicke_basis[2, 0] = 0.3
882        true_dicke_basis = Qobj(true_dicke_basis)
883        jmm1_1 = {(N / 2, 0, 0): 0.5}
884        jmm1_2 = {(0, 0, 0): 0.5}
885        jmm1_3 = {(N / 2, N / 2, N / 2 - 2): 0.3}
886        jmm1_4 = {(N / 2, N / 2 - 2, N / 2): 0.3}
887        db1 = dicke_basis(2, jmm1_1)
888        db2 = dicke_basis(2, jmm1_2)
889        db3 = dicke_basis(2, jmm1_3)
890        db4 = dicke_basis(2, jmm1_4)
891        test_dicke_basis = db1 + db2 + db3 + db4
892        assert_equal(test_dicke_basis, true_dicke_basis)
893
894        # error
895        assert_raises(AttributeError, dicke_basis, N)
896
897    def test_dicke(self):
898        """
899        PIQS: Test the calculation of the Dicke state as a pure state.
900        """
901        true_excited = np.zeros((4, 4))
902        true_excited[0, 0] = 1
903
904        true_superradiant = np.zeros((4, 4))
905        true_superradiant[1, 1] = 1
906
907        true_subradiant = np.zeros((4, 4))
908        true_subradiant[-1, -1] = 1
909
910        test_excited = dicke(2, 1, 1)
911        test_superradiant = dicke(2, 1, 0)
912        test_subradiant = dicke(2, 0, 0)
913
914        assert_equal(test_excited, Qobj(true_excited))
915        assert_equal(test_superradiant, Qobj(true_superradiant))
916        assert_equal(test_subradiant, Qobj(true_subradiant))
917
918    def test_excited(self):
919        """
920        PIQS: Test the calculation of the totally excited state density matrix.
921
922        The matrix has size (O(N^2), O(N^2)) in Dicke basis ('dicke').
923        The matrix has size (2^N, 2^N) in the uncoupled basis ('uncoupled').
924        """
925        N = 3
926        true_state = np.zeros((6, 6))
927        true_state[0, 0] = 1
928        true_state = Qobj(true_state)
929
930        test_state = excited(N)
931        assert_equal(test_state, true_state)
932
933        N = 4
934        true_state = np.zeros((9, 9))
935        true_state[0, 0] = 1
936        true_state = Qobj(true_state)
937
938        test_state = excited(N)
939        assert_equal(test_state, true_state)
940
941        # uncoupled
942        test_state_uncoupled = excited(2, basis="uncoupled")
943        assert_array_equal(test_state_uncoupled.dims, [[2, 2], [2, 2]])
944        assert_array_equal(test_state_uncoupled.shape, (4, 4))
945        assert_almost_equal(test_state_uncoupled.full()[0, 0], 1 + 0j)
946
947    def test_superradiant(self):
948        """
949        PIQS: Test the calculation of the superradiant state density matrix.
950
951        The state is |N/2, 0> for N even and |N/2, 0.5> for N odd.
952        The matrix has size (O(N^2), O(N^2)) in Dicke basis ('dicke').
953        The matrix has size (2^N, 2^N) in the uncoupled basis ('uncoupled').
954        """
955        N = 3
956        true_state = np.zeros((6, 6))
957        true_state[1, 1] = 1
958        true_state = Qobj(true_state)
959        test_state = superradiant(N)
960        assert_equal(test_state, true_state)
961        N = 4
962        true_state = np.zeros((9, 9))
963        true_state[2, 2] = 1
964        true_state = Qobj(true_state)
965        test_state = superradiant(N)
966        assert_equal(test_state, true_state)
967
968        # uncoupled
969        test_state_uncoupled = superradiant(2, basis="uncoupled")
970        assert_array_equal(test_state_uncoupled.dims, [[2, 2], [2, 2]])
971        assert_array_equal(test_state_uncoupled.shape, (4, 4))
972        assert_almost_equal(test_state_uncoupled.full()[1, 1], 1 + 0j)
973
974    def test_ghz(self):
975        """
976        PIQS: Test the calculation of the density matrix of the GHZ state.
977
978        PIQS: Test for N = 2 in the 'dicke' and in the 'uncoupled' basis.
979        """
980        ghz_dicke = Qobj(
981            [[0.5, 0, 0.5, 0], [0, 0, 0, 0], [0.5, 0, 0.5, 0], [0, 0, 0, 0]]
982        )
983        ghz_uncoupled = Qobj(
984            [[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]]
985        )
986        ghz_uncoupled.dims = [[2, 2], [2, 2]]
987        assert_equal(ghz(2), ghz_dicke)
988        assert_equal(ghz(2, "uncoupled"), ghz_uncoupled)
989
990    def test_ground(self):
991        """
992        PIQS: Test the calculation of the density matrix of the ground state.
993
994        PIQS: Test for N = 2 in the 'dicke' and in the 'uncoupled' basis.
995        """
996        zeros = np.zeros((4, 4), dtype=np.complex128)
997        gdicke = zeros.copy()
998        guncoupled = zeros.copy()
999        gdicke[2, 2] = 1
1000        guncoupled[3, 3] = 1
1001
1002        dim_dicke = [[4], [4]]
1003        dim_uncoupled = [[2, 2], [2, 2]]
1004
1005        test_ground_dicke = ground(2)
1006        test_ground_uncoupled = ground(2, "uncoupled")
1007
1008        assert_array_equal(test_ground_dicke.full(), gdicke)
1009        assert_array_equal(test_ground_dicke.dims, dim_dicke)
1010        assert_array_equal(test_ground_uncoupled.full(), guncoupled)
1011        assert_array_equal(test_ground_uncoupled.dims, dim_uncoupled)
1012
1013    def test_identity_uncoupled(self):
1014        """
1015        PIQS: Test the calculation of the identity in a 2^N dim Hilbert space.
1016        """
1017        test_identity = identity_uncoupled(4)
1018        assert_equal(test_identity.dims, [[2, 2, 2, 2], [2, 2, 2, 2]])
1019        assert_array_equal(
1020            np.diag(test_identity.full()), np.ones(16, np.complex128)
1021        )
1022
1023    def test_css(self):
1024        """
1025        PIQS: Test the calculation of the CSS state.
1026        """
1027        test_css_uncoupled = css(2, basis="uncoupled")
1028        test_css_dicke = css(2)
1029        css_uncoupled = 0.25 * np.ones((4, 4), dtype=np.complex128)
1030        css_dicke = np.array(
1031            [
1032                [
1033                    0.25000000 + 0.0j,
1034                    0.35355339 + 0.0j,
1035                    0.25000000 + 0.0j,
1036                    0.00000000 + 0.0j,
1037                ],
1038                [
1039                    0.35355339 + 0.0j,
1040                    0.50000000 + 0.0j,
1041                    0.35355339 + 0.0j,
1042                    0.00000000 + 0.0j,
1043                ],
1044                [
1045                    0.25000000 + 0.0j,
1046                    0.35355339 + 0.0j,
1047                    0.25000000 + 0.0j,
1048                    0.00000000 + 0.0j,
1049                ],
1050                [
1051                    0.00000000 + 0.0j,
1052                    0.00000000 + 0.0j,
1053                    0.00000000 + 0.0j,
1054                    0.00000000 + 0.0j,
1055                ],
1056            ]
1057        )
1058        assert_array_almost_equal(test_css_uncoupled.full(), css_uncoupled)
1059        assert_array_almost_equal(test_css_dicke.full(), css_dicke)
1060
1061    def test_collapse_uncoupled(self):
1062        """
1063        PIQS: Test the calculation of the correct collapse operators (c_ops) list.
1064
1065        In the "uncoupled" basis of N two-level system (TLS).
1066        The test is performed for N = 2 and emission = 1.
1067        """
1068        c1 = Qobj(
1069            [[0, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 0], [0, 1, 0, 0]],
1070            dims=[[2, 2], [2, 2]],
1071        )
1072        c2 = Qobj(
1073            [[0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]],
1074            dims=[[2, 2], [2, 2]],
1075        )
1076        true_c_ops = [c1, c2]
1077        assert_equal(true_c_ops, collapse_uncoupled(N=2, emission=1))
1078        system = Dicke(N=2, emission=1)
1079        assert_equal(true_c_ops, system.c_ops())
1080
1081    def test_get_blocks(self):
1082        """
1083        PIQS: Test the calculation of list of cumulative elements at each block.
1084
1085        For N = 4
1086
1087        1 1 1 1 1
1088        1 1 1 1 1
1089        1 1 1 1 1
1090        1 1 1 1 1
1091        1 1 1 1 1
1092                1 1 1
1093                1 1 1
1094                1 1 1
1095                     1
1096        Thus, the blocks are [5, 8, 9] denoting that after the first block 5
1097        elements have been accounted for and so on.
1098        """
1099        trueb1 = [2]
1100        trueb2 = [3, 4]
1101        trueb3 = [4, 6]
1102        trueb4 = [5, 8, 9]
1103
1104        test_b1 = get_blocks(1)
1105        test_b2 = get_blocks(2)
1106        test_b3 = get_blocks(3)
1107        test_b4 = get_blocks(4)
1108
1109        assert_equal(test_b1, trueb1)
1110        assert_equal(test_b2, trueb2)
1111        assert_equal(test_b3, trueb3)
1112
1113    def test_lindbladian_dims(self):
1114        """
1115        PIQS: Test the calculation of the lindbladian matrix.
1116        """
1117        true_L = [
1118            [-4, 0, 0, 3],
1119            [0, -3.54999995, 0, 0],
1120            [0, 0, -3.54999995, 0],
1121            [4, 0, 0, -3],
1122        ]
1123        true_L = Qobj(true_L)
1124        true_L.dims = [[[2], [2]], [[2], [2]]]
1125        N = 1
1126        test_dicke = _Dicke(
1127            N=N,
1128            pumping=1,
1129            collective_pumping=2,
1130            emission=1,
1131            collective_emission=3,
1132            dephasing=0.1,
1133        )
1134        test_L = test_dicke.lindbladian()
1135        assert_array_almost_equal(test_L.full(), true_L.full())
1136        assert_array_equal(test_L.dims, true_L.dims)
1137
1138    def test_liouvillian(self):
1139        """
1140        PIQS: Test the calculation of the liouvillian matrix.
1141        """
1142        true_L = [
1143            [-4, 0, 0, 3],
1144            [0, -3.54999995, 0, 0],
1145            [0, 0, -3.54999995, 0],
1146            [4, 0, 0, -3],
1147        ]
1148        true_L = Qobj(true_L)
1149        true_L.dims = [[[2], [2]], [[2], [2]]]
1150        true_H = [[1.0 + 0.0j, 1.0 + 0.0j], [1.0 + 0.0j, -1.0 + 0.0j]]
1151        true_H = Qobj(true_H)
1152        true_H.dims = [[[2], [2]]]
1153        true_liouvillian = [
1154            [-4, -1.0j, 1.0j, 3],
1155            [-1.0j, -3.54999995 + 2.0j, 0, 1.0j],
1156            [1.0j, 0, -3.54999995 - 2.0j, -1.0j],
1157            [4, +1.0j, -1.0j, -3],
1158        ]
1159        true_liouvillian = Qobj(true_liouvillian)
1160        true_liouvillian.dims = [[[2], [2]], [[2], [2]]]
1161        N = 1
1162        test_piqs = Dicke(
1163            hamiltonian=sigmaz() + sigmax(),
1164            N=N,
1165            pumping=1,
1166            collective_pumping=2,
1167            emission=1,
1168            collective_emission=3,
1169            dephasing=0.1,
1170        )
1171        test_liouvillian = test_piqs.liouvillian()
1172        test_hamiltonian = test_piqs.hamiltonian
1173        assert_array_almost_equal(
1174            test_liouvillian.full(), true_liouvillian.full()
1175        )
1176        assert_array_almost_equal(test_hamiltonian.full(), true_H.full())
1177        assert_array_equal(test_liouvillian.dims, test_liouvillian.dims)
1178
1179        # no Hamiltonian
1180        test_piqs = Dicke(
1181            N=N,
1182            pumping=1,
1183            collective_pumping=2,
1184            emission=1,
1185            collective_emission=3,
1186            dephasing=0.1,
1187        )
1188        liouv = test_piqs.liouvillian()
1189        lindblad = test_piqs.lindbladian()
1190        assert_equal(liouv, lindblad)
1191
1192    def test_gamma1(self):
1193        """
1194        PIQS: Test the calculation of gamma1.
1195        """
1196        true_gamma_1 = -2
1197        true_gamma_2 = -3
1198        true_gamma_3 = -7
1199        true_gamma_4 = -1
1200        true_gamma_5 = 0
1201        true_gamma_6 = 0
1202
1203        N = 4
1204        test_dicke = _Dicke(N=N, collective_emission=1)
1205        test_gamma_1 = test_dicke.gamma1((1, 1, 1))
1206        test_dicke = _Dicke(N=N, emission=1)
1207        test_gamma_2 = test_dicke.gamma1((1, 1, 1))
1208        test_dicke = _Dicke(N=N, emission=1, collective_emission=2)
1209        test_gamma_3 = test_dicke.gamma1((1, 1, 1))
1210        test_dicke = _Dicke(N=N, dephasing=4)
1211        test_gamma_4 = test_dicke.gamma1((1, 1, 1))
1212        test_dicke = _Dicke(N=N, collective_pumping=2)
1213        test_gamma_5 = test_dicke.gamma1((1, 1, 1))
1214        test_dicke = _Dicke(N=N, collective_dephasing=2)
1215        test_gamma_6 = test_dicke.gamma1((1, 1, 1))
1216
1217        assert_almost_equal(true_gamma_1, test_gamma_1)
1218        assert_almost_equal(true_gamma_2, test_gamma_2)
1219        assert_almost_equal(true_gamma_3, test_gamma_3)
1220        assert_almost_equal(true_gamma_4, test_gamma_4)
1221        assert_almost_equal(true_gamma_5, test_gamma_5)
1222        assert_almost_equal(true_gamma_6, test_gamma_6)
1223
1224    def test_gamma2(self):
1225        """
1226        PIQS: Test the calculation of gamma2. PIQS: Test performed for N = 4.
1227        """
1228        true_gamma_1 = 2
1229        true_gamma_2 = 1.5
1230        true_gamma_3 = 5.5
1231        true_gamma_4 = 0
1232        true_gamma_5 = 0
1233        true_gamma_6 = 0
1234
1235        N = 4
1236        test_dicke = _Dicke(N=N, collective_emission=1)
1237        test_gamma_1 = test_dicke.gamma2((1, 1, 1))
1238        test_dicke = _Dicke(N=N, emission=1)
1239        test_gamma_2 = test_dicke.gamma2((1, 1, 1))
1240        test_dicke = _Dicke(N=N, emission=1, collective_emission=2)
1241        test_gamma_3 = test_dicke.gamma2((1, 1, 1))
1242        test_dicke = _Dicke(N=N, dephasing=4)
1243        test_gamma_4 = test_dicke.gamma2((1, 1, 1))
1244        test_dicke = _Dicke(N=N, collective_pumping=2)
1245        test_gamma_5 = test_dicke.gamma2((1, 1, 1))
1246        test_dicke = _Dicke(N=N, collective_dephasing=2)
1247        test_gamma_6 = test_dicke.gamma2((1, 1, 1))
1248
1249        assert_almost_equal(true_gamma_1, test_gamma_1)
1250        assert_almost_equal(true_gamma_2, test_gamma_2)
1251        assert_almost_equal(true_gamma_3, test_gamma_3)
1252        assert_almost_equal(true_gamma_4, test_gamma_4)
1253        assert_almost_equal(true_gamma_5, test_gamma_5)
1254        assert_almost_equal(true_gamma_6, test_gamma_6)
1255
1256    def test_gamma3(self):
1257        """
1258        PIQS: Test the calculation of gamma3. PIQS: Test performed for N = 4.
1259        """
1260        true_gamma_1 = 0
1261        true_gamma_2 = 1.3333333730697632
1262        true_gamma_3 = 1.3333333730697632
1263        true_gamma_4 = 0
1264        true_gamma_5 = 0
1265        true_gamma_6 = 0
1266
1267        N = 4
1268        test_dicke = _Dicke(N=N, collective_emission=1)
1269        test_gamma_1 = test_dicke.gamma3((1, 1, 1))
1270        test_dicke = _Dicke(N=N, emission=1)
1271        test_gamma_2 = test_dicke.gamma3((1, 1, 1))
1272        test_dicke = _Dicke(N=N, emission=1, collective_emission=2)
1273        test_gamma_3 = test_dicke.gamma3((1, 1, 1))
1274        test_dicke = _Dicke(N=N, dephasing=4)
1275        test_gamma_4 = test_dicke.gamma3((1, 1, 1))
1276        test_dicke = _Dicke(N=N, collective_pumping=2)
1277        test_gamma_5 = test_dicke.gamma3((1, 1, 1))
1278        test_dicke = _Dicke(N=N, collective_dephasing=2)
1279        test_gamma_6 = test_dicke.gamma3((1, 1, 1))
1280        #
1281        assert_almost_equal(true_gamma_1, test_gamma_1)
1282        assert_almost_equal(true_gamma_2, test_gamma_2)
1283        assert_almost_equal(true_gamma_3, test_gamma_3)
1284        assert_almost_equal(true_gamma_4, test_gamma_4)
1285        assert_almost_equal(true_gamma_5, test_gamma_5)
1286        assert_almost_equal(true_gamma_6, test_gamma_6)
1287
1288    def test_gamma4(self):
1289        """
1290        PIQS: Test the calculation of gamma4. PIQS: Test performed for N = 4.
1291        """
1292        true_gamma_1 = 0.1666666716337204
1293        true_gamma_2 = 2
1294        true_gamma_3 = 0
1295        true_gamma_4 = 0.40824830532073975
1296
1297        N = 4
1298        test_dicke = _Dicke(N=N, emission=1, collective_emission=2)
1299        test_gamma_1 = test_dicke.gamma4((1, 1, 1))
1300        test_gamma_2 = test_dicke.gamma4((0, 0, 0))
1301        test_gamma_3 = test_dicke.gamma4((2, 1, 1))
1302        test_gamma_4 = test_dicke.gamma4((1, -1, 1))
1303
1304        assert_almost_equal(true_gamma_1, test_gamma_1)
1305        assert_almost_equal(true_gamma_2, test_gamma_2)
1306        assert_almost_equal(true_gamma_3, test_gamma_3)
1307        assert_almost_equal(true_gamma_4, test_gamma_4)
1308
1309    def test_gamma5(self):
1310        """
1311        PIQS: Test the calculation of gamma5. PIQS: Test performed for N = 4.
1312        """
1313        true_gamma_1 = 0
1314        true_gamma_2 = 0
1315        true_gamma_3 = 0.75
1316        true_gamma_4 = 0
1317
1318        N = 4
1319        test_dicke = _Dicke(N=N, dephasing=1)
1320        test_gamma_1 = test_dicke.gamma5((1, 1, 1))
1321        test_gamma_2 = test_dicke.gamma5((0, 0, 0))
1322        test_gamma_3 = test_dicke.gamma5((2, 1, 1))
1323        test_gamma_4 = test_dicke.gamma5((1, -1, 1))
1324
1325        assert_almost_equal(true_gamma_1, test_gamma_1)
1326        assert_almost_equal(true_gamma_2, test_gamma_2)
1327        assert_almost_equal(true_gamma_3, test_gamma_3)
1328        assert_almost_equal(true_gamma_4, test_gamma_4)
1329
1330    def test_gamma6(self):
1331        """
1332        PIQS: Test the calculation of gamma6. PIQS: Test performed for N = 4.
1333        """
1334        true_gamma_1 = 0.25
1335        true_gamma_2 = 1
1336        true_gamma_3 = 0
1337        true_gamma_4 = 0.25
1338
1339        N = 4
1340        test_dicke = _Dicke(N=N, dephasing=1)
1341        test_gamma_1 = test_dicke.gamma6((1, 1, 1))
1342        test_gamma_2 = test_dicke.gamma6((0, 0, 0))
1343        test_gamma_3 = test_dicke.gamma6((2, 1, 1))
1344        test_gamma_4 = test_dicke.gamma6((1, -1, 1))
1345
1346        assert_almost_equal(true_gamma_1, test_gamma_1)
1347        assert_almost_equal(true_gamma_2, test_gamma_2)
1348        assert_almost_equal(true_gamma_3, test_gamma_3)
1349        assert_almost_equal(true_gamma_4, test_gamma_4)
1350
1351    def test_gamma7(self):
1352        """
1353        PIQS: Test the calculation of gamma7. PIQS: Test performed for N = 4.
1354        """
1355        true_gamma_1 = 0
1356        true_gamma_2 = 0.5
1357        true_gamma_3 = 0
1358        true_gamma_4 = 1.5
1359
1360        N = 4
1361        test_dicke = _Dicke(N=N, pumping=1)
1362        test_gamma_1 = test_dicke.gamma7((1, 1, 1))
1363        test_gamma_2 = test_dicke.gamma7((2, 0, 0))
1364        test_gamma_3 = test_dicke.gamma7((1, 0, 0))
1365        test_gamma_4 = test_dicke.gamma7((2, -1, -1))
1366
1367        assert_almost_equal(true_gamma_1, test_gamma_1)
1368        assert_almost_equal(true_gamma_2, test_gamma_2)
1369        assert_almost_equal(true_gamma_3, test_gamma_3)
1370        assert_almost_equal(true_gamma_4, test_gamma_4)
1371
1372    def test_gamma8(self):
1373        """
1374        PIQS: Test the calculation of gamma8. PIQS: Test performed for N = 4.
1375        """
1376        true_gamma_1 = 0
1377        true_gamma_2 = 13.5
1378        true_gamma_3 = 5.5
1379        true_gamma_4 = 13.5
1380
1381        N = 4
1382        test_dicke = _Dicke(N=N, pumping=1, collective_pumping=2)
1383        test_gamma_1 = test_dicke.gamma8((1, 1, 1))
1384        test_gamma_2 = test_dicke.gamma8((2, 0, 0))
1385        test_gamma_3 = test_dicke.gamma8((1, 0, 0))
1386        test_gamma_4 = test_dicke.gamma8((2, -1, -1))
1387
1388        assert_almost_equal(true_gamma_1, test_gamma_1)
1389        assert_almost_equal(true_gamma_2, test_gamma_2)
1390        assert_almost_equal(true_gamma_3, test_gamma_3)
1391        assert_almost_equal(true_gamma_4, test_gamma_4)
1392
1393    def test_gamma9(self):
1394        """
1395        PIQS: Test the calculation of gamma9. PIQS: Test performed for N = 4.
1396        """
1397        true_gamma_1 = 1
1398        true_gamma_2 = 0
1399        true_gamma_3 = 0.5
1400        true_gamma_4 = 0
1401        N = 4
1402        test_dicke = _Dicke(N=N, pumping=1, collective_pumping=2)
1403        test_gamma_1 = test_dicke.gamma9((1, 1, 1))
1404        test_gamma_2 = test_dicke.gamma9((2, 0, 0))
1405        test_gamma_3 = test_dicke.gamma9((1, 0, 0))
1406        test_gamma_4 = test_dicke.gamma9((2, -1, -1))
1407        assert_almost_equal(true_gamma_1, test_gamma_1)
1408        assert_almost_equal(true_gamma_2, test_gamma_2)
1409        assert_almost_equal(true_gamma_3, test_gamma_3)
1410        assert_almost_equal(true_gamma_4, test_gamma_4)
1411
1412
1413class TestPim:
1414    """
1415    Tests for the `qutip.piqs.Pim` class.
1416    """
1417
1418    def test_gamma(self):
1419        """
1420        PIQS: Test the calculation of various gamma values for diagonal system.
1421
1422        For N = 6 |j, m> would be :
1423
1424        | 3, 3>
1425        | 3, 2> | 2, 2>
1426        | 3, 1> | 2, 1> | 1, 1>
1427        | 3, 0> | 2, 0> | 1, 0> |0, 0>
1428        | 3,-1> | 2,-1> | 1,-1>
1429        | 3,-2> | 2,-2>
1430        | 3,-3>
1431        """
1432        N = 6
1433        collective_emission = 1.0
1434        emission = 1.0
1435        dephasing = 1.0
1436        pumping = 1.0
1437        collective_pumping = 1.0
1438        model = Pim(
1439            N,
1440            collective_emission=collective_emission,
1441            emission=emission,
1442            dephasing=dephasing,
1443            pumping=pumping,
1444            collective_pumping=collective_pumping,
1445        )
1446        tau_calculated = [
1447            model.tau3(3, 1),
1448            model.tau2(2, 1),
1449            model.tau4(1, 1),
1450            model.tau5(3, 0),
1451            model.tau1(2, 0),
1452            model.tau6(1, 0),
1453            model.tau7(3, -1),
1454            model.tau8(2, -1),
1455            model.tau9(1, -1),
1456        ]
1457        tau_real = [
1458            2.0,
1459            8.0,
1460            0.333333,
1461            1.5,
1462            -19.5,
1463            0.666667,
1464            2.0,
1465            8.0,
1466            0.333333,
1467        ]
1468        assert_array_almost_equal(tau_calculated, tau_real)
1469
1470    def test_isdicke(self):
1471        """
1472        PIQS: Test the `isdicke` function
1473        """
1474        N1 = 1
1475        g0 = 0.01
1476        nth = 0.01
1477        gP = g0 * nth
1478        gL = g0 * (0.1 + nth)
1479        gS = 0.1
1480        gD = 0.1
1481
1482        pim1 = Pim(N1, gS, gL, gD, gP)
1483
1484        test_indices1 = [(0, 0), (0, 1), (1, 0), (-1, -1), (2, -1)]
1485        dicke_labels = [pim1.isdicke(x[0], x[1]) for x in test_indices1]
1486
1487        N2 = 4
1488
1489        pim2 = Pim(N2, gS, gL, gD, gP)
1490        test_indices2 = [(0, 0), (4, 4), (2, 0), (1, 3), (2, 2)]
1491        dicke_labels = [pim2.isdicke(x[0], x[1]) for x in test_indices2]
1492
1493        assert_array_equal(dicke_labels, [True, False, True, False, True])
1494
1495    def test_isdiagonal(self):
1496        """
1497        PIQS: Test the isdiagonal function.
1498        """
1499        mat1 = np.array([[1, 2], [3, 4]])
1500        mat2 = np.array([[1, 0.0], [0.0, 2]])
1501        mat3 = np.array([[1 + 1j, 0.0], [0.0 - 2j, 2 - 2j]])
1502        mat4 = np.array([[1 + 1j, 0.0], [0.0, 2 - 2j]])
1503        assert_equal(isdiagonal(mat1), False)
1504        assert_equal(isdiagonal(mat2), True)
1505        assert_equal(isdiagonal(mat3), False)
1506        assert_equal(isdiagonal(mat4), True)
1507
1508    def test_pisolve(self):
1509        """
1510        PIQS: Test the warning for diagonal Hamiltonians to use internal solver
1511        """
1512        jx, jy, jz = jspin(4)
1513        jp, jm = jspin(4, "+"), jspin(4, "-")
1514
1515    def test_coefficient_matrix(self):
1516        """
1517        PIQS: Test the coefficient matrix used by 'pisolve' for diagonal problems.
1518        """
1519        N = 2
1520        ensemble = Pim(N, emission=1)
1521        test_matrix = ensemble.coefficient_matrix().toarray()
1522        ensemble2 = Dicke(N, emission=1)
1523        test_matrix2 = ensemble.coefficient_matrix().toarray()
1524        true_matrix = [
1525            [-2, 0, 0, 0],
1526            [1, -1, 0, 0],
1527            [0, 1, 0, 1.0],
1528            [1, 0, 0, -1.0],
1529        ]
1530
1531        assert_array_almost_equal(test_matrix, true_matrix)
1532        assert_array_almost_equal(test_matrix2, true_matrix)
1533
1534    def test_pisolve(self):
1535        """
1536        PIQS: Test the warning for diagonal Hamiltonians to use internal solver.
1537        """
1538        jx, jy, jz = jspin(4)
1539        jp, jm = jspin(4, "+"), jspin(4, "-")
1540        non_diag_hamiltonian = jx
1541        diag_hamiltonian = jz
1542
1543        non_diag_system = Dicke(4, non_diag_hamiltonian, emission=0.1)
1544        diag_system = Dicke(4, diag_hamiltonian, emission=0.1)
1545
1546        diag_initial_state = dicke(4, 1, 0)
1547        non_diag_initial_state = ghz(4)
1548        tlist = np.linspace(0, 10, 100)
1549
1550        assert_raises(
1551            ValueError, non_diag_system.pisolve, diag_initial_state, tlist
1552        )
1553        assert_raises(
1554            ValueError, non_diag_system.pisolve, non_diag_initial_state, tlist
1555        )
1556        assert_raises(
1557            ValueError, diag_system.pisolve, non_diag_initial_state, tlist
1558        )
1559
1560        non_dicke_initial_state = excited(4, basis="uncoupled")
1561        assert_raises(
1562            ValueError, diag_system.pisolve, non_dicke_initial_state, tlist
1563        )
1564
1565        # no Hamiltonian
1566        no_hamiltonian_system = Dicke(4, emission=0.1)
1567        result = no_hamiltonian_system.pisolve(diag_initial_state, tlist)
1568        assert_equal(True, len(result.states) > 0)
1569
1570
1571if __name__ == "__main__":
1572    run_module_suite()
1573