Wednesday, 13 October 2021

Slicing a NumPy array with another set of arrays having beginning and ending set of indices

I am writing a code to simulate Continuous Time Random Walk phenomena with a function in python. My code works correctly so far, but I would like to exploit the indexing abilities of NumPy arrays and improve the speed. In the code, below I am generating an ensemble of trajectories, so I have to loop over each of them while generating it. Is it somehow possible to index the NumPy array x in such a way that I can get rid of the loop on Nens (the for loop in the below code snippet)

    for k in range(Nens):
            
        #start building the trajectory
        stop = 0
        i = 0
        while (stop < Nt):

            #reset the the stop time
            start = stop

            #increment the stop till the next waiting time
            stop += int(trand[i,k]) # A 2D numpy array

   
            #update the trajectory
            x[start:stop,k] = x[start-1,k] \ #x is also a 2D array
                            + (1-int(abs((x[start-1,k]+xrand[i,k])/(xmax))))* xrand[i,k] \
                            - int(abs((x[start-1,k]+xrand[i,k])/(xmax)))*np.sign(x[start-1,k]/xrand[i,k])* xrand[i,k] 
   
            i += 1

    print i
    return T, x

A plausible method that I can look to is as follows.

In this code, start and stop are scalar integers. However, I would like to index this is in way in which both start and stop are 1D Numpy integer arrays. But I have seen that if I can use only stop/start to slice the numpy array, but using slicing from a beginning to ending tuple of indices is not possible .

EDIT 1 (MWE): The following is the function that I have written, which produces random walk trajectory if given the appropriate parameters,

def ctrw_ens2d(sig,tau,sig2,tau2,alpha,xmax,Nens,Nt=1000,dt=1.0):
    #first define the array for time
    T = np.arange(0,Nt,1)*dt
    
    #generate at least Nt random time increments based on Poisson 
    # distribution (you will use only less than that)
    trand = np.random.exponential(tau, (2*Nt,Nens,1))
    xrand = np.random.normal(0.0,sig,(2*Nt,Nens,2))
    Xdist = np.random.lognormal(-1,0.9,(Nens))
    Xdist = np.clip(Xdist,2*sig,12*sig)
    
    trand2 = np.random.exponential(tau2, (2*Nt,Nens,1))
    xrand2 = np.random.normal(0.0,sig2,(2*Nt,Nens,2))
    
    
    #make a zero array of trajectory
    x1 = np.zeros((Nt,Nens))
    x2 = np.zeros((Nt,Nens))
    y1 = np.zeros((Nt,Nens))
    y2 = np.zeros((Nt,Nens))
    
    for k in range(Nens):
        #start building the trajectory
        stop = 0
        i = 0

        while (stop < Nt):

            #reset the the stop time
            start = stop

            #increment the stop till the next waiting time
            stop += int(trand[i,k,0])           

            #update the trajectory

            x1[start:stop,k] = x1[start-1,k] \
                               + (1-int(abs((r1+rr)/(Xdist[k]))))* xrand[i,k,0] \
                               - int(abs((r1+rr)/(Xdist[k])))* \
                               np.sign(x1[start-1,k]/xrand[i,k,0])* xrand[i,k,0] 
            
            y1[start:stop,k] = y1[start-1,k] \
                                 + (1-int(abs((r1+rr)/(Xdist[k]))))* xrand[i,k,1] \
                                - int(abs((r1+rr)/(Xdist[k])))* \
                                 np.sign(y1[start-1,k]/xrand[i,k,1])* xrand[i,k,1] 
            i += 1
        
        #randomly add jumps in between, at later stage
    
        stop = 1
        i = 0
        while (stop < Nt):

            #reset the the stop time
            start = stop

            #increment the stop till the next waiting time
            stop += int(trand2[i,k,0])

            #update the trajectory
            x2[start:stop,k] = x2[start-1,k] + xrand2[i,k,0]
            y2[start:stop,k] = y2[start-1,k] + xrand2[i,k,1]
            i += 1

        
    return T, (x1+x2), (y1+y2)


A simple run of the above function is given below,

``` lang-py
    Tmin = 0.61 # in ps
    Tmax = 1000 # in ps
    NT = int(Tmax/Tmin)*10
    delt = (Tmax-0.0)/NT
    print "Delta T, No. of timesteps:",delt,NT
    
    Dint = 0.21 #give it Ang^2/ps
    sig = 0.3 #in Ang
    xmax = 5.*sig
    tau = sig**2/(2*Dint)/delt  # from ps, convert it into the required units according to delt
    print "Waiting time for confined motion (in Delta T units)",tau
    
    Dj = 0.03 # in Ang^2/ps
    tau2 = 10 # in ps
    sig2 = np.sqrt(2*Dj*tau2)
    print "Sigma 2:", sig2
    tau2 = tau2/delt
    alpha  = 1
    
    tim, xtall, ytall = ctrw_ens2d(sig,tau,sig2,tau2,alpha,xmax,100,Nt=NT,dt=delt)

The generated trajectories can be plotted as follows,

    rall = np.stack((xtall,ytall),axis=-1)
    print rall.shape
    print xtall.shape
    print rall[:,99,:].shape
    k = 19
    plt.plot(xtall[:,k],ytall[:,k])

Random walk in 2D



from Slicing a NumPy array with another set of arrays having beginning and ending set of indices

No comments:

Post a Comment