CODEWITHSUNDEEP
pythonmatlabnumpy
Joe Bowen
Joe Bowen

Reputation: 49

Index error - Python, Numpy, MatLab

I have converted a section of MatLab code to Python using the numpy and scipy libraries. I am however stuck on the following index error;

IndexError: index 698 is out of bounds for axis 3 with size 2

698 is the size of the time list.

it occurs in the last section of code, on this line;

exp_Pes[indx,jndx,j,i]=np.trace(rhotemp[:,:,indx,jndx] * Pe)

the rest is included for completeness.

the following is the code;

import numpy as np
import math
import time
#from scipy.linalg import expm
tic = time.clock()

dt = 1e-2;                      # time step
t = np.arange(0, 7-dt, dt)         # t


gam=1
mt=2
nt=2
delta=1




ensemble_size=6

RKflag=0


itype = 'em'




#========================================================================
def _min(x):
    return [np.amin(x), np.argmin(x)]
#========================================================================
###############
#Wavepacket envelope
#rising exponential

sig1 = 1;
xi_t = np.zeros(len(t));
[min_dif, array_pos_start] = _min(abs(t -t[0] ));
[min_dif, array_pos_stop ] = _min(abs(t -t[-1]/2.0));
step = t[1]-t[0]
p_len = np.arange(0, t[array_pos_stop] - t[array_pos_start] + step, step)
xi_t[array_pos_start:array_pos_stop+1] = math.sqrt(sig1) * np.exp((sig1/2.0)*p_len);
norm = np.trapz(t, xi_t * np.conj(xi_t));
xi = xi_t/np.sqrt(abs(norm));


#Pauli matrices
#========================
Pe=np.array([[1,0],[0,0]])
Sm=np.array([[0,0],[1,0]])
Sz=np.array([[1,0],[0,- 1]])
Sy=np.array([[0,- 1j],[1j,0]])
Sx=np.array([[0,1],[1,0]])


#S,L,H coefficients
#=========================
S=np.eye(2)
L=np.sqrt(gam) * Sm
H=np.eye(2)
#=========================




psi=np.array([[0],[1]])
rhoi=psi * psi.T

#rho=np.zeros(2,2,2,2)
rho=np.zeros((2,2,2,2))
rhotot=np.copy(rho)
rhoblank=np.copy(rho)
rhos=np.copy(rho)
rhotots=np.copy(rho)
rhon=np.copy(rho)
rhototN=np.copy(rho)
rhov=np.copy(rhoi)

exp_Pes=np.zeros((2,2,2,2))
exp_Pes2=np.zeros((2,2,2,2))
#initial conditions into rho
#==================

rho[:,:,0,0]=rhoi
#rho[:,:,2,2]=rhoi
rhosi=np.copy(rho)



num_bad=0
avg_val=0
num_badRK=0
avg_valRK=0

#np.zeros(len(t));

exp_Sz=np.zeros((2,2,len(t)))
exp_Sy=np.copy(exp_Sz)
exp_Sx=np.copy(exp_Sz)
exp_Pe=np.copy(exp_Sz)



##########################3
#functions





#========================================================================
#D[X]rho = X*rho* ctranspose(X) -0.5*(ctranspose(X)*X*rho + rho*ctranspose(X)*X)
#
# To call write curlyD(X,rho)

def curlyD(X,rho,nargout=1):
    y=X * rho * X.conj().T - 0.5 * (X.conj().T * X * rho + rho * X.conj().T * X)
    return y

#========================================================================
def curlyC(A,B,nargout=1):
    y=A * B - B * A
    return y
#========================================================================
def calc_expect(rhotots,Pe,nargout=1):
    for indx in (1,2):
        for jndx in (1,2):
            exp_Pes[indx,jndx]=np.trace(rhotots[:,:,indx,jndx] * Pe)
            exp_Pes2[indx,jndx]=np.trace(rhotots[:,:,indx,jndx] * Pe) ** 2
    return exp_Pes,exp_Pes2
#========================================================================

#========================================================================




#========================================================================

#========================================================================









#========================================================================

def single_photon_liouvillian(S,L,H,rho,xi,nargout=1):
    rhotot[:,:,1,1]=curlyD(L,rho[:,:,1,1]) + xi * curlyC(S * rho[:,:,0,1],L.T) + xi.T * curlyC(L,rho[:,:,1,0] * S.T) + xi.T * xi * (S * rho[:,:,0,0] * S.T - rho[:,:,0,0])
    rhotot[:,:,1,0]=curlyD(L,rho[:,:,1,0]) + xi * curlyC(S * rho[:,:,0,0],L.T)
    rhotot[:,:,0,1]=curlyD(L,rho[:,:,0,1]) + xi.T * curlyC(L,rho[:,:,0,0] * S.T)
    rhotot[:,:,0,0]=curlyD(L,rho[:,:,0,0])
    A=np.copy(rhotot)
    return A
