#!/usr/bin/env python
# coding: utf-8

# # Rabi Cycling in the Two-Level-System — From Scratch (Variation for Pauli Matrices)

# This notebook uses sympy to symbolically solve the Schrödinger equation for a two-level system, resulting in the exact analytic expressions for the complex amplitudes and the population dynamics of Rabi cycling.
# 
# The particular variation in this notebook is for a Hamiltonian written in Terms of Pauli matrices, specifically a drift Hamiltonian of $\Delta \hat{S}_z$ where $\hat{S}_z = \frac{1}{2} \hat{\sigma}_z$, so that the energy levels are symmetric around zero. This symmetry eliminates the global phase induced otherwise, and thus leads to a simpler expression.

# In[1]:


from sympy import Function, symbols, sqrt, Eq, Abs, exp, cos, sin, Matrix, dsolve, solve


# In[2]:


from sympy import I as 𝕚


# In[3]:


g = Function('g')
e = Function('e')
t = symbols('t', positive=True)
Δ = symbols('Delta', real=True)
Ω0 = symbols('Ω_0', real=True)
Ω = symbols('Omega', real=True)


# In[4]:


Ĥ = Matrix([
    [ -Δ/2, -Ω0/2],
    [-Ω0/2,   Δ/2]
])
Ĥ


# In[5]:


TDSE_system = [
    g(t).diff(t) + 𝕚 * (Ĥ[0,0] * g(t) + Ĥ[0,1] * e(t)),
    e(t).diff(t) + 𝕚 * (Ĥ[1,0]*  g(t) + Ĥ[1,1] * e(t)),
]


# In[6]:


sols_gen = dsolve(TDSE_system, [g(t), e(t)])


# In[7]:


effective_rabi_freq = {
    Ω: sqrt(Δ**2 + Ω0**2)
}


# In[8]:


Eq(Ω, effective_rabi_freq[Ω])


# In[9]:


find_Ω = {
    effective_rabi_freq[Ω]: Ω
}


# In[10]:


sols_gen[0].subs(find_Ω)


# In[11]:


sols_gen[1].subs(find_Ω)


# In[12]:


boundary_conditions = {
    t: 0,
    g(t): 1,
    e(t) : 0,
}


# In[13]:


find_Ω0sq = {
    Ω**2 - Δ**2: Ω0**2
}


# In[14]:


C1, C2 = symbols("C1, C2")


# In[15]:


_integration_constants = solve(
    [sol.subs(find_Ω).subs(boundary_conditions) for sol in sols_gen],
    [C1, C2]
)
integration_constants = {
    k: v.subs(find_Ω0sq)
    for (k, v) in _integration_constants.items()
}


# In[16]:


Eq(C1, integration_constants[C1])


# In[17]:


Eq(C2, integration_constants[C2])


# In[18]:


def simplify_sol_g(sol_g):
    rhs = (
        sol_g
        .rhs
        .rewrite(sin)
        .expand()
        .collect(𝕚)
        .subs(effective_rabi_freq)
        .simplify()
        .subs(find_Ω)
    )
    return Eq(sol_g.lhs, rhs)


# In[19]:


sol_g = simplify_sol_g(sols_gen[0].subs(find_Ω).subs(integration_constants))
sol_g


# In[20]:


sol_e = (
    sols_gen[1]
    .subs(find_Ω)
    .subs(integration_constants)
    .expand()
    .rewrite(sin)
    .expand()
)
sol_e


# In[21]:


pop_e = (sol_e.rhs * sol_e.rhs.conjugate()).rewrite(sin).expand()
Eq(Abs(e(t))**2, pop_e)