jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def lyapunov_svd2(A, C, rtol=1e-4, eps=1e-7, use_svd=False):
    """Solve AX+XA=C using SVD"""

    # This method doesn't work for singular matrices, so regularize it
    # TODO: can optimize performance by reusing eigenvalues from regularization computations
    # A = regularize_mat(A, eps)
    # C = regularize_mat(C, eps)

    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    if use_svd:
        U, S, V = robust_svd(A)
    else:
        S, U = torch.linalg.eigh(A)
    S = S.diag() @ torch.ones_like(A)
    factor = (S + S.t())
    cutoff = max(S) * get_condition(S)
    factor = torch.where(factor > cutoff, 1 / factor, factor)
    X = U @ ((U.t() @ C @ U) * factor) @ U.t()
    error = A @ X + X @ A - C
    relative_error = torch.max(torch.abs(error)) / torch.max(torch.abs(A))
    if relative_error > rtol:
        # TODO(y): currently spams with errors, implement another method based on Newton iteration
        pass
        print(f"Warning, error {relative_error} encountered in lyapunov_svd")

    return X
//  [python]
def truncated_lyapunov_rho(A, C):
    """Returns quantities related to spectrum of solution to AX+XA=2C

    rho: measure of misfit
    erank: effective rank of X
    ."""

    C = 2 * C  # to center spectrum at 1
    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    cond = get_condition(A.dtype)

    S, U = torch.symeig(A, eigenvectors=True)
    S = torch.flip(S, [0])
    U = torch.flip(U, [1])
    cutoff = torch.max(S) * cond
    rank = torch.sum(S > cutoff)
    top_k = rank

    S = S[:top_k].diag()
    S = S @ torch.ones_like(S)

    U = U[:, :top_k]
    projected = U.t() @ C @ U
    divider = S + S.t()
    divider = divider[:top_k, :top_k]

    divided = torch.where(S > 0, projected / divider, torch.zeros_like(projected))
    X = U @ divided @ U.t()
    nan_check(X)

    U, S, V = torch.svd(X)  # can we use symeig here?
    # S = torch.symeig(X).eigenvalues
    # S = u.filter_evals(S)

    erank = torch.sum(S) / torch.max(S)
    rho = A.shape[0] / erank
    spectrum = filter_evals(S)

    return rho, erank, spectrum
//  [python]
def truncated_lyapunov_rho(A, C):
    """Returns quantities related to spectrum of solution to AX+XA=2C

    rho: measure of misfit
    erank: effective rank of X
    ."""

    C = 2 * C  # to center spectrum at 1
    assert A.shape[0] == A.shape[1]
    assert len(A.shape) == 2
    cond = get_condition(A.dtype)

    S, U = torch.linalg.eigh(A)
    S = torch.flip(S, [0])
    U = torch.flip(U, [1])
    cutoff = torch.max(S) * cond
    rank = torch.sum(S > cutoff)
    top_k = rank

    S = S[:top_k].diag()
    S = S @ torch.ones_like(S)

    U = U[:, :top_k]
    projected = U.t() @ C @ U
    divider = S + S.t()
    divider = divider[:top_k, :top_k]

    divided = torch.where(S > 0, projected / divider, torch.zeros_like(projected))
    X = U @ divided @ U.t()
    nan_check(X)

    U, S, V = torch.linalg.svd(X)  # can we use symeig here?
    # S = torch.symeig(X).eigenvalues
    # S = u.filter_evals(S)

    erank = torch.sum(S) / torch.max(S)
    rho = A.shape[0] / erank
    spectrum = filter_evals(S)

    return rho, erank, spectrum
//  [python]
def kl_div_cov(mat1, mat2, eps=1e-3):
    """KL divergence between two zero centered Gaussian's with given covariance matrices."""

    evals1 = torch.symeig(mat1).eigenvalues
    evals2 = torch.symeig(mat2).eigenvalues
    k = mat1.shape[0]
    # scale regularizer in proportion to achievable numerical precision (taken from scipy.pinv2)
    l1 = torch.max(evals1) * k
    l2 = torch.max(evals2) * k
    l = max(l1, l2)
    reg = torch.eye(mat1.shape[0]) * l * eps
    mat1 = mat1 + reg
    mat2 = mat2 + reg

    div = torch.trace(mat1 @ torch.inverse(mat2)) - (torch.logdet(mat1) - torch.logdet(mat2)) - k
    return div
//  [python]
def kl_div_cov(mat1, mat2, eps=1e-3):
    """KL divergence between two zero centered Gaussian's with given covariance matrices."""

    evals1 = torch.linalg.eigh(mat1).eigenvalues
    evals2 = torch.linalg.eigh(mat2).eigenvalues
    k = mat1.shape[0]
    # scale regularizer in proportion to achievable numerical precision (taken from scipy.pinv2)
    l1 = torch.max(evals1) * k
    l2 = torch.max(evals2) * k
    l = max(l1, l2)
    reg = torch.eye(mat1.shape[0]) * l * eps
    mat1 = mat1 + reg
    mat2 = mat2 + reg

    div = torch.trace(mat1 @ torch.inverse(mat2)) - (torch.logdet(mat1) - torch.logdet(mat2)) - k
    return div
