jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.eigh(
            self.b_velocities_cell
        )
        # 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
//  [python]
def forward(self, real, fake):
        S_b = torch.mm(real.transpose(0, 1), real)
        D_b = torch.mm(fake.transpose(0, 1), fake)

        lambda_real, v_real = torch.symeig(S_b, eigenvectors=True)
        lambda_fake, v_fake = torch.symeig(D_b, eigenvectors=True)

        # throw away the imaginary values RIP
        # lambda_real = lambda_real[:, 0]
        # lambda_fake = lambda_fake[:, 0]

        lambda_real_norm = lambda_real / lambda_real.sum(0).expand_as(lambda_real)

        L_m = torch.norm((lambda_real - lambda_fake), 2).sum()

        L_s = - (lambda_real_norm * F.cosine_similarity(v_real, v_fake)).sum()

        return L_m + L_s
//  [python]
def forward(self, real, fake):
        S_b = torch.mm(real.transpose(0, 1), real)
        D_b = torch.mm(fake.transpose(0, 1), fake)

        lambda_real, v_real = torch.linalg.eigh(S_b)
        lambda_fake, v_fake = torch.linalg.eigh(D_b)

        # throw away the imaginary values RIP
        # lambda_real = lambda_real[:, 0]
        # lambda_fake = lambda_fake[:, 0]

        lambda_real_norm = lambda_real / lambda_real.sum(0).expand_as(lambda_real)

        L_m = torch.norm((lambda_real - lambda_fake), 2).sum()

        L_s = - (lambda_real_norm * F.cosine_similarity(v_real, v_fake)).sum()

        return L_m + L_s
//  [python]
def _null_space(m, kernel_dim):
    """Finds the null space (kernel) basis of the matrix
    Args:
        m: the batch of input matrices, B x N x 12
        kernel_dim: number of dimensions to approximate the kernel
    Returns:
        * a batch of null space basis vectors
            of size B x 4 x 3 x kernel_dim
        * a batch of spectral values where near-0s correspond to actual
            kernel vectors, of size B x kernel_dim
    """
    mTm = torch.bmm(m.transpose(1, 2), m)

    s, v = torch.symeig(mTm, eigenvectors=True)
    return v[:, :, :kernel_dim].reshape(-1, 4, 3, kernel_dim), s[:, :kernel_dim]
//  [python]
def _null_space(m, kernel_dim):
    """Finds the null space (kernel) basis of the matrix
    Args:
        m: the batch of input matrices, B x N x 12
        kernel_dim: number of dimensions to approximate the kernel
    Returns:
        * a batch of null space basis vectors
            of size B x 4 x 3 x kernel_dim
        * a batch of spectral values where near-0s correspond to actual
            kernel vectors, of size B x kernel_dim
    """
    mTm = torch.bmm(m.transpose(1, 2), m)

    s, v = torch.linalg.eigh(mTm)
    return v[:, :, :kernel_dim].reshape(-1, 4, 3, kernel_dim), s[:, :kernel_dim]
//  [python]
def compute_gdpp(phi_real, phi_fake):
    def compute_diversity(phi):
        phi = f.normalize(phi, p=2, dim=1)
        S_B = torch.mm(phi, phi.t())
        eig_vals, eig_vecs = torch.symeig(S_B, eigenvectors=True)
        # eig_vals, eig_vecs = torch.eig(S_B, eigenvectors=True)
        return eig_vals, eig_vecs
        # return eig_vals[:, 0], eig_vecs

    def normalize_min_max(eig_vals):
        min_v, max_v = torch.min(eig_vals), torch.max(eig_vals)
        return (eig_vals - min_v) / (max_v - min_v)

    fake_eig_vals, fake_eig_vecs = compute_diversity(phi_fake)
    real_eig_vals, real_eig_vecs = compute_diversity(phi_real)
    # Scaling factor to make the two losses operating in comparable ranges.
    magnitude_loss = 0.0001 * f.mse_loss(target=real_eig_vals, input=fake_eig_vals)
    structure_loss = -torch.sum(torch.mul(fake_eig_vecs, real_eig_vecs), 0)
    normalized_real_eig_vals = normalize_min_max(real_eig_vals)
    weighted_structure_loss = torch.sum(torch.mul(normalized_real_eig_vals, structure_loss))
    return magnitude_loss + weighted_structure_loss
//  [python]
def compute_gdpp(phi_real, phi_fake):
    def compute_diversity(phi):
        phi = f.normalize(phi, p=2, dim=1)
        S_B = torch.mm(phi, phi.t())
        eig_vals, eig_vecs = torch.linalg.eigh(S_B)
        return eig_vals, eig_vecs

    def normalize_min_max(eig_vals):
        min_v, max_v = torch.min(eig_vals), torch.max(eig_vals)
        return (eig_vals - min_v) / (max_v - min_v)

    fake_eig_vals, fake_eig_vecs = compute_diversity(phi_fake)
    real_eig_vals, real_eig_vecs = compute_diversity(phi_real)
    magnitude_loss = 0.0001 * f.mse_loss(target=real_eig_vals, input=fake_eig_vals)
    structure_loss = -torch.sum(torch.mul(fake_eig_vecs, real_eig_vecs), 0)
    normalized_real_eig_vals = normalize_min_max(real_eig_vals)
    weighted_structure_loss = torch.sum(torch.mul(normalized_real_eig_vals, structure_loss))
    return magnitude_loss + weighted_structure_loss
//  [python]
def forward(self, x):
        """
        This function computes the forward pass throughout the network

        Parameters
        ----------
        :param x: input data

        Returns
        -------
        :return: output of the FullyConnectedNet given x
        """
        ldata = self.model(x)
        ldata = nn.Tanh()(ldata)

        # Build lower triangular matrix
        L = fill_tril(ldata)
        # Symmetric matrix
        Q = L + L.transpose(-1, -2)
        # SPD matrix using Cholesky decomposition
        # Q = torch.matmul(L, torch.transpose(L, -1, -2))

        # Compute the eigenvalue decomposition
        eigenvalues, eigenvectors = torch.symeig(Q, eigenvectors=True)
        # Softmax to constraint the eigenvalue of the matrix, i.e, norm(Q, fro) = 1
        eigenvalues = torch.sqrt(nn.Softmax(-1)(torch.exp(eigenvalues)**2))
        Q = torch.matmul(eigenvectors, torch.matmul(torch.diag_embed(eigenvalues), torch.inverse(eigenvectors)))

        return Q
