CPT_plotter_20231226.py

import h5py
import matplotlib.pyplot as plt
import numpy as np
import sys
import re
import ast
from scipy.optimize import curve_fit
import os
from scipy import interpolate

"""
Primero tengo mediciones de espectros cpt de un ion variando la tension dc_A
"""

#C:\Users\Usuario\Documents\artiq\artiq_experiments\analisis\plots\20220106_CPT_DosLaseres_v08_TISA_DR\Data

os.chdir('/home/nico/Documents/artiq_experiments/analisis/plots/20231226_CPTconmicromocion4/Data/')

CPT_FILES = """000016531-IR_Scan_withcal_optimized
000016532-IR_Scan_withcal_optimized
"""


def SeeKeys(files):
    for i, fname in enumerate(files.split()):
        data = h5py.File(fname+'.h5', 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'
        print(fname)
        print(list(data['datasets'].keys()))

print(SeeKeys(CPT_FILES))

#carpeta pc nico labo escritorio:
#C:\Users\Usuario\Documents\artiq\artiq_experiments\analisis\plots\20211101_CPT_DosLaseres_v03\Data

Counts = []
Freqs = []

AmpTisa = []
UVCPTAmp = []
No_measures = []
Voltages = []

for i, fname in enumerate(CPT_FILES.split()):
    print(str(i) + ' - ' + fname)
    #print(fname)
    data = h5py.File(fname+'.h5', 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'

    # Aca hago algo repugnante para poder levantar los strings que dejamos
    # que además tenian un error de tipeo al final. Esto no deberá ser necesario
    # cuando se solucione el error este del guardado.
    Freqs.append(np.array(data['datasets']['IR1_Frequencies']))
    Counts.append(np.array(data['datasets']['data_array']))
    #AmpTisa.append(np.array(data['datasets']['TISA_CPT_amp']))
    UVCPTAmp.append(np.array(data['datasets']['UV_CPT_amp']))
    No_measures.append(np.array(data['datasets']['no_measures']))
    Voltages.append(np.array(data['datasets']['scanning_voltages']))

def Split(array,n):
    length=len(array)/n
    splitlist = []
    jj = 0
    while jj<length:
        partial = []
        ii = 0
        while ii < n:
            partial.append(array[jj*n+ii])
            ii = ii + 1
        splitlist.append(partial)
        jj = jj + 1
    return splitlist


CountsSplit = []
CountsSplit.append(Split(Counts[0],len(Freqs[0])))
CountsSplit.append(Split(Counts[1],len(Freqs[1])))

#%%

"""
Ploteo la cpt de referencia / plotting the reference CPT
"""

medic = 1 #puede ser 0 o 1

jvec = [28] # de la 1 a la 9 vale la pena, despues no

drs = [390.5, 399.5, 406, 413.5]

drive=22.1

Frequencies = Freqs[medic]

plt.figure()
i = 0
for j in jvec:
    plt.errorbar([2*f*1e-6 for f in Frequencies], CountsSplit[medic][j], yerr=np.sqrt(CountsSplit[0][j]), fmt='o', capsize=2, markersize=2)
    i = i + 1
plt.xlabel('Frecuencia (MHz)')
plt.ylabel('counts')
plt.grid()
#for dr in drs:
#    plt.axvline(dr)
    #plt.axvline(dr+drive)
plt.legend()

#%%
"""
SUPERAJUSTE: TESTEO FITEOS
"""


from EITfit.lolo_modelo_full_8niveles import PerformExperiment_8levels_MM
from scipy.optimize import curve_fit
import time

"""
SUPER AJUSTE (SA)
"""

phidoppler, titadoppler = 0, 90
phirepump, titarepump = 0,  0
phiprobe = 0
titaprobe = 90

Temp = 0.5e-3

sg = 0.544
sp = 4.5
sr = 0
DetRepump = 0


lw = 0.1
DopplerLaserLinewidth, RepumpLaserLinewidth, ProbeLaserLinewidth = lw, lw, lw #ancho de linea de los laseres


u = 32.5e6

#B = (u/(2*np.pi))/c

# correccion = 13

#DetDoppler = -11.5-correccion

gPS, gPD, = 2*np.pi*21.58e6, 2*np.pi*1.35e6
alpha = 0


drivefreq = 2*np.pi*22.135*1e6

measurement = 1
SelectedCurveVec = [30]


for selectedcurve in SelectedCurveVec:

    FreqsDR = Freqs[measurement]
    CountsDR = CountsSplit[measurement][selectedcurve]

    freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))
    CircPr = 1
    alpha = 0


    def FitEIT_MM_single(Freqs, offset, DetDoppler, SG, SP, SCALE1, OFFSET, BETA1, TEMP, U, DL, PL, plot=False):
    #def FitEIT_MM(freqs, SG, SP, SCALE1, OFFSET, BETA1):
        #BETA = 1.8
        # SG = 0.6
        # SP = 8.1
        # TEMP = 0.2e-3

        freqs = [2*f*1e-6-offset for f in Freqs]

        #Detunings, Fluorescence1 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, U, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA1, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)
        Detunings, Fluorescence1 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, U, DL, PL, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA1, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)

        ScaledFluo1 = np.array([f*SCALE1 + OFFSET for f in Fluorescence1])
        if plot:
            return ScaledFluo1, Detunings
        else:
            return ScaledFluo1
        #return ScaledFluo1

    do_fit = True
    if do_fit:
        t1 = time.time()
        print(f'Beginning {selectedcurve}')
        popt_3_SA, pcov_3_SA = curve_fit(FitEIT_MM_single, FreqsDR, CountsDR, p0=[446, -36, 0.67, 8.5, 6e4, 250, 1, 0.3e-3, 32e6, 0.1, 0.15], bounds=((0, -50, 0, 0, 0, 0, 0, 0, 31e6, 0, 0), (1000, -20, 2, 20, 5e5, 5e4, 10, 100e-3, 34e6, 10, 10)))
        print(f'Ended {selectedcurve}, time elapsed: {round(time.time()-t1)} s')

        FittedEITpi_3_SA_short, Detunings_3_SA_short = FitEIT_MM_single(FreqsDR, *popt_3_SA, plot=True)
        freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))
        FittedEITpi_3_SA_long, Detunings_3_SA_long = FitEIT_MM_single(freqslong, *popt_3_SA, plot=True)


    plt.figure()
    plt.errorbar(Detunings_3_SA_short, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
    plt.plot(Detunings_3_SA_long, FittedEITpi_3_SA_long, color='darkolivegreen', linewidth=3, label=f'med {selectedcurve}')
    #plt.title(f'Sdop: {round(popt[0], 2)}, Spr: {round(popt[1], 2)}, T: {round(popt[2]*1e3, 2)} mK, detDop: {DetDoppler} MHz')
    plt.xlabel('Detuning (MHz)')
    plt.ylabel('Counts')
    plt.legend(loc='upper left', fontsize=20)
    plt.grid()

    print(popt_3_SA)

    print(f'Beta: {popt_3_SA[6]} +- {np.sqrt(pcov_3_SA[6,6])}')

#%%
"""
AHORA INTENTO SUPER AJUSTES O SEA CON OFFSETXPI Y DETDOPPLER INCLUIDOS
"""


from EITfit.lolo_modelo_full_8niveles import PerformExperiment_8levels_MM
from scipy.optimize import curve_fit
import time

"""
SUPER AJUSTE (SA)
"""

phidoppler, titadoppler = 0, 90
phirepump, titarepump = 0,  0
phiprobe = 0
titaprobe = 90

Temp = 0.5e-3

sg = 0.544
sp = 4.5
sr = 0
DetRepump = 0


lw = 0.1
DopplerLaserLinewidth, RepumpLaserLinewidth, ProbeLaserLinewidth = lw, lw, lw #ancho de linea de los laseres


u = 32.5e6

#B = (u/(2*np.pi))/c

correccion = 13

#DetDoppler = -11.5-correccion

gPS, gPD, = 2*np.pi*21.58e6, 2*np.pi*1.35e6
alpha = 0


drivefreq = 2*np.pi*22.135*1e6

measurement = 1
SelectedCurveVec = np.arange(0,31,1)
#SelectedCurveVec = [17]

popt_SA_vec = []
pcov_SA_vec = []
Detuningsshort_vec = []
Counts_vec = []
Detuningslong_vec = []
FittedCounts_vec = []

Betas_vec = []
ErrorBetas_vec = []
Temp_vec = []
ErrorTemp_vec = []

DetuningsUV_vec = []
ErrorDetuningsUV_vec = []

for selectedcurve in SelectedCurveVec:

    FreqsDR = Freqs[measurement]
    CountsDR = CountsSplit[measurement][selectedcurve]

    freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))
    CircPr = 1
    alpha = 0


    def FitEIT_MM_single(Freqs, offset, DetDoppler, SG, SP, SCALE1, OFFSET, BETA1, TEMP, U, plot=False):
    #def FitEIT_MM(freqs, SG, SP, SCALE1, OFFSET, BETA1):
        #BETA = 1.8
        # SG = 0.6
        # SP = 8.1
        # TEMP = 0.2e-3

        freqs = [2*f*1e-6-offset for f in Freqs]

        Detunings, Fluorescence1 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, U, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA1, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)

        ScaledFluo1 = np.array([f*SCALE1 + OFFSET for f in Fluorescence1])
        if plot:
            return ScaledFluo1, Detunings
        else:
            return ScaledFluo1
        #return ScaledFluo1

    do_fit = True
    if do_fit:
        t1 = time.time()
        print(f'Beginning {selectedcurve}')
        popt_3_SA, pcov_3_SA = curve_fit(FitEIT_MM_single, FreqsDR, CountsDR, p0=[443, -32, 0.67, 8, 4e4, 350, 1, 0.5e-3, 32e6], bounds=((0, -50, 0, 0, 0, 0, 0, 0, 31e6), (1000, -20, 2, 20, 5e5, 5e4, 10, 100e-3, 34e6)))
        print(f'Ended {selectedcurve}, time elapsed: {round(time.time()-t1)} s')

        popt_SA_vec.append(popt_3_SA)
        pcov_SA_vec.append(pcov_3_SA)

        FittedEITpi_3_SA_short, Detunings_3_SA_short = FitEIT_MM_single(FreqsDR, *popt_3_SA, plot=True)
        freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))
        FittedEITpi_3_SA_long, Detunings_3_SA_long = FitEIT_MM_single(freqslong, *popt_3_SA, plot=True)

    DetuningsUV_vec.append(popt_3_SA[1])
    ErrorDetuningsUV_vec.append(np.sqrt(pcov_3_SA[1,1]))

    Betas_vec.append(popt_3_SA[6])
    ErrorBetas_vec.append(np.sqrt(pcov_3_SA[6,6]))
    Temp_vec.append(popt_3_SA[7])
    ErrorTemp_vec.append(np.sqrt(pcov_3_SA[7,7]))

    Detuningsshort_vec.append(Detunings_3_SA_short)
    Counts_vec.append(CountsDR)
    Detuningslong_vec.append(Detunings_3_SA_long)
    FittedCounts_vec.append(FittedEITpi_3_SA_long)


    plt.figure()
    plt.errorbar(Detunings_3_SA_short, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
    plt.plot(Detunings_3_SA_long, FittedEITpi_3_SA_long, color='darkolivegreen', linewidth=3, label=f'med {selectedcurve}')
    #plt.title(f'Sdop: {round(popt[0], 2)}, Spr: {round(popt[1], 2)}, T: {round(popt[2]*1e3, 2)} mK, detDop: {DetDoppler} MHz')
    plt.xlabel('Detuning (MHz)')
    plt.ylabel('Counts')
    plt.legend(loc='upper left', fontsize=20)
    plt.grid()

    print(f'listo med {selectedcurve}')
    print(popt_3_SA)


#%%
"""
Grafico distintas variables que salieron del SUper ajuste
"""

import seaborn as sns
paleta = sns.color_palette("rocket")

medfin = 19


voltages_dcA = Voltages[0][0:medfin]

def lineal(x,a,b):
    return a*x+b

def hiperbola(x,a,b,c,x0):
    return a*np.sqrt(((x-x0)**2+c**2))+b

hiperbola_or_linear = True

if hiperbola_or_linear:
    #popthip,pcovhip = curve_fit(hiperbola,voltages_dcA[0:20],Betas_vec[0:20],p0=(100,0.1,1,-0.15))
    
    #xhip = np.linspace(-0.23,0.055,200)
    
    popthip,pcovhip = curve_fit(hiperbola,voltages_dcA,Betas_vec[:medfin],p0=(100,0.1,1,-0.15))

    xhip = np.linspace(-0.16,0.05,200)

    plt.figure()
    plt.errorbar(voltages_dcA,Betas_vec[0:medfin],yerr=ErrorBetas_vec[:medfin],fmt='o',capsize=5,markersize=5,color=paleta[1])
    plt.plot(xhip,hiperbola(xhip,*popthip))
    plt.xlabel('Endcap voltage (V)')
    plt.ylabel('Modulation factor')
    plt.ylim(-1,4)
    plt.grid()

else:
    poptini,pcovini = curve_fit(lineal,voltages_dcA[0:3],Betas_vec[0:3])
    poptfin,pcovfin = curve_fit(lineal,voltages_dcA[4:],Betas_vec[4:])

    minimum_voltage = -(poptini[1]-poptfin[1])/(poptini[0]-poptfin[0]) #voltaje donde se intersectan las rectas, es decir, donde deberia estar el minimo de micromocion
    minimum_modulationfactor = lineal(minimum_voltage,*poptini) #es lo mismo si pongo *poptfin

    xini = np.linspace(-0.23,-0.13,100)
    xfin = np.linspace(-0.15,0.005,100)

    plt.figure()
    plt.errorbar(voltages_dcA,Betas_vec,yerr=ErrorBetas_vec,fmt='o',capsize=5,markersize=5,color=paleta[1])
    plt.plot(xini,lineal(xini,*poptini))
    plt.plot(xfin,lineal(xfin,*poptfin))
    plt.y
    plt.axvline(minimum_voltage,linestyle='dashed',color='grey')
    plt.xlabel('Endcap voltage (V)')
    plt.ylabel('Modulation factor')
    plt.grid()


print([t*1e3 for t in Temp_vec])

plt.figure()
plt.errorbar(voltages_dcA,[t*1e3 for t in Temp_vec[:medfin]],yerr=[t*1e3 for t in ErrorTemp_vec[:medfin]],fmt='o',capsize=5,markersize=5,color=paleta[3])
#plt.axvline(minimum_voltage,linestyle='dashed',color='grey')
plt.axhline(0.538)
#plt.yscale('log')
plt.xlabel('Endcap voltage (V)')
plt.ylabel('Temperature (mK)')
plt.grid()
#plt.ylim(0,2)

#%%
"""
Ahora hago un ajuste con una hiperbola porque tiene mas sentido, por el hecho
de que en el punto optimo el ion no esta en el centro de la trampa
sino que esta a una distancia d
"""
def hiperbola(x,a,b,c,x0):
    return a*np.sqrt(((x-x0)**2+c**2))+b

medfin = 11
voltages_dcA = Voltages[0][1:medfin]


popthip,pcovhip = curve_fit(hiperbola,voltages_dcA[:10],Betas_vec[:10],p0=(100,0.1,1,-0.15))

xhip = np.linspace(-0.23,0.055,200)

plt.figure()
plt.errorbar(voltages_dcA,Betas_vec[:medfin-1],yerr=ErrorBetas_vec[:medfin-1],fmt='o',capsize=5,markersize=5,color=paleta[1])
plt.plot(xhip,hiperbola(xhip,*popthip))
plt.xlabel('Endcap voltage (V)')
plt.ylabel('Modulation factor')
#plt.yscale('log')
plt.grid()




#%%
from scipy.special import jv


def expo(x,tau,A,B):
    return A*np.exp(x/tau)+B

def cuadratica(x,a,c):
    return a*(x**2)+c

def InverseMicromotionSpectra(beta, A, det, x0, gamma, B):
    ftrap=22.1
    #gamma=30
    P = ((jv(0, beta)**2)/((((det-x0)**2)+(0.5*gamma)**2)**2))*(-2*(det-x0))
    i = 1
    #print(P)
    while i <= 5:
        P = P + (-2*(det-x0))*((jv(i, beta))**2)/(((((det-x0)+i*ftrap)**2)+(0.5*gamma)**2)**2) + (-2*(det-x0))*(((jv(-i, beta))**2)/((((det-x0)-i*ftrap)**2)+(0.5*gamma)**2)**2)
        i = i + 1
        #print(P)
    #return 1/(A*P+B)
    return 1/(A*P+B)


def MicromotionSpectra(beta,det, gamma):
    ftrap=22.1
    #gamma=23
    P = (jv(0, beta)**2)/(((det)**2)+(0.5*gamma)**2)
    i = 1
    #print(P)
    while i <= 5:
        P = P + ((jv(i, beta))**2)/((((det)+i*ftrap)**2)+(0.5*gamma)**2) + ((jv(-i, beta))**2)/((((det)-i*ftrap)**2)+(0.5*gamma)**2) 
        i = i + 1
        #print(P)
    return P


def polynomial(x,a,b,c,d,e):
    b=0
    d=0
    return a+b*x+c*x*x+d*x*x*x+e*x*x*x*x

def InverseDerivMicromotionSpectra(beta, det, gamma):
    ftrap=22.1
    #gamma=23
    #det = -gamma/2
    P = ((jv(0, beta)**2)/((((det)**2)+(0.5*gamma)**2)**2))*(-2*(det))
    i = 1
    #print(P)
    while i <= 5:
        P = P + (-2*(det))*((jv(i, beta))**2)/(((((det)+i*ftrap)**2)+(0.5*gamma)**2)**2) + (-2*(det))*(((jv(-i, beta))**2)/((((det)-i*ftrap)**2)+(0.5*gamma)**2)**2)
        i = i + 1
        #print(P)
    return 1/P


def FinalTemp(beta,det, C,D):
    gamma = 21
    #det=-11
    #D=-0.8
    #C = 1.68656122e-03  
    #D = 6.64227010e-02
    #C=0
    
    #print(MicromotionSpectra(beta,det,gamma))
    return (C*MicromotionSpectra(beta,det,gamma)+D*beta**2)*InverseDerivMicromotionSpectra(beta, det, gamma)
    #return (C*MicromotionSpectra(beta,det,gamma))*InverseDerivMicromotionSpectra(beta, det, gamma)

"""
Temperatura vs beta con un ajuste exponencial
"""

medfin = 10

popt_exp, pcov_exp = curve_fit(expo,Betas_vec[:medfin],[t*1e3 for t in Temp_vec[:medfin]])
#popt_quad, pcov_quad = curve_fit(cuadratica,Betas_vec[:11],[t*1e3 for t in Temp_vec[:11]],p0=(1,10))
#popt_rho22, pcov_rho22 = curve_fit(InverseMicromotionSpectra,Betas_vec,[t*1e3 for t in Temp_vec],p0=(10,10,-10,1,20)) #esto ajusta muy bien
#popt_rho22, pcov_rho22 = curve_fit(InverseMicromotionSpectra,Betas_vec, [t*1e3 for t in Temp_vec],p0=(-10,-10,10,1,20)) #esto ajusta muy bien
#popt_rho22_raw, pcov_rho22_raw = curve_fit(InverseMicromotionSpectra_raw,Betas_vec[:7], [t*1e3 for t in Temp_vec[:7]],p0=(-0.1, -10, 1)) #esto ajusta muy bien
popt_rho22_balance, pcov_rho22_balance = curve_fit(FinalTemp,Betas_vec[:medfin], [t*1e3 for t in Temp_vec[:medfin]],p0=(-10, 10,1)) #esto ajusta muy bien
popt_rho22_poly, pcov_rho22_poly = curve_fit(polynomial,Betas_vec[:medfin], [t*1e3 for t in Temp_vec[:medfin]],p0=(1,2,3,4,10)) #esto ajusta muy bien


print(popt_rho22_balance)

betaslong = np.arange(0,2.8,0.01)

print(f'Min temp predicted: {FinalTemp(betaslong,*popt_rho22_balance)[0]}')

print(f'Detuning: {popt_rho22_balance[0]} MHz')
print(f'rho22 coeff: {popt_rho22_balance[1]}')
print(f'betasquared coeff: {popt_rho22_balance[2]}')

print(f'cociente de los coeff: {popt_rho22_balance[2]/popt_rho22_balance[1]}')


print(f'params: {popt_rho22_balance}')
print(f'errores: {np.sqrt(np.diag(pcov_rho22_balance))}')

k_plot = medfin

plt.figure()
plt.errorbar(Betas_vec[:k_plot],[t*1e3 for t in Temp_vec[:k_plot]],xerr=ErrorBetas_vec[:k_plot], yerr=[t*1e3 for t in ErrorTemp_vec[:k_plot]],fmt='o',capsize=5,markersize=5,color=paleta[3])
#plt.plot(betaslong,polynomial(betaslong,*popt_rho22_poly),label='Ajuste exponencial')
#plt.plot(betaslong,FinalTemp(betaslong,popt_rho22_balance[0],popt_rho22_balance[1],popt_rho22_balance[2]*1),label='Ajuste con espectro modulado')
plt.xlim(-0,3)
plt.ylim(0,1)

#plt.axvline(minimum_voltage,linestyle='dashed',color='grey')
#plt.axhline(0.538)
plt.xlabel('Modulation factor')
plt.ylabel('Temperature (mK)')
plt.legend()
plt.grid()

#%%

hbar=1.05e-34
gammita = 22.1e6
kx = 2*np.pi/(0.397e-6)
kb = 1.38e-23
masita = 6.6e-26

coeff = gammita*hbar*hbar*kx*kx/(6*kb*masita)
print(coeff)

rfheatrate = coeff*popt_rho22_balance[2]/popt_rho22_balance[1]

print(f'heating rate due to rf heating: {rfheatrate*1e3} mK/s')



#%%
"""
Esto no es del super ajuste sino de los ajustes anteriores en donde DetDoppler y offset son puestos a mano
Aca grafico los betas con su error en funcion de la tension variada.
Ademas, hago ajuste lineal para primeros y ultimos puntos, ya que espero que
si la tension hace que la posicion del ion varie linealmente, el beta varia proporcional a dicha posicion.
"""

import seaborn as sns

def lineal(x,a,b):
    return a*x+b

paleta = sns.color_palette("rocket")

betavector = [beta1,beta2,beta3,beta4,beta5,beta6,beta7,beta8,beta9]
errorbetavector = [errorbeta1,errorbeta2,errorbeta3,errorbeta4,errorbeta5,errorbeta6,errorbeta7,errorbeta8,errorbeta9]
voltages_dcA = Voltages[0][1:10]

poptini,pcovini = curve_fit(lineal,voltages_dcA[0:3],betavector[0:3])
poptfin,pcovfin = curve_fit(lineal,voltages_dcA[4:],betavector[4:])

minimum_voltage = -(poptini[1]-poptfin[1])/(poptini[0]-poptfin[0]) #voltaje donde se intersectan las rectas, es decir, donde deberia estar el minimo de micromocion
minimum_modulationfactor = lineal(minimum_voltage,*poptini) #es lo mismo si pongo *poptfin

xini = np.linspace(-0.23,-0.13,100)
xfin = np.linspace(-0.15,0.005,100)

plt.figure()
plt.errorbar(voltages_dcA,betavector,yerr=errorbetavector,fmt='o',capsize=5,markersize=5,color=paleta[1])
plt.plot(xini,lineal(xini,*poptini))
plt.plot(xfin,lineal(xfin,*poptfin))
plt.axvline(minimum_voltage,linestyle='dashed',color='grey')
plt.xlabel('Endcap voltage (V)')
plt.ylabel('Modulation factor')
plt.grid()


#%%
"""
Aca veo la temperatura del ion en funcion del voltaje del endcap, ya que
al cambiar la cantidad de micromocion, cambia la calidad del enfriado
"""
tempvector = np.array([temp1,temp2,temp3,temp4,temp5,temp6,temp7,temp8,temp9])*1e3
errortempvector = np.array([errortemp1,errortemp2,errortemp3,errortemp4,errortemp5,errortemp6,errortemp7,errortemp8,errortemp9])*1e3

voltages_dcA = Voltages[0][1:10]

plt.figure()
plt.errorbar(voltages_dcA,tempvector,yerr=errortempvector,fmt='o',capsize=5,markersize=5,color=paleta[3])
plt.axvline(minimum_voltage,linestyle='dashed',color='grey')
plt.xlabel('Endcap voltage (V)')
plt.ylabel('Temperature (mK)')
plt.grid()
plt.ylim(0,2)


#%%
"""
Por las dudas, temperatura en funcion de beta
"""
plt.figure()
plt.errorbar(betavector,tempvector,yerr=errortempvector,xerr=errorbetavector,fmt='o',capsize=5,markersize=5)
plt.xlabel('Modulation factor')
plt.ylabel('Temperature (mK)')
plt.grid()



#%%
"""
Si quiero ver algun parametro del ajuste puntual. el orden es: 0:SG, 1:SP, 2:SCALE1, 3:OFFSET
"""
ki=2
plt.errorbar(np.arange(0,9,1),[popt_1[ki],popt_2[ki],popt_3[ki],popt_4[ki],popt_5[ki],popt_6[ki],popt_7[ki],popt_8[ki],popt_9[ki]],yerr=[np.sqrt(pcov_1[ki,ki]),np.sqrt(pcov_2[ki,ki]),np.sqrt(pcov_3[ki,ki]),np.sqrt(pcov_4[ki,ki]),np.sqrt(pcov_5[ki,ki]),np.sqrt(pcov_6[ki,ki]),np.sqrt(pcov_7[ki,ki]),np.sqrt(pcov_8[ki,ki]),np.sqrt(pcov_9[ki,ki])], fmt='o',capsize=3,markersize=3)

#%%
"""
AHORA VAMOS A MEDICIONES CON MAS DE UN ION!!!
"""

"""
Ploteo la cpt de referencia / plotting the reference CPT

1: 2 iones, -100 mV dcA
2: 2 iones, -150 mV dcA
3: 2 iones, -50 mV dcA
4: 2 iones, 5 voltajes (el ion se va en la 4ta medicion y en la 5ta ni esta)
5, 6 y 7: 3 iones en donde el scaneo esta centrado en distintos puntos

"""

jvec = [3] # desde la 1, pero la 4 no porque es un merge de curvitas

plt.figure()
i = 0
for j in jvec:
    plt.errorbar([2*f*1e-6 for f in Freqs[j]], Counts[j], yerr=np.sqrt(Counts[j]), fmt='o', capsize=2, markersize=2)
    i = i + 1
plt.xlabel('Frecuencia (MHz)')
plt.ylabel('counts')
plt.grid()
#for dr in drs:
#    plt.axvline(dr)
    #plt.axvline(dr+drive)
plt.legend()


#%%
"""
Mergeo la 5, 6 y 7
"""

Freqs5 = [2*f*1e-6 for f in Freqs[5]]
Freqs6 = [2*f*1e-6 for f in Freqs[6]]
Freqs7 = [2*f*1e-6 for f in Freqs[7]]

Counts5 = Counts[5]
Counts6 = Counts[6]
Counts7 = Counts[7]

i_1_ini = 0
i_1 = 36
i_2_ini = 0
i_2 = 24

f_1 = 18
f_2 = 30


scale_1 = 0.92
scale_2 = 0.98

#Merged_freqs_test = [f-f_2 for f in Freqs6[i_2_ini:i_2]]+[f-f_1 for f in Freqs5[i_1_ini:i_1]]+Freqs7

#plt.plot(Merged_freqs_test,'o')



Merged_freqs = [f-f_2 for f in Freqs6[0:i_2]]+[f-f_1 for f in Freqs5[0:i_1]]+Freqs7
Merged_counts = [scale_2*c for c in Counts6[0:i_2]]+[scale_1*c for c in Counts5[0:i_1]]+list(Counts7)

Merged_freqs_rescaled = np.linspace(np.min(Merged_freqs),np.max(Merged_freqs),len(Merged_freqs))

#drs = [391.5, 399.5, 405.5, 414]
drs = [370,379,385,391.5]

plt.figure()
i = 0
for j in jvec:
    plt.plot([f-f_1 for f in Freqs5[0:i_1]], [scale_1*c for c in Counts5[0:i_1]],'o')
    plt.plot([f-f_2 for f in Freqs6[0:i_2]], [scale_2*c for c in Counts6[0:i_2]],'o')
    plt.plot(Freqs7, Counts7,'o')
    plt.errorbar(Merged_freqs, Merged_counts, yerr=np.sqrt(Merged_counts), fmt='o', capsize=2, markersize=2)


    i = i + 1
plt.xlabel('Frecuencia (MHz)')
plt.ylabel('counts')
plt.grid()
for dr in drs:
    plt.axvline(dr)
    plt.axvline(dr+drive, color='red', linestyle='dashed', alpha=0.3)
    plt.axvline(dr-drive, color='red', linestyle='dashed', alpha=0.3)

plt.legend()

#%%
"""
ajusto la mergeada de 3 iones
"""

phidoppler, titadoppler = 0, 90
phirepump, titarepump = 0,  0
phiprobe = 0
titaprobe = 90

Temp = 0.5e-3

sg = 0.544
sp = 4.5
sr = 0
DetRepump = 0


lw = 0.1
DopplerLaserLinewidth, RepumpLaserLinewidth, ProbeLaserLinewidth = lw, lw, lw #ancho de linea de los laseres


u = 32.5e6

#B = (u/(2*np.pi))/c

correccion = -20

offsetxpi = 438+correccion
DetDoppler = -35-correccion-22


gPS, gPD, = 2*np.pi*21.58e6, 2*np.pi*1.35e6
alpha = 0


drivefreq = 2*np.pi*22.135*1e6

FreqsDR = [f-offsetxpi for f in Merged_freqs]
CountsDR = Merged_counts

freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))

