def compute(N, wc, wa, glist, use_rwa):

        # Pre-compute operators for the hamiltonian
        a  = tensor(destroy(N), qeye(2))
        sm = tensor(qeye(N), destroy(2))
        nc = a.dag() * a
        na = sm.dag() * sm

        idx = 0
        na_expt = zeros(shape(glist))
        nc_expt = zeros(shape(glist))
        for g in glist:

            # recalculate the hamiltonian for each value of g
            if use_rwa:
                H = wc * nc + wa * na + g * (a.dag() * sm + a * sm.dag())
            else:
                H = wc * nc + wa * na + g * (a.dag() + a) * (sm + sm.dag())

            # find the groundstate of the composite system
            evals, ekets = H.eigenstates()
            psi_gnd = ekets[0]
            na_expt[idx] = expect(na, psi_gnd)
            nc_expt[idx] = expect(nc, psi_gnd)

            idx += 1

        return nc_expt, na_expt, ket2dm(psi_gnd)

def test_diagHamiltonian2():
    """
    Diagonalization of composite systems
    """

    H1 = scipy.rand() * sigmax() + scipy.rand() * sigmay() +\
        scipy.rand() * sigmaz()
    H2 = scipy.rand() * sigmax() + scipy.rand() * sigmay() +\
        scipy.rand() * sigmaz()

    H = tensor(H1, H2)

    evals, ekets = H.eigenstates()

    for n in range(len(evals)):
        # assert that max(H * ket - e * ket) is small
        assert_equal(amax(
            abs((H * ekets[n] - evals[n] * ekets[n]).full())) < 1e-10, True)

    N1 = 10
    N2 = 2

    a1 = tensor(destroy(N1), qeye(N2))
    a2 = tensor(qeye(N1), destroy(N2))
    H = scipy.rand() * a1.dag() * a1 + scipy.rand() * a2.dag() * a2 + \
        scipy.rand() * (a1 + a1.dag()) * (a2 + a2.dag())
    evals, ekets = H.eigenstates()

    for n in range(len(evals)):
        # assert that max(H * ket - e * ket) is small
        assert_equal(amax(
            abs((H * ekets[n] - evals[n] * ekets[n]).full())) < 1e-10, True)

def carrier_flop(rho0, W, eta, delta, theta, phi, c_op_list = [], return_op_list = []):
    ''' Return values of atom and motion populations during carrier Rabi flop 
    for rotation angles theta. Calls numerical solution of master equation.
    @ var rho0: initial density matrix
    @ var W: bare Rabi frequency
    @ var eta: Lamb-Dicke parameter
    @ var delta: detuning between atom and motion
    @ var theta: list of Rabi rotation angles (i.e. theta, or g*time)
    @ var phi: phase of the input laser pulse
    @ var c_op_list: list of collapse operators for the master equation treatment
    @ var return_op_list: list of population operators the values of which will be returned 
    
    returns: time, populations of motional mode and atom
    '''    
    N = shape(rho0.data)[0]/2 # assume N Fock states and two atom states
    a = tensor(destroy(N), qeye(2))
    Wc = qeye(N)
    Wc.data = csr_matrix( qeye(N).data.dot( np.diag(rabi_coupling(N,0,eta) ) ) )    
    sm = tensor( Wc, destroy(2))

    # use the rotating wave approxiation
    H = delta * a.dag() * a + \
         (1./2.)* W * (sm.dag()*exp(1j*phi) + sm*exp(-1j*phi))
    
    if hasattr(theta, '__len__'):
        if len(theta)>1: # I need to be able to pass a list of length zero and not get an error
            time = theta/W
    else:
            time = theta/W        

    output = mesolve(H, rho0, time, c_op_list, return_op_list)
    return time, output

def collapse_operators(N, n_th_a, gamma_motion, gamma_motion_phi, gamma_atom):
    '''Collapse operators for the master equation of a single atom and a harmonic oscillator
    @ var N: size of the harmonic oscillator Hilbert space
    @ var n_th: temperature of the noise bath in quanta
    @ var gamma_motion: heating rate of the motion
    @ var gamma_motion_phi: dephasing rate of the motion
    @ var gamma_atom: decay rate of the atom
    
    returns: list of collapse operators for master equation solution of atom + harmonic oscillator
    '''
    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))  
    c_op_list = []

    rate = gamma_motion * (1 + n_th_a)
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a)

    rate = gamma_motion * n_th_a
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a.dag())
    
    rate = gamma_motion_phi
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a.dag() * a)

    rate = gamma_atom
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * sm)
    return c_op_list

