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#    and the QuTiP Developers.
5#    All rights reserved.
6#
7#    Redistribution and use in source and binary forms, with or without
8#    modification, are permitted provided that the following conditions are
9#    met:
10#
11#    1. Redistributions of source code must retain the above copyright notice,
12#       this list of conditions and the following disclaimer.
13#
14#    2. Redistributions in binary form must reproduce the above copyright
15#       notice, this list of conditions and the following disclaimer in the
16#       documentation and/or other materials provided with the distribution.
17#
18#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
19#       of its contributors may be used to endorse or promote products derived
20#       from this software without specific prior written permission.
21#
22#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
23#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
24#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
25#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
26#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
27#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
28#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
29#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
30#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
31#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
32#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
33###############################################################################
34"""Permutational Invariant Quantum Solver (PIQS)
35
36This module calculates the Liouvillian for the dynamics of ensembles of
37identical two-level systems (TLS) in the presence of local and collective
38processes by exploiting permutational symmetry and using the Dicke basis.
39It also allows to characterize nonlinear functions of the density matrix.
40"""
41
42# Authors: Nathan Shammah, Shahnawaz Ahmed
43# Contact: nathan.shammah@gmail.com, shahnawaz.ahmed95@gmail.com
44
45from math import factorial
46from decimal import Decimal
47
48import numpy as np
49from scipy.integrate import odeint
50from scipy.linalg import eigvalsh
51from scipy.special import entr
52from scipy.sparse import dok_matrix, block_diag, lil_matrix
53from qutip.solver import Options, Result
54from qutip import (
55    Qobj,
56    spre,
57    spost,
58    tensor,
59    identity,
60    ket2dm,
61)
62from qutip import sigmax, sigmay, sigmaz, sigmap, sigmam
63from qutip.cy.piqs import Dicke as _Dicke
64from qutip.cy.piqs import (
65    jmm1_dictionary,
66    _num_dicke_states,
67    _num_dicke_ladders,
68    get_blocks,
69    j_min,
70    j_vals,
71)
72
73__all__ = [
74    "num_dicke_states",
75    "num_dicke_ladders",
76    "num_tls",
77    "isdiagonal",
78    "dicke_blocks",
79    "dicke_blocks_full",
80    "dicke_function_trace",
81    "purity_dicke",
82    "entropy_vn_dicke",
83    "Dicke",
84    "state_degeneracy",
85    "m_degeneracy",
86    "energy_degeneracy",
87    "ap",
88    "am",
89    "spin_algebra",
90    "jspin",
91    "collapse_uncoupled",
92    "dicke_basis",
93    "dicke",
94    "excited",
95    "superradiant",
96    "css",
97    "ghz",
98    "ground",
99    "identity_uncoupled",
100    "block_matrix",
101    "tau_column",
102    "Pim",
103]
104
105
106def _ensure_int(x):
107    """
108    Ensure that a floating-point value `x` is exactly an integer, and return it
109    as an int.
110    """
111    out = int(x)
112    if out != x:
113        raise ValueError(f"{x} is not an integral value")
114    return out
115
116
117# Functions necessary to generate the Lindbladian/Liouvillian
118def num_dicke_states(N):
119    """Calculate the number of Dicke states.
120
121    Parameters
122    ----------
123    N: int
124        The number of two-level systems.
125
126    Returns
127    -------
128    nds: int
129        The number of Dicke states.
130    """
131    return _num_dicke_states(N)
132
133
134def num_dicke_ladders(N):
135    """Calculate the total number of ladders in the Dicke space.
136
137    For a collection of N two-level systems it counts how many different
138    "j" exist or the number of blocks in the block-diagonal matrix.
139
140    Parameters
141    ----------
142    N: int
143        The number of two-level systems.
144
145    Returns
146    -------
147    Nj: int
148        The number of Dicke ladders.
149    """
150    return _num_dicke_ladders(N)
151
152
153def num_tls(nds):
154    """Calculate the number of two-level systems.
155
156    Parameters
157    ----------
158    nds: int
159         The number of Dicke states.
160
161    Returns
162    -------
163    N: int
164        The number of two-level systems.
165    """
166    if np.sqrt(nds).is_integer():
167        # N is even
168        N = 2 * (np.sqrt(nds) - 1)
169    else:
170        # N is odd
171        N = 2 * (np.sqrt(nds + 1 / 4) - 1)
172    return int(N)
173
174
175def isdiagonal(mat):
176    """
177    Check if the input matrix is diagonal.
178
179    Parameters
180    ==========
181    mat: ndarray/Qobj
182        A 2D numpy array
183
184    Returns
185    =======
186    diag: bool
187        True/False depending on whether the input matrix is diagonal.
188    """
189    if isinstance(mat, Qobj):
190        mat = mat.full()
191
192    return np.all(mat == np.diag(np.diagonal(mat)))
193
194
195# nonlinear functions of the density matrix
196def dicke_blocks(rho):
197    """Create the list of blocks for block-diagonal density matrix in the Dicke basis.
198
199    Parameters
200    ----------
201    rho : :class:`qutip.Qobj`
202        A 2D block-diagonal matrix of ones with dimension (nds,nds),
203        where nds is the number of Dicke states for N two-level
204        systems.
205
206    Returns
207    -------
208    square_blocks: list of np.array
209        Give back the blocks list.
210
211    """
212    shape_dimension = rho.shape[0]
213    N = num_tls(shape_dimension)
214    ladders = num_dicke_ladders(N)
215    # create a list with the sizes of the blocks, in order
216    blocks_dimensions = int(N / 2 + 1 - 0.5 * (N % 2))
217    blocks_list = [
218        (2 * (i + 1 * (N % 2)) + 1 * ((N + 1) % 2))
219        for i in range(blocks_dimensions)
220    ]
221    blocks_list = np.flip(blocks_list, 0)
222    # create a list with each block matrix as element
223    square_blocks = []
224    block_position = 0
225    for block_size in blocks_list:
226        start_m = block_position
227        end_m = block_position + block_size
228        square_block = rho[start_m:end_m, start_m:end_m]
229        block_position = block_position + block_size
230        square_blocks.append(square_block)
231
232    return square_blocks
233
234
235def dicke_blocks_full(rho):
236    """Give the full (2^N-dimensional) list of blocks for a Dicke-basis matrix.
237
238    Parameters
239    ----------
240    rho : :class:`qutip.Qobj`
241        A 2D block-diagonal matrix of ones with dimension (nds,nds),
242        where nds is the number of Dicke states for N two-level
243        systems.
244
245    Returns
246    -------
247    full_blocks : list
248        The list of blocks expanded in the 2^N space for N qubits.
249
250    """
251    shape_dimension = rho.shape[0]
252    N = num_tls(shape_dimension)
253    ladders = num_dicke_ladders(N)
254    # create a list with the sizes of the blocks, in order
255    blocks_dimensions = int(N / 2 + 1 - 0.5 * (N % 2))
256    blocks_list = [
257        (2 * (i + 1 * (N % 2)) + 1 * ((N + 1) % 2))
258        for i in range(blocks_dimensions)
259    ]
260    blocks_list = np.flip(blocks_list, 0)
261    # create a list with each block matrix as element
262    full_blocks = []
263    k = 0
264    block_position = 0
265    for block_size in blocks_list:
266        start_m = block_position
267        end_m = block_position + block_size
268        square_block = rho[start_m:end_m, start_m:end_m]
269        block_position = block_position + block_size
270        j = N / 2 - k
271        djn = state_degeneracy(N, j)
272        for block_counter in range(0, djn):
273            full_blocks.append(square_block / djn)  # preserve trace
274        k = k + 1
275    return full_blocks
276
277
278def dicke_function_trace(f, rho):
279    """Calculate the trace of a function on a Dicke density matrix.
280    Parameters
281    ----------
282    f : function
283        A Taylor-expandable function of `rho`.
284
285    rho : :class:`qutip.Qobj`
286        A density matrix in the Dicke basis.
287
288    Returns
289    -------
290    res : float
291        Trace of a nonlinear function on `rho`.
292    """
293    N = num_tls(rho.shape[0])
294    blocks = dicke_blocks(rho)
295    block_f = []
296    degen_blocks = np.flip(j_vals(N))
297    state_degeneracies = []
298    for j in degen_blocks:
299        dj = state_degeneracy(N, j)
300        state_degeneracies.append(dj)
301    eigenvals_degeneracy = []
302    deg = []
303    for i, block in enumerate(blocks):
304        dj = state_degeneracies[i]
305        normalized_block = block / dj
306        eigenvals_block = eigvalsh(normalized_block)
307        for val in eigenvals_block:
308            eigenvals_degeneracy.append(val)
309            deg.append(dj)
310
311    eigenvalue = np.array(eigenvals_degeneracy)
312    function_array = f(eigenvalue) * deg
313    return sum(function_array)
314
315
316def entropy_vn_dicke(rho):
317    """Von Neumann Entropy of a Dicke-basis density matrix.
318
319    Parameters
320    ----------
321    rho : :class:`qutip.Qobj`
322        A 2D block-diagonal matrix of ones with dimension (nds,nds),
323        where nds is the number of Dicke states for N two-level
324        systems.
325
326    Returns
327    -------
328    entropy_dm: float
329        Entropy. Use degeneracy to multiply each block.
330
331    """
332    return dicke_function_trace(entr, rho)
333
334
335def purity_dicke(rho):
336    """Calculate purity of a density matrix in the Dicke basis.
337    It accounts for the degenerate blocks in the density matrix.
338
339    Parameters
340    ----------
341    rho : :class:`qutip.Qobj`
342        Density matrix in the Dicke basis of qutip.piqs.jspin(N), for N spins.
343
344    Returns
345    -------
346    purity : float
347        The purity of the quantum state.
348        It's 1 for pure states, 0<=purity<1 for mixed states.
349    """
350    f = lambda x: x * x
351    return dicke_function_trace(f, rho)
352
353
354class Dicke(object):
355    """The Dicke class which builds the Lindbladian and Liouvillian matrix.
356
357    Examples
358    --------
359    >>> from piqs import Dicke, jspin
360    >>> N = 2
361    >>> jx, jy, jz = jspin(N)
362    >>> jp = jspin(N, "+")
363    >>> jm = jspin(N, "-")
364    >>> ensemble = Dicke(N, emission=1.)
365    >>> L = ensemble.liouvillian()
366
367    Parameters
368    ----------
369    N: int
370        The number of two-level systems.
371
372    hamiltonian : :class:`qutip.Qobj`
373        A Hamiltonian in the Dicke basis.
374
375        The matrix dimensions are (nds, nds),
376        with nds being the number of Dicke states.
377        The Hamiltonian can be built with the operators
378        given by the `jspin` functions.
379
380    emission: float
381        Incoherent emission coefficient (also nonradiative emission).
382        default: 0.0
383
384    dephasing: float
385        Local dephasing coefficient.
386        default: 0.0
387
388    pumping: float
389        Incoherent pumping coefficient.
390        default: 0.0
391
392    collective_emission: float
393        Collective (superradiant) emmission coefficient.
394        default: 0.0
395
396    collective_pumping: float
397        Collective pumping coefficient.
398        default: 0.0
399
400    collective_dephasing: float
401        Collective dephasing coefficient.
402        default: 0.0
403
404    Attributes
405    ----------
406    N: int
407        The number of two-level systems.
408
409    hamiltonian : :class:`qutip.Qobj`
410        A Hamiltonian in the Dicke basis.
411
412        The matrix dimensions are (nds, nds),
413        with nds being the number of Dicke states.
414        The Hamiltonian can be built with the operators given
415        by the `jspin` function in the "dicke" basis.
416
417    emission: float
418        Incoherent emission coefficient (also nonradiative emission).
419        default: 0.0
420
421    dephasing: float
422        Local dephasing coefficient.
423        default: 0.0
424
425    pumping: float
426        Incoherent pumping coefficient.
427        default: 0.0
428
429    collective_emission: float
430        Collective (superradiant) emmission coefficient.
431        default: 0.0
432
433    collective_dephasing: float
434        Collective dephasing coefficient.
435        default: 0.0
436
437    collective_pumping: float
438        Collective pumping coefficient.
439        default: 0.0
440
441    nds: int
442        The number of Dicke states.
443
444    dshape: tuple
445        The shape of the Hilbert space in the Dicke or uncoupled basis.
446        default: (nds, nds).
447    """
448
449    def __init__(
450        self,
451        N,
452        hamiltonian=None,
453        emission=0.0,
454        dephasing=0.0,
455        pumping=0.0,
456        collective_emission=0.0,
457        collective_dephasing=0.0,
458        collective_pumping=0.0,
459    ):
460        self.N = N
461        self.hamiltonian = hamiltonian
462        self.emission = emission
463        self.dephasing = dephasing
464        self.pumping = pumping
465        self.collective_emission = collective_emission
466        self.collective_dephasing = collective_dephasing
467        self.collective_pumping = collective_pumping
468        self.nds = num_dicke_states(self.N)
469        self.dshape = (num_dicke_states(self.N), num_dicke_states(self.N))
470
471    def __repr__(self):
472        """Print the current parameters of the system."""
473        string = []
474        string.append("N = {}".format(self.N))
475        string.append("Hilbert space dim = {}".format(self.dshape))
476        string.append("Number of Dicke states = {}".format(self.nds))
477        string.append(
478            "Liouvillian space dim = {}".format((self.nds ** 2, self.nds ** 2))
479        )
480        if self.emission != 0:
481            string.append("emission = {}".format(self.emission))
482        if self.dephasing != 0:
483            string.append("dephasing = {}".format(self.dephasing))
484        if self.pumping != 0:
485            string.append("pumping = {}".format(self.pumping))
486        if self.collective_emission != 0:
487            string.append(
488                "collective_emission = {}".format(self.collective_emission)
489            )
490        if self.collective_dephasing != 0:
491            string.append(
492                "collective_dephasing = {}".format(self.collective_dephasing)
493            )
494        if self.collective_pumping != 0:
495            string.append(
496                "collective_pumping = {}".format(self.collective_pumping)
497            )
498        return "\n".join(string)
499
500    def lindbladian(self):
501        """Build the Lindbladian superoperator of the dissipative dynamics.
502
503        Returns
504        -------
505        lindbladian : :class:`qutip.Qobj`
506            The Lindbladian matrix as a `qutip.Qobj`.
507        """
508        cythonized_dicke = _Dicke(
509            int(self.N),
510            float(self.emission),
511            float(self.dephasing),
512            float(self.pumping),
513            float(self.collective_emission),
514            float(self.collective_dephasing),
515            float(self.collective_pumping),
516        )
517        return cythonized_dicke.lindbladian()
518
519    def liouvillian(self):
520        """Build the total Liouvillian using the Dicke basis.
521
522        Returns
523        -------
524        liouv : :class:`qutip.Qobj`
525            The Liouvillian matrix for the system.
526        """
527        lindblad = self.lindbladian()
528        if self.hamiltonian is None:
529            liouv = lindblad
530
531        else:
532            hamiltonian = self.hamiltonian
533            hamiltonian_superoperator = -1j * spre(hamiltonian) + 1j * spost(
534                hamiltonian
535            )
536            liouv = lindblad + hamiltonian_superoperator
537        return liouv
538
539    def pisolve(self, initial_state, tlist, options=None):
540        """
541        Solve for diagonal Hamiltonians and initial states faster.
542
543        Parameters
544        ==========
545        initial_state : :class:`qutip.Qobj`
546            An initial state specified as a density matrix of
547            `qutip.Qbj` type.
548
549        tlist: ndarray
550            A 1D numpy array of list of timesteps to integrate
551
552        options : :class:`qutip.solver.Options`
553            The options for the solver.
554
555        Returns
556        =======
557        result: list
558            A dictionary of the type `qutip.solver.Result` which holds the
559            results of the evolution.
560        """
561        if isdiagonal(initial_state) == False:
562            msg = "`pisolve` requires a diagonal initial density matrix. "
563            msg += "In general construct the Liouvillian using "
564            msg += "`piqs.liouvillian` and use qutip.mesolve."
565            raise ValueError(msg)
566
567        if self.hamiltonian and isdiagonal(self.hamiltonian) == False:
568            msg = "`pisolve` should only be used for diagonal Hamiltonians. "
569            msg += "Construct the Liouvillian using `piqs.liouvillian` and"
570            msg += "  use `qutip.mesolve`."
571            raise ValueError(msg)
572
573        if initial_state.full().shape != self.dshape:
574            msg = "Initial density matrix should be diagonal."
575            raise ValueError(msg)
576
577        pim = Pim(
578            self.N,
579            self.emission,
580            self.dephasing,
581            self.pumping,
582            self.collective_emission,
583            self.collective_pumping,
584            self.collective_dephasing,
585        )
586        result = pim.solve(initial_state, tlist, options=None)
587        return result
588
589    def c_ops(self):
590        """Build collapse operators in the full Hilbert space 2^N.
591
592        Returns
593        -------
594        c_ops_list: list
595            The list with the collapse operators in the 2^N Hilbert space.
596        """
597        ce = self.collective_emission
598        cd = self.collective_dephasing
599        cp = self.collective_pumping
600        c_ops_list = collapse_uncoupled(
601            N=self.N,
602            emission=self.emission,
603            dephasing=self.dephasing,
604            pumping=self.pumping,
605            collective_emission=ce,
606            collective_dephasing=cd,
607            collective_pumping=cp,
608        )
609        return c_ops_list
610
611    def coefficient_matrix(self):
612        """Build coefficient matrix for ODE for a diagonal problem.
613
614        Returns
615        -------
616        M: ndarray
617            The matrix M of the coefficients for the ODE dp/dt = Mp.
618            p is the vector of the diagonal matrix elements
619            of the density matrix rho in the Dicke basis.
620        """
621        diagonal_system = Pim(
622            N=self.N,
623            emission=self.emission,
624            dephasing=self.dephasing,
625            pumping=self.pumping,
626            collective_emission=self.collective_emission,
627            collective_dephasing=self.collective_dephasing,
628            collective_pumping=self.collective_pumping,
629        )
630        coef_matrix = diagonal_system.coefficient_matrix()
631        return coef_matrix
632
633
634# Utility functions for properties of the Dicke space
635def energy_degeneracy(N, m):
636    """Calculate the number of Dicke states with same energy.
637
638    The use of the `Decimals` class allows to explore N > 1000,
639    unlike the built-in function `scipy.special.binom`
640
641    Parameters
642    ----------
643    N: int
644        The number of two-level systems.
645
646    m: float
647        Total spin z-axis projection eigenvalue.
648        This is proportional to the total energy.
649
650    Returns
651    -------
652    degeneracy: int
653        The energy degeneracy
654    """
655    numerator = Decimal(factorial(N))
656    d1 = Decimal(factorial(_ensure_int(N / 2 + m)))
657    d2 = Decimal(factorial(_ensure_int(N / 2 - m)))
658    degeneracy = numerator / (d1 * d2)
659    return int(degeneracy)
660
661
662def state_degeneracy(N, j):
663    r"""Calculate the degeneracy of the Dicke state.
664
665    Each state :math:`\lvert j, m\rangle` includes D(N,j) irreducible
666    representations :math:`\lvert j, m, \alpha\rangle`.
667
668    Uses Decimals to calculate higher numerator and denominators numbers.
669
670    Parameters
671    ----------
672    N: int
673        The number of two-level systems.
674
675    j: float
676        Total spin eigenvalue (cooperativity).
677
678    Returns
679    -------
680    degeneracy: int
681        The state degeneracy.
682    """
683    if j < 0:
684        raise ValueError("j value should be >= 0")
685    numerator = Decimal(factorial(N)) * Decimal(2 * j + 1)
686    denominator_1 = Decimal(factorial(_ensure_int(N / 2 + j + 1)))
687    denominator_2 = Decimal(factorial(_ensure_int(N / 2 - j)))
688    degeneracy = numerator / (denominator_1 * denominator_2)
689    degeneracy = int(np.round(float(degeneracy)))
690    return degeneracy
691
692
693def m_degeneracy(N, m):
694    r"""Calculate the number of Dicke states :math:`\lvert j, m\rangle` with
695    same energy.
696
697    Parameters
698    ----------
699    N: int
700        The number of two-level systems.
701
702    m: float
703        Total spin z-axis projection eigenvalue (proportional to the total
704        energy).
705
706    Returns
707    -------
708    degeneracy: int
709        The m-degeneracy.
710    """
711    jvals = j_vals(N)
712    maxj = np.max(jvals)
713    if m < -maxj:
714        e = "m value is incorrect for this N."
715        e += " Minimum m value can be {}".format(-maxj)
716        raise ValueError(e)
717    degeneracy = N / 2 + 1 - abs(m)
718    return int(degeneracy)
719
720
721def ap(j, m):
722    r"""
723    Calculate the coefficient ``ap`` by applying :math:`J_+\lvert j,m\rangle`.
724
725    The action of ap is given by:
726    :math:`J_{+}\lvert j, m\rangle = A_{+}(j, m) \lvert j, m+1\rangle`
727
728    Parameters
729    ----------
730    j, m: float
731        The value for j and m in the dicke basis :math:`\lvert j, m\rangle`.
732
733    Returns
734    -------
735    a_plus: float
736        The value of :math:`a_{+}`.
737    """
738    a_plus = np.sqrt((j - m) * (j + m + 1))
739    return a_plus
740
741
742def am(j, m):
743    r"""Calculate the operator ``am`` used later.
744
745    The action of ``ap`` is given by:
746    :math:`J_{-}\lvert j,m\rangle = A_{-}(jm)\lvert j,m-1\rangle`
747
748    Parameters
749    ----------
750    j: float
751        The value for j.
752
753    m: float
754        The value for m.
755
756    Returns
757    -------
758    a_minus: float
759        The value of :math:`a_{-}`.
760    """
761    a_minus = np.sqrt((j + m) * (j - m + 1))
762    return a_minus
763
764
765def spin_algebra(N, op=None):
766    """Create the list [sx, sy, sz] with the spin operators.
767
768    The operators are constructed for a collection of N two-level systems
769    (TLSs). Each element of the list, i.e., sx, is a vector of `qutip.Qobj`
770    objects (spin matrices), as it cointains the list of the SU(2) Pauli
771    matrices for the N TLSs. Each TLS operator sx[i], with i = 0, ..., (N-1),
772    is placed in a :math:`2^N`-dimensional Hilbert space.
773
774    Notes
775    -----
776    sx[i] is :math:`\\frac{\\sigma_x}{2}` in the composite Hilbert space.
777
778    Parameters
779    ----------
780    N: int
781        The number of two-level systems.
782
783    Returns
784    -------
785    spin_operators: list or :class: qutip.Qobj
786        A list of `qutip.Qobj` operators - [sx, sy, sz] or the
787        requested operator.
788    """
789    # 1. Define N TLS spin-1/2 matrices in the uncoupled basis
790    N = int(N)
791    sx = [0 for i in range(N)]
792    sy = [0 for i in range(N)]
793    sz = [0 for i in range(N)]
794    sp = [0 for i in range(N)]
795    sm = [0 for i in range(N)]
796
797    sx[0] = 0.5 * sigmax()
798    sy[0] = 0.5 * sigmay()
799    sz[0] = 0.5 * sigmaz()
800    sp[0] = sigmap()
801    sm[0] = sigmam()
802
803    # 2. Place operators in total Hilbert space
804    for k in range(N - 1):
805        sx[0] = tensor(sx[0], identity(2))
806        sy[0] = tensor(sy[0], identity(2))
807        sz[0] = tensor(sz[0], identity(2))
808        sp[0] = tensor(sp[0], identity(2))
809        sm[0] = tensor(sm[0], identity(2))
810
811    # 3. Cyclic sequence to create all N operators
812    a = [i for i in range(N)]
813    b = [[a[i - i2] for i in range(N)] for i2 in range(N)]
814
815    # 4. Create N operators
816    for i in range(1, N):
817        sx[i] = sx[0].permute(b[i])
818        sy[i] = sy[0].permute(b[i])
819        sz[i] = sz[0].permute(b[i])
820        sp[i] = sp[0].permute(b[i])
821        sm[i] = sm[0].permute(b[i])
822
823    spin_operators = [sx, sy, sz]
824
825    if not op:
826        return spin_operators
827    elif op == "x":
828        return sx
829    elif op == "y":
830        return sy
831    elif op == "z":
832        return sz
833    elif op == "+":
834        return sp
835    elif op == "-":
836        return sm
837    else:
838        raise TypeError("Invalid type")
839
840
841def _jspin_uncoupled(N, op=None):
842    """Construct the the collective spin algebra in the uncoupled basis.
843
844    jx, jy, jz, jp, jm are constructed in the uncoupled basis of the
845    two-level system (TLS). Each collective operator is placed in a
846    Hilbert space of dimension 2^N.
847
848    Parameters
849    ----------
850    N: int
851        The number of two-level systems.
852
853    op: str
854        The operator to return 'x','y','z','+','-'.
855        If no operator given, then output is the list of operators
856        for ['x','y','z',].
857
858    Returns
859    -------
860    collective_operators: list or :class: qutip.Qobj
861        A list of `qutip.Qobj` representing all the operators in
862        uncoupled" basis or a single operator requested.
863    """
864    # 1. Define N TLS spin-1/2 matrices in the uncoupled basis
865    N = int(N)
866
867    sx, sy, sz = spin_algebra(N)
868    sp, sm = spin_algebra(N, "+"), spin_algebra(N, "-")
869
870    jx = sum(sx)
871    jy = sum(sy)
872    jz = sum(sz)
873    jp = sum(sp)
874    jm = sum(sm)
875
876    collective_operators = [jx, jy, jz]
877
878    if not op:
879        return collective_operators
880    elif op == "x":
881        return jx
882    elif op == "y":
883        return jy
884    elif op == "z":
885        return jz
886    elif op == "+":
887        return jp
888    elif op == "-":
889        return jm
890    else:
891        raise TypeError("Invalid type")
892
893
894def jspin(N, op=None, basis="dicke"):
895    r"""
896    Calculate the list of collective operators of the total algebra.
897
898    The Dicke basis :math:`\lvert j,m\rangle\langle j,m'\rvert` is used by
899    default. Otherwise with "uncoupled" the operators are in a
900    :math:`2^N` space.
901
902    Parameters
903    ----------
904    N: int
905        Number of two-level systems.
906
907    op: str
908        The operator to return 'x','y','z','+','-'.
909        If no operator given, then output is the list of operators
910        for ['x','y','z'].
911
912    basis: str
913        The basis of the operators - "dicke" or "uncoupled"
914        default: "dicke".
915
916    Returns
917    -------
918    j_alg: list or :class: qutip.Qobj
919        A list of `qutip.Qobj` representing all the operators in
920        the "dicke" or "uncoupled" basis or a single operator requested.
921    """
922    if basis == "uncoupled":
923        return _jspin_uncoupled(N, op)
924
925    nds = num_dicke_states(N)
926    num_ladders = num_dicke_ladders(N)
927    jz_operator = dok_matrix((nds, nds), dtype=np.complex128)
928    jp_operator = dok_matrix((nds, nds), dtype=np.complex128)
929    jm_operator = dok_matrix((nds, nds), dtype=np.complex128)
930    s = 0
931
932    for k in range(0, num_ladders):
933        j = 0.5 * N - k
934        mmax = int(2 * j + 1)
935        for i in range(0, mmax):
936            m = j - i
937            jz_operator[s, s] = m
938            if (s + 1) in range(0, nds):
939                jp_operator[s, s + 1] = ap(j, m - 1)
940            if (s - 1) in range(0, nds):
941                jm_operator[s, s - 1] = am(j, m + 1)
942            s = s + 1
943
944    jx_operator = 1 / 2 * (jp_operator + jm_operator)
945    jy_operator = 1j / 2 * (jm_operator - jp_operator)
946    jx = Qobj(jx_operator)
947    jy = Qobj(jy_operator)
948    jz = Qobj(jz_operator)
949    jp = Qobj(jp_operator)
950    jm = Qobj(jm_operator)
951
952    if not op:
953        return [jx, jy, jz]
954    if op == "+":
955        return jp
956    elif op == "-":
957        return jm
958    elif op == "x":
959        return jx
960    elif op == "y":
961        return jy
962    elif op == "z":
963        return jz
964    else:
965        raise TypeError("Invalid type")
966
967
968def collapse_uncoupled(
969    N,
970    emission=0.0,
971    dephasing=0.0,
972    pumping=0.0,
973    collective_emission=0.0,
974    collective_dephasing=0.0,
975    collective_pumping=0.0,
976):
977    """
978    Create the collapse operators (c_ops) of the Lindbladian in the
979    uncoupled basis
980
981    These operators are in the uncoupled basis of the two-level
982    system (TLS) SU(2) Pauli matrices.
983
984    Notes
985    -----
986    The collapse operator list can be given to `qutip.mesolve`.
987    Notice that the operators are placed in a Hilbert space of
988    dimension :math:`2^N`. Thus the method is suitable only for
989    small N (of the order of 10).
990
991    Parameters
992    ----------
993    N: int
994        The number of two-level systems.
995
996    emission: float
997        Incoherent emission coefficient (also nonradiative emission).
998        default: 0.0
999
1000    dephasing: float
1001        Local dephasing coefficient.
1002        default: 0.0
1003
1004    pumping: float
1005        Incoherent pumping coefficient.
1006        default: 0.0
1007
1008    collective_emission: float
1009        Collective (superradiant) emmission coefficient.
1010        default: 0.0
1011
1012    collective_pumping: float
1013        Collective pumping coefficient.
1014        default: 0.0
1015
1016    collective_dephasing: float
1017        Collective dephasing coefficient.
1018        default: 0.0
1019
1020    Returns
1021    -------
1022    c_ops: list
1023        The list of collapse operators as `qutip.Qobj` for the system.
1024    """
1025    N = int(N)
1026
1027    if N > 10:
1028        msg = "N > 10. dim(H) = 2^N. "
1029        msg += "Better use `piqs.lindbladian` to reduce Hilbert space "
1030        msg += "dimension and exploit permutational symmetry."
1031        raise Warning(msg)
1032
1033    [sx, sy, sz] = spin_algebra(N)
1034    sp, sm = spin_algebra(N, "+"), spin_algebra(N, "-")
1035    [jx, jy, jz] = jspin(N, basis="uncoupled")
1036    jp, jm = (
1037        jspin(N, "+", basis="uncoupled"),
1038        jspin(N, "-", basis="uncoupled"),
1039    )
1040
1041    c_ops = []
1042
1043    if emission != 0:
1044        for i in range(0, N):
1045            c_ops.append(np.sqrt(emission) * sm[i])
1046
1047    if dephasing != 0:
1048        for i in range(0, N):
1049            c_ops.append(np.sqrt(dephasing) * sz[i])
1050
1051    if pumping != 0:
1052        for i in range(0, N):
1053            c_ops.append(np.sqrt(pumping) * sp[i])
1054
1055    if collective_emission != 0:
1056        c_ops.append(np.sqrt(collective_emission) * jm)
1057
1058    if collective_dephasing != 0:
1059        c_ops.append(np.sqrt(collective_dephasing) * jz)
1060
1061    if collective_pumping != 0:
1062        c_ops.append(np.sqrt(collective_pumping) * jp)
1063
1064    return c_ops
1065
1066
1067# State definitions in the Dicke basis with an option for basis transformation
1068def dicke_basis(N, jmm1=None):
1069    r"""
1070    Initialize the density matrix of a Dicke state for several (j, m, m1).
1071
1072    This function can be used to build arbitrary states in the Dicke basis
1073    :math:`\lvert j, m\rangle\langle j, m'\rvert`. We create coefficients for
1074    each (j, m, m1) value in the dictionary jmm1. The mapping for the (i, k)
1075    index of the density matrix to the :math:`\lvert j, m\rangle` values is
1076    given by the cythonized function `jmm1_dictionary`. A density matrix is
1077    created from the given dictionary of coefficients for each (j, m, m1).
1078
1079    Parameters
1080    ----------
1081    N: int
1082        The number of two-level systems.
1083
1084    jmm1: dict
1085        A dictionary of {(j, m, m1): p} that gives a density p for the
1086        (j, m, m1) matrix element.
1087
1088    Returns
1089    -------
1090    rho: :class: qutip.Qobj
1091        The density matrix in the Dicke basis.
1092    """
1093    if jmm1 is None:
1094        msg = "Please specify the jmm1 values as a dictionary"
1095        msg += "or use the `excited(N)` function to create an"
1096        msg += "excited state where jmm1 = {(N/2, N/2, N/2): 1}"
1097        raise AttributeError(msg)
1098
1099    nds = _num_dicke_states(N)
1100    rho = np.zeros((nds, nds))
1101    jmm1_dict = jmm1_dictionary(N)[1]
1102    for key in jmm1:
1103        i, k = jmm1_dict[key]
1104        rho[i, k] = jmm1[key]
1105    return Qobj(rho)
1106
1107
1108def dicke(N, j, m):
1109    r"""
1110    Generate a Dicke state as a pure density matrix in the Dicke basis.
1111
1112    For instance, the superradiant state given by
1113    :math:`\lvert  j, m\rangle = \lvert 1, 0\rangle` for N = 2,
1114    and the state is represented as a density matrix of size (nds, nds) or
1115    (4, 4), with the (1, 1) element set to 1.
1116
1117
1118    Parameters
1119    ----------
1120    N: int
1121        The number of two-level systems.
1122
1123    j: float
1124        The eigenvalue j of the Dicke state (j, m).
1125
1126    m: float
1127        The eigenvalue m of the Dicke state (j, m).
1128
1129    Returns
1130    -------
1131    rho: :class: qutip.Qobj
1132        The density matrix.
1133    """
1134    nds = num_dicke_states(N)
1135    rho = np.zeros((nds, nds))
1136
1137    jmm1_dict = jmm1_dictionary(N)[1]
1138
1139    i, k = jmm1_dict[(j, m, m)]
1140    rho[i, k] = 1.0
1141    return Qobj(rho)
1142
1143
1144# Uncoupled states in the full Hilbert space. These are returned with the
1145# choice of the keyword argument `basis="uncoupled"` in the state functions.
1146def _uncoupled_excited(N):
1147    """
1148    Generate the density matrix of the excited Dicke state in the full
1149    :math:`2^N` dimensional Hilbert space.
1150
1151    Parameters
1152    ----------
1153    N: int
1154        The number of two-level systems.
1155
1156    Returns
1157    -------
1158    psi0: :class: qutip.Qobj
1159        The density matrix for the excited state in the uncoupled basis.
1160    """
1161    N = int(N)
1162    jz = jspin(N, "z", basis="uncoupled")
1163    en, vn = jz.eigenstates()
1164    psi0 = vn[2 ** N - 1]
1165    return ket2dm(psi0)
1166
1167
1168def _uncoupled_superradiant(N):
1169    """
1170    Generate the density matrix of a superradiant state in the full
1171    :math:`2^N`-dimensional Hilbert space.
1172
1173    Parameters
1174    ----------
1175    N: int
1176        The number of two-level systems.
1177
1178    Returns
1179    -------
1180    psi0: :class: qutip.Qobj
1181        The density matrix for the superradiant state in the full Hilbert
1182        space.
1183    """
1184    N = int(N)
1185    jz = jspin(N, "z", basis="uncoupled")
1186    en, vn = jz.eigenstates()
1187    psi0 = vn[2 ** N - (N + 1)]
1188    return ket2dm(psi0)
1189
1190
1191def _uncoupled_ground(N):
1192    """
1193    Generate the density matrix of the ground state in the full\
1194    :math:`2^N`-dimensional Hilbert space.
1195
1196    Parameters
1197    ----------
1198    N: int
1199        The number of two-level systems.
1200
1201    Returns
1202    -------
1203    psi0: :class: qutip.Qobj
1204        The density matrix for the ground state in the full Hilbert space.
1205    """
1206    N = int(N)
1207    jz = jspin(N, "z", basis="uncoupled")
1208    en, vn = jz.eigenstates()
1209    psi0 = vn[0]
1210    return ket2dm(psi0)
1211
1212
1213def _uncoupled_ghz(N):
1214    """
1215    Generate the density matrix of the GHZ state in the full 2^N
1216    dimensional Hilbert space.
1217
1218    Parameters
1219    ----------
1220    N: int
1221        The number of two-level systems.
1222
1223    Returns
1224    -------
1225    ghz: :class: qutip.Qobj
1226        The density matrix for the GHZ state in the full Hilbert space.
1227    """
1228    N = int(N)
1229    rho = np.zeros((2 ** N, 2 ** N))
1230    rho[0, 0] = 1 / 2
1231    rho[2 ** N - 1, 0] = 1 / 2
1232    rho[0, 2 ** N - 1] = 1 / 2
1233    rho[2 ** N - 1, 2 ** N - 1] = 1 / 2
1234    spin_dim = [2 for i in range(0, N)]
1235    spins_dims = list((spin_dim, spin_dim))
1236    rho = Qobj(rho, dims=spins_dims)
1237    return rho
1238
1239
1240def _uncoupled_css(N, a, b):
1241    r"""
1242    Generate the density matrix of the CSS state in the full 2^N
1243    dimensional Hilbert space.
1244
1245    The CSS states are non-entangled states given by
1246    :math:`\lvert a,b\rangle = \prod_i(a\lvert1\rangle_i + b\lvert0\rangle_i)`.
1247
1248    Parameters
1249    ----------
1250    N: int
1251        The number of two-level systems.
1252
1253    a: complex
1254        The coefficient of the :math:`\lvert1_i\rangle` state.
1255
1256    b: complex
1257        The coefficient of the :math:`\lvert0_i\rangle` state.
1258
1259    Returns
1260    -------
1261    css: :class: qutip.Qobj
1262        The density matrix for the CSS state in the full Hilbert space.
1263    """
1264    N = int(N)
1265    # 1. Define i_th factorized density matrix in the uncoupled basis
1266    rho_i = np.zeros((2, 2), dtype=complex)
1267    rho_i[0, 0] = a * np.conj(a)
1268    rho_i[1, 1] = b * np.conj(b)
1269    rho_i[0, 1] = a * np.conj(b)
1270    rho_i[1, 0] = b * np.conj(a)
1271    rho_i = Qobj(rho_i)
1272    rho = [0 for i in range(N)]
1273    rho[0] = rho_i
1274    # 2. Place single-two-level-system density matrices in total Hilbert space
1275    for k in range(N - 1):
1276        rho[0] = tensor(rho[0], identity(2))
1277    # 3. Cyclic sequence to create all N factorized density matrices
1278    # |CSS>_i<CSS|_i
1279    a = [i for i in range(N)]
1280    b = [[a[i - i2] for i in range(N)] for i2 in range(N)]
1281    # 4. Create all other N-1 factorized density matrices
1282    # |+><+| = Prod_(i=1)^N |CSS>_i<CSS|_i
1283    for i in range(1, N):
1284        rho[i] = rho[0].permute(b[i])
1285    identity_i = Qobj(np.eye(2 ** N), dims=rho[0].dims, shape=rho[0].shape)
1286    rho_tot = identity_i
1287    for i in range(0, N):
1288        rho_tot = rho_tot * rho[i]
1289    return rho_tot
1290
1291
1292def excited(N, basis="dicke"):
1293    """
1294    Generate the density matrix for the excited state.
1295
1296    This state is given by (N/2, N/2) in the default Dicke basis. If the
1297    argument `basis` is "uncoupled" then it generates the state in a
1298    2**N dim Hilbert space.
1299
1300    Parameters
1301    ----------
1302    N: int
1303        The number of two-level systems.
1304
1305    basis: str
1306        The basis to use. Either "dicke" or "uncoupled".
1307
1308    Returns
1309    -------
1310    state: :class: qutip.Qobj
1311        The excited state density matrix in the requested basis.
1312    """
1313    if basis == "uncoupled":
1314        state = _uncoupled_excited(N)
1315        return state
1316
1317    jmm1 = {(N / 2, N / 2, N / 2): 1}
1318    return dicke_basis(N, jmm1)
1319
1320
1321def superradiant(N, basis="dicke"):
1322    """
1323    Generate the density matrix of the superradiant state.
1324
1325    This state is given by (N/2, 0) or (N/2, 0.5) in the Dicke basis.
1326    If the argument `basis` is "uncoupled" then it generates the state
1327    in a 2**N dim Hilbert space.
1328
1329    Parameters
1330    ----------
1331    N: int
1332        The number of two-level systems.
1333
1334    basis: str
1335        The basis to use. Either "dicke" or "uncoupled".
1336
1337    Returns
1338    -------
1339    state: :class: qutip.Qobj
1340        The superradiant state density matrix in the requested basis.
1341    """
1342    if basis == "uncoupled":
1343        state = _uncoupled_superradiant(N)
1344        return state
1345
1346    if N % 2 == 0:
1347        jmm1 = {(N / 2, 0, 0): 1.0}
1348        return dicke_basis(N, jmm1)
1349    else:
1350        jmm1 = {(N / 2, 0.5, 0.5): 1.0}
1351    return dicke_basis(N, jmm1)
1352
1353
1354def css(
1355    N,
1356    x=1 / np.sqrt(2),
1357    y=1 / np.sqrt(2),
1358    basis="dicke",
1359    coordinates="cartesian",
1360):
1361    r"""
1362    Generate the density matrix of the Coherent Spin State (CSS).
1363
1364    It can be defined as,
1365    :math:`\lvert CSS\rangle = \prod_i^N(a\lvert1\rangle_i+b\lvert0\rangle_i)`
1366    with :math:`a = sin(\frac{\theta}{2})`,
1367    :math:`b = e^{i \phi}\cos(\frac{\theta}{2})`.
1368    The default basis is that of Dicke space
1369    :math:`\lvert j, m\rangle \langle j, m'\rvert`.
1370    The default state is the symmetric CSS,
1371    :math:`\lvert CSS\rangle = \lvert+\rangle`.
1372
1373    Parameters
1374    ----------
1375    N: int
1376        The number of two-level systems.
1377
1378    x, y: float
1379        The coefficients of the CSS state.
1380
1381    basis: str
1382        The basis to use. Either "dicke" or "uncoupled".
1383
1384    coordinates: str
1385        Either "cartesian" or "polar". If polar then the coefficients
1386        are constructed as sin(x/2), cos(x/2)e^(iy).
1387
1388    Returns
1389    -------
1390    rho: :class: qutip.Qobj
1391        The CSS state density matrix.
1392    """
1393    if coordinates == "polar":
1394        a = np.cos(0.5 * x) * np.exp(1j * y)
1395        b = np.sin(0.5 * x)
1396    else:
1397        a = x
1398        b = y
1399    if basis == "uncoupled":
1400        return _uncoupled_css(N, a, b)
1401    nds = num_dicke_states(N)
1402    num_ladders = num_dicke_ladders(N)
1403    rho = dok_matrix((nds, nds), dtype=np.complex128)
1404
1405    # loop in the allowed matrix elements
1406    jmm1_dict = jmm1_dictionary(N)[1]
1407
1408    j = 0.5 * N
1409    mmax = int(2 * j + 1)
1410    for i in range(0, mmax):
1411        m = j - i
1412        psi_m = (
1413            np.sqrt(float(energy_degeneracy(N, m)))
1414            * a ** (N * 0.5 + m)
1415            * b ** (N * 0.5 - m)
1416        )
1417        for i1 in range(0, mmax):
1418            m1 = j - i1
1419            row_column = jmm1_dict[(j, m, m1)]
1420            psi_m1 = (
1421                np.sqrt(float(energy_degeneracy(N, m1)))
1422                * np.conj(a) ** (N * 0.5 + m1)
1423                * np.conj(b) ** (N * 0.5 - m1)
1424            )
1425            rho[row_column] = psi_m * psi_m1
1426    return Qobj(rho)
1427
1428
1429def ghz(N, basis="dicke"):
1430    """
1431    Generate the density matrix of the GHZ state.
1432
1433    If the argument `basis` is "uncoupled" then it generates the state
1434    in a :math:`2^N`-dimensional Hilbert space.
1435
1436    Parameters
1437    ----------
1438    N: int
1439        The number of two-level systems.
1440
1441    basis: str
1442        The basis to use. Either "dicke" or "uncoupled".
1443
1444    Returns
1445    -------
1446    state: :class: qutip.Qobj
1447        The GHZ state density matrix in the requested basis.
1448    """
1449    if basis == "uncoupled":
1450        return _uncoupled_ghz(N)
1451    nds = _num_dicke_states(N)
1452    rho = dok_matrix((nds, nds), dtype=np.complex128)
1453    rho[0, 0] = 1 / 2
1454    rho[N, N] = 1 / 2
1455    rho[N, 0] = 1 / 2
1456    rho[0, N] = 1 / 2
1457    return Qobj(rho)
1458
1459
1460def ground(N, basis="dicke"):
1461    """
1462    Generate the density matrix of the ground state.
1463
1464    This state is given by (N/2, -N/2) in the Dicke basis. If the argument
1465    `basis` is "uncoupled" then it generates the state in a
1466    :math:`2^N`-dimensional Hilbert space.
1467
1468    Parameters
1469    ----------
1470    N: int
1471        The number of two-level systems.
1472
1473    basis: str
1474        The basis to use. Either "dicke" or "uncoupled"
1475
1476    Returns
1477    -------
1478    state: :class: qutip.Qobj
1479        The ground state density matrix in the requested basis.
1480    """
1481    if basis == "uncoupled":
1482        state = _uncoupled_ground(N)
1483        return state
1484    nds = _num_dicke_states(N)
1485    rho = dok_matrix((nds, nds), dtype=np.complex128)
1486    rho[N, N] = 1
1487    return Qobj(rho)
1488
1489
1490def identity_uncoupled(N):
1491    """
1492    Generate the identity in a :math:`2^N`-dimensional Hilbert space.
1493
1494    The identity matrix is formed from the tensor product of N TLSs.
1495
1496    Parameters
1497    ----------
1498    N: int
1499        The number of two-level systems.
1500
1501    Returns
1502    -------
1503    identity: :class: qutip.Qobj
1504        The identity matrix.
1505    """
1506    N = int(N)
1507    rho = np.zeros((2 ** N, 2 ** N))
1508    for i in range(0, 2 ** N):
1509        rho[i, i] = 1
1510    spin_dim = [2 for i in range(0, N)]
1511    spins_dims = list((spin_dim, spin_dim))
1512    identity = Qobj(rho, dims=spins_dims)
1513    return identity
1514
1515
1516def block_matrix(N, elements="ones"):
1517    """Construct the block-diagonal matrix for the Dicke basis.
1518
1519    Parameters
1520    ----------
1521    N : int
1522        Number of two-level systems.
1523    elements : str {'ones' (default),'degeneracy'}
1524
1525    Returns
1526    -------
1527    block_matr : ndarray
1528        A 2D block-diagonal matrix with dimension (nds,nds),
1529        where nds is the number of Dicke states for N two-level
1530        systems. Filled with ones or the value of degeneracy
1531        at each matrix element.
1532    """
1533    # create a list with the sizes of the blocks, in order
1534    blocks_dimensions = int(N / 2 + 1 - 0.5 * (N % 2))
1535    blocks_list = [
1536        (2 * (i + 1 * (N % 2)) + 1 * ((N + 1) % 2))
1537        for i in range(blocks_dimensions)
1538    ]
1539    blocks_list = np.flip(blocks_list, 0)
1540    # create a list with each block matrix as element
1541    square_blocks = []
1542    k = 0
1543    for i in blocks_list:
1544        if elements == "ones":
1545            square_blocks.append(np.ones((i, i)))
1546        elif elements == "degeneracy":
1547            j = N / 2 - k
1548            dj = state_degeneracy(N, j)
1549            square_blocks.append(dj * np.ones((i, i)))
1550        k = k + 1
1551    return block_diag(square_blocks)
1552
1553
1554# ============================================================================
1555# Adding a faster version to make a Permutational Invariant matrix
1556# ============================================================================
1557def tau_column(tau, k, j):
1558    """
1559    Determine the column index for the non-zero elements of the matrix for a
1560    particular row `k` and the value of `j` from the Dicke space.
1561
1562    Parameters
1563    ----------
1564    tau: str
1565        The tau function to check for this `k` and `j`.
1566
1567    k: int
1568        The row of the matrix M for which the non zero elements have
1569        to be calculated.
1570
1571    j: float
1572        The value of `j` for this row.
1573    """
1574    # In the notes, we indexed from k = 1, here we do it from k = 0
1575    k = k + 1
1576    mapping = {
1577        "tau3": k - (2 * j + 3),
1578        "tau2": k - 1,
1579        "tau4": k + (2 * j - 1),
1580        "tau5": k - (2 * j + 2),
1581        "tau1": k,
1582        "tau6": k + (2 * j),
1583        "tau7": k - (2 * j + 1),
1584        "tau8": k + 1,
1585        "tau9": k + (2 * j + 1),
1586    }
1587    # we need to decrement k again as indexing is from 0
1588    return int(mapping[tau] - 1)
1589
1590
1591class Pim(object):
1592    """
1593    The Permutation Invariant Matrix class.
1594
1595    Initialize the class with the parameters for generating a Permutation
1596    Invariant matrix which evolves a given diagonal initial state `p` as:
1597
1598                                dp/dt = Mp
1599
1600    Parameters
1601    ----------
1602    N: int
1603        The number of two-level systems.
1604
1605    emission: float
1606        Incoherent emission coefficient (also nonradiative emission).
1607        default: 0.0
1608
1609    dephasing: float
1610        Local dephasing coefficient.
1611        default: 0.0
1612
1613    pumping: float
1614        Incoherent pumping coefficient.
1615        default: 0.0
1616
1617    collective_emission: float
1618        Collective (superradiant) emmission coefficient.
1619        default: 0.0
1620
1621    collective_pumping: float
1622        Collective pumping coefficient.
1623        default: 0.0
1624
1625    collective_dephasing: float
1626        Collective dephasing coefficient.
1627        default: 0.0
1628
1629    Attributes
1630    ----------
1631    N: int
1632        The number of two-level systems.
1633
1634    emission: float
1635        Incoherent emission coefficient (also nonradiative emission).
1636        default: 0.0
1637
1638    dephasing: float
1639        Local dephasing coefficient.
1640        default: 0.0
1641
1642    pumping: float
1643        Incoherent pumping coefficient.
1644        default: 0.0
1645
1646    collective_emission: float
1647        Collective (superradiant) emmission coefficient.
1648        default: 0.0
1649
1650    collective_dephasing: float
1651        Collective dephasing coefficient.
1652        default: 0.0
1653
1654    collective_pumping: float
1655        Collective pumping coefficient.
1656        default: 0.0
1657
1658    M: dict
1659        A nested dictionary of the structure {row: {col: val}} which holds
1660        non zero elements of the matrix M
1661    """
1662
1663    def __init__(
1664        self,
1665        N,
1666        emission=0.0,
1667        dephasing=0,
1668        pumping=0,
1669        collective_emission=0,
1670        collective_pumping=0,
1671        collective_dephasing=0,
1672    ):
1673        self.N = N
1674        self.emission = emission
1675        self.dephasing = dephasing
1676        self.pumping = pumping
1677        self.collective_pumping = collective_pumping
1678        self.collective_dephasing = collective_dephasing
1679        self.collective_emission = collective_emission
1680        self.M = {}
1681
1682    def isdicke(self, dicke_row, dicke_col):
1683        """
1684        Check if an element in a matrix is a valid element in the Dicke space.
1685        Dicke row: j value index. Dicke column: m value index.
1686        The function returns True if the element exists in the Dicke space and
1687        False otherwise.
1688
1689        Parameters
1690        ----------
1691        dicke_row, dicke_col : int
1692            Index of the element in Dicke space which needs to be checked
1693        """
1694        rows = self.N + 1
1695        cols = 0
1696
1697        if (self.N % 2) == 0:
1698            cols = int(self.N / 2 + 1)
1699        else:
1700            cols = int(self.N / 2 + 1 / 2)
1701        if (dicke_row > rows) or (dicke_row < 0):
1702            return False
1703        if (dicke_col > cols) or (dicke_col < 0):
1704            return False
1705        if (dicke_row < int(rows / 2)) and (dicke_col > dicke_row):
1706            return False
1707        if (dicke_row >= int(rows / 2)) and (rows - dicke_row <= dicke_col):
1708            return False
1709        else:
1710            return True
1711
1712    def tau_valid(self, dicke_row, dicke_col):
1713        """
1714        Find the Tau functions which are valid for this value of (dicke_row,
1715        dicke_col) given the number of TLS. This calculates the valid tau
1716        values and reurns a dictionary specifying the tau function name and
1717        the value.
1718
1719        Parameters
1720        ----------
1721        dicke_row, dicke_col : int
1722            Index of the element in Dicke space which needs to be checked.
1723
1724        Returns
1725        -------
1726        taus: dict
1727            A dictionary of key, val as {tau: value} consisting of the valid
1728            taus for this row and column of the Dicke space element.
1729        """
1730        tau_functions = [
1731            self.tau3,
1732            self.tau2,
1733            self.tau4,
1734            self.tau5,
1735            self.tau1,
1736            self.tau6,
1737            self.tau7,
1738            self.tau8,
1739            self.tau9,
1740        ]
1741        N = self.N
1742        if self.isdicke(dicke_row, dicke_col) is False:
1743            return False
1744        # The 3x3 sub matrix surrounding the Dicke space element to
1745        # run the tau functions
1746        indices = [
1747            (dicke_row + x, dicke_col + y)
1748            for x in range(-1, 2)
1749            for y in range(-1, 2)
1750        ]
1751        taus = {}
1752        for idx, tau in zip(indices, tau_functions):
1753            if self.isdicke(idx[0], idx[1]):
1754                j, m = self.calculate_j_m(idx[0], idx[1])
1755                taus[tau.__name__] = tau(j, m)
1756        return taus
1757
1758    def calculate_j_m(self, dicke_row, dicke_col):
1759        """
1760        Get the value of j and m for the particular Dicke space element.
1761
1762        Parameters
1763        ----------
1764        dicke_row, dicke_col: int
1765            The row and column from the Dicke space matrix
1766
1767        Returns
1768        -------
1769        j, m: float
1770            The j and m values.
1771        """
1772        N = self.N
1773        j = N / 2 - dicke_col
1774        m = N / 2 - dicke_row
1775        return (j, m)
1776
1777    def calculate_k(self, dicke_row, dicke_col):
1778        """
1779        Get k value from the current row and column element in the Dicke space.
1780
1781        Parameters
1782        ----------
1783        dicke_row, dicke_col: int
1784            The row and column from the Dicke space matrix.
1785        Returns
1786        -------
1787        k: int
1788            The row index for the matrix M for given Dicke space
1789            element.
1790        """
1791        N = self.N
1792        if dicke_row == 0:
1793            k = dicke_col
1794        else:
1795            k = int(
1796                ((dicke_col) / 2) * (2 * (N + 1) - 2 * (dicke_col - 1))
1797                + (dicke_row - (dicke_col))
1798            )
1799        return k
1800
1801    def coefficient_matrix(self):
1802        """
1803        Generate the matrix M governing the dynamics for diagonal cases.
1804
1805        If the initial density matrix and the Hamiltonian is diagonal, the
1806        evolution of the system is given by the simple ODE: dp/dt = Mp.
1807        """
1808        N = self.N
1809        nds = num_dicke_states(N)
1810        rows = self.N + 1
1811        cols = 0
1812
1813        sparse_M = lil_matrix((nds, nds), dtype=float)
1814        if (self.N % 2) == 0:
1815            cols = int(self.N / 2 + 1)
1816        else:
1817            cols = int(self.N / 2 + 1 / 2)
1818        for (dicke_row, dicke_col) in np.ndindex(rows, cols):
1819            if self.isdicke(dicke_row, dicke_col):
1820                k = int(self.calculate_k(dicke_row, dicke_col))
1821                row = {}
1822                taus = self.tau_valid(dicke_row, dicke_col)
1823                for tau in taus:
1824                    j, m = self.calculate_j_m(dicke_row, dicke_col)
1825                    current_col = tau_column(tau, k, j)
1826                    sparse_M[k, int(current_col)] = taus[tau]
1827        return sparse_M.tocsr()
1828
1829    def solve(self, rho0, tlist, options=None):
1830        """
1831        Solve the ODE for the evolution of diagonal states and Hamiltonians.
1832        """
1833        if options is None:
1834            options = Options()
1835        output = Result()
1836        output.solver = "pisolve"
1837        output.times = tlist
1838        output.states = []
1839        output.states.append(Qobj(rho0))
1840        rhs_generate = lambda y, tt, M: M.dot(y)
1841        rho0_flat = np.diag(np.real(rho0.full()))
1842        L = self.coefficient_matrix()
1843        rho_t = odeint(rhs_generate, rho0_flat, tlist, args=(L,))
1844        for r in rho_t[1:]:
1845            diag = np.diag(r)
1846            output.states.append(Qobj(diag))
1847        return output
1848
1849    def tau1(self, j, m):
1850        """
1851        Calculate coefficient matrix element relative to (j, m, m).
1852        """
1853        yS = self.collective_emission
1854        yL = self.emission
1855        yD = self.dephasing
1856        yP = self.pumping
1857        yCP = self.collective_pumping
1858        N = float(self.N)
1859        spontaneous = yS * (1 + j - m) * (j + m)
1860        losses = yL * (N / 2 + m)
1861        pump = yP * (N / 2 - m)
1862        collective_pump = yCP * (1 + j + m) * (j - m)
1863        if j == 0:
1864            dephase = yD * N / 4
1865        else:
1866            dephase = yD * (N / 4 - m ** 2 * ((1 + N / 2) / (2 * j * (j + 1))))
1867        t1 = spontaneous + losses + pump + dephase + collective_pump
1868        return -t1
1869
1870    def tau2(self, j, m):
1871        """
1872        Calculate coefficient matrix element relative to (j, m+1, m+1).
1873        """
1874        yS = self.collective_emission
1875        yL = self.emission
1876        N = float(self.N)
1877        spontaneous = yS * (1 + j - m) * (j + m)
1878        losses = yL * (
1879            ((N / 2 + 1) * (j - m + 1) * (j + m)) / (2 * j * (j + 1))
1880        )
1881        t2 = spontaneous + losses
1882        return t2
1883
1884    def tau3(self, j, m):
1885        """
1886        Calculate coefficient matrix element relative to (j+1, m+1, m+1).
1887        """
1888        yL = self.emission
1889        N = float(self.N)
1890        num = (j + m - 1) * (j + m) * (j + 1 + N / 2)
1891        den = 2 * j * (2 * j + 1)
1892        t3 = yL * (num / den)
1893        return t3
1894
1895    def tau4(self, j, m):
1896        """
1897        Calculate coefficient matrix element relative to (j-1, m+1, m+1).
1898        """
1899        yL = self.emission
1900        N = float(self.N)
1901        num = (j - m + 1) * (j - m + 2) * (N / 2 - j)
1902        den = 2 * (j + 1) * (2 * j + 1)
1903        t4 = yL * (num / den)
1904        return t4
1905
1906    def tau5(self, j, m):
1907        """
1908        Calculate coefficient matrix element relative to (j+1, m, m).
1909        """
1910        yD = self.dephasing
1911        N = float(self.N)
1912        num = (j - m) * (j + m) * (j + 1 + N / 2)
1913        den = 2 * j * (2 * j + 1)
1914        t5 = yD * (num / den)
1915        return t5
1916
1917    def tau6(self, j, m):
1918        """
1919        Calculate coefficient matrix element relative to (j-1, m, m).
1920        """
1921        yD = self.dephasing
1922        N = float(self.N)
1923        num = (j - m + 1) * (j + m + 1) * (N / 2 - j)
1924        den = 2 * (j + 1) * (2 * j + 1)
1925        t6 = yD * (num / den)
1926        return t6
1927
1928    def tau7(self, j, m):
1929        """
1930        Calculate coefficient matrix element relative to (j+1, m-1, m-1).
1931        """
1932        yP = self.pumping
1933        N = float(self.N)
1934        num = (j - m - 1) * (j - m) * (j + 1 + N / 2)
1935        den = 2 * j * (2 * j + 1)
1936        t7 = yP * (float(num) / den)
1937        return t7
1938
1939    def tau8(self, j, m):
1940        """
1941        Calculate coefficient matrix element relative to (j, m-1, m-1).
1942        """
1943        yP = self.pumping
1944        yCP = self.collective_pumping
1945        N = float(self.N)
1946
1947        num = (1 + N / 2) * (j - m) * (j + m + 1)
1948        den = 2 * j * (j + 1)
1949        pump = yP * (float(num) / den)
1950        collective_pump = yCP * (j - m) * (j + m + 1)
1951        t8 = pump + collective_pump
1952        return t8
1953
1954    def tau9(self, j, m):
1955        """
1956        Calculate coefficient matrix element relative to (j-1, m-1, m-1).
1957        """
1958        yP = self.pumping
1959        N = float(self.N)
1960        num = (j + m + 1) * (j + m + 2) * (N / 2 - j)
1961        den = 2 * (j + 1) * (2 * j + 1)
1962        t9 = yP * (float(num) / den)
1963        return t9
1964