Index error - Python, Numpy, MatLab
Question: 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!


Answer: 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 be`min_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.

