import numpy as np
import matplotlib.pyplot as plt
import qutip as qt
from qutip import Bloch
from mpl_toolkits.mplot3d import Axes3D
# --- Set up plot styling ---
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 14
fig = plt.figure(figsize=(7, 4))
ax = fig.add_subplot(121, projection='3d')
b = Bloch(axes=ax)
b.add_states([qt.basis(2, 0), 1.3* (1 * qt.basis(2, 0) + 1*qt.basis(2, 1)).unit(), (1 * qt.basis(2, 0) - 1j*qt.basis(2, 1)).unit()])
b.vector_color = ['r','b', '#FF9999']
b.vector_width = 2
# Using QuTiP's built-in state evolution
# Rabi oscillation with detuning
H = - 1* qt.sigmax() # Hamiltonian (Pauli X)
psi0 = (1 * qt.basis(2, 0) - 1j*qt.basis(2, 1)).unit() # Initial state |0⟩
times = np.linspace(0, np.pi/2, 15)
result = qt.sesolve(H.unit(), psi0, times)
for i in range(len(result.states)-2):
b.add_arc(result.states[i], result.states[i+1], fmt="r", zorder=10, lw=1)
#b.add_states(result.states[i], colors='r')
#result2=qt.sesolve(H.unit(), result.states[-2], [0,1.5*np.pi/14])
#b.add_arc(result.states[-2], result2.states[1], fmt="r", zorder=10, lw=1)
#b.add_states(result.states[-2], colors='r')
b.xlabel = ['','']
b.ylabel = [r'$v$','']
b.zlabel = [r'$w$','']
b.render()
new_arrowhead = qt.bloch.Arrow3D(xs=[-0.1, 0], ys=[0, 0], zs=[0.99, 1],
mutation_scale=b.vector_mutation,
lw=1, arrowstyle=b.vector_style, color='r', zorder=15)
ax.add_artist(new_arrowhead)
ax.text(-.15,-1.1,0,r'$u$',ha='center', va='center', fontsize='20')
ax.text(-.2,-1.5,0,r'$\vec{\Omega}$',ha='center', va='center', fontsize='20')
ax.set_title(r'Case $1$:')
ax2 = fig.add_subplot(122, projection='3d')
b2 = Bloch(axes=ax2)
b2.add_states([qt.basis(2, 1), 1.3* (1 * qt.basis(2, 0) + 1*qt.basis(2, 1)).unit(), (1 * qt.basis(2, 0) + 1j*qt.basis(2, 1)).unit()])
b2.vector_color = ['r','b', '#FF9999']
b2.vector_width = 2
times2 = np.linspace(0, np.pi/2, 15)
result3 = qt.sesolve(H.unit(), (1 * qt.basis(2, 0) + 1j*qt.basis(2, 1)).unit(), times2)
for i in range(len(result3.states)-2):
b2.add_arc(result3.states[i], result3.states[i+1], fmt="r", zorder=10, lw=1)
b2.xlabel = ['','']
b2.ylabel = [r'$v$','']
b2.zlabel = [r'$w$','']
b2.render()
new_arrowhead2 = qt.bloch.Arrow3D(xs=[0.1, 0], ys=[0, 0], zs=[-0.99, -1],
mutation_scale=b.vector_mutation,
lw=1, arrowstyle=b.vector_style, color='r', zorder=15)
ax2.add_artist(new_arrowhead2)
ax2.text(-.15,-1.1,0,r'$u$',ha='center', va='center', fontsize='20')
ax2.text(-.2,-1.5,0,r'$\vec{\Omega}$',ha='center', va='center', fontsize='20')
ax2.set_title(r'Case $2$:')
plt.show()