【发布时间】：2021-11-28 03:14:43
【问题描述】：
我正在编写一个代码来使用 python 中的函数来模拟连续时间随机游走现象。到目前为止，我的代码工作正常，但我想利用 NumPy 数组的索引能力并提高速度。在下面的代码中，我正在生成一组轨迹，因此我必须在生成它时遍历它们中的每一个。是否有可能以某种方式索引 NumPy 数组 x，以便我可以摆脱 Nens 上的循环（下面的代码 sn-p 中的 for 循环）
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
我可以参考的一个合理的方法如下。
在此代码中，start 和 stop 是标量整数。但是，我想索引这是 start 和 stop 都是一维 Numpy 整数数组的方式。
但是我已经看到，如果我只能使用 stop/start 来切片 numpy 数组，但使用从索引元组的开始到结束的切片是不可能的。
编辑 1（MWE）： 以下是我写的函数，如果给定适当的参数，它会产生随机游走轨迹，
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
r1 = np.sqrt(x1[start-1,k]**2 + y1[start-1,k]**2)
rr = np.linalg.norm(xrand[i,k])
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)
下面给出上述函数的简单运行，
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)
生成的轨迹可以绘制如下，
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])
【问题讨论】：
-
似乎每个切片的长度可以变化？然后你就可以分别使用它们了。
-
确实，这是这个问题的主要瓶颈。如果不在 NumPy 中，是否还有其他数据对象可以帮助进行这种切片？
-
并非每个循环都可以进行 NumPy 向量化。由于您的动机是速度/效率，也许您可以尝试 Numba 编译循环。您没有发布 MWE，因此我们无法为您试用。
-
@BatWannaBe 到目前为止我还没有使用 Numba，请看一下我添加的 MWE 并告诉我它的可行性。
-
您可以（并且应该）自己尝试：
import numba as nb并使用@nb.njit装饰函数。请务必使用njit：它是jit(nopython=True)的别名，使用“no Python”进行编译是提高速度的方法。但这很棘手：Numba 本身并不支持所有 Python 或 NumPy 函数；它必须重新实现编译的东西。@nb.njit很方便，因为如果 Numba 无法编译某些东西，它会抛出错误。因此，您要做的是尝试划分出与njit兼容的函数部分并测试其速度（不保证加速）。
标签： python arrays numpy for-loop indexing