//  [python]
def forward(self, x):
        """
        This function computes the forward pass throughout the network

        Parameters
        ----------
        :param x: input data

        Returns
        -------
        :return: output of the FullyConnectedNet given x
        """
        ldata = self.model(x)
        ldata = nn.Tanh()(ldata)

        # Build lower triangular matrix
        L = fill_tril(ldata)
        # Symmetric matrix
        Q = L + L.transpose(-1, -2)
        # SPD matrix using Cholesky decomposition
        # Q = torch.matmul(L, torch.transpose(L, -1, -2))

        # Compute the eigenvalue decomposition
        eigenvalues, eigenvectors = torch.linalg.eigh(Q)
        # Softmax to constraint the eigenvalue of the matrix, i.e, norm(Q, fro) = 1
        eigenvalues = torch.sqrt(nn.Softmax(-1)(torch.exp(eigenvalues)**2))
        Q = torch.matmul(eigenvectors, torch.matmul(torch.diag_embed(eigenvalues), torch.inverse(eigenvectors)))

        return Q
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    # print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        # pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  # (B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  # (B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  # (B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask > 1] = 1

        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.symeig(L, eigenvectors=True)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    # print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        # pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  # (B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  # (B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  # (B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask > 1] = 1

        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.linalg.eigh(L)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def GDPPLoss(phiFake, phiReal, backward=True):
    r"""
    Implementation of the GDPP loss. Can be used with any kind of GAN
    architecture.
    Args:
        phiFake (tensor) : last feature layer of the discriminator on real data
        phiReal (tensor) : last feature layer of the discriminator on fake data
        backward (bool)  : should we perform the backward operation ?
    Returns:
        Loss's value. The backward operation in performed within this operator
    """
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.symeig(SB, eigenvectors=True)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phiFake)
    realEigVals, realEigVecs = compute_diversity(phiReal)

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()
//  [python]
def GDPPLoss(phiFake, phiReal, backward=True):
    r"""
    Implementation of the GDPP loss. Can be used with any kind of GAN
    architecture.
    Args:
        phiFake (tensor) : last feature layer of the discriminator on real data
        phiReal (tensor) : last feature layer of the discriminator on fake data
        backward (bool)  : should we perform the backward operation ?
    Returns:
        Loss's value. The backward operation in performed within this operator
    """
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.linalg.eigh(SB)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phiFake)
    realEigVals, realEigVecs = compute_diversity(phiReal)

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()
//  [python]
def test_custom_top_eigenpair():
    from data.network_diffusion import DiffusionDataset
    from data.brain_data import RealDataModule
    import time
    scs = DiffusionDataset(num_vertices=500, num_samples=2, fc_norm=None).full_ds()[1]
    #scs = RealDataModule(fc_norm=None).full_ds()[1]

    x = scs
    slice_shape = x.shape[-2:]
    x_slices = x.view(-1, *slice_shape).detach()

    evals_truth, evecs_truth = torch.symeig(x_slices, eigenvectors=False)
    largest_evals_truth = torch.max(torch.abs(evals_truth), dim=1)[0].view(-1, 1)

    for n_pow_iters in range(5, 25, 5):
        #largest_evals_approx, _ = TopEigenPair(A=x_slices, n_power_iterations=n_pow_iters)
        largest_evals_approx, evecs = torch.lobpcg(A=x_slices, k=1, largest=True, niter=n_pow_iters)  # , tol=10**-2)
        largest_error = torch.max(torch.abs(largest_evals_truth - largest_evals_approx))
        print(f'n_power_iterations: {n_pow_iters}, max error: {largest_error}')
//  [python]
def test_custom_top_eigenpair():
    from data.network_diffusion import DiffusionDataset
    from data.brain_data import RealDataModule
    import time
    scs = DiffusionDataset(num_vertices=500, num_samples=2, fc_norm=None).full_ds()[1]
    #scs = RealDataModule(fc_norm=None).full_ds()[1]

    x = scs
    slice_shape = x.shape[-2:]
    x_slices = x.view(-1, *slice_shape).detach()

    evals_truth, evecs_truth = torch.linalg.eigh(x_slices)
    largest_evals_truth = torch.max(torch.abs(evals_truth), dim=1)[0].view(-1, 1)

    for n_pow_iters in range(5, 25, 5):
        #largest_evals_approx, _ = TopEigenPair(A=x_slices, n_power_iterations=n_pow_iters)
        largest_evals_approx, evecs = torch.lobpcg(A=x_slices, k=1, largest=True, niter=n_pow_iters)  # , tol=10**-2)
        largest_error = torch.max(torch.abs(largest_evals_truth - largest_evals_approx))
        print(f'n_power_iterations: {n_pow_iters}, max error: {largest_error}')
//  [python]
def train(cls, vectors, processing_instruction, device):
        """[summary]
        
        Arguments:
            vectors {[type]} -- [description]
            processing_instruction {String} -- [Contains characters 'c', 'w', 'l'. For example 'cwlc' performs centering, whitening, length normalization, and centering (2nd time) in this order.]
        """
        print('Training vector processor ...')

        c_count = processing_instruction.count('c')
        w_count = processing_instruction.count('w')

        vec_size = vectors.size()[1]

        whitening_matrices = torch.zeros(w_count, vec_size, vec_size, device=device)
        centering_vectors = torch.zeros(c_count, vec_size, device=device)

        vectors = vectors.to(device)

        c_count = 0
        w_count = 0
        for c in processing_instruction:
            if c == 'c':
                print('Centering...')
                centering_vectors[c_count, :] = torch.mean(vectors, dim=0)
                vectors = vectors - centering_vectors[c_count, :]
                c_count += 1
            elif c == 'w':
                print('Whitening...')
                l, U = torch.symeig(torch.matmul(vectors.t(), vectors) / vectors.size()[0], eigenvectors=True)
                l = torch.clamp(l, min=1e-10)
                whitening_matrices[w_count, :, :] = torch.rsqrt(l) * U  # transposed
                vectors = torch.matmul(vectors, whitening_matrices[w_count, :, :])
                w_count += 1
            elif c == 'l':
                print('Normalizing length...')
                vectors = unit_len_norm(vectors)

        return VectorProcessor(centering_vectors, whitening_matrices, processing_instruction)