//  [python]
def check_eigen_decompose(num_dim, max_iter):
    # random matrix (positive)
    matrix_np = np.abs(np.random.randn(num_dim,num_dim))
    # make symmetric
    matrix_np = matrix_np + matrix_np.T
    matrix_torch = torch.autograd.Variable(torch.from_numpy(matrix_np), requires_grad=True)
    matrix_pi = torch.autograd.Variable(torch.from_numpy(matrix_np), requires_grad=True)

    
    # compute eigenvalue decompositions using numpy
    eigvals_np, eigvecs_np = np.linalg.eigh(matrix_np)
    # compute eigenvalue decompositions using PyTorch
    stime = time.time()
    eigvals_torch, eigvecs_torch = torch.symeig(matrix_torch, eigenvectors=True, upper=True)
    etime = time.time()
    time_eig = etime - stime
    leading_eig_torch = eigvecs_torch[:,-1]
    # print(f"Torch Leading EigenVector {etime-stime:.10f}s:", eigvecs_torch[:,-1].data.numpy())
    assert np.allclose(eigvals_np, eigvals_torch.data.numpy())

    # Power Iteration Algorithm 
    stime = time.time()
    solution = torch.ones_like(matrix_pi[:, 0:1])
    stop_iteration = max_iter
    for i in range(max_iter):
        last_solution = solution
        solution = matrix_pi @ solution
        solution = solution / torch.norm(solution)
        error = float(torch.norm(leading_eig_torch - solution.squeeze(), p=2))
        error2 = float(torch.norm(leading_eig_torch + solution.squeeze(), p=2))
        if torch.norm(solution - last_solution) < 1e-10:
            stop_iteration = i
            etime = time.time()
            time_pi = etime - stime 
            assert np.abs(error) < 1e-6 or np.abs(error2) < 1e-6               
            print(f"Converge at iter-{stop_iteration}")
            break
    if i == max_iter-1:
        etime = time.time()
        time_pi = etime - stime
    solution = solution.squeeze(-1)

    
    ## compute derivative of first eigenvalue with respect to the matrix
    # # Np gradient, see "On differentiating eigenvalues and eigenvectors" by Jan R. Magnus
    # grad_analytic = np.outer(eigvecs_np[:, 0], eigvecs_np[:, 0])
    # # PyTorch gradient
    # eigvals_torch[0].backward()
    # grad_torch = matrix_torch.grad.numpy()
    # assert np.allclose(grad_analytic, grad_torch)

    # comparing the gradient of Power iteration and torch.symeig
    leading_eig_torch[0].backward()
    grad_eig = matrix_torch.grad.numpy()
    solution[0].backward()
    grad_pi = matrix_pi.grad.numpy()
    # print(np.linalg.norm(grad_eig - grad_pi))
    # assert np.allclose(grad_eig, grad_pi)
    return time_eig, time_pi, stop_iteration, float(i!=max_iter-1)
//  [python]
def check_eigen_decompose(num_dim, max_iter):
    # random matrix (positive)
    matrix_np = np.abs(np.random.randn(num_dim,num_dim))
    # make symmetric
    matrix_np = matrix_np + matrix_np.T
    matrix_torch = torch.autograd.Variable(torch.from_numpy(matrix_np), requires_grad=True)
    matrix_pi = torch.autograd.Variable(torch.from_numpy(matrix_np), requires_grad=True)

    
    # compute eigenvalue decompositions using numpy
    eigvals_np, eigvecs_np = np.linalg.eigh(matrix_np)
    # compute eigenvalue decompositions using PyTorch
    stime = time.time()
    eigvals_torch, eigvecs_torch = torch.linalg.eigh(matrix_torch)
    etime = time.time()
    time_eig = etime - stime
    leading_eig_torch = eigvecs_torch[:,-1]
    # print(f"Torch Leading EigenVector {etime-stime:.10f}s:", eigvecs_torch[:,-1].data.numpy())
    assert np.allclose(eigvals_np, eigvals_torch.data.numpy())

    # Power Iteration Algorithm 
    stime = time.time()
    solution = torch.ones_like(matrix_pi[:, 0:1])
    stop_iteration = max_iter
    for i in range(max_iter):
        last_solution = solution
        solution = matrix_pi @ solution
        solution = solution / torch.norm(solution)
        error = float(torch.norm(leading_eig_torch - solution.squeeze(), p=2))
        error2 = float(torch.norm(leading_eig_torch + solution.squeeze(), p=2))
        if torch.norm(solution - last_solution) < 1e-10:
            stop_iteration = i
            etime = time.time()
            time_pi = etime - stime 
            assert np.abs(error) < 1e-6 or np.abs(error2) < 1e-6               
            print(f"Converge at iter-{stop_iteration}")
            break
    if i == max_iter-1:
        etime = time.time()
        time_pi = etime - stime
    solution = solution.squeeze(-1)

    
    ## compute derivative of first eigenvalue with respect to the matrix
    # # Np gradient, see "On differentiating eigenvalues and eigenvectors" by Jan R. Magnus
    # grad_analytic = np.outer(eigvecs_np[:, 0], eigvecs_np[:, 0])
    # # PyTorch gradient
    # eigvals_torch[0].backward()
    # grad_torch = matrix_torch.grad.numpy()
    # assert np.allclose(grad_analytic, grad_torch)

    # comparing the gradient of Power iteration and torch.symeig
    leading_eig_torch[0].backward()
    grad_eig = matrix_torch.grad.numpy()
    solution[0].backward()
    grad_pi = matrix_pi.grad.numpy()
    # print(np.linalg.norm(grad_eig - grad_pi))
    # assert np.allclose(grad_eig, grad_pi)
    return time_eig, time_pi, stop_iteration, float(i!=max_iter-1)
//  [python]
def symsqrt(mat, cond=None, return_rank=False, inverse=False):
    """Computes the symmetric square root of a symmetric matrix. Throws away small and negative eigenvalues."""

    nan_check(mat)
    s, u = torch.symeig(mat, eigenvectors=True)

    # check_symmetric(mat)

    # todo(y): dedupe with getcond
    cond_dict = {torch.float32: 1e3 * 1.1920929e-07, torch.float64: 1E6 * 2.220446049250313e-16}

    if cond in [None, -1]:
        cond = cond_dict[mat.dtype]

    # Note, this can include negative values, see https://github.com/pytorch/pytorch/issues/25972
    above_cutoff = (s > cond * torch.max(abs(s)))

    if torch.sum(above_cutoff) == 0:
        return torch.zeros_like(mat)

    sigma_diag = torch.sqrt(s[above_cutoff])
    if inverse:
        sigma_diag = 1 / sigma_diag
    u = u[:, above_cutoff]

    B = u @ torch.diag(sigma_diag) @ u.t()

    if torch.sum(torch.isnan(B)) > 0:
        if gl.debug_linalg_crashes:
            dump(mat, '/tmp/symsqrt.txt')
            assert False

    if return_rank:
        return B, len(sigma_diag)
    else:
        return B