CircPr = 1
alpha = 0


def FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3):
#def FitEIT_MM(freqs, SG, SP, SCALE1, OFFSET, BETA1):
    #BETA = 1.8
    # SG = 0.6
    # SP = 8.1
    TEMP = 0.1e-3


    #BETA1, BETA2, BETA3 = 0, 0, 2

    Detunings, Fluorescence1 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, u, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA1, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)
    Detunings, Fluorescence2 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, u, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA2, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)
    Detunings, Fluorescence3 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, u, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA3, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)


    ScaledFluo1 = np.array([f*SCALE1 for f in Fluorescence1])
    ScaledFluo2 = np.array([f*SCALE2 for f in Fluorescence2])
    ScaledFluo3 = np.array([f*SCALE3 for f in Fluorescence3])
    return ScaledFluo1+ScaledFluo2+ScaledFluo3+OFFSET
    #return ScaledFluo1


do_fit = True
if do_fit:
    popt_3ions, pcov_3ions = curve_fit(FitEIT_MM, FreqsDR, CountsDR, p0=[0.6, 6.2, 3.5e5, 3.5e5, 3.5e5, 2e3, 1, 1, 1], bounds=((0, 0, 0, 0, 0, 0, 0, 0, 0), (2, 20, 5e8, 5e8, 5e8, 7e3, 10, 10, 10)))
