lolo_CPT_plotter_20231123_multi.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Ploteo de datos y ajustes
@author: lolo
"""


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





#%% Funciones auxiliares

from scipy.stats.distributions import  t,chi2

def estadistica(datos_x,datos_y,modelo,pcov,parametros,nombres=None,alpha=0.05):
    
    if nombres is None:
        nombres = [ f'{j}' for j in range(len(parametros)) ]
    
    # Cantidad de parámetros
    P = len(parametros)
    
    # Número de datos
    N = len(datos_x)
    
    # Grados de libertas (Degrees Of Freedom)
    dof = N-P-1

    # Cauculamos coordenadas del modelo
    # modelo_x    = datos_x if modelo_x_arr is None else modelo_x_arr
    # modelo_y    = modelo( modelo_x, *parametros )
    
    # Predicción del modelo para los datos_x medidos
    prediccion_modelo = modelo( datos_x, *parametros )
    
    # Calculos de cantidades estadísticas relevantes
    COV       = pcov                                      # Matriz de Covarianza
    SE        = np.sqrt(np.diag( COV  ))                        # Standar Error / Error estandar de los parámetros
    residuos  = datos_y - prediccion_modelo               # diferencia enrte el modelo y los datos
    
    SSE       = sum(( residuos )**2 )                     # Resitual Sum of Squares
    SST       = sum(( datos_y - np.mean(datos_y))**2)        # Total Sum of Squares
    
    # http://en.wikipedia.org/wiki/Coefficient_of_determination
    # Expresa el porcentaje de la varianza que logra explicar el modelos propuesto
    Rsq       =  1 - SSE/SST                               # Coeficiente de determinación
    Rsq_adj   = 1-(1-Rsq) * (N-1)/(N-P-1)                  # Coeficiente de determinación Ajustado   
    
    # https://en.wikipedia.org/wiki/Pearson_correlation_coefficient#In_least_squares_regression_analysis
    # Expresa la correlación que hay entre los datos y la predicción del modelo
    r_pearson = np.corrcoef( datos_y ,  prediccion_modelo )[0,1]
    
    # Reduced chi squared
    # https://en.wikipedia.org/wiki/Reduced_chi-squared_statistic
    chi2_     = sum( residuos**2 )/N
    chi2_red  = sum( residuos**2 )/(N-P)
    
    # Chi squared test
    chi2_test = sum( residuos**2 / abs(prediccion_modelo) )
    # p-value del ajuste
    p_val  = chi2(dof).cdf( chi2_test )
    
    
    sT = t.ppf(1.0 - alpha/2.0, N - P ) # student T multiplier
    CI = sT * SE                        # Confidence Interval
    
    print('R-squared    ',Rsq)
    print('R-sq_adjusted',Rsq_adj)
    print('chi2         ',chi2_)
    print('chi2_reduced ',chi2_red)
    print('chi2_test    ',chi2_test)
    print('r-pearson    ',r_pearson)
    print('p-value      ',p_val)
    print('')
    print('Error Estandard (SE):')
    for i in range(P):
        print(f'parametro[{nombres[i]:>5s}]: ' , parametros[i], ' ± ' , SE[i])
    print('')
    print('Intervalo de confianza al '+str((1-alpha)*100)+'%:')
    for i in range(P):
        print(f'parametro[{nombres[i]:>5s}]: ' , parametros[i], ' ± ' , CI[i])
    
    return dict(R2=Rsq,R2_adj=Rsq_adj,chi2=chi2_,chi2_red=chi2_red,
                chi2_test=chi2_test,r=r_pearson,pvalue=p_val,
                SE=SE,CI=CI)





#%% Importaciones extra

# /home/lolo/Dropbox/marce/LIAF/Trampa_anular/artiq_experiments/analisis/plots/20231123_CPTconmicromocion3/Data/EITfit/MM_eightLevel_2repumps_AnalysisFunctions.py

from Data.EITfit.lolo_modelo_full_8niveles import PerformExperiment_8levels_MM


PARAMETROS = np.load('PARAMETROS.npz',allow_pickle=True)
for var_name in PARAMETROS.keys():
    globals()[var_name] = PARAMETROS[var_name]
    print(f'loaded: {var_name}')


#%%

"""
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('../20231123_CPTconmicromocion3/Data/')

folder = '../20231123_CPTconmicromocion3/Data/'
CPT_FILES = f"""
{folder}/000016262-IR_Scan_withcal_optimized
{folder}/000016239-IR_Scan_withcal_optimized
{folder}/000016240-IR_Scan_withcal_optimized
{folder}/000016241-IR_Scan_withcal_optimized
{folder}/000016244-IR_Scan_withcal_optimized
{folder}/000016255-IR_Scan_withcal_optimized
{folder}/000016256-IR_Scan_withcal_optimized
{folder}/000016257-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_2ions = []
CountsSplit_2ions.append(Split(Counts[4],len(Freqs[4])))

#%%

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

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

drs = [390.5, 399.5, 406, 413.5]

drive=22.1

Frequencies = Freqs[0]

plt.figure()
i = 0
for j in jvec:
    plt.errorbar([2*f*1e-6 for f in Frequencies], CountsSplit[0][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()

#%%


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


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

Las mediciones estan buenas, habria que ver de ajustarlas bien, yo no lo logre.

"""

