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) sparse matrix
        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 : (M, M) sparse matrix
        Expected value for transition counts after 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) sparse matrix
        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 : (M, M) sparse matrix
        Expected value for transition counts after N steps.

    """
    if (n <= 0):
        EC = coo_matrix(T.shape, dtype=float)
        return EC
    else:
        if mu is None:
            mu = stationary_distribution(T)
        D_mu = diags(mu, 0)
        EC = n * D_mu.dot(T)
        return EC