//  [python]
def symsqrt(mat, cond=None, return_rank=False, inverse=False):
    """Computes the symmetric square root of a symmetric matrix. Throws away small and negative eigenvalues."""

    nan_check(mat)
    s, u = torch.linalg.eigh(mat)

    # check_symmetric(mat)

    # todo(y): dedupe with getcond
    cond_dict = {torch.float32: 1e3 * 1.1920929e-07, torch.float64: 1E6 * 2.220446049250313e-16}

    if cond in [None, -1]:
        cond = cond_dict[mat.dtype]

    # Note, this can include negative values, see https://github.com/pytorch/pytorch/issues/25972
    above_cutoff = (s > cond * torch.max(abs(s)))

    if torch.sum(above_cutoff) == 0:
        return torch.zeros_like(mat)

    sigma_diag = torch.sqrt(s[above_cutoff])
    if inverse:
        sigma_diag = 1 / sigma_diag
    u = u[:, above_cutoff]

    B = u @ torch.diag(sigma_diag) @ u.t()

    if torch.sum(torch.isnan(B)) > 0:
        if gl.debug_linalg_crashes:
            dump(mat, '/tmp/symsqrt.txt')
            assert False

    if return_rank:
        return B, len(sigma_diag)
    else:
        return B
//  [python]
def compute_inv_covariance_v2(centroids, train_features, train_labels, jitters=None):
    jitter = 0
    jitter_eps = None
    if jitters is None:
        jitters = JITTERS

    cov = torch.zeros(
        len(centroids),
        centroids.shape[1],
        centroids.shape[1],
        device=centroids.device,
        dtype=torch.float32,
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        features_of_class_c = train_features[train_labels == c]
        for x in features_of_class_c:
            d = (x - mu_c).unsqueeze(1)
            cov[c] += d @ d.T
        cov[c] /= len(features_of_class_c)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(cov.shape[1], device=cov.device)
        cov_update = cov + jitter
        eigenvalues = torch.symeig(cov_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov = cov + jitter
    cov_inv = torch.empty(cov.shape, dtype=torch.float32, device=cov.device)
    for i in range(len(centroids)):
        cov_inv[i].copy_(torch.inverse(cov[i].to(torch.float64)).float())

    return cov_inv, jitter_eps
//  [python]
def compute_inv_covariance_v2(centroids, train_features, train_labels, jitters=None):
    jitter = 0
    jitter_eps = None
    if jitters is None:
        jitters = JITTERS

    cov = torch.zeros(
        len(centroids),
        centroids.shape[1],
        centroids.shape[1],
        device=centroids.device,
        dtype=torch.float32,
    ).float()
    for c, mu_c in tqdm(enumerate(centroids)):
        features_of_class_c = train_features[train_labels == c]
        for x in features_of_class_c:
            d = (x - mu_c).unsqueeze(1)
            cov[c] += d @ d.T
        cov[c] /= len(features_of_class_c)

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(cov.shape[1], device=cov.device)
        cov_update = cov + jitter
        eigenvalues = torch.linalg.eigh(cov_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov = cov + jitter
    cov_inv = torch.empty(cov.shape, dtype=torch.float32, device=cov.device)
    for i in range(len(centroids)):
        cov_inv[i].copy_(torch.inverse(cov[i].to(torch.float64)).float())

    return cov_inv, jitter_eps
//  [python]
def cpu_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    _,v = torch.symeig(X, True)
    v = v.cuda()   
    return v
//  [python]
def cpu_symeig(X):
    X = X.cpu()
    _,v = torch.linalg.eigh(X)
    v = v.cuda()
    return v
//  [python]
def quaternion_mean_robust(q, w=None):
    r"""
    Calculate the (weighted) average quaternion. (torch)

    This is more robust than quaternion_mean(), but slower.
    Reference: https://ntrs.nasa.gov/api/citations/20070017872/downloads/20070017872.pdf

    :param q: Tensor [N, 4].
    :param w: Tensor [N] for quaternion weights.
    :return: Tensor [4].
    """
    q = q.view(-1, 4)
    w = torch.ones_like(q[:, 0]) / q.shape[0] if w is None else w.view(-1) / w.sum()
    m = (q.unsqueeze(-1).bmm(q.unsqueeze(-2)) * w.view(-1, 1, 1)).sum(dim=0)
    q = torch.symeig(m, eigenvectors=True).eigenvectors[:, -1]
    return q
//  [python]
def quaternion_mean_robust(q, w=None):
    r"""
    Calculate the (weighted) average quaternion. (torch)

    This is more robust than quaternion_mean(), but slower.
    Reference: https://ntrs.nasa.gov/api/citations/20070017872/downloads/20070017872.pdf

    :param q: Tensor [N, 4].
    :param w: Tensor [N] for quaternion weights.
    :return: Tensor [4].
    """
    q = q.view(-1, 4)
    w = torch.ones_like(q[:, 0]) / q.shape[0] if w is None else w.view(-1) / w.sum()
    m = (q.unsqueeze(-1).bmm(q.unsqueeze(-2)) * w.view(-1, 1, 1)).sum(dim=0)
    q = torch.linalg.eigh(m).eigenvectors[:, -1]
    return q
//  [python]
def log_coral_loss(input1, input2, gamma=1e-3):
    ''' Implementation of eig version coral covariances (which is regularized)
   
    Note:
    - Return the covariances of all features. That's [f*f]
    - Send gradients to both input1 and input2
    '''
    #First: subtract the mean from the data matrix
    batch_size = input1.shape[0].double()
    h_src = input1 - torch.mean(input1, dim=0)
    h_trg = input2 - torch.mean(input2, dim=0)

    cov_src = (1./(batch_size-1)) * torch.mm(torch.transpose(h_src), h_src)
    cov_trg = (1./(batch_size-1)) * torch.mm(torch.transpose(h_trg), h_trg)
   
    #eigen decomposition
    e_src, v_src = torch.symeig(cov_src, eigenvectors=True)
    e_trg, v_trg = torch.symeig(cov_trg, eigenvectors=True)
   
	# Returns the Frobenius norm
    log_cov_src = torch.mm(e_src,  torch.mm(torch.diag(torch.log(e_src+0.0001)), torch.transpose(v_src)))
    log_cov_trg = torch.mm(e_trg,  torch.mm(torch.diag(torch.log(e_trg+0.0001)), torch.transpose(v_trg)))
   
    return torch.mean(torch.pow(log_cov_src-log_cov_trg, 2))
//  [python]
def log_coral_loss(input1, input2, gamma=1e-3):
    ''' Implementation of eig version coral covariances (which is regularized)
   
    Note:
    - Return the covariances of all features. That's [f*f]
    - Send gradients to both input1 and input2
    '''
    #First: subtract the mean from the data matrix
    batch_size = input1.shape[0].double()
    h_src = input1 - torch.mean(input1, dim=0)
    h_trg = input2 - torch.mean(input2, dim=0)

    cov_src = (1./(batch_size-1)) * torch.mm(torch.transpose(h_src), h_src)
    cov_trg = (1./(batch_size-1)) * torch.mm(torch.transpose(h_trg), h_trg)
   
    #eigen decomposition
    e_src, v_src = torch.linalg.eigh(cov_src)
    e_trg, v_trg = torch.linalg.eigh(cov_trg)
   
	# Returns the Frobenius norm
    log_cov_src = torch.mm(e_src,  torch.mm(torch.diag(torch.log(e_src+0.0001)), torch.transpose(v_src)))
    log_cov_trg = torch.mm(e_trg,  torch.mm(torch.diag(torch.log(e_trg+0.0001)), torch.transpose(v_trg)))
   
    return torch.mean(torch.pow(log_cov_src-log_cov_trg, 2))
//  [python]
def diversity_metrics(matrix):
    #find covariance matrix
    new_mat = matrix.view(opt.batchSize, -1)
#    print('SIZE: ', new_mat.size(), new_mat.min(), new_mat.max())
    covariance_mat = _cov(new_mat, rowvar=True)
#    print('SIZE 2: ', covariance_mat.size())

    eigvals_cov, vecs = torch.symeig(covariance_mat, eigenvectors=True, upper=False)
#    print('SIZE 2: ', eigvals_cov.size())
#    determinant_cov = torch.cholesky(covariance_mat_for_det).diag().prod()
#    vals=[]
#    for _ in range(opt.batchSize):
#        vals.append(matrix[_, :, :, :].reshape)
#    np.cov(vals)
    max_eig = torch.max(eigvals_cov)
    condition_number = torch.max(eigvals_cov)/torch.min(eigvals_cov)
    return torch.tensor(100000000000)/max_eig, torch.abs(condition_number-1)/torch.tensor(1000000000000000)
//  [python]
def diversity_metrics(matrix):
    #find covariance matrix
    new_mat = matrix.view(opt.batchSize, -1)
#    print('SIZE: ', new_mat.size(), new_mat.min(), new_mat.max())
    covariance_mat = _cov(new_mat, rowvar=True)
#    print('SIZE 2: ', covariance_mat.size())

    eigvals_cov, vecs = torch.linalg.eigh(covariance_mat)
#    print('SIZE 2: ', eigvals_cov.size())
#    determinant_cov = torch.cholesky(covariance_mat_for_det).diag().prod()
#    vals=[]
#    for _ in range(opt.batchSize):
#        vals.append(matrix[_, :, :, :].reshape)
#    np.cov(vals)
    max_eig = torch.max(eigvals_cov)
    condition_number = torch.max(eigvals_cov)/torch.min(eigvals_cov)
    return torch.tensor(100000000000)/max_eig, torch.abs(condition_number-1)/torch.tensor(1000000000000000)
//  [python]
def step(self, objective_f, rank_order='max'):

        # sample parameters
        params, z = sample(self.samples + floor(self.oversample), 1.0, self.mean, self.b, self.d)
        self.oversample = self.oversample * 0.8

        # rank by fitness
        f = objective_f(params)
        results = [{'parameters': params[i], 'z': z[i], 'fitness': f.item()} for i, f in enumerate(f)]
        ranked_results = self._rank(results, rank_order)

        selected_results = ranked_results[0:self.mu]
        g = torch.stack([g['parameters'] for g in selected_results])
        z = torch.stack([g['z'] for g in selected_results])

        self.mean = g.mean(0)
        bdz = self.b.matmul(self.d).matmul(z.t())
        c_mu = torch.matmul(bdz, torch.eye(self.mu) / self.mu)
        c_mu = c_mu.matmul(bdz.t())

        self.c = (1 - self.cmu) * self.c + self.cmu * c_mu

        self.d, self.b = torch.symeig(self.c, eigenvectors=True)
        self.d = self.d.sqrt().diag_embed()

        info = {'fitness_max': f.max(), 'fitness_mean': f.mean(), 'c_norm': self.c.norm(), 'max_eigv': self.d.max()}
        self.gen_count += 1

        return ranked_results, info
//  [python]
def step(self, objective_f, rank_order='max'):

        # sample parameters
        params, z = sample(self.samples + floor(self.oversample), 1.0, self.mean, self.b, self.d)
        self.oversample = self.oversample * 0.8

        # rank by fitness
        f = objective_f(params)
        results = [{'parameters': params[i], 'z': z[i], 'fitness': f.item()} for i, f in enumerate(f)]
        ranked_results = self._rank(results, rank_order)

        selected_results = ranked_results[0:self.mu]
        g = torch.stack([g['parameters'] for g in selected_results])
        z = torch.stack([g['z'] for g in selected_results])

        self.mean = g.mean(0)
        bdz = self.b.matmul(self.d).matmul(z.t())
        c_mu = torch.matmul(bdz, torch.eye(self.mu) / self.mu)
        c_mu = c_mu.matmul(bdz.t())

        self.c = (1 - self.cmu) * self.c + self.cmu * c_mu

        self.d, self.b = torch.linalg.eigh(self.c)
        self.d = self.d.sqrt().diag_embed()

        info = {'fitness_max': f.max(), 'fitness_mean': f.mean(), 'c_norm': self.c.norm(), 'max_eigv': self.d.max()}
        self.gen_count += 1

        return ranked_results, info
//  [python]
def estimate_normals_torch(inputpc, max_nn):
    knn = knn_graph(inputpc[:, :3], max_nn, loop=False)
    knn = knn.view(2, inputpc.shape[0], max_nn)[0]
    x = inputpc[knn][:, :, :3]
    temp = x[:, :, :3] - x.mean(dim=1)[:, None, :3]
    cov = temp.transpose(1, 2) @ temp / x.shape[0]
    e, v = torch.symeig(cov, eigenvectors=True)
    n = v[:, :, 0]
    return torch.cat([inputpc[:, :3], n], dim=-1)
//  [python]
def estimate_normals_torch(inputpc, max_nn):
    knn = knn_graph(inputpc[:, :3], max_nn, loop=False)
    knn = knn.view(2, inputpc.shape[0], max_nn)[0]
    x = inputpc[knn][:, :, :3]
    temp = x[:, :, :3] - x.mean(dim=1)[:, None, :3]
    cov = temp.transpose(1, 2) @ temp / x.shape[0]
    e, v = torch.linalg.eigh(cov)
    n = v[:, :, 0]
    return torch.cat([inputpc[:, :3], n], dim=-1)
//  [python]
def estimate_normals(points,k):
    batch_size = points.size(0)
    points_mean = points.mean(1,keepdim=True) 
    points_centered = points - points_mean
    knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k)
    knn_indices = torch.LongTensor(knn_indices).to(points.device)
    knn_points = []
    for b in range(batch_size):   
        knn_points.append(points[b,knn_indices[b]].unsqueeze(0))
    knn_points = torch.cat(knn_points,0)
    pt_mean = knn_points.mean(2, keepdim=True)
    central_diff = knn_points - pt_mean
    per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3)
    covariances = per_pt_cov.mean(2)    
    print(covariances.device)
    curvatures, local_coord_frames = torch.symeig(covariances, eigenvectors=True)    
    print(curvatures.shape)
    normals = local_coord_frames[:, :, :, 0]
    print(normals.shape)
    return normals, knn_indices
//  [python]
def estimate_normals(points,k):
    batch_size = points.size(0)
    points_mean = points.mean(1,keepdim=True) 
    points_centered = points - points_mean
    knn_indices, _ = get_nearest_neighbors_indices_batch(points_centered.detach().cpu().numpy(),points_centered.detach().cpu().numpy(),k)
    knn_indices = torch.LongTensor(knn_indices).to(points.device)
    knn_points = []
    for b in range(batch_size):   
        knn_points.append(points[b,knn_indices[b]].unsqueeze(0))
    knn_points = torch.cat(knn_points,0)
    pt_mean = knn_points.mean(2, keepdim=True)
    central_diff = knn_points - pt_mean
    per_pt_cov = central_diff.unsqueeze(4) * central_diff.unsqueeze(3)
    covariances = per_pt_cov.mean(2)    
    print(covariances.device)
    curvatures, local_coord_frames = torch.linalg.eigh(covariances)    
    print(curvatures.shape)
    normals = local_coord_frames[:, :, :, 0]
    print(normals.shape)
    return normals, knn_indices
//  [python]
def _simplegraph_cora(self, edges):
        graph = defaultdict(list)

        for edge in edges:  # 这里的key和value均存的是node_name(str)，不是node_index(按照content的顺序从0开始编号，int)
            graph[edge[0]].append(edge[1])
            graph[edge[1]].append(edge[0])
        # 去重复，和loop,求D_v
        node_degree_flat = torch.zeros(self.nodes_number)
        A = torch.zeros((self.nodes_number, self.nodes_number))
        for item in graph.items():
            if item[0] in item[1]:
                item[1].remove(item[0])
            graph[item[0]] = np.unique(item[1])
            node_degree_flat[self.nodes_names_map[item[0]]] = len(graph[item[0]])
            for node in graph[item[0]]:
                A[self.nodes_names_map[item[0]], self.nodes_names_map[node]] = 1

        # laplacian
        node_degree_flat_pow = torch.pow(node_degree_flat, -0.5)
        node_degree_flat_pow[node_degree_flat_pow == float("Inf")] = 0  # replace inf with 0
        node_degree_flat_pow[node_degree_flat_pow != node_degree_flat_pow] = 0  # replace nan with 0

        L = torch.eye(self.nodes_number) - torch.diag(node_degree_flat_pow) @ A @ torch.diag(node_degree_flat_pow)
        fourier_e, fourier_v = torch.symeig(L, eigenvectors=True)

        # wavelets
        wavelets = fourier_v @ torch.diag(torch.exp(-1.0 * fourier_e * self.s)) @ torch.transpose(fourier_v, 0, 1)
        wavelets_inv = fourier_v @ torch.diag(torch.exp(fourier_e * self.s)) @ torch.transpose(fourier_v, 0, 1)
        wavelets[wavelets < 0.00001] = 0
        wavelets_inv[wavelets_inv < 0.00001] = 0

        # metapath_list
        # metapath_list=[[0,0]]
        # node_type_map
        node_type_list = [0]
        node_type_map = defaultdict(int)  # {'1':0}
        type_node_map = defaultdict(list)  # {0:['1','3','4']}
        for node, _ in graph.items():
            type_node_map[0].append(node)

        for item in graph.items():
            node_type_map[item[0]] = 0

        simple_graph = {"graph": graph,  # {'1':['2','6','10'],}
                        "edges": edges,  # numpy array
                        "node_type_list": node_type_list,
                        # "metapath_list":metapath_list,
                        "node_degree_flat": node_degree_flat,
                        "node_type_map": node_type_map,
                        "type_node_map": type_node_map,
                        "adj": A,
                        "laplacian": L,
                        "fourier_e": fourier_e,
                        "fourier_v": fourier_v,
                        "wavelets": wavelets,
                        "wavelets_inv": wavelets_inv}
        return simple_graph
//  [python]
def _simplegraph_cora(self, edges):
        graph = defaultdict(list)

        for edge in edges:  # 这里的key和value均存的是node_name(str)，不是node_index(按照content的顺序从0开始编号，int)
            graph[edge[0]].append(edge[1])
            graph[edge[1]].append(edge[0])
        # 去重复，和loop,求D_v
        node_degree_flat = torch.zeros(self.nodes_number)
        A = torch.zeros((self.nodes_number, self.nodes_number))
        for item in graph.items():
            if item[0] in item[1]:
                item[1].remove(item[0])
            graph[item[0]] = np.unique(item[1])
            node_degree_flat[self.nodes_names_map[item[0]]] = len(graph[item[0]])
            for node in graph[item[0]]:
                A[self.nodes_names_map[item[0]], self.nodes_names_map[node]] = 1

        # laplacian
        node_degree_flat_pow = torch.pow(node_degree_flat, -0.5)
        node_degree_flat_pow[node_degree_flat_pow == float("Inf")] = 0  # replace inf with 0
        node_degree_flat_pow[node_degree_flat_pow != node_degree_flat_pow] = 0  # replace nan with 0

        L = torch.eye(self.nodes_number) - torch.diag(node_degree_flat_pow) @ A @ torch.diag(node_degree_flat_pow)
        fourier_e, fourier_v = torch.linalg.eigh(L)

        # wavelets
        wavelets = fourier_v @ torch.diag(torch.exp(-1.0 * fourier_e * self.s)) @ torch.transpose(fourier_v, 0, 1)
        wavelets_inv = fourier_v @ torch.diag(torch.exp(fourier_e * self.s)) @ torch.transpose(fourier_v, 0, 1)
        wavelets[wavelets < 0.00001] = 0
        wavelets_inv[wavelets_inv < 0.00001] = 0

        # metapath_list
        # metapath_list=[[0,0]]
        # node_type_map
        node_type_list = [0]
        node_type_map = defaultdict(int)  # {'1':0}
        type_node_map = defaultdict(list)  # {0:['1','3','4']}
        for node, _ in graph.items():
            type_node_map[0].append(node)

        for item in graph.items():
            node_type_map[item[0]] = 0

        simple_graph = {"graph": graph,  # {'1':['2','6','10'],}
                        "edges": edges,  # numpy array
                        "node_type_list": node_type_list,
                        # "metapath_list":metapath_list,
                        "node_degree_flat": node_degree_flat,
                        "node_type_map": node_type_map,
                        "type_node_map": type_node_map,
                        "adj": A,
                        "laplacian": L,
                        "fourier_e": fourier_e,
                        "fourier_v": fourier_v,
                        "wavelets": wavelets,
                        "wavelets_inv": wavelets_inv}
        return simple_graph
//  [python]
def closure():
    optimizer.zero_grad()
    Asymm = symmetrize(A)

    T1 = torch.einsum('xa,xby,yc',(Asymm,T,Asymm)).view(D,D,d,d,D,D).permute(0,2,4,1,3,5).contiguous().view(D**2*d, D**2*d)
    T2 = torch.einsum('xa,xb',(Asymm,Asymm)).view(D,D,D,D).permute(0,2,1,3).contiguous().view(D**2,D**2)

    eigval_1,_ = torch.symeig(T1,eigenvectors=True)
    eigval_2,_ = torch.symeig(T2,eigenvectors=True)

    lnT = torch.log(eigval_1.max())
    lnZ = torch.log(eigval_2.max())
    loss = -lnT + lnZ
    loss.backward()
    return loss
//  [python]
def closure():
    optimizer.zero_grad()
    Asymm = symmetrize(A)

    T1 = torch.einsum('xa,xby,yc',(Asymm,T,Asymm)).view(D,D,d,d,D,D).permute(0,2,4,1,3,5).contiguous().view(D**2*d, D**2*d)
    T2 = torch.einsum('xa,xb',(Asymm,Asymm)).view(D,D,D,D).permute(0,2,1,3).contiguous().view(D**2,D**2)

    eigval_1,_ = torch.linalg.eigh(T1)
    eigval_2,_ = torch.linalg.eigh(T2)

    lnT = torch.log(eigval_1.max())
    lnZ = torch.log(eigval_2.max())
    loss = -lnT + lnZ
    loss.backward()
    return loss
//  [python]
def power_by_eig_symmetric(tensor: torch.Tensor, damping: torch.Tensor, exp: float) -> torch.Tensor:
    # TODO: If we have multiple exponents and/or dampings we could save the eigenvalues
    eigenvalues, eigenvectors = torch.symeig(tensor, eigenvectors=True)
    return (eigenvectors * (eigenvalues + damping)**exp) @ eigenvectors.T
//  [python]
def power_by_eig_symmetric(tensor: torch.Tensor, damping: torch.Tensor, exp: float) -> torch.Tensor:
    # TODO: If we have multiple exponents and/or dampings we could save the eigenvalues
    eigenvalues, eigenvectors = torch.linalg.eigh(tensor)
    return (eigenvectors * (eigenvalues + damping)**exp) @ eigenvectors.T
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment, has " \
                               "{}".format(len(list(m.parameters())))
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            if p not in updates: continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            if p not in updates: continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment, has " \
                               "{}".format(len(list(m.parameters())))
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay

            if self.steps % self.Tf == 0:
                # My asynchronous implementation exists, I will add it later.
                # Experimenting with different ways to this in PyTorch.
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(
                    self.m_aa[m])
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
                    self.m_gg[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()

            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            if p not in updates: continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()

        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            if p not in updates: continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def test_fast_symeig3x3(seed, rand_sym, n):
    x = rand_sym(n, 3)
    assert_allclose(fast.symeig3x3(x), torch.symeig(x).eigenvalues, atol=1e-4)
//  [python]
def test_fast_symeig3x3(seed, rand_sym, n):
    x = rand_sym(n, 3)
    assert_allclose(fast.symeig3x3(x), torch.linalg.eigh(x).eigenvalues, atol=1e-4)
//  [python]
def test_fast_symeig2x2(seed, rand_sym, n):
    x = rand_sym(n, 2)
    assert_allclose(fast.symeig2x2(x), torch.symeig(x).eigenvalues, atol=1e-4)
//  [python]
def test_fast_symeig2x2(seed, rand_sym, n):
    x = rand_sym(n, 2)
    assert_allclose(fast.symeig2x2(x), torch.linalg.eigh(x).eigenvalues, atol=1e-4)
//  [python]
def propagate_main_step(self, system: System):
        """
        Main routine for propagating the system positions and cells. Since this is modified, no conventional velocity
        verlet integrator can be used.

        Args:
            system (schnetpack.md.System): System class containing all molecules and their
                             replicas.
        """
        eigval_b_velocities, eigvec_b_velocities = torch.symeig(
            self.b_velocities_cell, eigenvectors=True
        )

        scaled_velocity = eigval_b_velocities * self.time_step
        coeff_a = torch.exp(0.5 * scaled_velocity)[..., None] * self.aux_eye
        coeff_b = coeff_a * self.sinhdx.f(0.5 * scaled_velocity)[..., None]

        # Construct matrix operators and update positions using positions and momenta
        operator_a = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_a**2, eigvec_b_velocities.transpose(2, 3)),
        )
        operator_b = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_b, eigvec_b_velocities.transpose(2, 3)),
        )

        update_positions = torch.sum(
            system.expand_atoms(operator_a) * system.positions[..., None], dim=2
        )
        update_momenta = torch.sum(
            system.expand_atoms(operator_b)
            * (system.momenta / system.masses)[..., None],
            dim=2,
        )

        system.positions = update_positions + update_momenta * self.time_step

        # Update cells using first operator
        system.cells = torch.matmul(system.cells, operator_a)
