Markov transition matrix

transition_matrix.py

import numpy as np

def transition_matrix(states):
    """Create Markov transition matrix from 1D array of states.
    
    Parameters
    ----------
    states : array_like
        One-dimensional array of state. The dtype of ``x``
        must be integer (e.g. [0, 2, 2, 1, 1, 1, ...])
        
    Returns
    -------
    counts : array
        Counts transition matrix (number of
        transitions from state X to state Y, etc.)
    probs : array
        Conditional probability transition matrix, i.e. 
        given that current state is X, what is the probability that
        the next state is Y. See  more details at:
        https://en.wikipedia.org/wiki/Doubly_stochastic_matrix
        
    Examples
    --------
    >>> a = [1, 1, 1, 0, 0, 2, 2, 0, 2, 0, 1, 1, 0, 0]
    >>> counts, probs = transition_matrix(a)
    >>> counts
    array([[2, 1, 2],
           [2, 3, 0],
           [2, 0, 1]])
    
    >>> probs
    array([[0.4       , 0.2       , 0.4       ],
           [0.4       , 0.6       , 0.        ],
           [0.66666667, 0.        , 0.33333333]])
    """
    x = np.asarray(states, dtype=int)
    unique = np.unique(x)
    n = unique.size
    # Integer transition counts
    counts = np.zeros((n, n), dtype=int)
    np.add.at(counts, (x[:-1], x[1:]), 1)
    # Conditional probabilities
    probs = counts / counts.sum(axis=-1, keepdims=True)
    # Optional, convert to Pandas
    # counts = pd.DataFrame(counts, index=unique, columns=unique)
    # probs = pd.DataFrame(probs, index=unique, columns=unique)
    return counts, probs