[Python] Generate random paths using a 2d transition matrix

I have the following working function that generates n paths of n_steps steps using a predefined transition matrix T:

def simulate_paths(T: np.ndarray, n: int, n_steps: int) -> np.ndarray:
n_states = len(T)
states = np.zeros(n, dtype=np.int64) # every path starts at 0
paths = [states]

for _ in range(n_steps - 1):
c = np.array([
rng.choice(n_states, p=T[state]) for state in states
])
paths.append(c)
states = c

return np.array(paths).T


However, when T is defined as a 2D-matrix, the function doesn't work anymore. This is because parameter p of numpy.random.Generator.choice() must be a 1D-array.

Let's assume that T is defined as follows:

T = np.array([
[[0.0, 0.6, 0.4, 0.0, 0.0],
[0.4, 0.0, 0.0, 0.6, 0.0],
[0.6, 0.0, 0.0, 0.0, 0.4],
[0.0, 0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 1.0]],
[[0.0, 0.7, 0.2, 0.1, 0.0],
[0.4, 0.0, 0.0, 0.6, 0.0],
[0.6, 0.2, 0.0, 0.2, 0.0],
[0.0, 0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 1.0]],
])


I would like to change my function so it returns as fast as possible a numpy.ndarray with a shape of

(T.shape[0], n, n_steps)
​

Please advice

Continue reading...