content: """
  Tests for Permutational Invariant Quantum solver (PIQS).
  """
  import numpy as np
  from numpy.testing import (
      assert_raises,
      assert_array_equal,
      assert_array_almost_equal,
      assert_almost_equal,
      assert_equal,
  )
  from scipy.sparse import block_diag
  from qutip import Qobj, entropy_vn, sigmax, sigmaz
  from qutip.piqs._piqs import (
      get_blocks,
      j_min,
      j_vals,
      m_vals,
      _num_dicke_states,
      _num_dicke_ladders,
      get_index,
      jmm1_dictionary,
  )
  from qutip.piqs._piqs import Dicke as _Dicke
  from qutip.piqs.piqs import *
  
  import sys
  import unittest
  
  
  # Disable tests for python2 as qutip.piqs does not support python2.
  if sys.version_info[0] < 3:
      raise unittest.SkipTest("qutip.piqs module is not tested for Python 2")
  
  
  class TestDicke:
      """
      Tests for `qutip.piqs.Dicke` class and auxiliary functions.
      """
  
      def test_num_dicke_states(self):
          """
          PIQS: Test the `num_dicke_state` function.
          """
          N_list = [1, 2, 3, 4, 5, 6, 9, 10, 20, 100, 123]
          dicke_states = [num_dicke_states(i) for i in N_list]
          assert_array_equal(
              dicke_states, [2, 4, 6, 9, 12, 16, 30, 36, 121, 2601, 3906]
          )
          N = -1
          assert_raises(ValueError, num_dicke_states, N)
          N = 0.2
          assert_raises(ValueError, num_dicke_states, N)
  
      def test_num_tls(self):
          """
          PIQS: Test the `num_two_level` function.
          """
          N_dicke = [2, 4, 6, 9, 12, 16, 30, 36, 121, 2601, 3906]
          N = [1, 2, 3, 4, 5, 6, 9, 10, 20, 100, 123]
          calculated_N = [num_tls(i) for i in N_dicke]
          assert_array_equal(calculated_N, N)
  
      def test_num_dicke_ladders(self):
          """
          PIQS: Test the `_num_dicke_ladders` function.
          """
          ndl_true = [1, 2, 2, 3, 3, 4, 4, 5, 5]
          ndl = [num_dicke_ladders(N) for N in range(1, 10)]
          assert_array_equal(ndl, ndl_true)
  
      def test_get_blocks(self):
          """
          PIQS: Test the function to get blocks.
          """
          N_list = [1, 2, 5, 7]
          blocks = [
              np.array([2]),
              np.array([3, 4]),
              np.array([6, 10, 12]),
              np.array([8, 14, 18, 20]),
          ]
          calculated_blocks = [get_blocks(i) for i in N_list]
          for (i, j) in zip(calculated_blocks, blocks):
              assert_array_equal(i, j)
  
      def test_j_vals(self):
          """
          PIQS: Test calculation of j values for given N.
          """
          N_list = [1, 2, 3, 4, 7]
          j_vals_real = [
              np.array([0.5]),
              np.array([0.0, 1.0]),
              np.array([0.5, 1.5]),
              np.array([0.0, 1.0, 2.0]),
              np.array([0.5, 1.5, 2.5, 3.5]),
          ]
          j_vals_calc = [j_vals(i) for i in N_list]
  
          for (i, j) in zip(j_vals_calc, j_vals_real):
              assert_array_equal(i, j)
  
      def test_m_vals(self):
          """
          PIQS: Test calculation of m values for a particular j.
          """
          j_list = [0.5, 1, 1.5, 2, 2.5]
          m_real = [
              np.array([-0.5, 0.5]),
              np.array([-1, 0, 1]),
              np.array([-1.5, -0.5, 0.5, 1.5]),
              np.array([-2, -1, 0, 1, 2]),
              np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5]),
          ]
  
          m_calc = [m_vals(i) for i in j_list]
          for (i, j) in zip(m_real, m_calc):
              assert_array_equal(i, j)
  
      def test_dicke_blocks(self):
          """
          PIQS: Test the `dicke_blocks` function.
          """
          # test 1
          # mixed state with non-symmetrical block matrix elements
          N = 3
          true_matrix = (excited(N) + dicke(N, 0.5, 0.5)).unit()
          test_blocks = dicke_blocks(true_matrix)
          test_matrix = Qobj(block_diag(test_blocks))
          assert (test_matrix == true_matrix)
          # test 2
          # all elements in block-diagonal matrix
          N = 4
          true_matrix = Qobj(block_matrix(N)).unit()
          test_blocks = dicke_blocks(true_matrix)
          test_matrix = Qobj(block_diag(test_blocks))
          assert (test_matrix == true_matrix)
  
      def test_dicke_blocks_full(self):
          """
          PIQS: Test the `dicke_blocks_full` function.
          """
          N = 3
          test_blocks = dicke_blocks_full(excited(3))
          test_matrix = Qobj(block_diag(test_blocks))
          true_expanded = np.zeros((8, 8))
          true_expanded[0, 0] = 1.0
          assert (test_matrix == Qobj(true_expanded))
  
      def test_dicke_function_trace(self):
          """
          PIQS: Test the `dicke_function_trace` function.
          """
          ## test for N odd
          N = 3
          # test trace
          rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
          f = lambda x: x
          test_val = dicke_function_trace(f, rho_mixed)
          true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).tr()
          true_val2 = (rho_mixed).tr()
          # test with linear function (trace)
          assert_almost_equal(test_val, true_val)
          assert_almost_equal(test_val, true_val2)
          # test with nonlinear function
          f = lambda rho: rho ** 3
          test_val3 = dicke_function_trace(f, rho_mixed)
          true_val3 = (Qobj(block_diag(dicke_blocks_full(rho_mixed))) ** 3).tr()
          assert_almost_equal(test_val3, true_val3)
          ## test for N even
          N = 4
          # test trace
          rho_mixed = (excited(N) + dicke(N, 0, 0)).unit()
          f = lambda x: x
          test_val = dicke_function_trace(f, rho_mixed)
          true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).tr()
          true_val2 = (rho_mixed).tr()
          # test with linear function (trace)
          assert_almost_equal(test_val, true_val)
          assert_almost_equal(test_val, true_val2)
          # test with nonlinear function
          f = lambda rho: rho ** 3
          test_val3 = dicke_function_trace(f, rho_mixed)
          true_val3 = (Qobj(block_diag(dicke_blocks_full(rho_mixed))) ** 3).tr()
          assert np.allclose(test_val3, true_val3)
  
      def test_entropy_vn_dicke(self):
          """
          PIQS: Test the `entropy_vn_dicke` function.
          """
          N = 3
          rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
          test_val = entropy_vn_dicke(rho_mixed)
          true_val = entropy_vn(Qobj(block_diag(dicke_blocks_full(rho_mixed))))
          assert np.allclose(test_val, true_val)
  
      def test_purity_dicke(self):
          """
          PIQS: Test the `purity_dicke` function.
          """
          N = 3
          rho_mixed = (excited(N) + dicke(N, 0.5, 0.5)).unit()
          test_val = purity_dicke(rho_mixed)
          true_val = (Qobj(block_diag(dicke_blocks_full(rho_mixed)))).purity()
          assert np.allclose(test_val, true_val)
  
      def test_get_index(self):
          """
          PIQS: Test the index fetching function for given j, m, m1 value.
          """
          N = 1
          jmm1_list = [
              (0.5, 0.5, 0.5),
              (0.5, 0.5, -0.5),
              (0.5, -0.5, 0.5),
              (0.5, -0.5, -0.5),
          ]
          indices = [(0, 0), (0, 1), (1, 0), (1, 1)]
  
          blocks = get_blocks(N)
          calculated_indices = [
              get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
              for jmm1 in jmm1_list
          ]
          assert_array_almost_equal(calculated_indices, indices)
          N = 2
          blocks = get_blocks(N)
          jmm1_list = [
              (1, 1, 1),
              (1, 1, 0),
              (1, 1, -1),
              (1, 0, 1),
              (1, 0, 0),
              (1, 0, -1),
              (1, -1, 1),
              (1, -1, 0),
              (1, -1, -1),
              (0, 0, 0),
          ]
          indices = [
              (0, 0),
              (0, 1),
              (0, 2),
              (1, 0),
              (1, 1),
              (1, 2),
              (2, 0),
              (2, 1),
              (2, 2),
              (3, 3),
          ]
          calculated_indices = [
              get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
              for jmm1 in jmm1_list
          ]
          assert_array_almost_equal(calculated_indices, indices)
          N = 3
          blocks = get_blocks(N)
          jmm1_list = [
              (1.5, 1.5, 1.5),
              (1.5, 1.5, 0.5),
              (1.5, 1.5, -0.5),
              (1.5, 1.5, -1.5),
              (1.5, 0.5, 0.5),
              (1.5, -0.5, -0.5),
              (1.5, -1.5, -1.5),
              (1.5, -1.5, 1.5),
              (0.5, 0.5, 0.5),
              (0.5, 0.5, -0.5),
              (0.5, -0.5, 0.5),
              (0.5, -0.5, -0.5),
          ]
  
          indices = [
              (0, 0),
              (0, 1),
              (0, 2),
              (0, 3),
              (1, 1),
              (2, 2),
              (3, 3),
              (3, 0),
              (4, 4),
              (4, 5),
              (5, 4),
              (5, 5),
          ]
  
          calculated_indices = [
              get_index(N, jmm1[0], jmm1[1], jmm1[2], blocks)
              for jmm1 in jmm1_list
          ]
          assert_array_almost_equal(calculated_indices, indices)
  
      def test_jmm1_dictionary(self):
          """
          PIQS: Test the function to generate the mapping from jmm1 to ik matrix.
          """
          d1, d2, d3, d4 = jmm1_dictionary(1)
  
          d1_correct = {
              (0, 0): (0.5, 0.5, 0.5),
              (0, 1): (0.5, 0.5, -0.5),
              (1, 0): (0.5, -0.5, 0.5),
              (1, 1): (0.5, -0.5, -0.5),
          }
  
          d2_correct = {
              (0.5, -0.5, -0.5): (1, 1),
              (0.5, -0.5, 0.5): (1, 0),
              (0.5, 0.5, -0.5): (0, 1),
              (0.5, 0.5, 0.5): (0, 0),
          }
  
          d3_correct = {
              0: (0.5, 0.5, 0.5),
              1: (0.5, 0.5, -0.5),
              2: (0.5, -0.5, 0.5),
              3: (0.5, -0.5, -0.5),
          }
  
          d4_correct = {
              (0.5, -0.5, -0.5): 3,
              (0.5, -0.5, 0.5): 2,
              (0.5, 0.5, -0.5): 1,
              (0.5, 0.5, 0.5): 0,
          }
  
          assert_equal(d1, d1_correct)
          assert_equal(d2, d2_correct)
          assert_equal(d3, d3_correct)
          assert_equal(d4, d4_correct)
  
          d1, d2, d3, d4 = jmm1_dictionary(2)
  
          d1_correct = {
              (3, 3): (0.0, -0.0, -0.0),
              (2, 2): (1.0, -1.0, -1.0),
              (2, 1): (1.0, -1.0, 0.0),
              (2, 0): (1.0, -1.0, 1.0),
              (1, 2): (1.0, 0.0, -1.0),
              (1, 1): (1.0, 0.0, 0.0),
              (1, 0): (1.0, 0.0, 1.0),
              (0, 2): (1.0, 1.0, -1.0),
              (0, 1): (1.0, 1.0, 0.0),
              (0, 0): (1.0, 1.0, 1.0),
          }
  
          d2_correct = {
              (0.0, -0.0, -0.0): (3, 3),
              (1.0, -1.0, -1.0): (2, 2),
              (1.0, -1.0, 0.0): (2, 1),
              (1.0, -1.0, 1.0): (2, 0),
              (1.0, 0.0, -1.0): (1, 2),
              (1.0, 0.0, 0.0): (1, 1),
              (1.0, 0.0, 1.0): (1, 0),
              (1.0, 1.0, -1.0): (0, 2),
              (1.0, 1.0, 0.0): (0, 1),
              (1.0, 1.0, 1.0): (0, 0),
          }
  
          d3_correct = {
              15: (0.0, -0.0, -0.0),
              10: (1.0, -1.0, -1.0),
              9: (1.0, -1.0, 0.0),
              8: (1.0, -1.0, 1.0),
              6: (1.0, 0.0, -1.0),
              5: (1.0, 0.0, 0.0),
              4: (1.0, 0.0, 1.0),
              2: (1.0, 1.0, -1.0),
              1: (1.0, 1.0, 0.0),
              0: (1.0, 1.0, 1.0),
          }
  
          d4_correct = {
              (0.0, -0.0, -0.0): 15,
              (1.0, -1.0, -1.0): 10,
              (1.0, -1.0, 0.0): 9,
              (1.0, -1.0, 1.0): 8,
              (1.0, 0.0, -1.0): 6,
              (1.0, 0.0, 0.0): 5,
              (1.0, 0.0, 1.0): 4,
              (1.0, 1.0, -1.0): 2,
              (1.0, 1.0, 0.0): 1,
              (1.0, 1.0, 1.0): 0,
          }
  
          assert_equal(d1, d1_correct)
          assert_equal(d2, d2_correct)
          assert_equal(d3, d3_correct)
          assert_equal(d4, d4_correct)
  
      def test_lindbladian(self):
          """
          PIQS: Test the generation of the Lindbladian matrix.
          """
          N = 1
          gCE = 0.5
          gCD = 0.5
          gCP = 0.5
          gE = 0.1
          gD = 0.1
          gP = 0.1
  
          system = Dicke(
              N=N,
              emission=gE,
              pumping=gP,
              dephasing=gD,
              collective_emission=gCE,
              collective_pumping=gCP,
              collective_dephasing=gCD,
          )
  
          lindbladian = system.lindbladian()
          Ldata = np.zeros((4, 4), dtype="complex")
          Ldata[0] = [-0.6, 0, 0, 0.6]
          Ldata[1] = [0, -0.9, 0, 0]
          Ldata[2] = [0, 0, -0.9, 0]
          Ldata[3] = [0.6, 0, 0, -0.6]
  
          lindbladian_correct = Qobj(Ldata, dims=[[[2], [2]], [[2], [2]]])
          assert_array_almost_equal(lindbladian.full(), Ldata)
          N = 2
          gCE = 0.5
          gCD = 0.5
          gCP = 0.5
          gE = 0.1
          gD = 0.1
          gP = 0.1
          system = Dicke(
              N=N,
              emission=gE,
              pumping=gP,
              dephasing=gD,
              collective_emission=gCE,
              collective_pumping=gCP,
              collective_dephasing=gCD,
          )
  
          lindbladian = system.lindbladian()
  
          Ldata = np.zeros((16, 16), dtype="complex")
          Ldata[0][0], Ldata[0][5], Ldata[0][15] = -1.2, 1.1, 0.1
          Ldata[1, 1], Ldata[1, 6] = -2, 1.1
          Ldata[2, 2] = -2.2999999999999998
          Ldata[4, 4], Ldata[4, 9] = -2, 1.1
          Ldata[5, 0], Ldata[5, 5], Ldata[5, 10], Ldata[5, 15] = (
              1.1,
              -2.25,
              1.1,
              0.05,
          )
          Ldata[6, 1], Ldata[6, 6] = 1.1, -2
          Ldata[8, 8] = -2.2999999999999998
          Ldata[9, 4], Ldata[9, 9] = 1.1, -2
          Ldata[10, 5], Ldata[10, 10], Ldata[10, 15] = 1.1, -1.2, 0.1
          Ldata[15, 0], Ldata[15, 5], Ldata[15, 10], Ldata[15, 15] = (
              0.1,
              0.05,
              0.1,
              -0.25,
          )
          lindbladian_correct = Qobj(Ldata, dims=[[[4], [4]], [[4], [4]]])
          assert_array_almost_equal(lindbladian.full(), Ldata)
  
      def test_gamma(self):
          """
          PIQS: Test the calculation of various gamma values for diagonal system.
  
          For N = 6 |j, m> would be :
  
          | 3, 3>
          | 3, 2> | 2, 2>
          | 3, 1> | 2, 1> | 1, 1>
          | 3, 0> | 2, 0> | 1, 0> |0, 0>
          | 3,-1> | 2,-1> | 1,-1>
          | 3,-2> | 2,-2>
          | 3,-3>
          """
          N = 6
          collective_emission = 1.0
          emission = 1.0
          dephasing = 1.0
          pumping = 1.0
          collective_pumping = 1.0
          model = _Dicke(
              N,
              collective_emission=collective_emission,
              emission=emission,
              dephasing=dephasing,
              pumping=pumping,
              collective_pumping=collective_pumping,
          )
          tau_calculated = [
              model.gamma3((3, 1, 1)),
              model.gamma2((2, 1, 1)),
              model.gamma4((1, 1, 1)),
              model.gamma5((3, 0, 0)),
              model.gamma1((2, 0, 0)),
              model.gamma6((1, 0, 0)),
              model.gamma7((3, -1, -1)),
              model.gamma8((2, -1, -1)),
              model.gamma9((1, -1, -1)),
          ]
          tau_real = [
              2.0,
              8.0,
              0.333333,
              1.5,
              -19.5,
              0.666667,
              2.0,
              8.0,
              0.333333,
          ]
          assert_array_almost_equal(tau_calculated, tau_real)
  
      def test_jspin(self):
          """
          PIQS: Test calculation of the j algebra relation for the total operators.
  
          The jx, jy, jz, jp and jm for a given N in the (j, m, m1)
          basis should follow the following algebra
          [jx, jy] == 1j * jz, [jp, jm] == 2 * jz, jx^2 + jy^2 + jz^2 == j2^2.
          """
          N_list = [1, 2, 3, 4, 7]
  
          for nn in N_list:
              # tests 1
              [jx, jy, jz] = jspin(nn)
              jp, jm = jspin(nn, "+"), jspin(nn, "-")
              test_jxjy = jx * jy - jy * jx
              true_jxjy = 1j * jz
              test_jpjm = jp * jm - jm * jp
              true_jpjm = 2 * jz
  
              assert_array_almost_equal(test_jxjy.full(), true_jxjy.full())
              assert_array_almost_equal(test_jpjm.full(), true_jpjm.full())
  
              # tests 2
              [jx, jy, jz] = jspin(nn)
              jp, jm = jspin(nn, "+"), jspin(nn, "-")
              test_jxjy = jx * jy - jy * jx
              true_jxjy = 1j * jz
              test_jpjm = jp * jm - jm * jp
              true_jpjm = 2 * jz
  
              assert_array_almost_equal(test_jxjy.full(), true_jxjy.full())
              assert_array_almost_equal(test_jpjm.full(), true_jpjm.full())
  
              assert_array_almost_equal(jspin(nn, "x").full(), jx.full())
              assert_array_almost_equal(jspin(nn, "y").full(), jy.full())
              assert_array_almost_equal(jspin(nn, "z").full(), jz.full())
              assert_array_almost_equal(jspin(nn, "+").full(), jp.full())
              assert_array_almost_equal(jspin(nn, "-").full(), jm.full())
              assert_raises(TypeError, jspin, nn, "q")
  
      def test_j_min_(self):
          """
          PIQS: Test the `j_min` function.
          """
          even = [2, 4, 6, 8]
          odd = [1, 3, 5, 7]
  
          for i in even:
              assert (j_min(i) == 0)
  
          for i in odd:
              assert (j_min(i) == 0.5)
  
      def test_energy_degeneracy(self):
          """
          PIQS: Test the energy degeneracy (m) of Dicke state | j, m >.
          """
          true_en_deg = [1, 1, 1, 1, 1]
          true_en_deg_even = [2, 6, 20]
          true_en_deg_odd = [1, 1, 3, 3, 35, 35]
          test_en_deg = []
          test_en_deg_even = []
          test_en_deg_odd = []
  
          for nn in [1, 2, 3, 4, 7]:
              test_en_deg.append(energy_degeneracy(nn, nn / 2))
  
          for nn in [2, 4, 6]:
              test_en_deg_even.append(energy_degeneracy(nn, 0))
  
          for nn in [1, 3, 7]:
              test_en_deg_odd.append(energy_degeneracy(nn, 1 / 2))
              test_en_deg_odd.append(energy_degeneracy(nn, -1 / 2))
  
          assert_array_equal(test_en_deg, true_en_deg)
          assert_array_equal(test_en_deg_even, true_en_deg_even)
          assert_array_equal(test_en_deg_odd, true_en_deg_odd)
  
      def test_state_degeneracy(self):
          """
          PIQS: Test the calculation of the degeneracy of the Dicke state |j, m>,
          state_degeneracy(N, j).
          """
          true_state_deg = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 14, 14, 42, 42]
          state_deg = []
          state_deg = []
          for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
              state_deg.append(state_degeneracy(nn, nn / 2))
          for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
              state_deg.append(state_degeneracy(nn, (nn / 2) % 1))
          assert_array_equal(state_deg, true_state_deg)
  
          # check error
          assert_raises(ValueError, state_degeneracy, 2, -1)
  
      def test_m_degeneracy(self):
          """
          PIQS: Test the degeneracy of TLS states with same m eigenvalue.
          """
          true_m_deg = [1, 2, 2, 3, 4, 5, 5, 6]
          m_deg = []
          for nn in [1, 2, 3, 4, 7, 8, 9, 10]:
              m_deg.append(m_degeneracy(nn, -(nn / 2) % 1))
          assert_array_equal(m_deg, true_m_deg)
  
          # check error
          assert_raises(ValueError, m_degeneracy, 6, -6)
  
      def test_ap(self):
          """
          PIQS: Test the calculation of the real coefficient A_{+}(j,m).
  
          For given values of j, m. For a Dicke state,
          J_{+} |j, m> = A_{+}(j,m) |j, m + 1>.
          """
          true_ap_list = [110, 108, 104, 98, 90, 54, 38, 20, 0]
          ap_list = []
          for m in [0, 1, 2, 3, 4, 7, 8, 9, 10]:
              ap_list.append(ap(10, m) ** 2)
  
          assert_almost_equal(ap_list, true_ap_list)
  
      def test_am(self):
          """
          PIQS: Test the calculation of the real coefficient A_{-}(j,m).
  
          For a Dicke state,  J_{-} |j, m> = A_{+}(j,m) |j, m - 1>.
          """
          true_am_list = [110, 110, 108, 104, 98, 68, 54, 38, 20]
          am_list = []
          for m in [0, 1, 2, 3, 4, 7, 8, 9, 10]:
              am_list.append(am(10, m) ** 2)
  
          assert_almost_equal(am_list, true_am_list)
  
      def test_spin_algebra(self):
          """
          PIQS: Test the function that creates the SU2 algebra in uncoupled basis.
          The list [sx, sy, sz, sp, sm] is checked for N = 2.
          """
          sx1 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          sx2 = [
              [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
          ]
  
          sy1 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
              [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          sy2 = [
              [0.0 + 0.0j, 0.0 - 0.5j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.0j],
          ]
  
          sz1 = [
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j],
          ]
  
          sz2 = [
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, -0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -0.5 + 0.0j],
          ]
  
          sp1 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          sp2 = [
              [0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          sm1 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          sm2 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          assert_array_equal(spin_algebra(2, "x")[0].full(), sx1)
          assert_array_equal(spin_algebra(2, "x")[1].full(), sx2)
          assert_array_equal(spin_algebra(2, "y")[0].full(), sy1)
          assert_array_equal(spin_algebra(2, "y")[1].full(), sy2)
          assert_array_equal(spin_algebra(2, "z")[0].full(), sz1)
          assert_array_equal(spin_algebra(2, "z")[1].full(), sz2)
          assert_array_equal(spin_algebra(2, "+")[0].full(), sp1)
          assert_array_equal(spin_algebra(2, "+")[1].full(), sp2)
          assert_array_equal(spin_algebra(2, "-")[0].full(), sm1)
          assert_array_equal(spin_algebra(2, "-")[1].full(), sm2)
  
          # test error
          assert_raises(TypeError, spin_algebra, 2, "q")
  
      def test_collective_algebra(self):
          """
          PIQS: Test the generation of the collective algebra in uncoupled basis.
  
          The list [jx, jy, jz] created in the 2^N Hilbert space is
          checked for N = 2.
          """
  
          jx_n2 = [
              [0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
              [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
              [0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.0 + 0.0j],
          ]
  
          jy_n2 = [
              [0.0 + 0.0j, 0.0 - 0.5j, 0.0 - 0.5j, 0.0 + 0.0j],
              [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
              [0.0 + 0.5j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 - 0.5j],
              [0.0 + 0.0j, 0.0 + 0.5j, 0.0 + 0.5j, 0.0 + 0.0j],
          ]
  
          jz_n2 = [
              [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, -1.0 + 0.0j],
          ]
  
          jp_n2 = [
              [0.0 + 0.0j, 1.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 1.0 + 0.0j],
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          jm_n2 = [
              [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [1.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
              [0.0 + 0.0j, 1.0 + 0.0j, 1.0 + 0.0j, 0.0 + 0.0j],
          ]
  
          assert_array_equal(jspin(2, "x", basis="uncoupled").full(), jx_n2)
          assert_array_equal(jspin(2, "y", basis="uncoupled").full(), jy_n2)
          assert_array_equal(jspin(2, "z", basis="uncoupled").full(), jz_n2)
          assert_array_equal(jspin(2, "+", basis="uncoupled").full(), jp_n2)
          assert_array_equal(jspin(2, "-", basis="uncoupled").full(), jm_n2)
  
          # error
          assert_raises(TypeError, spin_algebra, 2, "q")
  
      def test_block_matrix(self):
          """
          PIQS: Test the calculation of the block-diagonal matrix for given N.
  
          If the matrix element |j,m><j,m'| is allowed it is 1, otherwise 0.
          """
          # N = 1 TLSs
          block_1 = [[1.0, 1.0], [1.0, 1.0]]
  
          # N = 2 TLSs
          block_2 = [
              [1.0, 1.0, 1.0, 0.0],
              [1.0, 1.0, 1.0, 0.0],
              [1.0, 1.0, 1.0, 0.0],
              [0.0, 0.0, 0.0, 1.0],
          ]
  
          # N = 3 TLSs
          block_3 = [
              [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
              [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
              [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
              [1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
              [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
          ]
  
          assert_equal(Qobj(block_1), Qobj(block_matrix(1)))
          assert_equal(Qobj(block_2), Qobj(block_matrix(2)))
          assert_equal(Qobj(block_3), Qobj(block_matrix(3)))
  
      def test_dicke_basis(self):
          """
          PIQS: Test if the Dicke basis (j, m, m') is constructed correctly.
  
          We test the state with for N = 2,
  
          0   0   0.3 0
          0   0.5 0   0
          0.3 0   0   0
          0   0   0   0.5
          """
          N = 2
          true_dicke_basis = np.zeros((4, 4))
          true_dicke_basis[1, 1] = 0.5
          true_dicke_basis[-1, -1] = 0.5
          true_dicke_basis[0, 2] = 0.3
          true_dicke_basis[2, 0] = 0.3
          true_dicke_basis = Qobj(true_dicke_basis)
          jmm1_1 = {(N / 2, 0, 0): 0.5}
          jmm1_2 = {(0, 0, 0): 0.5}
          jmm1_3 = {(N / 2, N / 2, N / 2 - 2): 0.3}
          jmm1_4 = {(N / 2, N / 2 - 2, N / 2): 0.3}
          db1 = dicke_basis(2, jmm1_1)
          db2 = dicke_basis(2, jmm1_2)
          db3 = dicke_basis(2, jmm1_3)
          db4 = dicke_basis(2, jmm1_4)
          test_dicke_basis = db1 + db2 + db3 + db4
          assert_equal(test_dicke_basis, true_dicke_basis)
  
          # error
          assert_raises(AttributeError, dicke_basis, N, None)
  
      def test_dicke(self):
          """
          PIQS: Test the calculation of the Dicke state as a pure state.
          """
          true_excited = np.zeros((4, 4))
          true_excited[0, 0] = 1
  
          true_superradiant = np.zeros((4, 4))
          true_superradiant[1, 1] = 1
  
          true_subradiant = np.zeros((4, 4))
          true_subradiant[-1, -1] = 1
  
          test_excited = dicke(2, 1, 1)
          test_superradiant = dicke(2, 1, 0)
          test_subradiant = dicke(2, 0, 0)
  
          assert_equal(test_excited, Qobj(true_excited))
          assert_equal(test_superradiant, Qobj(true_superradiant))
          assert_equal(test_subradiant, Qobj(true_subradiant))
  
      def test_excited(self):
          """
          PIQS: Test the calculation of the totally excited state density matrix.
  
          The matrix has size (O(N^2), O(N^2)) in Dicke basis ('dicke').
          The matrix has size (2^N, 2^N) in the uncoupled basis ('uncoupled').
          """
          N = 3
          true_state = np.zeros((6, 6))
          true_state[0, 0] = 1
          true_state = Qobj(true_state)
  
          test_state = excited(N)
          assert_equal(test_state, true_state)
  
          N = 4
          true_state = np.zeros((9, 9))
          true_state[0, 0] = 1
          true_state = Qobj(true_state)
  
          test_state = excited(N)
          assert_equal(test_state, true_state)
  
          # uncoupled
          test_state_uncoupled = excited(2, basis="uncoupled")
          assert_array_equal(test_state_uncoupled.dims, [[2, 2], [2, 2]])
          assert_array_equal(test_state_uncoupled.shape, (4, 4))
          assert_almost_equal(test_state_uncoupled.full()[0, 0], 1 + 0j)
  
      def test_superradiant(self):
          """
          PIQS: Test the calculation of the superradiant state density matrix.
  
          The state is |N/2, 0> for N even and |N/2, 0.5> for N odd.
          The matrix has size (O(N^2), O(N^2)) in Dicke basis ('dicke').
          The matrix has size (2^N, 2^N) in the uncoupled basis ('uncoupled').
          """
          N = 3
          true_state = np.zeros((6, 6))
          true_state[1, 1] = 1
          true_state = Qobj(true_state)
          test_state = superradiant(N)
          assert_equal(test_state, true_state)
          N = 4
          true_state = np.zeros((9, 9))
          true_state[2, 2] = 1
          true_state = Qobj(true_state)
          test_state = superradiant(N)
          assert_equal(test_state, true_state)
  
          # uncoupled
          test_state_uncoupled = superradiant(2, basis="uncoupled")
          assert_array_equal(test_state_uncoupled.dims, [[2, 2], [2, 2]])
          assert_array_equal(test_state_uncoupled.shape, (4, 4))
          assert_almost_equal(test_state_uncoupled.full()[1, 1], 1 + 0j)
  
      def test_ghz(self):
          """
          PIQS: Test the calculation of the density matrix of the GHZ state.
  
          PIQS: Test for N = 2 in the 'dicke' and in the 'uncoupled' basis.
          """
          ghz_dicke = Qobj(
              [[0.5, 0, 0.5, 0], [0, 0, 0, 0], [0.5, 0, 0.5, 0], [0, 0, 0, 0]]
          )
          ghz_uncoupled = Qobj(
              [[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]]
          )
          ghz_uncoupled.dims = [[2, 2], [2, 2]]
          assert_equal(ghz(2), ghz_dicke)
          assert_equal(ghz(2, "uncoupled"), ghz_uncoupled)
  
      def test_ground(self):
          """
          PIQS: Test the calculation of the density matrix of the ground state.
  
          PIQS: Test for N = 2 in the 'dicke' and in the 'uncoupled' basis.
          """
          zeros = np.zeros((4, 4), dtype=np.complex128)
          gdicke = zeros.copy()
          guncoupled = zeros.copy()
          gdicke[2, 2] = 1
          guncoupled[3, 3] = 1
  
          dim_dicke = [[4], [4]]
          dim_uncoupled = [[2, 2], [2, 2]]
  
          test_ground_dicke = ground(2)
          test_ground_uncoupled = ground(2, "uncoupled")
  
          assert_array_equal(test_ground_dicke.full(), gdicke)
          assert_array_equal(test_ground_dicke.dims, dim_dicke)
          assert_array_equal(test_ground_uncoupled.full(), guncoupled)
          assert_array_equal(test_ground_uncoupled.dims, dim_uncoupled)
  
      def test_identity_uncoupled(self):
          """
          PIQS: Test the calculation of the identity in a 2^N dim Hilbert space.
          """
          test_identity = identity_uncoupled(4)
          assert_equal(test_identity.dims, [[2, 2, 2, 2], [2, 2, 2, 2]])
          assert_array_equal(
              np.diag(test_identity.full()), np.ones(16, np.complex128)
          )
  
      def test_css(self):
          """
          PIQS: Test the calculation of the CSS state.
          """
          test_css_uncoupled = css(2, basis="uncoupled")
          test_css_dicke = css(2)
          css_uncoupled = 0.25 * np.ones((4, 4), dtype=np.complex128)
          css_dicke = np.array(
              [
                  [
                      0.25000000 + 0.0j,
                      0.35355339 + 0.0j,
                      0.25000000 + 0.0j,
                      0.00000000 + 0.0j,
                  ],
                  [
                      0.35355339 + 0.0j,
                      0.50000000 + 0.0j,
                      0.35355339 + 0.0j,
                      0.00000000 + 0.0j,
                  ],
                  [
                      0.25000000 + 0.0j,
  