使用取决于先前值的操作来向量化numpy代码

下面的代码模拟了一个系统,它可以随时采样3种不同的状态,这些状态之间的恒定转移概率由矩阵给出prob_nor.因此,每个点都trace取决于之前的状态.

n_states, n_frames = 3, 1000
state_val = np.linspace(0, 1, n_states)

prob = np.random.randint(1, 10, size=(n_states,)*2)
prob[np.diag_indices(n_states)] += 50

prob_nor = prob/prob.sum(1)[:,None] # transition probability matrix, 
                                    # row sum normalized to 1.0

state_idx = range(n_states) # states is a list of integers 0, 1, 2...
current_state = np.random.choice(state_idx)

trace = []      
sigma = 0.1     
for _ in range(n_frames):
    trace.append(np.random.normal(loc=state_val[current_state], scale=sigma))
    current_state = np.random.choice(state_idx, p=prob_nor[current_state, :])

上面代码中的循环使得它运行得非常慢,特别是当我必须建模数百万个数据点时.有没有办法矢量化/加速它？