//  [python]
def train(cls, vectors, processing_instruction, device):
        """[summary]
        
        Arguments:
            vectors {[type]} -- [description]
            processing_instruction {String} -- [Contains characters 'c', 'w', 'l'. For example 'cwlc' performs centering, whitening, length normalization, and centering (2nd time) in this order.]
        """
        print('Training vector processor ...')

        c_count = processing_instruction.count('c')
        w_count = processing_instruction.count('w')

        vec_size = vectors.size()[1]

        whitening_matrices = torch.zeros(w_count, vec_size, vec_size, device=device)
        centering_vectors = torch.zeros(c_count, vec_size, device=device)

        vectors = vectors.to(device)

        c_count = 0
        w_count = 0
        for c in processing_instruction:
            if c == 'c':
                print('Centering...')
                centering_vectors[c_count, :] = torch.mean(vectors, dim=0)
                vectors = vectors - centering_vectors[c_count, :]
                c_count += 1
            elif c == 'w':
                print('Whitening...')
                l, U = torch.linalg.eigh(torch.matmul(vectors.t(), vectors) / vectors.size()[0])
                l = torch.clamp(l, min=1e-10)
                whitening_matrices[w_count, :, :] = torch.rsqrt(l) * U  # transposed
                vectors = torch.matmul(vectors, whitening_matrices[w_count, :, :])
                w_count += 1
            elif c == 'l':
                print('Normalizing length...')
                vectors = unit_len_norm(vectors)

        return VectorProcessor(centering_vectors, whitening_matrices, processing_instruction)
//  [python]
def vec_to_quat(vec: torch.Tensor) -> torch.Tensor:
    """
    convert a vector to unit quaternion
    args:
    vec: 10 dim vector (default) from netwrok training
    return:
    4 dim unit quaternion
    """
    adj = vec_to_adj(vec)
    _, evs = torch.symeig(adj, eigenvectors=True)
    if evs.dim() < 3:
        evs.unsqueeze_(dim = 0)

    return evs[:,:,0]
//  [python]
def vec_to_quat(vec: torch.Tensor) -> torch.Tensor:
    """
    convert a vector to unit quaternion
    args:
    vec: 10 dim vector (default) from netwrok training
    return:
    4 dim unit quaternion
    """
    adj = vec_to_adj(vec)
    _, evs = torch.linalg.eigh(adj)
    if evs.dim() < 3:
        evs.unsqueeze_(dim = 0)

    return evs[:,:,0]
//  [python]
def compute_multiplets(env, eps_multiplet_gap=1.0e-10):
    D= torch.zeros(env.chi+1, dtype=env.dtype, device=env.device)
    if _torch_version_check("1.8.1"):
        D[:env.chi]= torch.linalg.eigvalsh(env.C[env.keyC])
    else:
        D[:env.chi], U= torch.symeig(env.C[env.keyC])
    D, p= torch.sort(torch.abs(D),descending=True)
    m=[]
    l=0
    for i in range(env.chi):
        l+=1
        g=D[i]-D[i+1]
        #print(f"{i} {D[i]} {g}", end=" ")
        if g>eps_multiplet_gap:
            #print(f"{l}", end=" ")
            m.append(l)
            l=0
    return m
//  [python]
def compute_multiplets(env, eps_multiplet_gap=1.0e-10):
    D= torch.zeros(env.chi+1, dtype=env.dtype, device=env.device)
    D[:env.chi]= torch.linalg.eigh(env.C[env.keyC]).eigenvalues
    D, p= torch.sort(torch.abs(D),descending=True)
    m=[]
    l=0
    for i in range(env.chi):
        l+=1
        g=D[i]-D[i+1]
        #print(f"{i} {D[i]} {g}", end=" ")
        if g>eps_multiplet_gap:
            #print(f"{l}", end=" ")
            m.append(l)
            l=0
    return m
//  [python]
def PCA(xyz, adjust=False, mode='random_v1'):
    """
    :param xyz: [B, N, 3]
    :param adjust:
    :param mode:
    :return: pca basis: [B, 3, 3]
    """
    assert xyz.size()[-1] == 3, "invalid input xyz shape"
    sample_cov = (xyz.transpose(2, 1) @ xyz) / (xyz.size()[1] - 1)  # [B, 3, 3]
    basis = torch.symeig(sample_cov.cpu(), eigenvectors=True)[1]  # [v1, v2, v3]
    basis = torch.flip(basis.to(xyz.device), [-1])  # [B, 3, 3]
    if adjust:
        basis = global_pca_adjust(basis, xyz, mode)
    xyz = xyz @ basis
    return xyz, basis
//  [python]
def PCA(xyz, adjust=False, mode='random_v1'):
    """
    :param xyz: [B, N, 3]
    :param adjust:
    :param mode:
    :return: pca basis: [B, 3, 3]
    """
    assert xyz.size()[-1] == 3, "invalid input xyz shape"
    sample_cov = (xyz.transpose(2, 1) @ xyz) / (xyz.size()[1] - 1)  # [B, 3, 3]
    basis = torch.linalg.eigh(sample_cov.cpu(), eigenvectors=True)[1]  # [v1, v2, v3]
    basis = torch.flip(basis.to(xyz.device), [-1])  # [B, 3, 3]
    if adjust:
        basis = global_pca_adjust(basis, xyz, mode)
    xyz = xyz @ basis
    return xyz, basis