//  [python]
def propagate_main_step(self, system: System):
        """
        Main routine for propagating the system positions and cells. Since this is modified, no conventional velocity
        verlet integrator can be used.

        Args:
            system (schnetpack.md.System): System class containing all molecules and their
                             replicas.
        """
        eigval_b_velocities, eigvec_b_velocities = torch.linalg.eigh(
            self.b_velocities_cell
        )

        scaled_velocity = eigval_b_velocities * self.time_step
        coeff_a = torch.exp(0.5 * scaled_velocity)[..., None] * self.aux_eye
        coeff_b = coeff_a * self.sinhdx.f(0.5 * scaled_velocity)[..., None]

        # Construct matrix operators and update positions using positions and momenta
        operator_a = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_a**2, eigvec_b_velocities.transpose(2, 3)),
        )
        operator_b = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_b, eigvec_b_velocities.transpose(2, 3)),
        )

        update_positions = torch.sum(
            system.expand_atoms(operator_a) * system.positions[..., None], dim=2
        )
        update_momenta = torch.sum(
            system.expand_atoms(operator_b)
            * (system.momenta / system.masses)[..., None],
            dim=2,
        )

        system.positions = update_positions + update_momenta * self.time_step

        # Update cells using first operator
        system.cells = torch.matmul(system.cells, operator_a)
