numpy-markov-chains.py

import numpy as np

# The transition matrix. Notice that each column sums to 1
T = np.array([[0.25, 0.33, 0.3 ],
              [0.5 , 0.66, 0.2 ],
              [0.25, 0.01, 0.5 ]])

# The state vector. We are initially in the first state with 100% probability
s = np.array([1.0, 0.0, 0.0])

s1 = np.matmul(T,s)
# array([0.25, 0.5 , 0.25])
# After one transition, these are the probabilites that we are in states 1, 2 and 3 respectively

np.matmul(T,s1)
# array([0.3025, 0.505 , 0.1925])
# After two transitions, these are the probabilities of each state

# We can also compute 2 transitions at once by applying the transition matrix to itself
# and then applying the result to our state vector

T2 = np.linalg.matrix_power(T,2)
# array([[0.3025, 0.3033, 0.291 ],
#        [0.505 , 0.6026, 0.382 ],
#        [0.1925, 0.0941, 0.327 ]])

np.matmul(T2, s)
# array([0.3025, 0.505 , 0.1925])
# As expected, we got the same result!


# Let's look at what happens after 5 transitions
T5 = np.linalg.matrix_power(T,5)
# array([[0.3008838 , 0.30147219, 0.30010961],
#        [0.53428791, 0.54460013, 0.52066214],
#        [0.16482829, 0.15392768, 0.17922825]])

# ... and 15...
np.linalg.matrix_power(A,15)
# array([[0.30107515, 0.30107551, 0.30107468],
#        [0.53763236, 0.53763866, 0.53762404],
#        [0.16129248, 0.16128583, 0.16130128]])
 
# ... and 50
T50 = np.linalg.matrix_power(T,50)
# array([[0.30107527, 0.30107527, 0.30107527],
#        [0.53763441, 0.53763441, 0.53763441],
#        [0.16129032, 0.16129032, 0.16129032]])