//  [python]
def graph_laplacian_embedds(edges, eigen_k, center_idx=1, norm=False):
    """ Returns the embeddings of points in the K 
        first eigenvectors of the graph Laplacian.
        Inputs:
        * edges: (2, N). long tensor or list. undirected edges are enough.
        * eigen_k: int. N of first eigenvectors to return embeddings for.
        * center_idx: int. index to take as center for the embeddings
        * norm: bool. Whether to use the normalized Laplacian. Not recommended.
        Output: (n_points, eigen_k)
    """
    if isinstance(edges, list):
        edges = torch.tensor(edges).long()
        # correct dims
        if edges.shape[0] != 2:
            edges = edges.t()
        # early stopping if empty entry
        if edges.shape[0] == 0:
            return torch.zeros(1, eigen_k)
    # get params
    size = torch.max(edges)+1
    device = edges.device
    # crate laplacian
    adj_mat = torch.eye(size, device=device) 
    for i,j in edges.t():
        adj_mat[i,j] = adj_mat[j,i] = 1.
        
    deg_mat = torch.eye(size) * adj_mat.sum(dim=-1, keepdim=True)
    laplace = deg_mat - adj_mat
    # use norm-laplace if arg passed
    if norm:
        for i,j in edges.t():
            laplace[i,j] = laplace[j,i] = -1 / (deg_mat[i,i] * deg_mat[j,j])**0.5
    # get laplacian basis - eigendecomposition - order importance by eigenvalue
    e, v = torch.symeig(laplace, eigenvectors=True)
    idxs = torch.sort( e.abs(), descending=True)[1]
    # take embedds and center
    embedds = v[:, idxs[:eigen_k]]
    embedds = embedds - embedds[center_idx].unsqueeze(-2)
    return embedds
//  [python]
def graph_laplacian_embedds(edges, eigen_k, center_idx=1, norm=False):
    """ Returns the embeddings of points in the K 
        first eigenvectors of the graph Laplacian.
        Inputs:
        * edges: (2, N). long tensor or list. undirected edges are enough.
        * eigen_k: int. N of first eigenvectors to return embeddings for.
        * center_idx: int. index to take as center for the embeddings
        * norm: bool. Whether to use the normalized Laplacian. Not recommended.
        Output: (n_points, eigen_k)
    """
    if isinstance(edges, list):
        edges = torch.tensor(edges).long()
        # correct dims
        if edges.shape[0] != 2:
            edges = edges.t()
        # early stopping if empty entry
        if edges.shape[0] == 0:
            return torch.zeros(1, eigen_k)
    # get params
    size = torch.max(edges)+1
    device = edges.device
    # crate laplacian
    adj_mat = torch.eye(size, device=device) 
    for i,j in edges.t():
        adj_mat[i,j] = adj_mat[j,i] = 1.
        
    deg_mat = torch.eye(size) * adj_mat.sum(dim=-1, keepdim=True)
    laplace = deg_mat - adj_mat
    # use norm-laplace if arg passed
    if norm:
        for i,j in edges.t():
            laplace[i,j] = laplace[j,i] = -1 / (deg_mat[i,i] * deg_mat[j,j])**0.5
    # get laplacian basis - eigendecomposition - order importance by eigenvalue
    e, v = torch.linalg.eigh(laplace)
    idxs = torch.sort( e.abs(), descending=True)[1]
    # take embedds and center
    embedds = v[:, idxs[:eigen_k]]
    embedds = embedds - embedds[center_idx].unsqueeze(-2)
    return embedds
//  [python]
def test_fast_symeig_against_torch(seed, rand_sym, n, d, eig):
    x = rand_sym(n, d)
    symeig_time = timeit(lambda: torch.symeig(x), number=100)
    eig_time = timeit(lambda: eig(x), number=100)
    logging.info('Fast symeig speedup (n={}, d={}): {:.2f}x'.format(
            n, d, symeig_time / eig_time))
    assert eig_time < symeig_time
//  [python]
def test_fast_symeig_against_torch(seed, rand_sym, n, d, eig):
    x = rand_sym(n, d)
    symeig_time = timeit(lambda: torch.linalg.eigh(x), number=100)
    eig_time = timeit(lambda: eig(x), number=100)
    logging.info('Fast symeig speedup (n={}, d={}): {:.2f}x'.format(
            n, d, symeig_time / eig_time))
    assert eig_time < symeig_time
//  [python]
def test_lowrank():
    n = 300
    A = torch.randn(n, n, dtype=torch.float64)
    A = A + A.T
    eigvalues, eigvectors = torch.symeig(A, eigenvectors=True)
    alpha = eigvalues[0]
    x = eigvectors[:, 0]
    Aprime = A - alpha * torch.eye(n, dtype=torch.float64)
    b = torch.randn(n, dtype=torch.float64)
    b = b - torch.matmul(x, b) * x
    initialx = torch.randn(n, dtype=torch.float64)
    initialx = initialx - torch.matmul(x, initialx) * x
    result = CG_torch(Aprime, b, initialx)
    assert torch.allclose(torch.matmul(Aprime, result) - b, 
                          torch.zeros(n, dtype=torch.float64), 
                          atol=1e-06)
    assert torch.allclose(torch.matmul(result, x)[None], 
                          torch.zeros(1, dtype=torch.float64), 
                          atol=1e-06)
//  [python]
def test_lowrank():
    n = 300
    A = torch.randn(n, n, dtype=torch.float64)
    A = A + A.T
    eigvalues, eigvectors = torch.linalg.eigh(A)
    alpha = eigvalues[0]
    x = eigvectors[:, 0]
    Aprime = A - alpha * torch.eye(n, dtype=torch.float64)
    b = torch.randn(n, dtype=torch.float64)
    b = b - torch.matmul(x, b) * x
    initialx = torch.randn(n, dtype=torch.float64)
    initialx = initialx - torch.matmul(x, initialx) * x
    result = CG_torch(Aprime, b, initialx)
    assert torch.allclose(torch.matmul(Aprime, result) - b, 
                          torch.zeros(n, dtype=torch.float64), 
                          atol=1e-06)
    assert torch.allclose(torch.matmul(result, x)[None], 
                          torch.zeros(1, dtype=torch.float64), 
                          atol=1e-06)