#========================================================================


def single_photon_stochastic(S,L,H,rho,xi,nargout=1):
    K=np.trace((L + L.T) * rho[:,:,1,1]) + np.trace(S * rho[:,:,0,1]) * xi + np.trace(S.T * rho[:,:,1,0]) * xi.T
    rhotot[:,:,1,1]=(L * rho[:,:,1,1] + rho[:,:,1,1] * L.T + S * rho[:,:,0,1] * xi + rho[:,:,1,0] * S.T * xi.T - K * rho[:,:,1,1])
    rhotot[:,:,1,0]=(L * rho[:,:,1,0] + rho[:,:,1,0] * L.T + S * rho[:,:,0,0] * xi - K * rho[:,:,1,0])
    rhotot[:,:,0,1]=(L * rho[:,:,0,1] + rho[:,:,0,1] * L.T + rho[:,:,0,0] * S.T * xi.T - K * rho[:,:,0,1])
    rhotot[:,:,0,0]=(L * rho[:,:,0,0] + rho[:,:,0,0] * L.T - K * rho[:,:,0,0])
    B=np.copy(rhotot)
    return B

#========================================================================
def sde_int_io2_rk(a,b,S,L,H,yn,xi,dt,dW,nargout=1):
    Gn=yn + a[S,L,H,yn,xi] * dt + b[S,L,H,yn,xi] * dW
    Gnp=yn + a[S,L,H,yn,xi] * dt + b[S,L,H,yn,xi] * np.sqrt(dt)
    Gnm=yn + a[S,L,H,yn,xi] * dt - b[S,L,H,yn,xi] * np.sqrt(dt)
    ynp1=yn + 0.5 * (a[S,L,H,Gn,xi] + a[S,L,H,yn,xi]) * dt + 0.25 * (b[S,L,H,Gnp,xi] + b[S,L,H,Gnm,xi] + 2 * b[S,L,H,yn,xi]) * dW + 0.25 * (b[S,L,H,Gnp,xi] - b[S,L,H,Gnm,xi]) * (dW ** 2 - dt) * (dt) ** (- 0.5)
    return ynp1
#======================================================================== 


#========================================================================
def sde_int_photon(itype,rhos,S,L,H,Pe,xi,t,nargout=1):
    dt=t[1] - t[0]
    rhoblank=np.zeros(len(rhos))
    Ax=single_photon_liouvillian
    Bx=single_photon_stochastic
    #if strcmp(itype,'em'):
    if itype == 'em':    
        for i in (1,(len(t)-1)):
            if i == 1:
                exp_Pes[:,:,i],exp_Pes2[:,:,i]=calc_expect(rhos,Pe,nargout=2)
                continue
            dW=np.sqrt(dt) * np.randn
            rhotots=rhos + dt * single_photon_liouvillian(S,L,H,rhos,xi[i]) + dW * single_photon_stochastic(S,L,H,rhos,xi[i])
            exp_Pes[:,:,i],exp_Pes2[:,:,i]=calc_expect(rhotots,Pe,nargout=2)
            rhos=np.copy(rhotots)
            rhotots=np.copy(rhoblank)
    if itype == 'rk': 
        for i in (1,(len(t)-1)):
            if i == 1:
                exp_Pes[:,:,i],exp_Pes2[:,:,i]=calc_expect(rhos,Pe,nargout=2)
                continue
            dW=np.sqrt(dt) * np.randn
            rhotots=sde_int_io2_rk(Ax,Bx,S,L,H,rhos,xi[i],dt,dW)
            exp_Pes[:,:,i],exp_Pes2[:,:,i]=calc_expect(rhotots,Pe,nargout=2)
            rhos=np.copy(rhotots)
            rhotots=np.copy(rhoblank)
    return exp_Pes,exp_Pes2
#========================================================================

"""
def md_expm(Ain,nargout=1):
    Aout=np.zeros(len(Ain))
    r,c,d1,d2=len(Ain,nargout=4)
    for indx in (1,d1):
        for jndx in (1,d2):
            Aout[:,:,indx,jndx]=expm(Ain[:,:,indx,jndx])
    return Aout
"""
#========================================================================

#========================================================================


Ax=single_photon_liouvillian
Bx=single_photon_stochastic





toc = time.clock()


for indx in (1,range(mt)):
    for jndx in (1,range(nt)):
        exp_Sz[indx,jndx,1]=np.trace(rho[:,:,indx,jndx] * Sz)
        exp_Sy[indx,jndx,1]=np.trace(rho[:,:,indx,jndx] * Sy)
        exp_Sx[indx,jndx,1]=np.trace(rho[:,:,indx,jndx] * Sx)
        exp_Pe[indx,jndx,1]=np.trace(rho[:,:,indx,jndx] * Pe)



