r"""Compute expected transition counts for Markov chain after n steps.

    Expected counts are computed according to ..math::

    E[C_{ij}^{(n)}]=\sum_{k=0}^{n-1} (p_0^T T^{k})_{i} p_{ij}

    For N<=M, the sum is computed via successive summation of the
    following matrix vector products, p_1^t=p_0^t
    T,...,p_n^t=P_{n-1}^t T. For n>M, the sum is computed using the
    eigenvalue decomposition of T and applying the expression for a
    finite geometric series to each of the eigenvalues.

    Parameters
    ----------
    p0 : (M,) ndarray
        Starting (probability) vector of the chain.
    T : (M, M) ndarray
        Transition matrix of the chain.
    n : int
        Number of steps to take from initial state.

    Returns
    --------
    EC : (M, M) ndarray
        Expected value for transition counts after n steps.

def expected_counts(p0, T, n):
    r"""Compute expected transition counts for Markov chain after n steps.

    Expected counts are computed according to ..math::

    E[C_{ij}^{(n)}]=\sum_{k=0}^{n-1} (p_0^T T^{k})_{i} p_{ij}

    For N<=M, the sum is computed via successive summation of the
    following matrix vector products, p_1^t=p_0^t
    T,...,p_n^t=P_{n-1}^t T. For n>M, the sum is computed using the
    eigenvalue decomposition of T and applying the expression for a
    finite geometric series to each of the eigenvalues.

    Parameters
    ----------
    p0 : (M,) ndarray
        Starting (probability) vector of the chain.
    T : (M, M) ndarray
        Transition matrix of the chain.
    n : int
        Number of steps to take from initial state.

    Returns
    --------
    EC : (M, M) ndarray
        Expected value for transition counts after n steps.

    """
    M = T.shape[0]
    if n <= M:
        return ec_matrix_vector(p0, T, n)
    else:
        return ec_geometric_series(p0, T, n)