from sympy import *
init_printing()

from sympsi import *
from sympsi.boson import *
from sympsi.operatorordering import *

omega_a, omega_b, g, A, Delta, t = symbols("omega_a, omega_b, g, A, Delta, t", positive=True)
Hsym, omega_d = symbols("H, omega_d")

a, b = BosonOp("a"), BosonOp("b")

H0 = omega_a * Dagger(a) * a + omega_b * Dagger(b) * b - g * Dagger(a) * a * (b + Dagger(b))

Hdrive = (A * exp(-I * omega_d * t) + conjugate(A) * exp(I * omega_d * t)) * (a + Dagger(a))

H = H0 + Hdrive

Eq(Hsym, H)

U = exp(I * Dagger(a) * a * omega_d * t)

U

H1 = hamiltonian_transformation(U, H, independent=True)

H1

H2 = drop_terms_containing(H1.expand(), [exp(-2*I*omega_d*t), exp(2*I*omega_d*t)])

Eq(Symbol("H_{rwa}"), H2)

H3 = H2.subs(omega_a, Delta + omega_d).expand()

H3

alpha = symbols("alpha")

UH = Dagger(a) * alpha - conjugate(alpha) * a
U = exp(UH)

U

H4 = hamiltonian_transformation(U, H3, independent=True)

H4

H5 = H4.expand().subs({A: alpha * Delta, conjugate(alpha): alpha}) 

H5 = collect(H5, [g * Dagger(a) * a, - alpha * g])

H5

H6 = drop_c_number_terms(H5)

H6

H7 = H6.subs(g * Dagger(a) * a, 0)

e = (a + Dagger(a)) * (b + Dagger(b))
H7 = H7.subs(e.expand(), e)

Eq(Hsym, H7)

H1 = H7
H1

U = exp(I * Dagger(a) * a * Delta * t)
U

H2 = hamiltonian_transformation(U, H1, independent=True)
H2

U = exp(I * Dagger(b) * b * omega_b * t)
U

H3 = hamiltonian_transformation(U, H2, independent=True)
H3

H4 = H3.expand().subs(Delta, omega_b)
H4

H5 = drop_terms_containing(H4, [exp(+2 * I * omega_b * t), exp(-2 * I * omega_b * t)])
H5

U = exp(-I * Dagger(a) * a * omega_b * t)  # Delta = omega_b
H6 = hamiltonian_transformation(U, H5, independent=True)
U = exp(-I * Dagger(b) * b * omega_b * t)
H7 = hamiltonian_transformation(U, H6, independent=True)
H7 = collect(H7, [alpha**2, g])
H7

H4 = H3.expand().subs(Delta, -omega_b)
H4

H5 = drop_terms_containing(H4, [exp(+2 * I * omega_b * t), exp(-2 * I * omega_b * t)])
H5

U = exp( I * Dagger(a) * a * omega_b * t)  # Delta = -omega_b
H6 = hamiltonian_transformation(U, H5, independent=True)
U = exp(-I * Dagger(b) * b * omega_b * t)
H7 = hamiltonian_transformation(U, H6, independent=True)
H7 = collect(H7, [alpha**2, g])
H7

H0

x = symbols("x")

U = exp(- x * Dagger(a) * a * (Dagger(b) - b))

U

H1 = H0.subs(A, 0)

H1

H2 = hamiltonian_transformation(U, H1, independent=True, expansion_search=False, N=2)

H3 = normal_ordered_form(H2.expand(), independent=True)

H3

H4 = H3.subs({x**2: 0, x**3: 0, x: g/omega_b})  # neglect higher order terms

H4 = collect(H4, Dagger(a)*a)

H4

%reload_ext version_information

%version_information sympy, sympsi