#popt, pcov = curve_fit(FitEIT_MM, FreqsDR, CountsDR, p0=[0.8, 8, 4e4, 3.5e3, 0], bounds=((0, 0, 0, 0, 0), (2, 15, 1e5, 1e5, 10)))

#array([7.12876797e-01, 7.92474752e+00, 4.29735308e+04, 1.74240582e+04,
       #1.53401696e+03, 1.17073206e-06, 2.53804151e+00])


FittedEITpi_3ions = FitEIT_MM(freqslong, *popt_3ions)
#FittedEITpi_3ions = FitEIT_MM(freqslong, popt_3ions[0],popt_3ions[1],popt_3ions[2],popt_3ions[3],popt_3ions[4],popt_3ions[5],4,2,0)

#FittedEITpi_3ions = FitEIT_MM(freqslong, *popt_3ions)


print(popt_3ions)

plt.figure()
plt.errorbar(FreqsDR, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
plt.plot(freqslong, FittedEITpi_3ions, color='darkgreen', linewidth=3)
#plt.title(f'Sdop: {round(popt[0], 2)}, Spr: {round(popt[1], 2)}, T: {round(popt[2]*1e3, 2)} mK, detDop: {DetDoppler} MHz')
plt.xlabel('Detuning (MHz)')
plt.ylabel('Counts')
plt.title(f'Corr:{correccion},DetD:{DetDoppler}')
plt.grid()





#%%
"""
Veo la medicion de varios voltajes uno atras de otro
Se va en medio de la medicion 4, y en la 5 ni esta
"""

jvec = [2] # desde la 1, pero la 4 no porque es un merge de curvitas

Freqs

plt.figure()
i = 0
for j in jvec:
    plt.errorbar([2*f*1e-6 for f in Freqs[4]], CountsSplit_2ions[0][j], yerr=np.sqrt(CountsSplit_2ions[0][j]), fmt='o', capsize=2, markersize=2)
    i = i + 1
plt.xlabel('Frecuencia (MHz)')
plt.ylabel('counts')
plt.grid()
#for dr in drs:
#    plt.axvline(dr)
    #plt.axvline(dr+drive)
plt.legend()


#%%
#from EITfit.MM_eightLevel_2repumps_AnalysisFunctions import PerformExperiment_8levels
from scipy.optimize import curve_fit
import time

"""
AJUSTO LA CPT DE 2 IONES CON UN MODELO EN DONDE SUMO DOS ESPECTROS CON BETAS DISTINTOS
"""

phidoppler, titadoppler = 0, 90
phirepump, titarepump = 0,  0
phiprobe = 0
titaprobe = 90

Temp = 0.5e-3

sg = 0.544
sp = 4.5
sr = 0
DetRepump = 0


lw = 0.1
DopplerLaserLinewidth, RepumpLaserLinewidth, ProbeLaserLinewidth = lw, lw, lw #ancho de linea de los laseres


u = 32.5e6

#B = (u/(2*np.pi))/c

correccion = 27

offsetxpi = 421+correccion
DetDoppler = -16-correccion+5


gPS, gPD, = 2*np.pi*21.58e6, 2*np.pi*1.35e6
alpha = 0


drivefreq = 2*np.pi*22.135*1e6

FreqsDR = [2*f*1e-6-offsetxpi for f in Freqs[1]]
CountsDR = Counts[1]

freqslong = np.arange(min(FreqsDR), max(FreqsDR)+FreqsDR[1]-FreqsDR[0], 0.1*(FreqsDR[1]-FreqsDR[0]))

CircPr = 1
alpha = 0


def FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, OFFSET):
#def FitEIT_MM(freqs, SG, SP, SCALE1, OFFSET, BETA1):
    #BETA = 1.8
    # SG = 0.6
    # SP = 8.1
    TEMP = 0.1e-3

    BETA1, BETA2 = 3, 0

    Detunings, Fluorescence1 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, u, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA1, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)
    Detunings, Fluorescence2 = PerformExperiment_8levels_MM(SG, SP, gPS, gPD, DetDoppler, u, DopplerLaserLinewidth, ProbeLaserLinewidth, TEMP, alpha, phidoppler, titadoppler, phiprobe, titaprobe,  BETA2, drivefreq, min(freqs), max(freqs)+(freqs[1]-freqs[0]), freqs[1]-freqs[0], circularityprobe=CircPr, plot=False, solvemode=1, detpvec=None)

    ScaledFluo1 = np.array([f*SCALE1 + OFFSET for f in Fluorescence1])
    ScaledFluo2 = np.array([f*SCALE2 + OFFSET for f in Fluorescence2])
    return ScaledFluo1+ScaledFluo2
    #return ScaledFluo1