"""
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

import numba



@numba.jit
def FitEIT_MM1(freqs, SG, SP, SCALE1, OFFSET, BETA1):
#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)

    ScaledFluo1 = np.array([f*SCALE1 for f in Fluorescence1])

    return ScaledFluo1+OFFSET
    #return ScaledFluo1


@numba.jit
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
    freqs = np.array(freqs)

    #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
    Fluorescence1 = np.array(Fluorescence1)
    Fluorescence2 = np.array(Fluorescence2)
    Fluorescence3 = np.array(Fluorescence3)
    return SCALE1*Fluorescence1+SCALE2*Fluorescence2+SCALE3*Fluorescence3+OFFSET




if not 'popt_1ions' in globals().keys():
    t0 = time.time()
    print("Arranamos FIT 1")

    par_ini = [0.65, 7.06, 86070, 3917,  1.64]
    
    popt_1ions, pcov_1ions = curve_fit(FitEIT_MM1, FreqsDR, CountsDR, p0=par_ini, bounds=((0, 0, 0, 0, 0), (2, 20, 5e8, 7e3, 10)))
    
    pp1 = estadistica(FreqsDR,CountsDR,FitEIT_MM1,pcov_1ions,popt_1ions,nombres=None,alpha=0.05)
    
    print(f"time: {round(time.time()-t0,1)} seg")



if not 'popt_3ions' in globals().keys():
    t0 = time.time()
    print("Arranamos FIT 1")
    
    par_ini = [0.65, 7.06, 86070, 3917,  1.64]
    popt_1ions, pcov_1ions = curve_fit(FitEIT_MM1, FreqsDR, CountsDR, p0=par_ini, bounds=((0, 0, 0, 0, 0), (2, 20, 5e8, 7e3, 10)))
    print(f"time: {round(time.time()-t0,1)} seg")
    
    print("Arranamos FIT 3")
    def fun(freqs, SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3):
        SCALE3 = max(1-SCALE2-SCALE1-OFFSET,0)
        Fluorescence3 = FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3)
        return Fluorescence3/sum(Fluorescence3)
    
    par_ini = popt_1ions.tolist()[:2]+[1/3]*2 + [0.1] +popt_1ions.tolist()[-1:]*3
    bounds  = ((0, 0,  0, 0, 0, 0, 0, 0), 
               (2, 20, 1, 1, 1, 10, 10, 10))
    
    popt_3ions, pcov_3ions = curve_fit(fun, FreqsDR, CountsDR, p0=par_ini,bounds=bounds )
   
    
    print(f"time: {round(time.time()-t0,1)} seg")
    
    freqs = np.array(FreqsDR)
    SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3 = popt_3ions
    SCALE3 = max(1-SCALE2-SCALE1-OFFSET,0)
    # Fluorescence= FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3)
    Fluorescence= fun(freqs, SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3)
    
    
    plt.figure()
    plt.errorbar(freqs, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
    
    plt.plot(freqs, Fluorescence*sum(CountsDR))
    
    plt.plot(freqs, fun(freqs, *par_ini )*sum(CountsDR))
    
    
#%% Acá hay un ajuste del de 3 iones que da razonable

FreqsDR, CountsDR = np.array(FreqsDR) , np.array(CountsDR)

t0 = time.time()
print("Arranamos FIT 1")

par_ini = [0.65, 7.06, 86070, 3917,  1.64]
popt_1ions, pcov_1ions = curve_fit(FitEIT_MM1, FreqsDR, CountsDR, p0=par_ini, bounds=((0, 0, 0, 0, 0), (2, 20, 5e8, 7e3, 10)))
print(f"time: {round(time.time()-t0,1)} seg")




par_ini = popt_1ions.tolist()[:2]+[1/3]*2 + [0.1] +popt_1ions.tolist()[-1:]*3
bounds  = ((0, 0,  0, 0, 0, 0, 0, 0), 
           (2, 20, 1, 1, 1, 10, 10, 10))

# bounds = (-np.inf, np.inf)

plt.figure()
plt.errorbar(freqs, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
l1, =plt.plot(freqs, fun(freqs, *par_ini )*sum(CountsDR))
plt.draw()

print("Arranamos FIT 3")
def fun(freqs, SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3):
    SCALE3 = max(1-SCALE2-SCALE1-OFFSET,0)
    Fluorescence3 = FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3)
    
    l1.set_ydata(Fluorescence3/sum(Fluorescence3)*sum(CountsDR))
    print(SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3)
    plt.pause(0.1)
    return Fluorescence3/sum(Fluorescence3)


popt_3ions, pcov_3ions = curve_fit(fun, FreqsDR, CountsDR/CountsDR.sum(), p0=par_ini,bounds=bounds )
   

print(f"time: {round(time.time()-t0,1)} seg")

freqs = np.array(FreqsDR)
SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3 = popt_3ions
SCALE3 = max(1-SCALE2-SCALE1-OFFSET,0)
# Fluorescence= FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3)
Fluorescence= fun(freqs, SG, SP, SCALE1, SCALE2, OFFSET, BETA1, BETA2, BETA3)

plt.plot(freqs, Fluorescence*sum(CountsDR))

if False:
    popt_3ions = np.array([0.70790618, 7.41165226, 0.1707309 , 0.13974759, 0.02322333,
                           3.69298275, 3.68847693, 1.36216489])
    

#%%
    
    
    
    
    
#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])


if False:
    # Ejemplo 3 iones
    t0 = time.time()
    freqs, SG, SP = np.array(FreqsDR), 0.65, 7.06
    SCALE1, SCALE2, SCALE3, OFFSET = 1,1,1,1
    BETA1, BETA2, BETA3 = 1.64, 2, 3
    TEMP = 0.1e-3
    Fluorescence3= FitEIT_MM(freqs, SG, SP, SCALE1, SCALE2, SCALE3, OFFSET, BETA1, BETA2, BETA3)
    print(f"time: {round(time.time()-t0,1)} seg")
    Detunings, Fluorescence = np.array(freqs), np.array(Fluorescence3)
    plt.figure()
    plt.errorbar(FreqsDR, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
    plt.plot(Detunings, Fluorescence*sum(CountsDR)/sum(Fluorescence))
    

if False:
    # Ejemplo 1 ion
    t0 = time.time()
    freqs, SG, SP, BETA1 = FreqsDR, 0.65, 7.06, 1.64
    TEMP = 0.1e-3
    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)
    print(f"time: {round(time.time()-t0,1)} seg")
    Detunings, Fluorescence1 = np.array(Detunings), np.array(Fluorescence1)
    plt.figure()
    plt.errorbar(FreqsDR, CountsDR, yerr=2*np.sqrt(CountsDR), fmt='o', color='darkgreen', alpha=0.5, capsize=2, markersize=2)
    plt.plot(Detunings, Fluorescence1*sum(CountsDR)/sum(Fluorescence1))
    

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


if not 'popt_2ions_1' in globals().keys():
    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


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


SelectedCurveVec = [3]


if not 'popt_SA_vec_2ions' in globals().keys():

    popt_SA_vec_2ions = []
    pcov_SA_vec_2ions = []

    for selectedcurve in SelectedCurveVec:

        FreqsDR = Freqs[selectedcurve]
        CountsDR = Counts[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, SCALE2, OFFSET, BETA1, BETA2, TEMP, 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, 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 for f in Fluorescence2])
            if plot:
                return ScaledFluo1+ScaledFluo2, Detunings
            else:
                return ScaledFluo1+ScaledFluo2
            #return ScaledFluo1

        if True:
            popt_3_SA_2ions, pcov_3_SA_2ions = curve_fit(FitEIT_MM_single, FreqsDR, CountsDR, p0=[448, -42, 0.6, 8.1, 4e4, 4e4, 6e3, 1, 1.2, 0.5e-3], bounds=((0, -100,0, 0, 0,0,0,0,0, 0), (1000, 0, 2, 20,5e6, 5e6,5e4, 10, 10,10e-3)))

            #popt_3_SA_2ions = [448, -42, 8e4, 6e3, 2, 0.5e-3]


            popt_SA_vec_2ions.append(popt_3_SA_2ions)
            pcov_SA_vec_2ions.append(pcov_3_SA_2ions)

            FittedEITpi_3_SA_short, Detunings_3_SA_short = FitEIT_MM_single(FreqsDR, *popt_3_SA_2ions, 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_2ions, plot=True)

raise ValueError('Acá tenes que levantar de nuevo los valores que van')


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_2ions)
#print(f'Detdop:{popt_3_SA[1]},popt_3_SA:{popt[0]}')


#%%
#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


if not 'popt_2ions_1' in globals().keys():
    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()


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


#%%

"""
SUPER AJUSTE (SA)


"""

if False:
    GUARDAR = {}
    for var in [ kk for kk in globals().keys() if kk.startswith('pop') ]:
        print(var)
        GUARDAR[var] = globals()[var]
    print('')
    for var in [ kk for kk in globals().keys() if kk.startswith('pcov') ]:
        print(var)
        GUARDAR[var] = globals()[var]

    print('')
    for var in [ kk for kk in globals().keys() if kk.startswith('Fitted') ]:
        print(var)
        GUARDAR[var] = globals()[var]
    print('')
    for var in [ kk for kk in globals().keys() if kk.endswith('_vec') ]:
        print(var)
        GUARDAR[var] = globals()[var]


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