r"""Compute time-correlation of obs1, or time-cross-correlation with obs2.

    The time-correlation at time=k is computed by the matrix-vector expression:
    cor(k) = obs1' diag(pi) P^k obs2


    Parameters
    ----------
    P : ndarray, shape=(n, n) or scipy.sparse matrix
        Transition matrix
    obs1 : ndarray, shape=(n)
        Vector representing observable 1 on discrete states
    obs2 : ndarray, shape=(n)
        Vector representing observable 2 on discrete states. If not given,
        the autocorrelation of obs1 will be computed
    mu : ndarray, shape=(n)
        stationary distribution vector.
    time : int
        time point at which the (auto)correlation will be evaluated.
    start_values : (time, ndarray <P, <P, obs2>>_t)
        start iteration of calculation of matrix power product, with this values.
        only useful when calling this function out of a loop over times.
    return_P_k_obs : bool
        if True, the dot product <P^time, obs2> will be returned for further
        calculations.

    Returns
    -------
    cor(k) : float
           correlation between observations

def time_correlation_direct_by_mtx_vec_prod(P, mu, obs1, obs2=None, time=1, start_values=None, return_P_k_obs=False):
    r"""Compute time-correlation of obs1, or time-cross-correlation with obs2.

    The time-correlation at time=k is computed by the matrix-vector expression:
    cor(k) = obs1' diag(pi) P^k obs2


    Parameters
    ----------
    P : ndarray, shape=(n, n) or scipy.sparse matrix
        Transition matrix
    obs1 : ndarray, shape=(n)
        Vector representing observable 1 on discrete states
    obs2 : ndarray, shape=(n)
        Vector representing observable 2 on discrete states. If not given,
        the autocorrelation of obs1 will be computed
    mu : ndarray, shape=(n)
        stationary distribution vector.
    time : int
        time point at which the (auto)correlation will be evaluated.
    start_values : (time, ndarray <P, <P, obs2>>_t)
        start iteration of calculation of matrix power product, with this values.
        only useful when calling this function out of a loop over times.
    return_P_k_obs : bool
        if True, the dot product <P^time, obs2> will be returned for further
        calculations.

    Returns
    -------
    cor(k) : float
           correlation between observations
    """
    # input checks
    if not (type(time) == int):
        if not (type(time) == np.int64):
            raise TypeError("given time (%s) is not an integer, but has type: %s"
                            % (str(time), type(time)))
    if obs1.shape[0] != P.shape[0]:
        raise ValueError("observable shape not compatible with given matrix")
    if obs2 is None:
        obs2 = obs1
    # multiply element-wise obs1 and pi. this is obs1' diag(pi)
    l = np.multiply(obs1, mu)
    # raise transition matrix to power of time by substituting dot product
    # <Pk, obs2> with something like <P, <P, obs2>>.
    # This saves a lot of matrix matrix multiplications.
    if start_values:  # begin with a previous calculated val
        P_i_obs = start_values[1]
        # calculate difference properly!
        time_prev = start_values[0]
        t_diff = time - time_prev
        r = range(t_diff)
    else:
        if time >= 2:
            P_i_obs = np.dot(P, np.dot(P, obs2))  # vector <P, <P, obs2> := P^2 * obs
            r = range(time - 2)
        elif time == 1:
            P_i_obs = np.dot(P, obs2)  # P^1 = P*obs
            r = range(0)
        elif time == 0:  # P^0 = I => I*obs2 = obs2
            P_i_obs = obs2
            r = range(0)

    for k in r:  # since we already substituted started with 0
        P_i_obs = np.dot(P, P_i_obs)
    corr = np.dot(l, P_i_obs)
    if return_P_k_obs:
        return corr, (time, P_i_obs)
    else:
        return corr