r"""Compute time-relaxations of obs with respect of given initial distribution.

    relaxation(k) = p0 P^k obs

    Parameters
    ----------
    P : ndarray, shape=(n, n) or scipy.sparse matrix
        Transition matrix
    p0 : ndarray, shape=(n)
        initial distribution
    obs : ndarray, shape=(n)
        Vector representing observable on discrete states.
    time : int or array like
        time point at which the (auto)correlation will be evaluated.

    start_values = (time,
    Returns
    -------
    relaxation : float

def time_relaxation_direct_by_mtx_vec_prod(P, p0, obs, time=1, start_values=None, return_pP_k=False):
    r"""Compute time-relaxations of obs with respect of given initial distribution.

    relaxation(k) = p0 P^k obs

    Parameters
    ----------
    P : ndarray, shape=(n, n) or scipy.sparse matrix
        Transition matrix
    p0 : ndarray, shape=(n)
        initial distribution
    obs : ndarray, shape=(n)
        Vector representing observable on discrete states.
    time : int or array like
        time point at which the (auto)correlation will be evaluated.

    start_values = (time,
    Returns
    -------
    relaxation : float
    """
    # 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 obs.shape[0] != P.shape[0]:
        raise ValueError("observable shape not compatible with given matrix")
    if p0.shape[0] != P.shape[0]:
        raise ValueError("shape of init dist p0 (%s) not compatible with given matrix (shape=%s)"
                         % (p0.shape[0], P.shape))
    # propagate in time
    if start_values:  # begin with a previous calculated val
        pk_i = start_values[1]
        time_prev = start_values[0]
        t_diff = time - time_prev
        r = range(t_diff)
    else:
        if time >= 2:
            pk_i = np.dot(np.dot(p0, P), P)  # pk_2
            r = range(time - 2)
        elif time == 1:
            pk_i = np.dot(p0, P)  # propagate once
            r = range(0)
        elif time == 0:  # P^0 = I => p0*I = p0
            pk_i = p0
            r = range(0)

    for k in r:  # perform the rest of the propagations p0 P^t_diff
        pk_i = np.dot(pk_i, P)

    # result
    l = np.dot(pk_i, obs)
    if return_pP_k:
        return l, (time, pk_i)
    else:
        return l