How to use the pyemma._ext.variational.solvers.direct.spd_inv_sqrt function in pyEMMA

def _svd_sym_koopman(K, C00_train, Ctt_train):
    """ Computes the SVD of the symmetrized Koopman operator in the empirical distribution.
    """
    from pyemma._ext.variational.solvers.direct import spd_inv_sqrt
    # reweight operator to empirical distribution
    C0t_re = mdot(C00_train, K)
    # symmetrized operator and SVD
    K_sym = mdot(spd_inv_sqrt(C00_train), C0t_re, spd_inv_sqrt(Ctt_train))
    U, S, Vt = np.linalg.svd(K_sym, compute_uv=True, full_matrices=False)
    # projects back to singular functions of K
    U = mdot(spd_inv_sqrt(C00_train), U)
    Vt = mdot(Vt,spd_inv_sqrt(Ctt_train))
    return U, S, Vt.T

Returns:
    --------
    vamp1 : float
        VAMP-1 score

    """
    from pyemma._ext.variational.solvers.direct import spd_inv_sqrt

    # SVD of symmetrized operator in empirical distribution
    U, S, V = _svd_sym_koopman(K, C00_train, Ctt_train)
    if k is not None:
        U = U[:, :k]
        # S = S[:k][:, :k]
        V = V[:, :k]
    A = spd_inv_sqrt(mdot(U.T, C00_test, U))
    B = mdot(U.T, C0t_test, V)
    C = spd_inv_sqrt(mdot(V.T, Ctt_test, V))

    # compute trace norm (nuclear norm), equal to the sum of singular values
    score = np.linalg.norm(mdot(A, B, C), ord='nuc')
    return score

References
        ----------
        .. [1] Wu, H. and Noe, F. 2017. Variational approach for learning Markov processes from time series data.
            arXiv:1707.04659v1
        .. [2] Noe, F. and Clementi, C. 2015. Kinetic distance and kinetic maps from molecular dynamics simulation.
            J. Chem. Theory. Comput. doi:10.1021/acs.jctc.5b00553
        """
        # TODO: implement for TICA too
        if test_model is None:
            test_model = self
        Uk = self.U[:, 0:self.dimension()]
        Vk = self.V[:, 0:self.dimension()]
        res = None
        if score_method == 'VAMP1' or score_method == 'VAMP2':
            A = spd_inv_sqrt(Uk.T.dot(test_model.C00).dot(Uk))
            B = Uk.T.dot(test_model.C0t).dot(Vk)
            C = spd_inv_sqrt(Vk.T.dot(test_model.Ctt).dot(Vk))
            ABC = mdot(A, B, C)
            if score_method == 'VAMP1':
                res = np.linalg.norm(ABC, ord='nuc')
            elif score_method == 'VAMP2':
                res = np.linalg.norm(ABC, ord='fro')**2
        elif score_method == 'VAMPE':
            Sk = np.diag(self.singular_values[0:self.dimension()])
            res = np.trace(2.0 * mdot(Vk, Sk, Uk.T, test_model.C0t) - mdot(Vk, Sk, Uk.T, test_model.C00, Uk, Sk, Vk.T, test_model.Ctt))
        else:
            raise ValueError('"score" should be one of VAMP1, VAMP2 or VAMPE')
        # add the contribution (+1) of the constant singular functions to the result
        assert res
        return res + 1

Returns:
    --------
    vamp2 : float
        VAMP-2 score

    """
    from pyemma._ext.variational.solvers.direct import spd_inv_sqrt

    # SVD of symmetrized operator in empirical distribution
    U, _, V = _svd_sym_koopman(K, C00_train, Ctt_train)
    if k is not None:
        U = U[:, :k]
        V = V[:, :k]
    A = spd_inv_sqrt(mdot(U.T, C00_test, U))
    B = mdot(U.T, C0t_test, V)
    C = spd_inv_sqrt(mdot(V.T, Ctt_test, V))

    # compute square frobenius, equal to the sum of squares of singular values
    score = np.linalg.norm(mdot(A, B, C), ord='fro') ** 2
    return score

def _svd_sym_koopman(K, C00_train, Ctt_train):
    """ Computes the SVD of the symmetrized Koopman operator in the empirical distribution.
    """
    from pyemma._ext.variational.solvers.direct import spd_inv_sqrt
    # reweight operator to empirical distribution
    C0t_re = mdot(C00_train, K)
    # symmetrized operator and SVD
    K_sym = mdot(spd_inv_sqrt(C00_train), C0t_re, spd_inv_sqrt(Ctt_train))
    U, S, Vt = np.linalg.svd(K_sym, compute_uv=True, full_matrices=False)
    # projects back to singular functions of K
    U = mdot(spd_inv_sqrt(C00_train), U)
    Vt = mdot(Vt,spd_inv_sqrt(Ctt_train))
    return U, S, Vt.T

----------
        .. [1] Wu, H. and Noe, F. 2017. Variational approach for learning Markov processes from time series data.
            arXiv:1707.04659v1
        .. [2] Noe, F. and Clementi, C. 2015. Kinetic distance and kinetic maps from molecular dynamics simulation.
            J. Chem. Theory. Comput. doi:10.1021/acs.jctc.5b00553
        """
        # TODO: implement for TICA too
        if test_model is None:
            test_model = self
        Uk = self.U[:, 0:self.dimension()]
        Vk = self.V[:, 0:self.dimension()]
        res = None
        if score_method == 'VAMP1' or score_method == 'VAMP2':
            A = spd_inv_sqrt(Uk.T.dot(test_model.C00).dot(Uk))
            B = Uk.T.dot(test_model.C0t).dot(Vk)
            C = spd_inv_sqrt(Vk.T.dot(test_model.Ctt).dot(Vk))
            ABC = mdot(A, B, C)
            if score_method == 'VAMP1':
                res = np.linalg.norm(ABC, ord='nuc')
            elif score_method == 'VAMP2':
                res = np.linalg.norm(ABC, ord='fro')**2
        elif score_method == 'VAMPE':
            Sk = np.diag(self.singular_values[0:self.dimension()])
            res = np.trace(2.0 * mdot(Vk, Sk, Uk.T, test_model.C0t) - mdot(Vk, Sk, Uk.T, test_model.C00, Uk, Sk, Vk.T, test_model.Ctt))
        else:
            raise ValueError('"score" should be one of VAMP1, VAMP2 or VAMPE')
        # add the contribution (+1) of the constant singular functions to the result
        assert res
        return res + 1