<jupyter_start><jupyter_text>Model quantum FROG (GS) Algorithm<jupyter_code>from whitepeaks import *<jupyter_output><empty_output><jupyter_text>Generate state<jupyter_code>wi0,sigma_i=w2f(0.822),0.01166
ws0,sigma_s=w2f(0.732),0.01275

rho=-0.995
Ai,As=20000,-10000.0
Ai3,As3=1e5,1e5
Ap=0

Wi,Ws,F=gaussian_state(wi0,ws0,sigma_i,sigma_s,rho,Ai=Ai,As=As,Ap=Ap,Nsigma=[1.5,1.5],Nsamples=[2**6,2**6])

F=F*np.exp(1j*As3*(Ws-ws0)**3+1j*Ai3*(Wi-wi0)**3)


print('Input state')
amplitude=np.abs(F)
phase=np.angle(F)
phase[np.where(amplitude**2/np.max(amplitude**2)<0.05)]=0

fig,ax=plt.subplots(1,4,figsize=(15,3))
ax.reshape(-1)
S1=ax[0].pcolormesh(Wi,Ws,amplitude**2)
S2=ax[1].pcolormesh(Wi,Ws,phase)
S3=ax[2].pcolormesh(Wi,Ws,F.real)
S4=ax[3].pcolormesh(Wi,Ws,F.imag)
[ax[i].set_title([r'$|F(\omega_i,\omega_s)|^2$',r'$\phi(\omega_i,\omega_s)$',r'$Re[F(\omega_i,\omega_s)]$',r'$Im[F(\omega_i,\omega_s)]$'][i]) for i in range(4)]
[fig.colorbar([S1,S2,S3,S4][i],ax=ax[i]) for i in range(4)]
plt.tight_layout()
plt.show()

print('\nTime-frequency data')
tf_data=create_tf_data((Wi,Ws,F),add_poisson_noise=True,Nphotons=3000)
Wi,Ws,Ti,Ts,Iww,Itw,Iwt,Itt=tf_data

plot_tf_data(tf_data,fit=True)<jupyter_output>Input state<jupyter_text>Run Algorithm<jupyter_code>output=run_gs_algorithm(tf_data,5000,error_model='FROG',method='full')<jupyter_output>Iterations 5000
Grid_size (64, 64)<jupyter_text>Fit phase<jupyter_code>i,Fk,_,_,fk,errors=output

c=polyfit2phase(Wi,Ws,Fk)
def complexgauss2D(xy,a,x0,y0,sigma_x,sigma_y,rho,k1,A1,A2,B1,B2,C1,C2,k2):
    
    X,Y=xy
    
    amplitude=gauss2d(xy,a,x0,y0,sigma_x,sigma_y,rho,k1)
    
    phase=np.exp(1j*(C1*(X-x0)**3+C2*(Y-y0)**3
           +B1*(X-x0)**2+B2*(Y-y0)**2 
           +A1*(X-x0) + A2*(Y-y0) +k2))
    
    return amplitude*phase
    
          
def complexgauss2D_cost(p,x,y,z):
    '''
    Cost function for poly fit. To be used with
    scipy.optimize.least_squares fitting function.
    '''
    #return np.sqrt((np.real(z)-np.real(complexgauss2D([x,y],*p)))**2+(np.imag(z)-np.imag(complexgauss2D([x,y],*p)))**2)
    return np.abs(np.real(z)-np.real(complexgauss2D([x,y],*p)))+np.abs(np.imag(z)-np.imag(complexgauss2D([x,y],*p)))


def fit_complexgauss2D(data,p0=None):

        X,Y,Z=data

        res=least_squares(complexgauss2D_cost,p0,args=(X.reshape(-1),Y.reshape(-1),Z.reshape(-1)))

        return res

pfit=fit_gaussian((Wi-wi0,Ws-ws0,np.abs(Fk))).x

print(pfit)

p0=np.concatenate((pfit,np.array([0,0,3e4,-1e4,1e4,1e4,0])))
p0=np.concatenate((pfit,np.random.rand(7)))
print(p0)


                          
fit=fit_complexgauss2D([Wi-wi0,Ws-ws0,Fk],p0=p0)

Ffit=complexgauss2D([Wi-wi0,Ws-ws0],*fit.x)

fig,ax=plt.subplots(2,2,figsize=(12,8))
ax=ax.reshape(-1)
S1=ax[0].pcolormesh(Wi-wi0,Ws-ws0,np.real(Fk))
S2=ax[1].pcolormesh(Wi-wi0,Ws-ws0,np.imag(Fk))
S3=ax[2].pcolormesh(Wi-wi0,Ws-ws0,np.real(Ffit))
S4=ax[3].pcolormesh(Wi-wi0,Ws-ws0,np.imag(Ffit))
[fig.colorbar([S1,S2,S3,S4][i],ax=ax[i]) for i in range(4)]
plt.tight_layout()
plt.show()

print(fit)
Fk.imag

amplitude=np.abs(Fk)
phase=unwrap_phase(np.angle(Fk))
wphase=amplitude*phase

#Create function
def surfit(x,A1,A2,B1,B2,C1,C2,k2):
    X,Y,A=x

    return A*(C1*X**3+C2*Y**3
        +B1*X**2+B2*Y**2 
        +A1*X + A2*Y +k2)

p0=np.random.rand(7,1)
print(p0)
p0=[0,0,2e4,-1e4,1e5,1e5,10]

out=curve_fit(surfit,[(Wi-wi0).reshape(-1),(Ws-ws0).reshape(-1),np.abs(Fk).reshape(-1)],wphase.reshape(-1),p0=p0)

fig,ax=plt.subplots(1,2,figsize=(12,4))
ax.reshape(-1)
S1=ax[0].pcolormesh(Wi-wi0,Ws-ws0,wphase)
S2=ax[1].pcolormesh(Wi-wi0,Ws-ws0,surfit([Wi-wi0,Ws-ws0,amplitude],*out[0]))
[fig.colorbar([S1,S2][i],ax=ax[i]) for i in range(2)]
plt.tight_layout()
plt.show()

print(out[0])<jupyter_output>[[0.47393124]
 [0.51678497]
 [0.54730157]
 [0.88613537]
 [0.99293373]
 [0.48395506]
 [0.89847827]]
{'fraction_non_alphanumeric': 0.23211946554886037, 'fraction_numerical': 0.0877652606759235, 'mean_word_length': 6.50098231827112, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '596557', 'license': 'MIT', 'n_tokens_mistral': 2146, 'n_tokens_neox': 2004, 'n_words': 145, 'path': 'examples/model_qFROG.ipynb', 'repo_name': 'SachaSchwarz/whitepeaks', 'size': '3817'}