//  [python]
d2E0s = np.empty(Npoints)
E0s_computation = np.empty(Npoints)
dE0s_computation = np.empty(Npoints)
d2E0s_computation = np.empty(Npoints)

for i in range(Npoints):
    g = ps[i]
    squareroot = np.sqrt((a - b)**2 + 4 * g**2)
    E0s[i] = 0.5 * (a + b - squareroot)
    dE0s[i] = - 2 * g / squareroot
    d2E0s[i] = - 2 * (a - b)**2 / squareroot**3

    dpsi0 = np.array([compute_dpsi0(a, b, g)])

    g = torch.Tensor([g]).to(torch.float64)
    g.requires_grad_(True)
    H = diag + g * offdiag
    Es, psis = torch.symeig(H, eigenvectors=True)

    psi0 = psis[:, 0]
    dpsi0_computation = torch.empty(2, dtype=torch.float64)
    I = torch.eye(2, dtype=torch.float64)
    for idx in range(2):
        dpsi0_computation[idx], = torch.autograd.grad(psi0, g, grad_outputs=I[:, idx], retain_graph=True)
    #dpsi0_computation, = torch.autograd.grad(psi0, g, grad_outputs=I, retain_graph=True)

    E0 = Es[0]
    dE0, = torch.autograd.grad(E0, g, create_graph=True)
    d2E0, = torch.autograd.grad(dE0, g)
    E0s_computation[i] = E0.item()
    dE0s_computation[i] = dE0.item()
    d2E0s_computation[i] = d2E0.item()

    print(ps[i], E0s[i], E0s_computation[i], 
          dE0s[i], dE0s_computation[i], 
          d2E0s[i], d2E0s_computation[i], 
          dpsi0, dpsi0_computation.numpy())
//  [python]
d2E0s = np.empty(Npoints)
E0s_computation = np.empty(Npoints)
dE0s_computation = np.empty(Npoints)
d2E0s_computation = np.empty(Npoints)

for i in range(Npoints):
    g = ps[i]
    squareroot = np.sqrt((a - b)**2 + 4 * g**2)
    E0s[i] = 0.5 * (a + b - squareroot)
    dE0s[i] = - 2 * g / squareroot
    d2E0s[i] = - 2 * (a - b)**2 / squareroot**3

    dpsi0 = np.array([compute_dpsi0(a, b, g)])

    g = torch.Tensor([g]).to(torch.float64)
    g.requires_grad_(True)
    H = diag + g * offdiag
    Es, psis = torch.linalg.eigh(H)

    psi0 = psis[:, 0]
    dpsi0_computation = torch.empty(2, dtype=torch.float64)
    I = torch.eye(2, dtype=torch.float64)
    for idx in range(2):
        dpsi0_computation[idx], = torch.autograd.grad(psi0, g, grad_outputs=I[:, idx], retain_graph=True)
    #dpsi0_computation, = torch.autograd.grad(psi0, g, grad_outputs=I, retain_graph=True)

    E0 = Es[0]
    dE0, = torch.autograd.grad(E0, g, create_graph=True)
    d2E0, = torch.autograd.grad(dE0, g)
    E0s_computation[i] = E0.item()
    dE0s_computation[i] = dE0.item()
    d2E0s_computation[i] = d2E0.item()

    print(ps[i], E0s[i], E0s_computation[i], 
          dE0s[i], dE0s_computation[i], 
          d2E0s[i], d2E0s_computation[i], 
          dpsi0, dpsi0_computation.numpy())
//  [python]
def test_lowrank_gpu():
    n = 300
    cuda = torch.device("cuda")
    dtype = torch.float64
    A = torch.randn(n, n, device=cuda, dtype=dtype)
    A = A + A.T
    eigvalues, eigvectors = torch.symeig(A, eigenvectors=True)
    alpha = eigvalues[0]
    x = eigvectors[:, 0]
    Aprime = A - alpha * torch.eye(n, device=cuda, dtype=dtype)
    b = torch.randn(n, device=cuda, dtype=dtype)
    b = b - torch.matmul(x, b) * x
    initialx = torch.randn(n, device=cuda, dtype=dtype)
    initialx = initialx - torch.matmul(x, initialx) * x
    result = CG_torch(Aprime, b, initialx)
    assert torch.allclose(torch.matmul(Aprime, result) - b, 
                          torch.zeros(n, device=cuda, dtype=dtype), 
                          atol=1e-06)
    assert torch.allclose(torch.matmul(result, x)[None], 
                          torch.zeros(1, device=cuda, dtype=dtype), 
                          atol=1e-06)
//  [python]
def test_lowrank_gpu():
    n = 300
    cuda = torch.device("cuda")
    dtype = torch.float64
    A = torch.randn(n, n, device=cuda, dtype=dtype)
    A = A + A.T
    eigvalues, eigvectors = torch.linalg.eigh(A)
    alpha = eigvalues[0]
    x = eigvectors[:, 0]
    Aprime = A - alpha * torch.eye(n, device=cuda, dtype=dtype)
    b = torch.randn(n, device=cuda, dtype=dtype)
    b = b - torch.matmul(x, b) * x
    initialx = torch.randn(n, device=cuda, dtype=dtype)
    initialx = initialx - torch.matmul(x, initialx) * x
    result = CG_torch(Aprime, b, initialx)
    assert torch.allclose(torch.matmul(Aprime, result) - b, 
                          torch.zeros(n, device=cuda, dtype=dtype), 
                          atol=1e-06)
    assert torch.allclose(torch.matmul(result, x)[None], 
                          torch.zeros(1, device=cuda, dtype=dtype), 
                          atol=1e-06)
//  [python]
def batch_symeig(X, device):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b, d, d)
    for batch_idx in range(X.shape[0]):
        e, v = torch.symeig(X[batch_idx, :, :].squeeze(), True)
        bv[batch_idx, :, :] = v
    bv = bv.to(device)
    return bv