do_fit = True
if do_fit:
    popt_2ions_1, pcov_2ions_1 = curve_fit(FitEIT_MM, FreqsDR, CountsDR, p0=[0.9, 6.2, 3.5e3, 2.9e3, 3e3], bounds=((0, 0, 0, 0, 0), (2, 20, 5e8, 5e8, 8e3)))
#popt, pcov = curve_fit(FitEIT_MM, FreqsDR, CountsDR, p0=[0.8, 8, 4e4, 3.5e3, 0], bounds=((0, 0, 0, 0, 0), (2, 15, 1e5, 1e5, 10)))

#array([7.12876797e-01, 7.92474752e+00, 4.29735308e+04, 1.74240582e+04,
       #1.53401696e+03, 1.17073206e-06, 2.53804151e+00])

FittedEITpi_2sp = FitEIT_MM(freqslong, *popt_2ions_1)
#FittedEITpi = FitEIT_MM(freqslong, 0.8, 8, 4e4, 3.5e3, 0)

# beta1_2ions = popt_2ions_1[5]
# beta2_2ions = popt_2ions_1[6]

# errbeta1_2ions = np.sqrt(pcov_2ions_1[5,5])
# errbeta2_2ions = np.sqrt(pcov_2ions_1[6,6])

"""
Estos params dan bien poniendo beta2=0 y correccion=0 y son SG, SP, SCALE1, SCALE2, OFFSET, BETA1
#array([9.03123248e-01, 6.25865542e+00, 3.47684055e+04, 2.92076804e+04, 1.34556420e+03, 3.55045904e+00])
"""