//  [python]
def make_symposdef(mat, eval_threshold=1e-1):
    '''Enforce a matrix to be symmetric and positive definite.

    Args:
        mat (``torch.Tensor[dim, dim]``): Input matrix.
        eval_threshold (float): Minimum value of the eigen values of
            the matrix.

    Returns:
        ``torch.Tensor[dim, dim]``

    '''
    sym_mat = symmetrize_matrix(mat)
    evals, evecs = torch.symeig(sym_mat, eigenvectors=True)

    threshold = torch.tensor(eval_threshold, dtype=sym_mat.dtype,
                             device=sym_mat.device)
    new_evals = torch.where(
        evals < threshold * evals[-1],
        threshold * evals[-1],
        evals
    )
    return (evecs @ torch.diag(new_evals) @ evecs.t()).view(*mat.shape)
//  [python]
def make_symposdef(mat, eval_threshold=1e-1):
    '''Enforce a matrix to be symmetric and positive definite.

    Args:
        mat (``torch.Tensor[dim, dim]``): Input matrix.
        eval_threshold (float): Minimum value of the eigen values of
            the matrix.

    Returns:
        ``torch.Tensor[dim, dim]``

    '''
    sym_mat = symmetrize_matrix(mat)
    evals, evecs = torch.linalg.eigh(sym_mat)

    threshold = torch.tensor(eval_threshold, dtype=sym_mat.dtype,
                             device=sym_mat.device)
    new_evals = torch.where(
        evals < threshold * evals[-1],
        threshold * evals[-1],
        evals
    )
    return (evecs @ torch.diag(new_evals) @ evecs.t()).view(*mat.shape)