//  [python]
def batch_symeig(X, device):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b, d, d)
    for batch_idx in range(X.shape[0]):
        e, v = torch.linalg.eigh(X[batch_idx, :, :].squeeze())
        bv[batch_idx, :, :] = v
    bv = bv.to(device)
    return bv
//  [python]
def entropy(self, x, n_n, alpha, sigma=None):

        eigv = th.abs(th.symeig(x)[0])
        eig_pow = eigv**alpha
        return (1/(1-alpha))*th.log2(eig_pow.sum())
//  [python]
def entropy(self, x, n_n, alpha, sigma=None):

        eigv = th.abs(th.linalg.eigh(x)[0])
        eig_pow = eigv**alpha
        return (1/(1-alpha))*th.log2(eig_pow.sum())
//  [python]
def GetBaseVector(H, label):
    N, C, d, _ = H.shape
    N_new = N * d * d
    H = H.permute(0, 2, 3, 1)
    H = torch.reshape(H, (N_new, C))
    H_bar = H - torch.mean(H, 0, True)
    label = label.view(N, 1)
    labels = torch.reshape(label * torch.Tensor().new_ones((N, d * d), device=device, dtype=torch.long), (N_new,))
    S_w = torch.Tensor().new_zeros((C, C), device = device, dtype = dtype)
    S_t = H_bar.t().matmul(H_bar) / (N_new - 1)
    for i in range(categ):
        H_i = H[torch.nonzero(labels == i).view(-1)]
        H_i_bar = H_i - torch.mean(H_i, 0, True)
        N_i = H_i.shape[0]
        if N_i == 0:
            continue
        S_w += H_i_bar.t().matmul(H_i_bar) / (N_i - 1) / categ
    temp = (S_w + lamb * torch.diag(torch.Tensor().new_ones((C), device = device, dtype = dtype))).pinverse().matmul(S_t - S_w)
    w, v = torch.symeig(temp, eigenvectors = True)
    w = w.detach()
    v = v.detach()
    return v[:, base_num:]
//  [python]
def GetBaseVector(H, label):
    N, C, d, _ = H.shape
    N_new = N * d * d
    H = H.permute(0, 2, 3, 1)
    H = torch.reshape(H, (N_new, C))
    H_bar = H - torch.mean(H, 0, True)
    label = label.view(N, 1)
    labels = torch.reshape(label * torch.Tensor().new_ones((N, d * d), device=device, dtype=torch.long), (N_new,))
    S_w = torch.Tensor().new_zeros((C, C), device = device, dtype = dtype)
    S_t = H_bar.t().matmul(H_bar) / (N_new - 1)
    for i in range(categ):
        H_i = H[torch.nonzero(labels == i).view(-1)]
        H_i_bar = H_i - torch.mean(H_i, 0, True)
        N_i = H_i.shape[0]
        if N_i == 0:
            continue
        S_w += H_i_bar.t().matmul(H_i_bar) / (N_i - 1) / categ
    temp = (S_w + lamb * torch.diag(torch.Tensor().new_ones((C), device = device, dtype = dtype))).pinverse().matmul(S_t - S_w)
    w, v = torch.linalg.eigh(temp)
    w = w.detach()
    v = v.detach()
    return v[:, base_num:]
//  [python]
def j_entropy(self, x, y, n_n, alpha, sigma=None):

        k = x*y / (th.trace(x*y))
        eigv = th.abs(th.symeig(k)[0])
        eig_pow = eigv**alpha
        return (1/(1-alpha))*th.log2(eig_pow.sum())
//  [python]
def j_entropy(self, x, y, n_n, alpha, sigma=None):

        k = x*y / (th.trace(x*y))
        eigv = th.abs(th.linalg.eigh(k)[0])
        eig_pow = eigv**alpha
        return (1/(1-alpha))*th.log2(eig_pow.sum())
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma):
        Lambda_F, Q_F = torch.symeig(g(F), eigenvectors=True)
        Lambda_F = Lambda_F.view(-1,1)
        G = get_G(Lambda_F, Lambda_S, gamma)
        Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ Q_S.t()
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def forward(ctx, X, F, S, Q_S, Lambda_S, gamma):
        Lambda_F, Q_F = torch.linalg.eigh(g(F))
        Lambda_F = Lambda_F.view(-1,1)
        G = get_G(Lambda_F, Lambda_S, gamma)
        Z = Q_F @ (G * (Q_F.t() @ X @ Q_S)) @ Q_S.t()
        ctx.save_for_backward(F, S, Q_F, Q_S, Z, G, X, gamma)
        return Z
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b, d, d)
    for batch_idx in range(X.shape[0]):
        e, v = torch.symeig(X[batch_idx, :, :].squeeze(), True)
        bv[batch_idx, :, :] = v
    bv = bv.cuda()
    return bv
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b, d, d)
    for batch_idx in range(X.shape[0]):
        e, v = torch.linalg.eigh(X[batch_idx, :, :].squeeze())
        bv[batch_idx, :, :] = v
    bv = bv.cuda()
    return bv