def testJCZeroTemperature():
    """
    brmesolve: Jaynes-Cummings model, zero temperature
    """

    N = 10
    a = tensor(destroy(N), identity(2))
    sm = tensor(identity(N), destroy(2))
    psi0 = ket2dm(tensor(basis(N, 1), basis(2, 0)))
    a_ops = [(a + a.dag())]
    e_ops = [a.dag() * a, sm.dag() * sm]

    w0 = 1.0 * 2 * np.pi
    g = 0.05 * 2 * np.pi
    kappa = 0.05
    times = np.linspace(0, 2 * 2 * np.pi / g, 1000)

    c_ops = [np.sqrt(kappa) * a]
    H = w0 * a.dag() * a + w0 * sm.dag() * sm + \
        g * (a + a.dag()) * (sm + sm.dag())

    res_me = mesolve(H, psi0, times, c_ops, e_ops)
    res_brme = brmesolve(H, psi0, times, a_ops, e_ops,
                         spectra_cb=[lambda w: kappa * (w >= 0)])

    for idx, e in enumerate(e_ops):
        diff = abs(res_me.expect[idx] - res_brme.expect[idx]).max()
        assert_(diff < 5e-2)  # accept 5% error

def test_spectrum_espi_legacy():
    """
    correlation: legacy spectrum from es and pi methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    wlist = 2 * np.pi * np.linspace(0.5, 1.5, 100)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        spec1 = spectrum_ss(H, wlist, c_ops, a.dag(), a)
        spec2 = spectrum_pi(H, wlist, c_ops, a.dag(), a)

    assert_(max(abs(spec1 - spec2)) < 1e-3)

def test_spectrum_esfft():
    """
    correlation: comparing spectrum from es and fft methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        tlist = np.linspace(0, 100, 2500)
        corr = correlation_ss(H, tlist, c_ops, a.dag(), a)
        wlist1, spec1 = spectrum_correlation_fft(tlist, corr)
        spec2 = spectrum_ss(H, wlist1, c_ops, a.dag(), a)

    assert_(max(abs(spec1 - spec2)) < 1e-3)

