OBE3levels_functions.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 15 18:11:00 2020

@author: oem
"""
import numpy as np
import matplotlib.pyplot as plt
from odeintw import odeintw
import time
#matplotlib.use('Qt5Agg')


"""
odeintw es un wrapper de la funcion odeint de scipy.integrate que sirve para integrar
sistemas de ecuaciones con parámetros complejos.
En este script se resuelven las ecuaciones
de Bloch de un sistema de 3 niveles en configuración lambda (1 y 3 son ground, 2 es excitado)
para ion de Bario o de Calcio.
Este script considera que hay dos láseres repump prendidos con sus respectivas frecuencias de
rabi y detuning.
Además, se usan dos distintos sistemas de referencia rotantes, condensados en las funciones
RepumpReferenceSystemModel 1 y 2. La 2 es más útil para prender de a poco el láser repump2, y
cuando su frec. de rabi es cero, no aparece en las ecuaciones.
https://github.com/WarrenWeckesser/odeintw


"""


#CHEQUEADOS LOS DAMPINGS


def EquationSystemModel(z, t, rabi12, rabi23, g21, g23, glg, glr, detr, detg, no_IR_lasers):
    """
    Se plantean acá las ecuaciones dinámicas de la matriz densidad (OBE: Optical Bloch Equations)
    para la interacción entre un átomo y un campo electromagnético.
        -rabi12, rabi23: frecuencias de rabi, proporcionales a la intensidad de los láseres
        -g21, g23: anchos de línea naturales de las transiciones
        -glg, glr: anchos de línea de los láseres doppler (UV) y repump (IR) respectivamente
        -detg, detr: detunings de transición doppler (UV) y repump (IR) respectivamente
        -no_IR_lasers: número de láseres IR aplicados
    """

    rho33, rho22, rho11, rho12, rho13, rho23 = z

    rho21 = np.conjugate(rho12)
    rho31 = np.conjugate(rho13)
    rho32 = np.conjugate(rho23)

    drho11dt = (1j)*0.5*rabi12*(rho12 - rho21) + rho22*g21
    drho22dt = (-1j)*0.5*rabi12*(rho12 - rho21) - 1j*0.5*(np.conjugate(rabi23)*rho32 - rabi23*rho23) - rho22*(g21 + g23)
    drho33dt = (-1j)*0.5*(rabi23*rho23 - np.conjugate(rabi23)*rho32) + g23*rho22

    drho12dt = (-1j)*0.5*rabi12*(rho22 - rho11) - (1j)*(detg*rho12 - 0.5*rabi23*rho13) - rho12*(0.5*g21 + 0.5*g23 + glg)
    drho13dt = (-1j)*(detg*rho13 + 0.5*rabi12*rho23 - 0.5*np.conjugate(rabi23)*rho12 - detr*rho13) - rho13*(glg + no_IR_lasers*glr)
    drho23dt = (-1j)*0.5*rabi12*rho13 - (1j)*0.5*np.conjugate(rabi23)*(rho33-rho22) + (1j)*detr*rho23 - rho23*(0.5*g21 + 0.5*g23 + no_IR_lasers*glr)

    dzdt = [drho33dt, drho22dt, drho11dt, drho12dt, drho13dt, drho23dt]

    return dzdt


def RepumpReferenceSystemModel(rabi23repump1, rabi23repump2, Detuningrepump1, Detuningrepump2, t):
    """
    Cuando hay más de un láser en una transición se producen batidos entre los mismos, haciendo que
    la frecuencia de rabi sea un parámetro que varía en el tiempo. Para resolver eso numéricamente se
    implementa de esta manera, calculando la frecuencia de rabi efectiva antes de resolver las OBE.
    Al resolver el sistema, se pasa a un marco de referencia rotante respecto a los láseres. Como los dos láseres
    IR excitan la misma transición, hay que elegir una de ambas frecuencias para rotar. Esto elige al repump1
    como tal referencia, pero es posible elegir a repump2, o incluso a un promedio entre ambas frecuencias (por eso se llama detuningmedio).
    Es decir, esto considera a la matriz de rotación como:
        U = exp(-i.wg.t)|1><1| + |2><2| + exp(-i.wr1.t)|3><3|
    """

    rabi23 = [rabi23repump1 + rabi23repump2*np.exp(-1*(Detuningrepump1-Detuningrepump2)*1j*tiemp) for tiemp in t]
    DetuningMedioRepump = Detuningrepump1
    return rabi23, DetuningMedioRepump


def GetTransitionLinewidths(species):
    """
    Devuelve los anchos de línea naturales según la especie iónica
    """


    if species == 'Calcium':
        g21, g23, = 2*np.pi*21.58e6, 2*np.pi*1.35e6 #anchos de linea de las transiciones
    elif species == 'Barium':
        g21, g23, = 2*np.pi*15.1e6, 2*np.pi*5.3e6 #anchos de linea de las transiciones
    else:
        print('Species not implemented')
        raise

    return g21, g23


def Solve3LevelBloch(EquationSystemModel, species, SatParDoppler, SatParRep1, SatParRep2, DetDoppler, DetRep1, DetRep2,
                     glg, glr, no_IR_lasers, initcond=1, tfin=10, paso=1e-3, printprogress=False):
    """
    Resuelve las ecuaciones de bloch del modelo dado. Los inputs son:
        glg, glr: anchos de línea de doppler y repump en MHz
        SatPar: parámetro de saturación de cada láser
        Det: detunings en MHz
        rabi12, rabi23 = sr*g21, sp*g23 #frecuencias de rabi
    """

    t = np.arange(0, tfin*1e-6, paso*1e-6)

    g21, g23 = GetTransitionLinewidths(species)


    #convierto todo a MHz
    DopplerLaserLinewidth = glg*2*np.pi*1e6
    RepumpLaserLinewidth = glr*2*np.pi*1e6

    DetuningDoppler = DetDoppler*2*np.pi*1e6
    Detuningrepump1 = DetRep1*2*np.pi*1e6
    Detuningrepump2 = DetRep2*2*np.pi*1e6

    rabi12 = np.sqrt(SatParDoppler*(4*(DetuningDoppler**2) + g21**2)/2)
    rabi23repump1 = np.sqrt(SatParRep1*(4*(Detuningrepump1**2) + g23**2)/2)
    rabi23repump2 = np.sqrt(SatParRep2*(4*(Detuningrepump2**2) + g23**2)/2)

    rabi23vst, _ = RepumpReferenceSystemModel(rabi23repump1, rabi23repump2, Detuningrepump1, Detuningrepump2, t)

    rho33sol = []
    rho22sol = []
    rho11sol = []

    rho12sol = []
    rho13sol = []
    rho23sol = []


    #condicion inicial: poblacion en el nivel 1
    if initcond==1:
        z0 = [0+0j, 0+0j, 1+0j, 0+0j, 0+0j, 0+0j] #formato: [rho33, rho22, rho11, rho12, rho13, rho23]
    elif initcond==2:
        z0 = [0+0j, 1+0j, 0+0j, 0+0j, 0+0j, 0+0j]
    elif initcond==3:
        z0 = [1+0j, 0+0j, 0+0j, 0+0j, 0+0j, 0+0j]
    else:
        print('Mala condicion inicial')
        raise

    #z0 = [0.5+0j, 0+0j, 0.5+0j, -0.5+0j, 0+0j, 0+0j] DARK STATE

    for i in range(1,len(t)):
        # span for next time step
        tspan = [t[i-1],t[i]]
        # solve for next step
        #z = odeintw(EquationSystemModel, z0, tspan, args=(rabi12, rabi23vst[i], g21, g23, DopplerLaserLinewidth, RepumpLaserLinewidth, DetuningMedioRepump, DetuningDoppler), full_output=True)
        z = odeintw(EquationSystemModel, z0, tspan, args=(rabi12, rabi23repump1, g21, g23, DopplerLaserLinewidth, RepumpLaserLinewidth, Detuningrepump1, DetuningDoppler, no_IR_lasers), full_output=True)
        # store solution for plotting
        rho22sol.append(z[0][1][1])
        rho33sol.append(z[0][1][0])
        rho11sol.append(z[0][1][2])

        rho12sol.append(z[0][1][3])
        rho13sol.append(z[0][1][4])
        rho23sol.append(z[0][1][5])

        # next initial condition
        z0 = z[0][1]
        if printprogress:
            print(i, '/', len(t), end="\r", flush=True)

    #Tomo la parte real de los diagonales y el módulo de las coherencias
    RealRho22 = [r.real for r in rho22sol]
    RealRho11 = [r.real for r in rho11sol]
    RealRho33 = [r.real for r in rho33sol]
    AbsRho12 = [np.abs(r) for r in rho12sol]
    AbsRho13 = [np.abs(r) for r in rho13sol]
    AbsRho23 = [np.abs(r) for r in rho23sol]

    return t, RealRho11, RealRho22, RealRho33, AbsRho12, AbsRho13, AbsRho23



def RepumpSpectrum(EquationSystemModel, species, SatParDoppler, SatParRep1, SatParRep2, DetDoppler, DetRep2,
                     glg, glr, tfinOBE=10, pasoOBE=1e-3, frecini=-100, frecfin=100, pasofrec=1, plotSpectrum=True, sweepingrepump=True):
    """
    Hace un barrido en la frecuencia del repump.
    Considerando que puede haber oscilaciones de la población por haber dos repumps,
    """

    g21, g23 = GetTransitionLinewidths(species)

    FluorescenceVector = []

    t = np.arange(0, tfinOBE*1e-6, pasoOBE*1e-6)

    DopplerLaserLinewidth = glg*2*np.pi*1e6
    RepumpLaserLinewidth = glr*2*np.pi*1e6

    DetuningDoppler = DetDoppler*2*np.pi*1e6
    Detuningrepump2 = DetRep2*2*np.pi*1e6

    rabi12 = np.sqrt(SatParDoppler*(4*(DetuningDoppler**2) + g21**2)/2)
    rabi23repump1 = np.sqrt(SatParRep1*(4*(Detuningrepump1**2) + g23**2)/2)
    rabi23repump2 = np.sqrt(SatParRep2*(4*(Detuningrepump2**2) + g23**2)/2)

    Repump1DetuningVector = np.arange(frecini*2*np.pi*1e6, frecfin*2*np.pi*1e6, pasofrec*2*np.pi*1e6)

    counter = 1

    FluorescenceVector = []

    for DetRep1 in Repump1DetuningVector:
        tini = time.time()

        rho33sol = []
        rho22sol = []
        rho11sol = []

        rho12sol = []
        rho13sol = []
        rho23sol = []

        rabi23vst, DetuningMedioRepump = RepumpReferenceSystemModel(rabi23repump1, rabi23repump2, DetRep1, Detuningrepump2, t)

        z0 = [0+0j, 0+0j, 1+0j, 0+0j, 0+0j, 0+0j] #formato: [rho33, rho22, rho11, rho12, rho13, rho23]


        for i in range(1,len(t)):
            # span for next time step
            tspan = [t[i-1],t[i]]
            # solve for next step
            z = odeintw(EquationSystemModel,z0,tspan, args=(rabi12, rabi23vst[i], g21, g23, DopplerLaserLinewidth, RepumpLaserLinewidth, DetuningMedioRepump, DetuningDoppler), full_output=True)
            # store solution for plotting

            rho22sol.append(z[0][1][1])
            rho33sol.append(z[0][1][0])
            rho11sol.append(z[0][1][2])

            rho12sol.append(z[0][1][3])
            rho13sol.append(z[0][1][4])
            rho23sol.append(z[0][1][5])

            # next initial condition
            z0 = z[0][1]
    #        print(i, '/', len(t))
        print

        FluorescenceVector.append(np.real(rho12sol[-1]))

        tfin = time.time()
        print(counter, '/', len(Repump1DetuningVector), ', elapsed time: ', round(tfin-tini, 2), ' s')
        counter = counter + 1

    if plotSpectrum:
        plt.plot([1e-6*d/(2*np.pi) for d in Repump1DetuningVector], [100*f for f in FluorescenceVector], 'o')
        plt.xlabel('Repump detuning (MHz')
        plt.ylabel('Fluorescence (A.U.)')
        plt.title('Doppler detuning: ' + str(DetDoppler) + ' MHz, Repump 2 detuning: ' + str(DetRep2) + ' MHz')

    return Repump1DetuningVector, FluorescenceVector


#Funciones auxiliares para el análisis
def FindFluoAtTime(fluo, time, timetarget, pasoti):
    """
    Encuentra la fluorescencia para determinado tiempo (en us).
    El pasoti es la tolerancia temporal para el cual busca esa fluorescencia
    """
    ti = timetarget*1e-6
    pasoti = pasoti*1e-6
    i = 0
    fluoattivec = []
    while i < len(fluo):
        if abs(time[i]-ti) < 2*pasoti:
            fluoattivec.append(fluo[i])
        i = i + 1
    return np.mean(fluoattivec)

def FindTimeForFluo(fluo, time, fluotarget, tolerancefluo=1e-3):
    i = 0
    timesforfluo = []
    while i < len(time)-1:
        if np.abs(fluo[i]-fluotarget) < 2*tolerancefluo:
            timesforfluo.append(time[i])
        i = i + 1
    if len(timesforfluo) == 0:
        print('La fluorescencia no toma ese valor')
        return 0
    elif len(timesforfluo) == 1:
        return timesforfluo[0]
    else:
#        return np.mean(timesforfluo[1:])
        return timesforfluo[-1]

def FindCharacteristicDecayTime(fluo, time):
    pass


def IntegrateWindow(fluo, time, ti, window):
    ti = ti*1e-6
    window = window*1e-6
    i = 0
    fluowindow = []
    while i<len(fluo):
        if time[i]>ti and time[i]<=ti+window:
            fluowindow.append(fluo[i])
        i = i + 1
        if time[i]>ti+window:
            return np.mean(fluowindow)


def Lorentz(x, x0, gamma, ymax):
    scale = ymax*np.pi*gamma*0.5
    return (scale/np.pi)*(0.5*gamma)/((x-x0)**2 + (0.5*gamma)**2)

def rho23model(rabip, gamma, gammas, detuningp, rho0, t):
    rho23 = ((1j)*rho0*rabip/(2*gamma))*(np.exp(-gamma*t) - np.exp(-(gammas+(rabip**2)/(4*gamma))*t -(1j)*detuningp*t))
    return [np.real(rho23), np.abs(np.imag(rho23))]


def dopplerBroadening(T, wldopp=0.397e-6, wlrep=0.866e-6, alpha=0, mcalcio = 6.655e-23*1e-3):
    """
    Calcula el broadening extra semiclásico por temperatura considerando que el ion atrapado se mueve.
    wlr es la longitud de onda doppler, wlp la longitud de onda repump, T la temperatura del ion en kelvin, y alpha (en rads) el ángulo
    que forman ambos láseres.
    """

    kboltzmann = 1.38e-23 #J/K

    gammaD = (2*np.pi)*np.sqrt((1/(wldopp*wldopp)) + (1/(wlrep*wlrep)) - 2*(1/(wldopp*wlrep))*np.cos(alpha))*np.sqrt(kboltzmann*T/(2*mcalcio))

    return gammaD