pythonnumpymarkov

How to efficiently lookup state changes in a transition matrix with numpy?

I am doing some work with Markov Chains and I need to lookup the transition probability from a transition matrix given a sequence of state changes. How does one do this efficiently in numpy?

For example:

import numpy as np

#here is the sequence that I need to look up in the transition matrix
sequence = [0, 1, 0, 1, 1]

#the transition matrix that gives the probability to change between each 
of the states
transition_matrix = np.array([[0.2, 0.8], [0.6, 0.4]])

#desired output
result = [0.8, 0.6, 0.8, 0.4]

So the result is just the probability value that was looked up in the transition matrix. How to do this efficiently when there are many states and the sequence is very long?

Thanks.

Solution

  • Just use zip:

    result = []
for (step, next_step) in zip(sequence[:-1], sequence[1:]):
    result.append(transition_matrix[step][next_step])

    Result:

    [0.8, 0.6, 0.8, 0.4]