def test_spectrum_espi():
    """
    correlation: comparing spectrum from es and pi methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    wlist = 2 * pi * np.linspace(0.5, 1.5, 100)
    spec1 = spectrum(H, wlist, c_ops, a.dag(), a, solver='es')
    spec2 = spectrum(H, wlist, c_ops, a.dag(), a, solver='pi')

    assert_(max(abs(spec1 - spec2)) < 1e-3)

def rsb_flop(rho0, W, eta, delta, theta, phi, c_op_list = [], return_op_list = []):
    ''' Return values of atom and motion populations during red sideband Rabi flop 
    for rotation angles theta. Calls numerical solution of master equation for the 
    Jaynes-Cummings Hamiltonian.
    @ var rho0: initial density matrix
    @ var W: bare Rabi frequency
    @ var delta: detuning between atom and motion
    @ var theta: list of Rabi rotation angle (i.e. theta, or g*time)
    @ var phi: phase of the input laser pulse
    @ var c_op_list: list of collapse operators for the master equation treatment
    @ var return_op_list: list of population operators the values of which will be returned 
    
    returns: time, populations of motional mode and atom
    '''
    N = shape(rho0.data)[0]/2 # assume N Fock states and two atom states
    a = tensor(destroy(N), qeye(2))
    sm = tensor( qeye(N), destroy(2))
    Wrsb = destroy(N)
    one_then_zero = ([float(x<1) for x in range(N)])
    Wrsb.data = csr_matrix( destroy(N).data.dot( np.diag( rabi_coupling(N,-1,eta) / np.sqrt(one_then_zero+np.linspace(0,N-1,N)) ) ) ) 
    Arsb = tensor(Wrsb, qeye(2))
    # use the rotating wave approxiation
    # Note that the regular a, a.dag() is used for the time evolution of the oscillator
    # Arsb is the destruction operator including the state dependent coupling strength
    H = delta * a.dag() * a + \
        (1./2.) * W * (Arsb.dag() * sm * exp(1j*phi) + Arsb * sm.dag() * exp(-1j*phi))
    
    if hasattr(theta, '__len__'):
        if len(theta)>1: # I need to be able to pass a list of length zero and not get an error
            time = theta/(eta*W)
    else:
            time = theta/(eta*W)        

    output = mesolve(H, rho0, time, c_op_list, return_op_list)
    return time, output

def test_enr_destory_full():
    "Excitation-number-restricted state-space: full state space"
    a1, a2 = enr_destroy([4, 4], 4**2)
    b1, b2 = tensor(destroy(4), identity(4)), tensor(identity(4), destroy(4))

    assert_(a1 == b1)
    assert_(a2 == b2)

 def __init__(self):
     N = 3
     self.t1 = QobjEvo([qeye(N)*(1.+0.1j),[create(N)*(1.-0.1j),f]])
     self.t2 = QobjEvo([destroy(N)*(1.-0.2j)])
     self.t3 = QobjEvo([[destroy(N)*create(N)*(1.+0.2j),f]])
     self.q1 = qeye(N)*(1.+0.3j)
     self.q2 = destroy(N)*(1.-0.3j)
     self.q3 = destroy(N)*create(N)*(1.+0.4j)

def population_operators(N):
    '''Population operators for the master equation 
    @ var N: size of the oscillator Hilbert space
    
    returns: list of population operators for the harmonic oscillator and the atom
    '''
    a = tensor(destroy(N), qeye(2))
    sm = tensor( qeye(N), destroy(2))
    return [a.dag()*a, sm.dag()*sm]

def test_destroy():
    "Destruction operator"
    b4 = basis(5, 4)
    d5 = destroy(5)
    test1 = d5 * b4
    assert_equal(np.allclose(test1.full(), 2.0 * basis(5, 3).full()), True)
    d3 = destroy(3)
    matrix3 = np.array([[0.00000000 + 0.j, 1.00000000 + 0.j, 0.00000000 + 0.j],
                        [0.00000000 + 0.j, 0.00000000 + 0.j, 1.41421356 + 0.j],
                        [0.00000000 + 0.j, 0.00000000 + 0.j, 0.00000000 + 0.j]])

    assert_equal(np.allclose(matrix3, d3.full()), True)

def full_hamiltonian(cav_dim, w_1, w_2, w_c, g_1, g_2):
    """Return a QObj denoting the full Hamiltonian including cavity
     and two qubits"""
    a = qt.destroy(cav_dim)
    num = a.dag() * a
    return (
            g_1 * qt.tensor(qt.create(cav_dim), SMinus, I) +
            g_1 * qt.tensor(qt.destroy(cav_dim), SPlus, I) +
            g_2 * qt.tensor(qt.create(cav_dim), I, SMinus) +
            g_2 * qt.tensor(qt.destroy(cav_dim), I, SPlus) +
            w_c * qt.tensor(num, I, I) +
            0.5 * w_1 * qt.tensor(qt.qeye(cav_dim), SZ, I) +
            0.5 * w_2 * qt.tensor(qt.qeye(cav_dim), I, SZ))

def test_smesolve_photocurrent():
    "Stochastic: photocurrent_mesolve"
    tol = 0.01

    N = 4
    gamma = 0.25
    ntraj = 20
    nsubsteps = 100
    a = destroy(N)

    H = [[a.dag() * a,f]]
    psi0 = coherent(N, 0.5)
    sc_ops = [np.sqrt(gamma) * a, np.sqrt(gamma) * a * 0.5]
    e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]

    times = np.linspace(0, 1.0, 21)
    res_ref = mesolve(H, psi0, times, sc_ops, e_ops, args={"a":2})
    res = photocurrent_mesolve(H, psi0, times, [], sc_ops, e_ops, args={"a":2},
                   ntraj=ntraj, nsubsteps=nsubsteps,  store_measurement=True,
                   map_func=parallel_map)

    assert_(all([np.mean(abs(res.expect[idx] - res_ref.expect[idx])) < tol
                 for idx in range(len(e_ops))]))
    assert_(len(res.measurement) == ntraj)
    assert_(all([m.shape == (len(times), len(sc_ops))
                 for m in res.measurement]))

def testHOZeroTemperature():
    """
    brmesolve: harmonic oscillator, zero temperature
    """

    N = 10
    w0 = 1.0 * 2 * np.pi
    g = 0.05 * w0
    kappa = 0.15

    times = np.linspace(0, 25, 1000)
    a = destroy(N)
    H = w0 * a.dag() * a + g * (a + a.dag())
    psi0 = ket2dm((basis(N, 4) + basis(N, 2) + basis(N, 0)).unit())

    c_ops = [np.sqrt(kappa) * a]
    a_ops = [a + a.dag()]
    e_ops = [a.dag() * a, a + a.dag()]

    res_me = mesolve(H, psi0, times, c_ops, e_ops)
    res_brme = brmesolve(H, psi0, times, a_ops, e_ops,
                         spectra_cb=[lambda w: kappa * (w >= 0)])

    for idx, e in enumerate(e_ops):
        diff = abs(res_me.expect[idx] - res_brme.expect[idx]).max()
        assert_(diff < 1e-2)

def test_qubit_power():
    "Steady state: Thermal qubit - power solver"
    # thermal steadystate of a qubit: compare numerics with analytical formula
    sz = sigmaz()
    sm = destroy(2)

    H = 0.5 * 2 * np.pi * sz
    gamma1 = 0.05

    wth_vec = np.linspace(0.1, 3, 20)
    p_ss = np.zeros(np.shape(wth_vec))

    for idx, wth in enumerate(wth_vec):

        n_th = 1.0 / (np.exp(1.0 / wth) - 1)  # bath temperature
        c_op_list = []
        rate = gamma1 * (1 + n_th)
        c_op_list.append(np.sqrt(rate) * sm)
        rate = gamma1 * n_th
        c_op_list.append(np.sqrt(rate) * sm.dag())
        rho_ss = steadystate(H, c_op_list, method='power')
        p_ss[idx] = expect(sm.dag() * sm, rho_ss)

    p_ss_analytic = np.exp(-1.0 / wth_vec) / (1 + np.exp(-1.0 / wth_vec))
    delta = sum(abs(p_ss_analytic - p_ss))
    assert_equal(delta < 1e-5, True)

def test_ho_lgmres():
    "Steady state: Thermal HO - iterative-lgmres solver"
    # thermal steadystate of an oscillator: compare numerics with analytical
    # formula
    a = destroy(40)
    H = 0.5 * 2 * np.pi * a.dag() * a
    gamma1 = 0.05

    wth_vec = np.linspace(0.1, 3, 20)
    p_ss = np.zeros(np.shape(wth_vec))

    for idx, wth in enumerate(wth_vec):

        n_th = 1.0 / (np.exp(1.0 / wth) - 1)  # bath temperature
        c_op_list = []
        rate = gamma1 * (1 + n_th)
        c_op_list.append(np.sqrt(rate) * a)
        rate = gamma1 * n_th
        c_op_list.append(np.sqrt(rate) * a.dag())
        rho_ss = steadystate(H, c_op_list, method='iterative-lgmres')
        p_ss[idx] = np.real(expect(a.dag() * a, rho_ss))

    p_ss_analytic = 1.0 / (np.exp(1.0 / wth_vec) - 1)
    delta = sum(abs(p_ss_analytic - p_ss))
    assert_equal(delta < 1e-3, True)

    def to_matrix(self, fd):
        n = num(fd)
        a = destroy(fd)
        ic = qeye(fd)
        sz = sigmaz()
        sm = sigmam()
        iq = qeye(2)

        ms = {
            "id": tensor(iq, ic),
            "a*ad" : tensor(iq, n),
            "a+hc" : tensor(iq, a),
            "sz" : tensor(sz, ic),
            "sm+hc" : tensor(sm, ic)
        }

        H0 = 0
        H1s = []
        for (p1, p2), v in self.coefs.items():
            h = ms[p1] * ms[p2]
            try:
                term = float(v) * h
                if not term.isherm:
                    term += term.dag()
                H0 += term
            except ValueError:
                H1s.append([h, v])
                if not h.isherm:
                    replacement = lambda m: '(-' + m.group() + ')'
                    conj_v = re.sub('[1-9]+j', replacement, v)
                    H1s.append([h.dag(), conj_v])
        if H1s:
            return [H0] + H1s
        else:
            return H0

    def steady(N = 20):  # number of basis states to consider
        n=num(N)
        a = destroy(N)
        H = a.dag() * a
        print H.eigenstates()
        #psi0 = basis(N, 10)  # initial state
        kappa = 0.1  # coupling to oscillator
        c_op_list = []
        n_th_a = 2  # temperature with average of 2 excitations
        rate = kappa * (1 + n_th_a)
        c_op_list.append(sqrt(rate) * a)  # decay operators
        rate = kappa * n_th_a
        c_op_list.append(sqrt(rate) * a.dag())  # excitation operators
        final_state = steadystate(H, c_op_list)
        fexpt = expect(a.dag() * a, final_state)

        #tlist = linspace(0, 100, 100)

        #mcdata = mcsolve(H, psi0, tlist, c_op_list, [a.dag() * a], ntraj=100)

        #medata = mesolve(H, psi0, tlist, c_op_list, [a.dag() * a])
        #plot(tlist, mcdata.expect[0],
        #plt.plot(tlist, medata.expect[0], lw=2)
        plt.axhline(y=fexpt, color='r', lw=1.5) # ss expt. value as horiz line (= 2)
        plt.ylim([0, 10])
        plt.show()