I couldn’t find a library to draw simple state transition diagrams for Markov Chains in Python – and had a couple of days off – so I made my own. The code only works with 2 to 4 state transition matrices, which was enough for what I needed, but feel free to contribute and improve upon it 🙂
UPDATE: An updated version of this post is now available >> here <<.
Installation
You can clone the repo, or simply copy the files node.py and markovchain.py in the same directory as your script.Â
Two-State Transition Diagram
import numpy as np import matplotlib.pyplot as plt # Import the MarkovChain class from markovchain.py from markovchain import MarkovChain P = np.array([[0.8, 0.2], [0.1, 0.9]]) # Transition matrix mc = MarkovChain(P, ['1', '2']) mc.draw()
Will plot this
Three-State Transition Diagram
P = np.array([ [0.8, 0.1, 0.1], [0.1, 0.7, 0.2], [0.1, 0.7, 0.2], ]) mc = MarkovChain(P, ['A', 'B', 'C']) mc.draw()
Four-State Transition Diagram
P = np.array([ [0.8, 0.1, 0.1, 0.0], [0.1, 0.7, 0.0, 0.2], [0.1, 0.0, 0.7, 0.2], [0.1, 0.0, 0.7, 0.2] ]) mc = MarkovChain(P, ['1', '2', '3', '4']) mc.draw("../img/markov-chain-four-states.png")
Â
Full demo available here.