atomo_2_niveles.py

from scipy.integrate import ode
import numpy as np
import pylab as pl

# Parámetros físicos
omega = np.pi

# Parámetros de simulación
t = np.linspace(0,4,101)

# La solución se guarda en el arreglo y que es de 2xlen(t).
# El 2 es necesario porque nuestro vector de estado tiene 2 componentes 
y = np.zeros((2,len(t)),dtype=np.complex)

# Condiciones iniciales [c_g, c_e]
y[:,0] = [1,0]

# Definir lado derecho de ecuación de la forma c'(t) = rhs(t,c)
# En este caso c es un vector con componentes [c_g, c_e]
def rhs(t,c):
    return [-0.5j*omega*c[1],-0.5j*omega*c[0]]

# Definir solucionador de ecuación diferencial, el parámetro 'zvode' 
# es para poder usar valores complejos
r = ode(rhs).set_integrator('zvode')

# Definir condiciones inciales
r.set_initial_value(y[:,0], t[0])

# Integrar la ecuación
i = 0
while r.successful() and i < len(t)-1:
    r.integrate(t[i+1])
    y[:,i+1] = r.y
    i = i+1


# Graficar la solución
pl.plot(t,np.abs(y[0,:])**2,label="$P_g$")
pl.plot(t,np.abs(y[1,:])**2,label="$P_e$")
pl.legend(loc=0)
pl.xlabel("A")
pl.ylabel("A")

pl.savefig("atomo_2_niveles.png",dpi=100)
pl.show()