"""
Ahora considerando ambos betas, con los parametros iniciales dados por los que se obtuvieron con beta2=0
y correccion=0 dan estos parametros que son los de antes pero con BETA2 incluido:
array([8.52685426e-01, 7.42939084e+00, 3.61998310e+04, 3.40160472e+04, 8.62651715e+02, 3.89756335e+00, 7.64867601e-01])
"""

#arreglito = np.array([8.52685426e-01, 7.42939084e+00, 3.61998310e+04, 3.40160472e+04, 8.62651715e+02, 3.89756335e+00, 7.64867601e-01])

FittedEITpi_2ions_1 = FitEIT_MM(freqslong, *popt_2ions_1)


print(popt_2ions_1)

plt.figure()
plt.errorbar(FreqsDR, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
plt.plot(freqslong, FittedEITpi_2ions_1, color='darkgreen', linewidth=3)
#plt.title(f'Sdop: {round(popt[0], 2)}, Spr: {round(popt[1], 2)}, T: {round(popt[2]*1e3, 2)} mK, detDop: {DetDoppler} MHz')
plt.xlabel('Detuning (MHz)')
plt.ylabel('Counts')
plt.title(f'Corr:{correccion},DetD:{DetDoppler}')
plt.grid()


#%%
"""
SUPER AJUSTE PARA MED DE 2 IONES
"""


#from EITfit.MM_eightLevel_2repumps_AnalysisFunctions import PerformExperiment_8levels
from scipy.optimize import curve_fit
import time


phidoppler, titadoppler = 0, 90
phirepump, titarepump = 0,  0
phiprobe = 0
titaprobe = 90

Temp = 0.5e-3

sg = 0.544
sp = 4.5
sr = 0
DetRepump = 0