//  [python]
def test_eig_gradients(seed, rand_sym, n, d, eig):
    x1 = rand_sym(n, d).requires_grad_()
    s1 = eig(x1).pow(2).sum()
    s1.backward()
    x2 = x1.detach().clone().requires_grad_()
    s2 = torch.symeig(x2, eigenvectors=True).eigenvalues.pow(2).sum()
    s2.backward()
    assert_allclose(tb.sym(x1.grad), tb.sym(x2.grad), atol=1e-4)
//  [python]
def test_eig_gradients(seed, rand_sym, n, d, eig):
    x1 = rand_sym(n, d).requires_grad_()
    s1 = eig(x1).pow(2).sum()
    s1.backward()
    x2 = x1.detach().clone().requires_grad_()
    s2 = torch.linalg.eigh(x2).eigenvalues.pow(2).sum()
    s2.backward()
    assert_allclose(tb.sym(x1.grad), tb.sym(x2.grad), atol=1e-4)
//  [python]
def pairloss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

    H1, H2 = H1.t(), H2.t()

    o1 = o2 = H1.size(0)

    m = H1.size(1)

    H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
    H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

    SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
    SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar, H1bar.t()) + r1 * torch.eye(o1, device=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

    # Calculating the root inverse of covariance matrices by using eigen decomposition
    [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
    [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

    # Added to increase stability
    posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
    D1 = D1[posInd1]
    V1 = V1[:, posInd1]
    posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
    D2 = D2[posInd2]
    V2 = V2[:, posInd2]

    SigmaHat11RootInv = torch.matmul(torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
    SigmaHat22RootInv = torch.matmul(torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

    Tval = torch.matmul(torch.matmul(SigmaHat11RootInv, SigmaHat12), SigmaHat22RootInv)

    if use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.symeig(trace_TT, eigenvectors=True)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def pairloss(H1, H2, outdim_size, use_all_singular_values=False):
    r1 = 1e-3
    r2 = 1e-3
    eps = 1e-9

    H1, H2 = H1.t(), H2.t()

    o1 = o2 = H1.size(0)

    m = H1.size(1)

    H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
    H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

    SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
    SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar, H1bar.t()) + r1 * torch.eye(o1, device=H1.device)
    SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar, H2bar.t()) + r2 * torch.eye(o2, device=H1.device)

    # Calculating the root inverse of covariance matrices by using eigen decomposition
    [D1, V1] = torch.linalg.eigh(SigmaHat11)
    [D2, V2] = torch.linalg.eigh(SigmaHat22)

    # Added to increase stability
    posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
    D1 = D1[posInd1]
    V1 = V1[:, posInd1]
    posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
    D2 = D2[posInd2]
    V2 = V2[:, posInd2]

    SigmaHat11RootInv = torch.matmul(torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
    SigmaHat22RootInv = torch.matmul(torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

    Tval = torch.matmul(torch.matmul(SigmaHat11RootInv, SigmaHat12), SigmaHat22RootInv)

    if use_all_singular_values:
        # all singular values are used to calculate the correlation
        tmp = torch.matmul(Tval.t(), Tval)
        corr = torch.trace(torch.sqrt(tmp))
        # assert torch.isnan(corr).item() == 0
    else:
        # just the top outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(H1.device)) # regularization for more stability
        U, V = torch.linalg.eigh(trace_TT)
        U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(H1.device))
        U = U.topk(outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))

    return - corr
//  [python]
def PCA(Z,n=3):
  zm = torch.mean(Z,1,keepdim=True)
  Z2 = Z - zm
  zm = torch.mean(Z2,0,keepdim=True)
  zc = torch.matmul(torch.transpose(zm, 0, -1),zm)
  e, v = torch.symeig(zc, eigenvectors=True)
  r=range(len(e)-1,len(e)-n-1,-1)
  v = v[r]
  return torch.matmul(Z ,torch.transpose(v, 0, -1))
//  [python]
def PCA(Z,n=3):
  zm = torch.mean(Z,1,keepdim=True)
  Z2 = Z - zm
  zm = torch.mean(Z2,0,keepdim=True)
  zc = torch.matmul(torch.transpose(zm, 0, -1),zm)
  e, v = torch.linalg.eigh(zc)
  r=range(len(e)-1,len(e)-n-1,-1)
  v = v[r]
  return torch.matmul(Z ,torch.transpose(v, 0, -1))
//  [python]
def propagate_half_step(self, system: System):
        """
        Main routine for propagating the system momenta. Since this is modified, no conventional velocity
        verlet integrator can be used.

        Args:
            system (schnetpack.md.System): System class containing all molecules and their
                             replicas.
        """
        eigval_b_velocities, eigvec_b_velocities = torch.symeig(
            self.b_velocities_cell, eigenvectors=True
        )
        # Trace of matrix is sum of eigenvalues
        trace_b_velocities = torch.sum(eigval_b_velocities, dim=2, keepdim=True)

        scaled_velocity = (
            (eigval_b_velocities + trace_b_velocities / self.degrees_of_freedom)
            * self.time_step
            * 0.5
        )

        coeff_a = torch.exp(-0.5 * scaled_velocity)[..., None] * self.aux_eye
        coeff_b = coeff_a * self.sinhdx.f(0.5 * scaled_velocity)[..., None]

        # Construct matrix operators and update positions using positions and momenta
        operator_a = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_a**2, eigvec_b_velocities.transpose(2, 3)),
        )
        operator_b = torch.matmul(
            eigvec_b_velocities,
            torch.matmul(coeff_b, eigvec_b_velocities.transpose(2, 3)),
        )

        update_momenta = torch.sum(
            system.expand_atoms(operator_a) * system.momenta[..., None], dim=2
        )
        update_forces = torch.sum(
            system.expand_atoms(operator_b) * system.forces[..., None], dim=2
        )

        system.momenta = update_momenta + 0.5 * self.time_step * update_forces