content: """Permutational Invariant Quantum Solver (PIQS)
  
  This module calculates the Liouvillian for the dynamics of ensembles of
  identical two-level systems (TLS) in the presence of local and collective
  processes by exploiting permutational symmetry and using the Dicke basis.
  It also allows to characterize nonlinear functions of the density matrix.
  """
  
  # Authors: Nathan Shammah, Shahnawaz Ahmed
  # Contact: nathan.shammah@gmail.com, shahnawaz.ahmed95@gmail.com
  
  from math import factorial
  from decimal import Decimal
  
  import numpy as np
  from scipy.integrate import odeint
  from scipy.linalg import eigvalsh
  from scipy.special import entr
  from scipy.sparse import dok_matrix, block_diag, lil_matrix
  from .. import (
      Qobj, spre, spost, tensor, identity, ket2dm, sigmax, sigmay, sigmaz,
      sigmap, sigmam,
  )
  from ..entropy import entropy_vn
  from ._piqs import Dicke as _Dicke
  from ._piqs import (
      jmm1_dictionary,
      _num_dicke_states,
      _num_dicke_ladders,
      get_blocks,
      j_min,
      j_vals,
  )
  
  __all__ = [
      "num_dicke_states",
      "num_dicke_ladders",
      "num_tls",
      "isdiagonal",
      "dicke_blocks",
      "dicke_blocks_full",
      "dicke_function_trace",
      "purity_dicke",
      "entropy_vn_dicke",
      "Dicke",
      "state_degeneracy",
      "m_degeneracy",
      "energy_degeneracy",
      "ap",
      "am",
      "spin_algebra",
      "jspin",
      "collapse_uncoupled",
      "dicke_basis",
      "dicke",
      "excited",
      "superradiant",
      "css",
      "ghz",
      "ground",
      "identity_uncoupled",
      "block_matrix",
      "tau_column",
      "Pim",
  ]
  
  
  def _ensure_int(x):
      """
      Ensure that a floating-point value `x` is exactly an integer, and return it
      as an int.
      """
      out = int(x)
      if out != x:
          raise ValueError(f"{x} is not an integral value")
      return out
  
  
  # Functions necessary to generate the Lindbladian/Liouvillian
  def num_dicke_states(N):
      """Calculate the number of Dicke states.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      nds: int
          The number of Dicke states.
      """
      return _num_dicke_states(N)
  
  
  def num_dicke_ladders(N):
      """Calculate the total number of ladders in the Dicke space.
  
      For a collection of N two-level systems it counts how many different
      "j" exist or the number of blocks in the block-diagonal matrix.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      Nj: int
          The number of Dicke ladders.
      """
      return _num_dicke_ladders(N)
  
  
  def num_tls(nds):
      """Calculate the number of two-level systems.
  
      Parameters
      ----------
      nds: int
           The number of Dicke states.
  
      Returns
      -------
      N: int
          The number of two-level systems.
      """
      if np.sqrt(nds).is_integer():
          # N is even
          N = 2 * (np.sqrt(nds) - 1)
      else:
          # N is odd
          N = 2 * (np.sqrt(nds + 1 / 4) - 1)
      return int(N)
  
  
  def isdiagonal(mat):
      """
      Check if the input matrix is diagonal.
  
      Parameters
      ==========
      mat: ndarray/Qobj
          A 2D numpy array
  
      Returns
      =======
      diag: bool
          True/False depending on whether the input matrix is diagonal.
      """
      if isinstance(mat, Qobj):
          mat = mat.full()
  
      return np.all(mat == np.diag(np.diagonal(mat)))
  
  
  # nonlinear functions of the density matrix
  def dicke_blocks(rho):
      """Create the list of blocks for block-diagonal density matrix in the Dicke
      basis.
  
      Parameters
      ----------
      rho : :class:`.Qobj`
          A 2D block-diagonal matrix of ones with dimension (nds,nds),
          where nds is the number of Dicke states for N two-level
          systems.
  
      Returns
      -------
      square_blocks: list of np.ndarray
          Give back the blocks list.
  
      """
      shape_dimension = rho.shape[0]
      N = num_tls(shape_dimension)
      ladders = num_dicke_ladders(N)
      # create a list with the sizes of the blocks, in order
      blocks_dimensions = int(N / 2 + 1 - 0.5 * (N % 2))
      blocks_list = [
          (2 * (i + 1 * (N % 2)) + 1 * ((N + 1) % 2))
          for i in range(blocks_dimensions)
      ]
      blocks_list = np.flip(blocks_list, 0)
      # create a list with each block matrix as element
      square_blocks = []
      block_position = 0
      for block_size in blocks_list:
          start_m = block_position
          end_m = block_position + block_size
          square_block = rho[start_m:end_m, start_m:end_m]
          block_position = block_position + block_size
          square_blocks.append(square_block)
  
      return square_blocks
  
  
  def dicke_blocks_full(rho):
      """Give the full (2^N-dimensional) list of blocks for a Dicke-basis matrix.
  
      Parameters
      ----------
      rho : :class:`.Qobj`
          A 2D block-diagonal matrix of ones with dimension (nds,nds),
          where nds is the number of Dicke states for N two-level
          systems.
  
      Returns
      -------
      full_blocks : list
          The list of blocks expanded in the 2^N space for N qubits.
  
      """
      shape_dimension = rho.shape[0]
      N = num_tls(shape_dimension)
      ladders = num_dicke_ladders(N)
      # create a list with the sizes of the blocks, in order
      blocks_dimensions = int(N / 2 + 1 - 0.5 * (N % 2))
      blocks_list = [
          (2 * (i + 1 * (N % 2)) + 1 * ((N + 1) % 2))
          for i in range(blocks_dimensions)
      ]
      blocks_list = np.flip(blocks_list, 0)
      # create a list with each block matrix as element
      full_blocks = []
      k = 0
      block_position = 0
      for block_size in blocks_list:
          start_m = block_position
          end_m = block_position + block_size
          square_block = rho[start_m:end_m, start_m:end_m]
          block_position = block_position + block_size
          j = N / 2 - k
          djn = state_degeneracy(N, j)
          for block_counter in range(0, djn):
              full_blocks.append(square_block / djn)  # preserve trace
          k = k + 1
      return full_blocks
  
  
  def dicke_function_trace(f, rho):
      """Calculate the trace of a function on a Dicke density matrix.
      Parameters
      ----------
      f : function
          A Taylor-expandable function of `rho`.
  
      rho : :class:`.Qobj`
          A density matrix in the Dicke basis.
  
      Returns
      -------
      res : float
          Trace of a nonlinear function on `rho`.
      """
      N = num_tls(rho.shape[0])
      blocks = dicke_blocks(rho)
      block_f = []
      degen_blocks = np.flip(j_vals(N))
      state_degeneracies = []
      for j in degen_blocks:
          dj = state_degeneracy(N, j)
          state_degeneracies.append(dj)
      eigenvals_degeneracy = []
      deg = []
      for i, block in enumerate(blocks):
          dj = state_degeneracies[i]
          normalized_block = block / dj
          eigenvals_block = eigvalsh(normalized_block)
          for val in eigenvals_block:
              eigenvals_degeneracy.append(abs(val))
              deg.append(dj)
  
      eigenvalue = np.array(eigenvals_degeneracy)
      function_array = f(eigenvalue) * deg
      return sum(function_array)
  
  
  def entropy_vn_dicke(rho):
      """Von Neumann Entropy of a Dicke-basis density matrix.
  
      Parameters
      ----------
      rho : :class:`.Qobj`
          A 2D block-diagonal matrix of ones with dimension (nds, nds),
          where nds is the number of Dicke states for N two-level
          systems.
  
      Returns
      -------
      entropy_dm: float
          Entropy. Use degeneracy to multiply each block.
  
      """
      return dicke_function_trace(entr, rho)
  
  
  def purity_dicke(rho):
      """Calculate purity of a density matrix in the Dicke basis.
      It accounts for the degenerate blocks in the density matrix.
  
      Parameters
      ----------
      rho : :class:`.Qobj`
          Density matrix in the Dicke basis of qutip.piqs.jspin(N), for N spins.
  
      Returns
      -------
      purity : float
          The purity of the quantum state.
          It's 1 for pure states, 0<=purity<1 for mixed states.
      """
      f = lambda x: x * x
      return dicke_function_trace(f, rho)
  
  
  class Result:
      pass
  
  
  class Dicke(object):
      """The Dicke class which builds the Lindbladian and Liouvillian matrix.
  
      Examples
      --------
      >>> from qutip.piqs import Dicke, jspin
      >>> N = 2
      >>> jx, jy, jz = jspin(N)
      >>> jp = jspin(N, "+")
      >>> jm = jspin(N, "-")
      >>> ensemble = Dicke(N, emission=1.)
      >>> L = ensemble.liouvillian()
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      hamiltonian : :class:`.Qobj`
          A Hamiltonian in the Dicke basis.
  
          The matrix dimensions are (nds, nds),
          with nds being the number of Dicke states.
          The Hamiltonian can be built with the operators
          given by the `jspin` functions.
  
      emission: float
          Incoherent emission coefficient (also nonradiative emission).
          default: 0.0
  
      dephasing: float
          Local dephasing coefficient.
          default: 0.0
  
      pumping: float
          Incoherent pumping coefficient.
          default: 0.0
  
      collective_emission: float
          Collective (superradiant) emmission coefficient.
          default: 0.0
  
      collective_pumping: float
          Collective pumping coefficient.
          default: 0.0
  
      collective_dephasing: float
          Collective dephasing coefficient.
          default: 0.0
  
      Attributes
      ----------
      N: int
          The number of two-level systems.
  
      hamiltonian : :class:`.Qobj`
          A Hamiltonian in the Dicke basis.
  
          The matrix dimensions are (nds, nds),
          with nds being the number of Dicke states.
          The Hamiltonian can be built with the operators given
          by the `jspin` function in the "dicke" basis.
  
      emission: float
          Incoherent emission coefficient (also nonradiative emission).
          default: 0.0
  
      dephasing: float
          Local dephasing coefficient.
          default: 0.0
  
      pumping: float
          Incoherent pumping coefficient.
          default: 0.0
  
      collective_emission: float
          Collective (superradiant) emmission coefficient.
          default: 0.0
  
      collective_dephasing: float
          Collective dephasing coefficient.
          default: 0.0
  
      collective_pumping: float
          Collective pumping coefficient.
          default: 0.0
  
      nds: int
          The number of Dicke states.
  
      dshape: tuple
          The shape of the Hilbert space in the Dicke or uncoupled basis.
          default: (nds, nds).
      """
  
      def __init__(
          self,
          N,
          hamiltonian=None,
          emission=0.0,
          dephasing=0.0,
          pumping=0.0,
          collective_emission=0.0,
          collective_dephasing=0.0,
          collective_pumping=0.0,
      ):
          self.N = N
          self.hamiltonian = hamiltonian
          self.emission = emission
          self.dephasing = dephasing
          self.pumping = pumping
          self.collective_emission = collective_emission
          self.collective_dephasing = collective_dephasing
          self.collective_pumping = collective_pumping
          self.nds = num_dicke_states(self.N)
          self.dshape = (num_dicke_states(self.N), num_dicke_states(self.N))
  
      def __repr__(self):
          """Print the current parameters of the system."""
          string = []
          string.append("N = {}".format(self.N))
          string.append("Hilbert space dim = {}".format(self.dshape))
          string.append("Number of Dicke states = {}".format(self.nds))
          string.append(
              "Liouvillian space dim = {}".format((self.nds ** 2, self.nds ** 2))
          )
          if self.emission != 0:
              string.append("emission = {}".format(self.emission))
          if self.dephasing != 0:
              string.append("dephasing = {}".format(self.dephasing))
          if self.pumping != 0:
              string.append("pumping = {}".format(self.pumping))
          if self.collective_emission != 0:
              string.append(
                  "collective_emission = {}".format(self.collective_emission)
              )
          if self.collective_dephasing != 0:
              string.append(
                  "collective_dephasing = {}".format(self.collective_dephasing)
              )
          if self.collective_pumping != 0:
              string.append(
                  "collective_pumping = {}".format(self.collective_pumping)
              )
          return "\n".join(string)
  
      def lindbladian(self):
          """Build the Lindbladian superoperator of the dissipative dynamics.
  
          Returns
          -------
          lindbladian : :class:`.Qobj`
              The Lindbladian matrix as a `qutip.Qobj`.
          """
          cythonized_dicke = _Dicke(
              int(self.N),
              float(self.emission),
              float(self.dephasing),
              float(self.pumping),
              float(self.collective_emission),
              float(self.collective_dephasing),
              float(self.collective_pumping),
          )
          return cythonized_dicke.lindbladian()
  
      def liouvillian(self):
          """Build the total Liouvillian using the Dicke basis.
  
          Returns
          -------
          liouv : :class:`.Qobj`
              The Liouvillian matrix for the system.
          """
          lindblad = self.lindbladian()
          if self.hamiltonian is None:
              liouv = lindblad
  
          else:
              hamiltonian = self.hamiltonian
              hamiltonian_superoperator = -1j * spre(hamiltonian) + 1j * spost(
                  hamiltonian
              )
              liouv = lindblad + hamiltonian_superoperator
          return liouv
  
      def pisolve(self, initial_state, tlist):
          """
          Solve for diagonal Hamiltonians and initial states faster.
  
          Parameters
          ==========
          initial_state : :class:`.Qobj`
              An initial state specified as a density matrix of
              `qutip.Qbj` type.
  
          tlist: ndarray
              A 1D numpy array of list of timesteps to integrate
  
          Returns
          =======
          result: list
              A dictionary of the type `qutip.piqs.Result` which holds the
              results of the evolution.
          """
          if isdiagonal(initial_state) == False:
              msg = "`pisolve` requires a diagonal initial density matrix. "
              msg += "In general construct the Liouvillian using "
              msg += "`piqs.liouvillian` and use qutip.mesolve."
              raise ValueError(msg)
  
          if self.hamiltonian and isdiagonal(self.hamiltonian) == False:
              msg = "`pisolve` should only be used for diagonal Hamiltonians. "
              msg += "Construct the Liouvillian using `piqs.liouvillian` and"
              msg += "  use `qutip.mesolve`."
              raise ValueError(msg)
  
          if initial_state.full().shape != self.dshape:
              msg = "Initial density matrix should be diagonal."
              raise ValueError(msg)
  
          pim = Pim(
              self.N,
              self.emission,
              self.dephasing,
              self.pumping,
              self.collective_emission,
              self.collective_pumping,
              self.collective_dephasing,
          )
          result = pim.solve(initial_state, tlist)
          return result
  
      def c_ops(self):
          """Build collapse operators in the full Hilbert space 2^N.
  
          Returns
          -------
          c_ops_list: list
              The list with the collapse operators in the 2^N Hilbert space.
          """
          ce = self.collective_emission
          cd = self.collective_dephasing
          cp = self.collective_pumping
          c_ops_list = collapse_uncoupled(
              N=self.N,
              emission=self.emission,
              dephasing=self.dephasing,
              pumping=self.pumping,
              collective_emission=ce,
              collective_dephasing=cd,
              collective_pumping=cp,
          )
          return c_ops_list
  
      def coefficient_matrix(self):
          """Build coefficient matrix for ODE for a diagonal problem.
  
          Returns
          -------
          M: ndarray
              The matrix M of the coefficients for the ODE dp/dt = Mp.
              p is the vector of the diagonal matrix elements
              of the density matrix rho in the Dicke basis.
          """
          diagonal_system = Pim(
              N=self.N,
              emission=self.emission,
              dephasing=self.dephasing,
              pumping=self.pumping,
              collective_emission=self.collective_emission,
              collective_dephasing=self.collective_dephasing,
              collective_pumping=self.collective_pumping,
          )
          coef_matrix = diagonal_system.coefficient_matrix()
          return coef_matrix
  
  
  # Utility functions for properties of the Dicke space
  def energy_degeneracy(N, m):
      """Calculate the number of Dicke states with same energy.
  
      The use of the ``Decimals`` class allows to explore N > 1000,
      unlike the built-in function ``scipy.special.binom``.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      m: float
          Total spin z-axis projection eigenvalue.
          This is proportional to the total energy.
  
      Returns
      -------
      degeneracy: int
          The energy degeneracy
      """
      numerator = Decimal(factorial(N))
      d1 = Decimal(factorial(_ensure_int(N / 2 + m)))
      d2 = Decimal(factorial(_ensure_int(N / 2 - m)))
      degeneracy = numerator / (d1 * d2)
      return int(degeneracy)
  
  
  def state_degeneracy(N, j):
      r"""Calculate the degeneracy of the Dicke state.
  
      Each state :math:`\lvert j, m\rangle` includes D(N,j) irreducible
      representations :math:`\lvert j, m, \alpha\rangle`.
  
      Uses Decimals to calculate higher numerator and denominators numbers.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      j: float
          Total spin eigenvalue (cooperativity).
  
      Returns
      -------
      degeneracy: int
          The state degeneracy.
      """
      if j < 0:
          raise ValueError("j value should be >= 0")
      numerator = Decimal(factorial(N)) * Decimal(2 * j + 1)
      denominator_1 = Decimal(factorial(_ensure_int(N / 2 + j + 1)))
      denominator_2 = Decimal(factorial(_ensure_int(N / 2 - j)))
      degeneracy = numerator / (denominator_1 * denominator_2)
      degeneracy = int(np.round(float(degeneracy)))
      return degeneracy
  
  
  def m_degeneracy(N, m):
      r"""Calculate the number of Dicke states :math:`\lvert j, m\rangle` with
      same energy.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      m: float
          Total spin z-axis projection eigenvalue (proportional to the total
          energy).
  
      Returns
      -------
      degeneracy: int
          The m-degeneracy.
      """
      jvals = j_vals(N)
      maxj = np.max(jvals)
      if m < -maxj:
          e = "m value is incorrect for this N."
          e += " Minimum m value can be {}".format(-maxj)
          raise ValueError(e)
      degeneracy = N / 2 + 1 - abs(m)
      return int(degeneracy)
  
  
  def ap(j, m):
      r"""
      Calculate the coefficient ``ap`` by applying :math:`J_+\lvert j,m\rangle`.
  
      The action of ap is given by:
      :math:`J_{+}\lvert j, m\rangle = A_{+}(j, m) \lvert j, m+1\rangle`
  
      Parameters
      ----------
      j, m: float
          The value for j and m in the dicke basis :math:`\lvert j, m\rangle`.
  
      Returns
      -------
      a_plus: float
          The value of :math:`a_{+}`.
      """
      a_plus = np.sqrt((j - m) * (j + m + 1))
      return a_plus
  
  
  def am(j, m):
      r"""Calculate the operator ``am`` used later.
  
      The action of ``ap`` is given by:
      :math:`J_{-}\lvert j,m\rangle = A_{-}(jm)\lvert j,m-1\rangle`
  
      Parameters
      ----------
      j: float
          The value for j.
  
      m: float
          The value for m.
  
      Returns
      -------
      a_minus: float
          The value of :math:`a_{-}`.
      """
      a_minus = np.sqrt((j + m) * (j - m + 1))
      return a_minus
  
  
  def spin_algebra(N, op=None):
      """Create the list [sx, sy, sz] with the spin operators.
  
      The operators are constructed for a collection of N two-level systems
      (TLSs). Each element of the list, i.e., sx, is a vector of `qutip.Qobj`
      objects (spin matrices), as it cointains the list of the SU(2) Pauli
      matrices for the N TLSs. Each TLS operator sx[i], with i = 0, ..., (N-1),
      is placed in a :math:`2^N`-dimensional Hilbert space.
  
      Notes
      -----
      sx[i] is :math:`\\frac{\\sigma_x}{2}` in the composite Hilbert space.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      spin_operators: list or :class:`.Qobj`
          A list of `qutip.Qobj` operators - [sx, sy, sz] or the
          requested operator.
      """
      # 1. Define N TLS spin-1/2 matrices in the uncoupled basis
      N = int(N)
      sx = [0 for i in range(N)]
      sy = [0 for i in range(N)]
      sz = [0 for i in range(N)]
      sp = [0 for i in range(N)]
      sm = [0 for i in range(N)]
  
      sx[0] = 0.5 * sigmax()
      sy[0] = 0.5 * sigmay()
      sz[0] = 0.5 * sigmaz()
      sp[0] = sigmap()
      sm[0] = sigmam()
  
      # 2. Place operators in total Hilbert space
      for k in range(N - 1):
          sx[0] = tensor(sx[0], identity(2))
          sy[0] = tensor(sy[0], identity(2))
          sz[0] = tensor(sz[0], identity(2))
          sp[0] = tensor(sp[0], identity(2))
          sm[0] = tensor(sm[0], identity(2))
  
      # 3. Cyclic sequence to create all N operators
      a = [i for i in range(N)]
      b = [[a[i - i2] for i in range(N)] for i2 in range(N)]
  
      # 4. Create N operators
      for i in range(1, N):
          sx[i] = sx[0].permute(b[i])
          sy[i] = sy[0].permute(b[i])
          sz[i] = sz[0].permute(b[i])
          sp[i] = sp[0].permute(b[i])
          sm[i] = sm[0].permute(b[i])
  
      spin_operators = [sx, sy, sz]
  
      if not op:
          return spin_operators
      elif op == "x":
          return sx
      elif op == "y":
          return sy
      elif op == "z":
          return sz
      elif op == "+":
          return sp
      elif op == "-":
          return sm
      else:
          raise TypeError("Invalid type")
  
  
  def _jspin_uncoupled(N, op=None):
      """Construct the the collective spin algebra in the uncoupled basis.
  
      jx, jy, jz, jp, jm are constructed in the uncoupled basis of the
      two-level system (TLS). Each collective operator is placed in a
      Hilbert space of dimension 2^N.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      op: str
          The operator to return 'x','y','z','+','-'.
          If no operator given, then output is the list of operators
          for ['x','y','z',].
  
      Returns
      -------
      collective_operators: list or :class:`.Qobj`
          A list of `qutip.Qobj` representing all the operators in
          uncoupled" basis or a single operator requested.
      """
      # 1. Define N TLS spin-1/2 matrices in the uncoupled basis
      N = int(N)
  
      sx, sy, sz = spin_algebra(N)
      sp, sm = spin_algebra(N, "+"), spin_algebra(N, "-")
  
      jx = sum(sx)
      jy = sum(sy)
      jz = sum(sz)
      jp = sum(sp)
      jm = sum(sm)
  
      collective_operators = [jx, jy, jz]
  
      if not op:
          return collective_operators
      elif op == "x":
          return jx
      elif op == "y":
          return jy
      elif op == "z":
          return jz
      elif op == "+":
          return jp
      elif op == "-":
          return jm
      else:
          raise TypeError("Invalid type")
  
  
  def jspin(N, op=None, basis="dicke"):
      r"""
      Calculate the list of collective operators of the total algebra.
  
      The Dicke basis :math:`\lvert j,m\rangle\langle j,m'\rvert` is used by
      default. Otherwise with "uncoupled" the operators are in a
      :math:`2^N` space.
  
      Parameters
      ----------
      N: int
          Number of two-level systems.
  
      op: str {'x', 'y', 'z', '+', '-'}, optional
          The operator to return 'x','y','z','+','-'.
          If no operator given, then output is the list of operators
          for ['x','y','z'].
  
      basis: str {"dicke", "uncoupled"}, default: "dicke"
          The basis of the operators.
  
      Returns
      -------
      j_alg: list or :class:`.Qobj`
          A list of `qutip.Qobj` representing all the operators in
          the "dicke" or "uncoupled" basis or a single operator requested.
      """
      if basis == "uncoupled":
          return _jspin_uncoupled(N, op)
  
      nds = num_dicke_states(N)
      num_ladders = num_dicke_ladders(N)
      jz_operator = dok_matrix((nds, nds), dtype=np.complex128)
      jp_operator = dok_matrix((nds, nds), dtype=np.complex128)
      jm_operator = dok_matrix((nds, nds), dtype=np.complex128)
      s = 0
  
      for k in range(0, num_ladders):
          j = 0.5 * N - k
          mmax = int(2 * j + 1)
          for i in range(0, mmax):
              m = j - i
              jz_operator[s, s] = m
              if (s + 1) in range(0, nds):
                  jp_operator[s, s + 1] = ap(j, m - 1)
              if (s - 1) in range(0, nds):
                  jm_operator[s, s - 1] = am(j, m + 1)
              s = s + 1
  
      jx_operator = 1 / 2 * (jp_operator + jm_operator)
      jy_operator = 1j / 2 * (jm_operator - jp_operator)
      jx = Qobj(jx_operator)
      jy = Qobj(jy_operator)
      jz = Qobj(jz_operator)
      jp = Qobj(jp_operator)
      jm = Qobj(jm_operator)
  
      if not op:
          return [jx, jy, jz]
      if op == "+":
          return jp
      elif op == "-":
          return jm
      elif op == "x":
          return jx
      elif op == "y":
          return jy
      elif op == "z":
          return jz
      else:
          raise TypeError("Invalid type")
  
  
  def collapse_uncoupled(
      N,
      emission=0.0,
      dephasing=0.0,
      pumping=0.0,
      collective_emission=0.0,
      collective_dephasing=0.0,
      collective_pumping=0.0,
  ):
      """
      Create the collapse operators (c_ops) of the Lindbladian in the
      uncoupled basis
  
      These operators are in the uncoupled basis of the two-level
      system (TLS) SU(2) Pauli matrices.
  
      Notes
      -----
      The collapse operator list can be given to `qutip.mesolve`.
      Notice that the operators are placed in a Hilbert space of
      dimension :math:`2^N`. Thus the method is suitable only for
      small N (of the order of 10).
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      emission: float, default: 0.0
          Incoherent emission coefficient (also nonradiative emission).
  
  
      dephasing: float, default: 0.0
          Local dephasing coefficient.
  
      pumping: float, default: 0.0
          Incoherent pumping coefficient.
  
      collective_emission: float, default: 0.0
          Collective (superradiant) emmission coefficient.
  
      collective_pumping: float, default: 0.0
          Collective pumping coefficient.
  
      collective_dephasing: float, default: 0.0
          Collective dephasing coefficient.
  
      Returns
      -------
      c_ops: list
          The list of collapse operators as :obj:`.Qobj` for the system.
      """
      N = int(N)
  
      if N > 10:
          msg = "N > 10. dim(H) = 2^N. "
          msg += "Better use `piqs.lindbladian` to reduce Hilbert space "
          msg += "dimension and exploit permutational symmetry."
          raise Warning(msg)
  
      [sx, sy, sz] = spin_algebra(N)
      sp, sm = spin_algebra(N, "+"), spin_algebra(N, "-")
      [jx, jy, jz] = jspin(N, basis="uncoupled")
      jp, jm = (
          jspin(N, "+", basis="uncoupled"),
          jspin(N, "-", basis="uncoupled"),
      )
  
      c_ops = []
  
      if emission != 0:
          for i in range(0, N):
              c_ops.append(np.sqrt(emission) * sm[i])
  
      if dephasing != 0:
          for i in range(0, N):
              c_ops.append(np.sqrt(dephasing) * sz[i])
  
      if pumping != 0:
          for i in range(0, N):
              c_ops.append(np.sqrt(pumping) * sp[i])
  
      if collective_emission != 0:
          c_ops.append(np.sqrt(collective_emission) * jm)
  
      if collective_dephasing != 0:
          c_ops.append(np.sqrt(collective_dephasing) * jz)
  
      if collective_pumping != 0:
          c_ops.append(np.sqrt(collective_pumping) * jp)
  
      return c_ops
  
  
  # State definitions in the Dicke basis with an option for basis transformation
  def dicke_basis(N, jmm1):
      r"""
      Initialize the density matrix of a Dicke state for several (j, m, m1).
  
      This function can be used to build arbitrary states in the Dicke basis
      :math:`\lvert j, m\rangle\langle j, m'\rvert`. We create coefficients for
      each (j, m, m1) value in the dictionary jmm1. The mapping for the (i, k)
      index of the density matrix to the :math:`\lvert j, m\rangle` values is
      given by the cythonized function `jmm1_dictionary`. A density matrix is
      created from the given dictionary of coefficients for each (j, m, m1).
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      jmm1: dict
          A dictionary of {(j, m, m1): p} that gives a density p for the
          (j, m, m1) matrix element.
  
      Returns
      -------
      rho: :class:`.Qobj`
          The density matrix in the Dicke basis.
      """
      if not isinstance(jmm1, dict):
          msg = "Please specify the jmm1 values as a dictionary"
          msg += "or use the `excited(N)` function to create an"
          msg += "excited state where jmm1 = {(N/2, N/2, N/2): 1}"
          raise AttributeError(msg)
  
      nds = _num_dicke_states(N)
      rho = np.zeros((nds, nds))
      jmm1_dict = jmm1_dictionary(N)[1]
      for key in jmm1:
          i, k = jmm1_dict[key]
          rho[i, k] = jmm1[key]
      return Qobj(rho)
  
  
  def dicke(N, j, m):
      r"""
      Generate a Dicke state as a pure density matrix in the Dicke basis.
  
      For instance, the superradiant state given by
      :math:`\lvert  j, m\rangle = \lvert 1, 0\rangle` for N = 2,
      and the state is represented as a density matrix of size (nds, nds) or
      (4, 4), with the (1, 1) element set to 1.
  
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      j: float
          The eigenvalue j of the Dicke state (j, m).
  
      m: float
          The eigenvalue m of the Dicke state (j, m).
  
      Returns
      -------
      rho: :class:`.Qobj`
          The density matrix.
      """
      nds = num_dicke_states(N)
      rho = np.zeros((nds, nds))
  
      jmm1_dict = jmm1_dictionary(N)[1]
  
      i, k = jmm1_dict[(j, m, m)]
      rho[i, k] = 1.0
      return Qobj(rho)
  
  
  # Uncoupled states in the full Hilbert space. These are returned with the
  # choice of the keyword argument `basis="uncoupled"` in the state functions.
  def _uncoupled_excited(N):
      """
      Generate the density matrix of the excited Dicke state in the full
      :math:`2^N` dimensional Hilbert space.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      psi0: :class:`.Qobj`
          The density matrix for the excited state in the uncoupled basis.
      """
      N = int(N)
      jz = jspin(N, "z", basis="uncoupled")
      en, vn = jz.eigenstates()
      psi0 = vn[2 ** N - 1]
      return ket2dm(psi0)
  
  
  def _uncoupled_superradiant(N):
      """
      Generate the density matrix of a superradiant state in the full
      :math:`2^N`-dimensional Hilbert space.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      psi0: :class:`.Qobj`
          The density matrix for the superradiant state in the full Hilbert
          space.
      """
      N = int(N)
      jz = jspin(N, "z", basis="uncoupled")
      en, vn = jz.eigenstates()
      psi0 = vn[2 ** N - (N + 1)]
      return ket2dm(psi0)
  
  
  def _uncoupled_ground(N):
      """
      Generate the density matrix of the ground state in the full\
      :math:`2^N`-dimensional Hilbert space.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      psi0: :class:`.Qobj`
          The density matrix for the ground state in the full Hilbert space.
      """
      N = int(N)
      jz = jspin(N, "z", basis="uncoupled")
      en, vn = jz.eigenstates()
      psi0 = vn[0]
      return ket2dm(psi0)
  
  
  def _uncoupled_ghz(N):
      """
      Generate the density matrix of the GHZ state in the full 2^N
      dimensional Hilbert space.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      Returns
      -------
      ghz: :class:`.Qobj`
          The density matrix for the GHZ state in the full Hilbert space.
      """
      N = int(N)
      rho = np.zeros((2 ** N, 2 ** N))
      rho[0, 0] = 1 / 2
      rho[2 ** N - 1, 0] = 1 / 2
      rho[0, 2 ** N - 1] = 1 / 2
      rho[2 ** N - 1, 2 ** N - 1] = 1 / 2
      spin_dim = [2 for i in range(0, N)]
      spins_dims = list((spin_dim, spin_dim))
      rho = Qobj(rho, dims=spins_dims)
      return rho
  
  
  def _uncoupled_css(N, a, b):
      r"""
      Generate the density matrix of the CSS state in the full 2^N
      dimensional Hilbert space.
  
      The CSS states are non-entangled states given by
      :math:`\lvert a,b\rangle = \prod_i(a\lvert1\rangle_i + b\lvert0\rangle_i)`.
  
      Parameters
      ----------
      N: int
          The number of two-level systems.
  
      a: complex
          The coefficient of the :math:`\lvert1_i\rangle` state.
  
      b: complex
  