for i in (2,len(t)-1):
    #Master equation
    rhotot=rho + dt * single_photon_liouvillian(S,L,H,rho,xi[i - 1])

    for indx in (1,range(mt)):
        for jndx in (1,range(nt)):
            exp_Sz[indx,jndx,i]=np.trace(rhotot[:,:,indx,jndx] * Sz)
            exp_Sy[indx,jndx,i]=np.trace(rhotot[:,:,indx,jndx] * Sy)
            exp_Sx[indx,jndx,i]=np.trace(rhotot[:,:,indx,jndx] * Sx)
            exp_Pe[indx,jndx,i]=np.trace(rhotot[:,:,indx,jndx] * Pe)
    rho=np.copy(rhotot)
    rhotot=np.copy(rhoblank)




for j in (1,range(ensemble_size)):
    psi1=np.array([[0],[1]])
    rho1=psi1 * psi1.T
    rhotemp = np.zeros((2,2,2,2))
    rhotemp[:,:,0,0]=rho1
    rhotemp[:,:,1,1]=rho1
    rhos=np.copy(rhotemp)
    for indx in (1,range(2)):
        for jndx in (1,range(2)):
            exp_Pes[indx,jndx,j,i]=np.trace(rhotemp[:,:,indx,jndx] * Pe)
            exp_Pes2[indx,jndx,j,i]=np.trace(rhotemp[:,:,indx,jndx] * Pe) ** 2
    for i in (2,(len(t)-1)):
        dW=np.sqrt(dt) * np.random.randn()
        rhotots=rhos + dt * single_photon_liouvillian(S,L,H,rhos,xi[i - 1]) + dW * single_photon_stochastic(S,L,H,rhos,xi[i - 1])
        for indx in (1,range(mt)):
            for jndx in (1,range(nt)):
                exp_Pes[indx,jndx,j,i]=np.trace(rhotots[:,:,indx,jndx] * Pe)
                exp_Pes2[indx,jndx,j,i]=np.trace(rhotots[:,:,indx,jndx] * Pe) ** 2
        rhos=np.copy(rhotots)
        rhotots=np.copy(rhoblank)
    Irow=np.where(np.squeeze(exp_Pes[2,2,j,:]) > 1)
    Val=np.squeeze(exp_Pes[2,2,j,Irow])
    if Irow:
        num_bad=num_bad + 1
        avg_val=avg_val + max(Val)

Any help would be appreciated as I have been stuck on this for a while.

Thanks!

Upvotes: 0

Views: 208

Answers (1)

TheBlackCat
TheBlackCat

Reputation: 10308

The problem is that you are not defining i in the loop. i is being left over from the last time it was used. Specifically, it is using the last value of i from the previous loop. This value is len(t)-1, which is much larger than the length of the 3rd dimension of exp_Pes which has a length of 2. You need to define a valid value for i somewhere.

If you don't want to loop over i, generally you could just define it once before the loop starts. In your case, however, you would need to set it every time you loop since it is also being defined inside this loop a few lines below the one causing the error.

However, it would probably be clearer to just do something like exp_Pes[indx,jndx,j,0] = ..., since people reading your code won't need to look back and find what value i was set to.

A couple more minor points of advice:

np.arange does not include the last value. So the case where you use 7-dt at the beginning, you are ending up with the start being 0 and the end being 7-dt-dt, which is probably not what you want. But anyway, you should never use np.arange (or the equivalent in MATLAB) with floats, since it can have floating-point errors. Use np.linspace instead.

Second, for your _min function, you don't need [], you can just do return np.amin(x), np.argmin(x).

However, it is usually easier for arrays to use the method version, which is the version attached to the array. This would be return x.min(), x.argmin(). Same with np.copy(rho), use rho.copy() instead.

However, this function could be further simplified to a lambda expression, the equivalent of a Matlab anonymous function, like so: _min = lambda x: [x.min(), x.argmin()] (you do need the [] there).

You should use the numpy version of functions with numpy arrays. They will be faster. So np.abs(t-t[0]) is faster than abs(t-t[0]). Same with using np.sqrt(sig1) instead of math.sqrt(sig1).

You also don't need the [] in returned values, so[min_dif, array_pos_start] = _min(abs(t -t[0] ))can bemin_dif, array_pos_start = _min(np.abs(t-t[0])). However, since you are never actually using min_dif, this could just be array_pos_start = np.abs(t-t[0]).argmax().

You don't need to end lines with ;.

You don't need to assign to a variable before returning. So you can just have, for example, return A*B-B*A.

You seem to unnecessarily use the nargout argument. Python doesn't use this, in large part because you can just index the results So say you have a function foo(x) that returns the min and max of x. You can just want the max, you can do foo(x)[1]. Matlab doesn't let you do that, so it relies on nargout to work around it.

Upvotes: 1

Related Questions