//  [python]
def eigh(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:  # pragma: no cover
    """
    Like torch.linalg.eigh, assuming the argument is a symmetric real matrix.
    """
    if hasattr(torch, "linalg") and hasattr(torch.linalg, "eigh"):
        return torch.linalg.eigh(A)
    return torch.symeig(A, eigenvectors=True)
//  [python]
def eigh(A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:  # pragma: no cover
    """
    Like torch.linalg.eigh, assuming the argument is a symmetric real matrix.
    """
    return torch.linalg.eigh(A)
//  [python]
def opt_loss(self, adj):
        Ly_inv_rt, Ly_inv = mx_inv_sqrt(adj)
        m = self.step_size
        P = self.P.abs()

        for _ in range(self.sinkhorn_iter):
            P = P / P.sum(dim=1, keepdim=True)
            P = P / P.sum(dim=0, keepdim=True)
        
        # if self.args.use_symeig:
        sqrt = torch.symeig(Ly_inv_rt @ self.P.t() @ self.Lx_inv @ self.P @ Ly_inv_rt, eigenvectors=True)
        loss = torch.abs(mx_tr(Ly_inv)*m - 2*torch.sqrt(sqrt[0].clamp(min=2e-20)).sum())
        return loss
//  [python]
def opt_loss(self, adj):
        Ly_inv_rt, Ly_inv = mx_inv_sqrt(adj)
        m = self.step_size
        P = self.P.abs()

        for _ in range(self.sinkhorn_iter):
            P = P / P.sum(dim=1, keepdim=True)
            P = P / P.sum(dim=0, keepdim=True)
        
        # if self.args.use_symeig:
        sqrt = torch.linalg.eigh(Ly_inv_rt @ self.P.t() @ self.Lx_inv @ self.P @ Ly_inv_rt, eigenvectors=True)
        loss = torch.abs(mx_tr(Ly_inv)*m - 2*torch.sqrt(sqrt[0].clamp(min=2e-20)).sum())
        return loss
//  [python]
def init_factors(X, rank, nonnegative, init='svd'):
    if init == 'random':
        factors = [torch.randn((dim_size, rank)) for dim_size in X.size()]
        if nonnegative:
            factors[0] = abs(factors[0])
            factors[1] = abs(factors[1])
            factors[2] = -abs(factors[2])
        return factors

    elif init == 'svd':
        # Initialize using left singular vectors
        # Fill with random values if any mode.size() < rank
        factors = []
        for mode in range(X.dim()):
            # Use eigendecomposition method on Gram matrix (XX^T) for SVD
            # mode-unfold of X
            G = unfold(X, mode)
            G = torch.mm(G, G.t())
            eigvals, eigvecs = torch.symeig(G, eigenvectors=True)
            # eigenvectors are reverse sorted (ascending)
            factor = torch.flip(eigvecs[:, -int(min(X.size(mode), rank)):],
                                [1])
            # Fill remaining columns with random vectors
            if rank > X.size(mode):
                factor = torch.cat(
                    (factor, torch.randn(X.size(mode), rank - X.size(mode)).to(
                        X.device)),
                    dim=1)
            if nonnegative:
                if mode in [2]:
                    factor = -abs(factor)
                else:
                    factor = abs(factor)
            factors.append(factor)
        return factors
    else:
        raise ValueError('Wrong value for init: {}'.format(init))
//  [python]
def init_factors(X, rank, nonnegative, init='svd'):
    if init == 'random':
        factors = [torch.randn((dim_size, rank)) for dim_size in X.size()]
        if nonnegative:
            factors[0] = abs(factors[0])
            factors[1] = abs(factors[1])
            factors[2] = -abs(factors[2])
        return factors

    elif init == 'svd':
        factors = []
        for mode in range(X.dim()):
            G = unfold(X, mode)
            G = torch.mm(G, G.t())
            eigvals, eigvecs = torch.linalg.eigh(G)
            factor = torch.flip(eigvecs[:, -int(min(X.size(mode), rank)):], [1])
            if rank > X.size(mode):
                factor = torch.cat(
                    (factor, torch.randn(X.size(mode), rank - X.size(mode)).to(
                        X.device)),
                    dim=1)
            if nonnegative:
                if mode in [2]:
                    factor = -abs(factor)
                else:
                    factor = abs(factor)
            factors.append(factor)
        return factors
    else:
        raise ValueError('Wrong value for init: {}'.format(init))
//  [python]
def loss_func_renyi_mod1(preds,Y,Lam):
	''' If numS <= numY ''' 
	pred = soft(preds).to(dev)
	Py = torch.mean(pred,0).to(dev)
	Q = torch.zeros((numS,numY)).to(dev)
	# tmp = 0
	for s in range(numS):
	    Q[s] = PS[s] * torch.mean(pred[S==s],0) / torch.sqrt(Py*PS[s])	
	e, v = torch.symeig(torch.matmul(Q,torch.transpose(Q,0,1)),eigenvectors=True)

	return F.cross_entropy(preds,Y,reduction='mean') + Lam * e[numS - 2]
//  [python]
def loss_func_renyi_mod1(preds,Y,Lam):
	''' If numS <= numY ''' 
	pred = soft(preds).to(dev)
	Py = torch.mean(pred,0).to(dev)
	Q = torch.zeros((numS,numY)).to(dev)
	# tmp = 0
	for s in range(numS):
	    Q[s] = PS[s] * torch.mean(pred[S==s],0) / torch.sqrt(Py*PS[s])	
	e, v = torch.linalg.eigh(torch.matmul(Q,torch.transpose(Q,0,1)))

	return F.cross_entropy(preds,Y,reduction='mean') + Lam * e[numS - 2]
//  [python]
def loss_func_renyi_mod2(preds,Y,Lam):
	''' If numS > numY ''' 
	pred = soft(preds).to(dev)
	Py = torch.mean(pred,0).to(dev)
	Q = torch.zeros((numS,numY)).to(dev)

	for s in range(numS):
	    Q[s] = PS[s] * torch.mean(pred[S==s],0) / torch.sqrt(Py*PS[s])	
	e, v = torch.symeig(torch.matmul(torch.transpose(Q,0,1),Q),eigenvectors=True)

	return F.cross_entropy(preds,Y,reduction='mean') + Lam * e[numY - 2]
//  [python]
def loss_func_renyi_mod2(preds,Y,Lam):
	''' If numS > numY ''' 
	pred = soft(preds).to(dev)
	Py = torch.mean(pred,0).to(dev)
	Q = torch.zeros((numS,numY)).to(dev)

	for s in range(numS):
	    Q[s] = PS[s] * torch.mean(pred[S==s],0) / torch.sqrt(Py*PS[s])	
	e, v = torch.linalg.eigh(torch.matmul(torch.transpose(Q,0,1),Q))

	return F.cross_entropy(preds,Y,reduction='mean') + Lam * e[numY - 2]
//  [python]
def covariance_eigen_decomposition(self, point_neighbors, point_locations):
        # Filter by range
        epsilon_range = self.config["epsilon_range"]
        range_deviation_indices = torch.abs(torch.norm(point_neighbors, dim=1)
                                            - torch.norm(point_locations, dim=1)) > epsilon_range
        # Annulate neighbors that have a range which deviates too much from point location
        point_neighbors.permute(0, 2, 1)[range_deviation_indices] = 0.0

        point_neighbors_permuted = point_neighbors.permute(2, 1, 0)
        covariance_matrices, number_neighbours = utility.linalg.cov(
            point_neighbors=point_neighbors_permuted)

        # Filter by amount of neighbors
        # Dropping eigenvectors and values where min number of points is not met
        where_enough_neighbors_bool = number_neighbours >= self.config[
            "min_num_points_in_neighborhood_to_determine_point_class"]
        covariance_matrices = covariance_matrices[where_enough_neighbors_bool]
        eigenvalues, eigenvectors = torch.symeig(covariance_matrices.to(torch.device("cpu")),
                                                 eigenvectors=True)
        eigenvalues = eigenvalues.to(self.config["device"])
        eigenvectors = eigenvectors.to(self.config["device"])

        # e-vectors corresponding to smallest e-value
        all_normal_vectors = eigenvectors[:, :, 0]

        point_locations = point_locations[0].permute(1, 0)
        dot_products = all_normal_vectors.view(-1, 1, 3).matmul(
            point_locations[where_enough_neighbors_bool].view(-1, 3, 1)).squeeze()
        all_normal_vectors[dot_products > 0] *= -1

        # Now for all points which do not have a normal we simply store (0, 0, 0)^T
        all_normal_vectors_with_zeros = torch.zeros_like(point_locations)
        all_normal_vectors_with_zeros[where_enough_neighbors_bool] = all_normal_vectors

        return all_normal_vectors_with_zeros, where_enough_neighbors_bool, point_locations
//  [python]
def covariance_eigen_decomposition(self, point_neighbors, point_locations):
        # Filter by range
        epsilon_range = self.config["epsilon_range"]
        range_deviation_indices = torch.abs(torch.norm(point_neighbors, dim=1)
                                            - torch.norm(point_locations, dim=1)) > epsilon_range
        # Annulate neighbors that have a range which deviates too much from point location
        point_neighbors.permute(0, 2, 1)[range_deviation_indices] = 0.0

        point_neighbors_permuted = point_neighbors.permute(2, 1, 0)
        covariance_matrices, number_neighbours = utility.linalg.cov(
            point_neighbors=point_neighbors_permuted)

        # Filter by amount of neighbors
        # Dropping eigenvectors and values where min number of points is not met
        where_enough_neighbors_bool = number_neighbours >= self.config[
            "min_num_points_in_neighborhood_to_determine_point_class"]
        covariance_matrices = covariance_matrices[where_enough_neighbors_bool]
        eigenvalues, eigenvectors = torch.linalg.eigh(covariance_matrices.to(torch.device("cpu")))
        eigenvalues = eigenvalues.to(self.config["device"])
        eigenvectors = eigenvectors.to(self.config["device"])

        # e-vectors corresponding to smallest e-value
        all_normal_vectors = eigenvectors[:, :, 0]

        point_locations = point_locations[0].permute(1, 0)
        dot_products = all_normal_vectors.view(-1, 1, 3).matmul(
            point_locations[where_enough_neighbors_bool].view(-1, 3, 1)).squeeze()
        all_normal_vectors[dot_products > 0] *= -1

        # Now for all points which do not have a normal we simply store (0, 0, 0)^T
        all_normal_vectors_with_zeros = torch.zeros_like(point_locations)
        all_normal_vectors_with_zeros[where_enough_neighbors_bool] = all_normal_vectors

        return all_normal_vectors_with_zeros, where_enough_neighbors_bool, point_locations
//  [python]
def get_local_frame(self, pnts):
        center = pnts.mean(1, True)
        pnts_centered = pnts - center
        R = torch.bmm(pnts_centered.transpose(1, 2), pnts_centered)
        lambdas, V_ = torch.symeig(R.detach().cpu(), True)
        F = V_.to(R)
        return F.detach(), center.detach()
//  [python]
def get_local_frame(self, pnts):
        center = pnts.mean(1, True)
        pnts_centered = pnts - center
        R = torch.bmm(pnts_centered.transpose(1, 2), pnts_centered)
        lambdas, V_ = torch.linalg.eigh(R.detach().cpu())
        F = V_.to(R)
        return F.detach(), center.detach()
//  [python]
def compute_gdpp(phi_fake, phi_real, backward=True):
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.symeig(SB, eigenvectors=True)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        if abs(minV - maxV) < 1e-10:
            return eigVals
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phi_fake.double())
    realEigVals, realEigVecs = compute_diversity(phi_real.double())

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()
//  [python]
def compute_gdpp(phi_fake, phi_real, backward=True):
    def compute_diversity(phi):
        phi = F.normalize(phi, p=2, dim=1)
        SB = torch.mm(phi, phi.t())
        eigVals, eigVecs = torch.linalg.eigh(SB)
        return eigVals, eigVecs

    def normalize_min_max(eigVals):
        minV, maxV = torch.min(eigVals), torch.max(eigVals)
        if abs(minV - maxV) < 1e-10:
            return eigVals
        return (eigVals - minV) / (maxV - minV)

    fakeEigVals, fakeEigVecs = compute_diversity(phi_fake.double())
    realEigVals, realEigVecs = compute_diversity(phi_real.double())

    # Scaling factor to make the two losses operating in comparable ranges.
    magnitudeLoss = 0.0001 * F.mse_loss(target=realEigVals, input=fakeEigVals)
    structureLoss = -torch.sum(torch.mul(fakeEigVecs, realEigVecs), 0)
    normalizedRealEigVals = normalize_min_max(realEigVals)
    weightedStructureLoss = torch.sum(
        torch.mul(normalizedRealEigVals, structureLoss))
    gdppLoss = magnitudeLoss + weightedStructureLoss

    if backward:
        gdppLoss.backward(retain_graph=True)

    return gdppLoss.item()