Convert pairs of states to a transition probability matrix.

state_pairs_to_tpm.py

# -*- coding: utf-8 -*-
# state_pairs_to_tpm.py

import numpy as np
import pandas as pd


def state_pairs_to_tpm(previous_states, next_states):
    """Return a TPM from observed state transitions.

    Arguments:
        previous_states (array-like): An array where the first dimension
            indexes states, such that ``previous_states[i]`` precedes
            ``next_states[i]``.
        next_states (array-like): An array where the first dimension indexes
            states, such that ``next_states[i]`` follows
            ``previous_states[i]``.

    Returns:
        pd.DataFrame: A TPM with unique ``previous_states`` as the index and
        unique ``next_states`` as the columns.

    Example:
        >>> sensors = np.array([
        ...     [0, 0],
        ...     [0, 1],
        ...     [1, 0],
        ...     [1, 0],
        ...     [1, 1],
        ...     [1, 1],
        ... ])
        >>> motors = np.array([
        ...     [1, 0],
        ...     [0, 1],
        ...     [0, 0],
        ...     [0, 1],
        ...     [0, 0],
        ...     [1, 1],
        ... ])
        >>> tpm = state_pairs_to_tpm(sensors, motors)
        >>> tpm
        next      (0, 0)  (0, 1)  (1, 0)  (1, 1)
        previous                                
        (0, 0)       0.0     0.0     1.0     0.0
        (0, 1)       0.0     1.0     0.0     0.0
        (1, 0)       0.5     0.5     0.0     0.0
        (1, 1)       0.5     0.0     0.0     0.5
        >>> # The underlying NumPy array can be accessed with the
        >>> # `.values` attribute:
        >>> tpm.values
        array([[0. , 0. , 1. , 0. ],
               [0. , 1. , 0. , 0. ],
               [0.5, 0.5, 0. , 0. ],
               [0.5, 0. , 0. , 0.5]])
    """
    if previous_states.ndim != 2 or next_states.ndim != 2:
        raise ValueError(
            f"input arrays must be 2D; got {previous_states.ndim}D and "
            f"{next_states.ndim}D"
        )
    previous_states = pd.Categorical([tuple(states) for states in previous_states])
    next_states = pd.Categorical([tuple(states) for states in next_states])
    tpm = pd.crosstab(previous_states, next_states, normalize="index")
    tpm.index.name = "previous"
    tpm.columns.name = "next"
    return tpm