from qutip import *
from pylab import *
%pylab inline

N=10 # wymiar przestrzenie stanów

chi=1.0 # amplituda nieliniowości
#omega=1.0

H=num(N)+chi*(num(N)*num(N)-num(N))

kk=H.eigenstates(sort='low')

E=kk[0] # wartości własne
W=kk[1] # wektory własne

def gk(J,gamma,E,W,t): # t : czas
    E=E-E[0] #skalowanie stanu podstawowego
    rho=1.0
    gk=W[0]
    for i1 in range(1,N):
        rho=rho*E[i1]
        gk=gk+(J**(i1/2.0)*exp(-1J*(gamma+t)*E[i1]))/(sqrt(rho))*W[i1]
    gk=gk.unit()
    return gk

def mandel(rho): # UWAGA: rho jest macierzą gęstości (QuTip quantum object)
    N=dims(rho)[0][0] # określenie wymiaru macierzy (przestrzeni stanów)
    Q=(num(N)*num(N)*rho).tr()-((num(N)*rho).tr())**2.0-(num(N)*rho).tr()
    Q=Q/((num(N)*rho).tr())
    return Q

chi_list=linspace(0.0,10.0,30) # lista od 0 do 1.0 o 20 elementach

mandel_list=zeros(len(chi_list)) # lista wyników $Q$ dla odpowiednich $\chi$
for i1 in range(0,len(chi_list)):
    chi=chi_list[i1] # wartość $\chi$
    HH=num(N)+chi*(num(N)*num(N)-num(N)) # postać hamiltonianiu
    kkk=HH.eigenstates(sort='low') # zagadnienie własne
    EE=kkk[0] # wartości własne
    WW=kkk[1] # wektory własne
    psi=gk(0.1,0,EE,WW,0) # GKCS dla HH
    mandel_list[i1]=real(mandel(psi*psi.dag())) # macierz gęstosci w agrumencie mandel 

plot(chi_list,mandel_list)
show() # wykres

tlist=linspace(0.0,10.0,10)# czas
psi0=gk(1.0,0,E,W,0) #stan poczatkowy

qq=mesolve(H, psi0, tlist, [], [])

print(((qq.states[9]).dag())*gk(1.0,0.0,E,W,10.0))