r"""Expected transition counts for Markov chain in equilibrium.

    Since mu is stationary for T we have

    .. math::

        E(C^{(N)})=N diag(mu)*T.

    Parameters
    ----------
    T : (M, M) ndarray
        Transition matrix.
    n : int
        Number of steps for chain.
    mu : (M,) ndarray (optional)
        Stationary distribution for T. If mu is not specified it will be
        computed via diagonalization of T.

    Returns
    -------
    EC : numpy array, shape=(n,n)
        Expected value for transition counts after a propagation of n steps.

def expected_counts_stationary(T, n, mu=None):
    r"""Expected transition counts for Markov chain in equilibrium.

    Since mu is stationary for T we have

    .. math::

        E(C^{(N)})=N diag(mu)*T.

    Parameters
    ----------
    T : (M, M) ndarray
        Transition matrix.
    n : int
        Number of steps for chain.
    mu : (M,) ndarray (optional)
        Stationary distribution for T. If mu is not specified it will be
        computed via diagonalization of T.

    Returns
    -------
    EC : numpy array, shape=(n,n)
        Expected value for transition counts after a propagation of n steps.

    """
    if n <= 0:
        EC = np.zeros(T.shape)
        return EC
    else:
        if mu is None:
            mu = stationary_distribution(T)
        EC = n * mu[:, np.newaxis] * T
        return EC