jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def pca_torch(self, X):
        covariance = torch.transpose(X, 1, 0) @ X
        S, U = torch.linalg.eig(covariance)
        return S, U
//  [python]
def fit(self, X):
        n = X.shape[0]
        self.mean = torch.mean(X, axis=0)
        X = X - self.mean
        covariance_matrix = 1 / n * torch.matmul(X.T, X)
        eigenvalues, eigenvectors = torch.eig(covariance_matrix, eigenvectors=True)
        eigenvalues = torch.norm(eigenvalues, dim=1)
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        self.proj_mat = eigenvectors[:, 0:self.n_components]
//  [python]
def fit(self, X):
        n = X.shape[0]
        self.mean = torch.mean(X, axis=0)
        X = X - self.mean
        covariance_matrix = 1 / n * torch.matmul(X.T, X)
        eigenvalues, eigenvectors = torch.linalg.eig(covariance_matrix)
        eigenvalues = torch.norm(eigenvalues, dim=1)
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        self.proj_mat = eigenvectors[:, 0:self.n_components]
//  [python]
def get_singular_gaussian_penalty(model):
    """Return scalar high when attention covariance get very singular
    """
    if config["attention_isotropic_gaussian"]:
        # TODO move at setup
        print("Singular gaussian penalty ignored as `attention_isotropic_gaussian` is True")
        return 0

    condition_numbers = []
    for layer in model.encoder.layer:
        for sigma_half_inv in layer.attention.self.attention_spreads:
            sigma_inv = sigma_half_inv.transpose(0, 1) @ sigma_half_inv
            eig_values = torch.eig(sigma_inv)[0][:, 0].abs()
            condition_number = eig_values.max() / eig_values.min()
            condition_numbers.append(condition_number)

    return torch.mean((torch.tensor(condition_numbers)  - 1) ** 2)
//  [python]
def get_singular_gaussian_penalty(model):
    """Return scalar high when attention covariance get very singular
    """
    if config["attention_isotropic_gaussian"]:
        # TODO move at setup
        print("Singular gaussian penalty ignored as `attention_isotropic_gaussian` is True")
        return 0

    condition_numbers = []
    for layer in model.encoder.layer:
        for sigma_half_inv in layer.attention.self.attention_spreads:
            sigma_inv = sigma_half_inv.transpose(0, 1) @ sigma_half_inv
            eig_values = torch.linalg.eigvals(sigma_inv).abs()
            condition_number = eig_values.max() / eig_values.min()
            condition_numbers.append(condition_number)

    return torch.mean((torch.tensor(condition_numbers)  - 1) ** 2)
//  [python]
def knl(self, x, xp, eigeps=EPS):
        f = self.gp
        if xp is x:
            xp = xp.detach().clone()

        grad_k_func = lambda xs, xt: torch.autograd.grad(
            f.knl(xs, xt), xs, create_graph=True)[0]
        if self.grad_check:
            old_dtype = self.dtype
            self.to(torch.float64)
            f_knl_func = lambda xt: f.knl(xt, xp.double())
            with variable_required_grad(x):
                torch.autograd.gradcheck(f_knl_func, x.double())
            torch.autograd.gradgradcheck(lambda x: f.knl(x, x), x.double())
            with variable_required_grad(x):
                with variable_required_grad(xp):
                    torch.autograd.gradcheck(
                        lambda xp: grad_k_func(x.double(), xp)[0], xp.double())
            self.to(dtype=old_dtype)

        analytical = self.analytical_hessian
        if analytical:
            Hxx_k = t_hessian(f.knl, x, xp)
        else:
            with variable_required_grad(x):
                with variable_required_grad(xp):
                    old_dtype = self.dtype
                    self.to(torch.float64)
                    Hxx_k = tgradcheck.get_numerical_jacobian(
                        partial(grad_k_func, x.double()), xp.double())
                    self.to(dtype=old_dtype)
                    Hxx_k = Hxx_k.to(old_dtype)
        if torch.allclose(x, xp):
            eigenvalues, eigenvectors = torch.eig(Hxx_k, eigenvectors=False)
            assert (eigenvalues[:, 0] > -eigeps).all(),  " Hessian must be positive definite"
            small_neg_eig = ((eigenvalues[:, 0] > -eigeps) & (eigenvalues[:, 0] < 0))
            if small_neg_eig.any():
                eigenvalues, eigenvectors = torch.eig(Hxx_k, eigenvectors=True)
                evalz = eigenvalues[:, 0]
                evalz[small_neg_eig] = 0
                Hxx_k = eigenvectors.T @ torch.diag(evalz) @ eigenvectors
        return Hxx_k
//  [python]
def knl(self, x, xp, eigeps=EPS):
        f = self.gp
        if xp is x:
            xp = xp.detach().clone()

        grad_k_func = lambda xs, xt: torch.autograd.grad(
            f.knl(xs, xt), xs, create_graph=True)[0]
        if self.grad_check:
            old_dtype = self.dtype
            self.to(torch.float64)
            f_knl_func = lambda xt: f.knl(xt, xp.double())
            with variable_required_grad(x):
                torch.autograd.gradcheck(f_knl_func, x.double())
            torch.autograd.gradgradcheck(lambda x: f.knl(x, x), x.double())
            with variable_required_grad(x):
                with variable_required_grad(xp):
                    torch.autograd.gradcheck(
                        lambda xp: grad_k_func(x.double(), xp)[0], xp.double())
            self.to(dtype=old_dtype)

        analytical = self.analytical_hessian
        if analytical:
            Hxx_k = t_hessian(f.knl, x, xp)
        else:
            with variable_required_grad(x):
                with variable_required_grad(xp):
                    old_dtype = self.dtype
                    self.to(torch.float64)
                    Hxx_k = tgradcheck.get_numerical_jacobian(
                        partial(grad_k_func, x.double()), xp.double())
                    self.to(dtype=old_dtype)
                    Hxx_k = Hxx_k.to(old_dtype)
        if torch.allclose(x, xp):
            eigenvalues, eigenvectors = torch.linalg.eig(Hxx_k)
            assert (eigenvalues[:, 0] > -eigeps).all(),  " Hessian must be positive definite"
            small_neg_eig = ((eigenvalues[:, 0] > -eigeps) & (eigenvalues[:, 0] < 0))
            if small_neg_eig.any():
                eigenvalues, eigenvectors = torch.linalg.eig(Hxx_k)
                evalz = eigenvalues[:, 0]
                evalz[small_neg_eig] = 0
                Hxx_k = eigenvectors @ torch.diag(evalz) @ eigenvectors.T
        return Hxx_k
//  [python]
def test_eig_backward_complex(self, device, dtype):
        # torch.eig's backward is not supported yet for complex types. We
        # should kill this test once it's implemented.
        a = torch.tensor([[1., 2], [3, 4]], device=device, dtype=dtype, requires_grad=True)
        with self.assertRaisesRegex(RuntimeError,
                                    "eig does not support automatic differentiation for outputs with complex dtype"):
            e, v = torch.eig(a, True)
//  [python]
def test_eig_backward_complex(self, device, dtype):
        # torch.linalg.eig's backward is not supported yet for complex types. We
        # should kill this test once it's implemented.
        a = torch.tensor([[1., 2], [3, 4]], device=device, dtype=dtype, requires_grad=True)
        with self.assertRaisesRegex(RuntimeError,
                                    "linalg.eig does not support automatic differentiation for outputs with complex dtype"):
            e, v = torch.linalg.eig(a)
//  [python]
def _matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    https://discuss.pytorch.org/t/pytorch-square-root-of-a-positive-semi-definite-matrix/100138/5
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: matrix
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.eig(matrix, eigenvectors=True)
    vals = torch.view_as_complex(vals.contiguous())
    vals_pow = vals.pow(p)
    vals_pow = torch.view_as_real(vals_pow)[:, 0]
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs + 1e-6 * torch.eye(vecs.shape[0], device=vecs.device))))
    return matrix_pow
//  [python]
def _matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    https://discuss.pytorch.org/t/pytorch-square-root-of-a-positive-semi-definite-matrix/100138/5
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: matrix
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.linalg.eig(matrix)
    vals = torch.view_as_complex(vals.contiguous())
    vals_pow = vals.pow(p)
    vals_pow = torch.view_as_real(vals_pow)[:, 0]
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs + 1e-6 * torch.eye(vecs.shape[0], device=vecs.device))))
    return matrix_pow
//  [python]
def eig(self, eigenvectors=False):
        """
        Returns the eigenvalues (optionally eigenvectors) of the linear map used in matrix multiplication.

        :param eigenvectors: (bool) Whether to return eigenvectors along with eigenvalues.
        :return: (torch.Tensor) Vector of eigenvalues, optionally a tuple including a matrix of eigenvectors.
        """
        return torch.eig(self.effective_W(), eigenvectors=eigenvectors)
//  [python]
def eig(self, eigenvectors=False):
        """
        Returns the eigenvalues (optionally eigenvectors) of the linear map used in matrix multiplication.

        :param eigenvectors: (bool) Whether to return eigenvectors along with eigenvalues.
        :return: (torch.Tensor) Vector of eigenvalues, optionally a tuple including a matrix of eigenvectors.
        """
        return torch.linalg.eig(self.effective_W())
//  [python]
def find_degenerated_heads(model):
    """
    returns a dict of degenerated head per layer like {layer_idx -> [head_idx, ...]}
    """
    model_params = dict(model.named_parameters())
    degenerated_heads = OrderedDict()
    degenerated_reasons = []

    for layer_idx in range(config["num_hidden_layers"]):
        prune_heads = []
        sigmas_half_inv = model_params["encoder.layer.{}.attention.self.attention_spreads".format(layer_idx)]

        for head_idx in range(config["num_attention_heads"]):
            head_is_degenerated = False


            if config["attention_isotropic_gaussian"]:
                sigma_inv = sigmas_half_inv[head_idx]
                if sigma_inv ** 2 < 1e-5:
                    degenerated_reasons.append("Sigma too low -> uniform attention: sigma**-2= {}".format(sigma_inv ** 2))
                    head_is_degenerated = True
            else:
                sigma_half_inv = sigmas_half_inv[head_idx]
                sigma_inv = sigma_half_inv.transpose(0, 1) @ sigma_half_inv
                eig_values = torch.eig(sigma_inv)[0][:, 0].abs()
                condition_number = eig_values.max() / eig_values.min()

                if condition_number > 1000:
                    degenerated_reasons.append("Covariance matrix is ill defined: condition number = {}".format(condition_number))
                    head_is_degenerated = True
                elif eig_values.max() < 1e-5:
                    degenerated_reasons.append("Covariance matrix is close to 0: largest eigen value = {}".format(eig_values.max()))
                    head_is_degenerated = True

            if head_is_degenerated:
                prune_heads.append(head_idx)

        if prune_heads:
            degenerated_heads[layer_idx] = prune_heads

    if degenerated_heads:
        print("Degenerated heads:")
        reasons = iter(degenerated_reasons)
        table = [(layer, head, next(reasons)) for layer, heads in degenerated_heads.items() for head in heads]
        print(tabulate.tabulate(table, headers=["layer", "head", "reason"]))

    return degenerated_heads
//  [python]
def find_degenerated_heads(model):
    """
    returns a dict of degenerated head per layer like {layer_idx -> [head_idx, ...]}
    """
    model_params = dict(model.named_parameters())
    degenerated_heads = OrderedDict()
    degenerated_reasons = []

    for layer_idx in range(config["num_hidden_layers"]):
        prune_heads = []
        sigmas_half_inv = model_params["encoder.layer.{}.attention.self.attention_spreads".format(layer_idx)]

        for head_idx in range(config["num_attention_heads"]):
            head_is_degenerated = False


            if config["attention_isotropic_gaussian"]:
                sigma_inv = sigmas_half_inv[head_idx]
                if sigma_inv ** 2 < 1e-5:
                    degenerated_reasons.append("Sigma too low -> uniform attention: sigma**-2= {}".format(sigma_inv ** 2))
                    head_is_degenerated = True
            else:
                sigma_half_inv = sigmas_half_inv[head_idx]
                sigma_inv = sigma_half_inv.transpose(0, 1) @ sigma_half_inv
                eig_values = torch.linalg.eig(sigma_inv)[0][:, 0].abs()
                condition_number = eig_values.max() / eig_values.min()

                if condition_number > 1000:
                    degenerated_reasons.append("Covariance matrix is ill defined: condition number = {}".format(condition_number))
                    head_is_degenerated = True
                elif eig_values.max() < 1e-5:
                    degenerated_reasons.append("Covariance matrix is close to 0: largest eigen value = {}".format(eig_values.max()))
                    head_is_degenerated = True

            if head_is_degenerated:
                prune_heads.append(head_idx)

        if prune_heads:
            degenerated_heads[layer_idx] = prune_heads

    if degenerated_heads:
        print("Degenerated heads:")
        reasons = iter(degenerated_reasons)
        table = [(layer, head, next(reasons)) for layer, heads in degenerated_heads.items() for head in heads]
        print(tabulate.tabulate(table, headers=["layer", "head", "reason"]))

    return degenerated_heads
//  [python]
def PCA(X, k, center=True, scale=False):
    n, p = X.size()
    ones = torch.ones(n).view([n, 1])
    h = ((1 / n) * torch.mm(ones, ones.t())) if center else torch.zeros(n * n).view([n, n])
    H = torch.eye(n) - h
    X_center = torch.mm(H.double(), X.double())
    covariance = 1 / (n - 1) * torch.mm(X_center.t(), X_center).view(p, p)
    scaling = torch.sqrt(1 / torch.diag(covariance)).double() if scale else torch.ones(p).double()
    scaled_covariance = torch.mm(torch.diag(scaling).view(p, p), covariance)
    eigenvalues, eigenvectors = torch.eig(scaled_covariance, True)
    components = (eigenvectors[:, :k]).t()
    explained_variance = eigenvalues[:k, 0]
    return {'X': X, 'k': k, 'components': components, 'explained_variance': explained_variance}
//  [python]
def PCA(X, k, center=True, scale=False):
    n, p = X.size()
    ones = torch.ones(n).view([n, 1])
    h = ((1 / n) * torch.mm(ones, ones.t())) if center else torch.zeros(n * n).view([n, n])
    H = torch.eye(n) - h
    X_center = torch.mm(H.double(), X.double())
    covariance = 1 / (n - 1) * torch.mm(X_center.t(), X_center).view(p, p)
    scaling = torch.sqrt(1 / torch.diag(covariance)).double() if scale else torch.ones(p).double()
    scaled_covariance = torch.mm(torch.diag(scaling).view(p, p), covariance)
    eigenvalues, eigenvectors = torch.linalg.eig(scaled_covariance)
    components = (eigenvectors[:, :k]).t()
    explained_variance = eigenvalues[:k, 0]
    return {'X': X, 'k': k, 'components': components, 'explained_variance': explained_variance}
//  [python]
def forward(self, adj, feat, lambda_max=None):
        r"""Compute (Dense) Chebyshev Spectral Graph Convolution layer

        Parameters
        ----------
        adj : torch.Tensor
            The adjacency matrix of the graph to apply Graph Convolution on,
            should be of shape :math:`(N, N)`, where a row represents the destination
            and a column represents the source.
        feat : torch.Tensor
            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
            is size of input feature, :math:`N` is the number of nodes.
        lambda_max : float or None, optional
            A float value indicates the largest eigenvalue of given graph.
            Default: None.

        Returns
        -------
        torch.Tensor
            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
            is size of output feature.
        """
        A = adj.to(feat)
        num_nodes = A.shape[0]

        in_degree = 1 / A.sum(dim=1).clamp(min=1).sqrt()
        D_invsqrt = th.diag(in_degree)
        I = th.eye(num_nodes).to(A)
        L = I - D_invsqrt @ A @ D_invsqrt

        if lambda_max is None:
            lambda_ = th.eig(L)[0][:, 0]
            lambda_max = lambda_.max()

        L_hat = 2 * L / lambda_max - I
        Z = [th.eye(num_nodes).to(A)]
        for i in range(1, self._k):
            if i == 1:
                Z.append(L_hat)
            else:
                Z.append(2 * L_hat @ Z[-1] - Z[-2])

        Zs = th.stack(Z, 0)  # (k, n, n)

        Zh = Zs @ feat.unsqueeze(0) @ self.W
        Zh = Zh.sum(0)

        if self.bias is not None:
            Zh = Zh + self.bias
        return Zh
//  [python]
def forward(self, adj, feat, lambda_max=None):
        r"""Compute (Dense) Chebyshev Spectral Graph Convolution layer

        Parameters
        ----------
        adj : torch.Tensor
            The adjacency matrix of the graph to apply Graph Convolution on,
            should be of shape :math:`(N, N)`, where a row represents the destination
            and a column represents the source.
        feat : torch.Tensor
            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
            is size of input feature, :math:`N` is the number of nodes.
        lambda_max : float or None, optional
            A float value indicates the largest eigenvalue of given graph.
            Default: None.

        Returns
        -------
        torch.Tensor
            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
            is size of output feature.
        """
        A = adj.to(feat)
        num_nodes = A.shape[0]

        in_degree = 1 / A.sum(dim=1).clamp(min=1).sqrt()
        D_invsqrt = th.diag(in_degree)
        I = th.eye(num_nodes).to(A)
        L = I - D_invsqrt @ A @ D_invsqrt

        if lambda_max is None:
            lambda_ = th.linalg.eig(L)[0][:, 0]
            lambda_max = lambda_.max()

        L_hat = 2 * L / lambda_max - I
        Z = [th.eye(num_nodes).to(A)]
        for i in range(1, self._k):
            if i == 1:
                Z.append(L_hat)
            else:
                Z.append(2 * L_hat @ Z[-1] - Z[-2])

        Zs = th.stack(Z, 0)  # (k, n, n)

        Zh = Zs @ feat.unsqueeze(0) @ self.W
        Zh = Zh.sum(0)

        if self.bias is not None:
            Zh = Zh + self.bias
        return Zh
//  [python]
def _eig_helper(H: Tensor):
    """
    wrapper around `torch.linalg.eig` that sorts eigenvalues / eigenvectors by
    ascending eigenvalues, like `torch.linalg.eigh`, and returns the real component
    (since `H` is never complex, there should never be a complex component. however,
    `torch.linalg.eig` always returns a complex tensor, which in this case would
    actually have no complex component)
    """
    version = _parse_version(torch.__version__)
    if version < (1, 9):
        ls, vs = torch.eig(H, eigenvectors=True)
        ls = ls[:, 0]
    else:
        ls, vs = torch.linalg.eig(H)
        ls, vs = ls.real, vs.real

    ls_argsort = torch.argsort(ls)
    vs = vs[:, ls_argsort]
    ls = ls[ls_argsort]
    return ls, vs
//  [python]
def _eig_helper(H: Tensor):
    """
    wrapper around `torch.linalg.eig` that sorts eigenvalues / eigenvectors by
    ascending eigenvalues, like `torch.linalg.eigh`, and returns the real component
    (since `H` is never complex, there should never be a complex component. however,
    `torch.linalg.eig` always returns a complex tensor, which in this case would
    actually have no complex component)
    """
    ls, vs = torch.linalg.eig(H)
    ls, vs = ls.real, vs.real

    ls_argsort = torch.argsort(ls)
    vs = vs[:, ls_argsort]
    ls = ls[ls_argsort]
    return ls, vs
//  [python]
def spectral_radius(m):
    """
    Compute spectral radius of a square 2-D tensor
    :param m: squared 2D tensor
    :return:
    """
    return torch.max(torch.abs(torch.eig(m)[0])).item()
//  [python]
def spectral_radius(m):
    """
    Compute spectral radius of a square 2-D tensor
    :param m: squared 2D tensor
    :return:
    """
    return torch.max(torch.abs(torch.linalg.eig(m)[0])).item()
//  [python]
def test_ARNOLDI_random():
    m= 50
    k= 5
    M= torch.rand(m, m, dtype=torch.float64)

    D0, U0= torch.eig(M, eigenvectors=True)
    D0abs= torch.norm(D0,dim=1)
    D0abs,p= torch.sort(D0abs,descending=True)
    D0= D0[p,:]
    U0= U0[:,p]

    D, U= ARNOLDI.apply(M,k)
    
    # verify spectrum
    assert( torch.norm(D0[:k,0]-D[:,0])/(D0abs[0]*k)<1.0e-14 )
    assert( torch.norm(torch.abs(D0[:k,1])-torch.abs(D[:,1]))/(D0abs[0]*k)<1.0e-14 )
//  [python]
def test_ARNOLDI_random():
    m= 50
    k= 5
    M= torch.rand(m, m, dtype=torch.float64)

    D0, U0= torch.linalg.eig(M)
    D0abs= torch.norm(D0,dim=1)
    D0abs,p= torch.sort(D0abs,descending=True)
    D0= D0[p,:]
    U0= U0[:,p]

    D, U= ARNOLDI.apply(M,k)
    
    # verify spectrum
    assert( torch.norm(D0[:k,0]-D[:,0])/(D0abs[0]*k)<1.0e-14 )
    assert( torch.norm(torch.abs(D0[:k,1])-torch.abs(D[:,1]))/(D0abs[0]*k)<1.0e-14 )
//  [python]
def test_eig_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.eig(a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            a.eig()
//  [python]
def test_eig_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.eig(a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            a.eig()
//  [python]
def eigvec_following(ev_atoms,
                     step,
                     maxstepsize,
                     device,
                     method,
                     hessian_old=None,
                     gradient_old=None,
                     h_old=None):

    Ndim = 3
    old_xyz = torch.Tensor(ev_atoms.get_positions()
                           ).reshape(1, -1, Ndim).to(device)

    if(method == 'NR' or step == 0):
        hessian = get_hessian(atoms=ev_atoms,
                              device=device)

    neural_energy_ase(ev_atoms)
    neural_gradient = -1 * torch.Tensor(neural_force_ase(ev_atoms)).to(device)
    grad = neural_gradient.reshape(-1)

    if (method == 'Powell' and step != 0):
        hessian = powell_update(hessian_old, h_old, gradient_old, grad)

    # eigenvectors are stored in a transposed form
    eigenvalues, eigvecs = torch.eig(hessian, eigenvectors=True)

    # Ordering eigenvalues and eigenvectors in ascending order
    idx = eigenvalues[:, 0].argsort()
    eigenvalues = eigenvalues[idx][:, 0]
    # print(eigenvalues)

    eigvecs = eigvecs[:, idx]
    eigvecs_t = torch.t(eigvecs)
    # print(eigvecs_t)

    F = torch.mv(eigvecs_t, grad).reshape(-1, 1)

    matrix_p = torch.Tensor([[eigenvalues[0], F[0]], [F[0], 0]])
    lambda_p = torch.symeig(matrix_p, eigenvectors=False)[0][1]

    matrix_n = torch.zeros(Ndim * len(old_xyz[0]),
                           Ndim * len(old_xyz[0]))

    for i in range(Ndim * len(old_xyz[0]) - 1):
        matrix_n[i][i] = eigenvalues[i + 1]
        matrix_n[i][Ndim * len(old_xyz[0]) - 1] = F[i + 1]
        matrix_n[Ndim * len(old_xyz[0]) - 1][i] = F[i + 1]

    lambda_n = torch.symeig(matrix_n, eigenvectors=False)[0][0]

    lambda_n = lambda_n.new_full(
        (Ndim * len(old_xyz[0]) - 1,), lambda_n.item()).to(device)

    h_p = -1.0 * F[0] * eigvecs_t[0] / (eigenvalues[0] - lambda_p)
    h_n = -1.0 * F[1:] * eigvecs_t[1:] / \
        ((eigenvalues[1:] - lambda_n).reshape(-1, 1))

    h = torch.add(h_p, torch.sum(h_n, dim=0)
                  ).reshape(-1, len(old_xyz[0]), Ndim)

    step_size = h.norm()
    if step_size <= maxstepsize:
        new_xyz = old_xyz + h
    else:
        new_xyz = old_xyz + (h / (step_size / maxstepsize))

    output = (new_xyz.detach(), grad.detach(),
              hessian.detach(), h.reshape(-1).detach())

    return output
//  [python]
def eigvec_following(ev_atoms,
                     step,
                     maxstepsize,
                     device,
                     method,
                     hessian_old=None,
                     gradient_old=None,
                     h_old=None):

    Ndim = 3
    old_xyz = torch.Tensor(ev_atoms.get_positions()
                           ).reshape(1, -1, Ndim).to(device)

    if(method == 'NR' or step == 0):
        hessian = get_hessian(atoms=ev_atoms,
                              device=device)

    neural_energy_ase(ev_atoms)
    neural_gradient = -1 * torch.Tensor(neural_force_ase(ev_atoms)).to(device)
    grad = neural_gradient.reshape(-1)

    if (method == 'Powell' and step != 0):
        hessian = powell_update(hessian_old, h_old, gradient_old, grad)

    # eigenvectors are stored in a transposed form
    eigenvalues, eigvecs = torch.linalg.eig(hessian)

    # Ordering eigenvalues and eigenvectors in ascending order
    idx = eigenvalues[:, 0].argsort()
    eigenvalues = eigenvalues[idx][:, 0]
    # print(eigenvalues)

    eigvecs = eigvecs[:, idx]
    eigvecs_t = torch.t(eigvecs)
    # print(eigvecs_t)

    F = torch.mv(eigvecs_t, grad).reshape(-1, 1)

    matrix_p = torch.Tensor([[eigenvalues[0], F[0]], [F[0], 0]])
    lambda_p = torch.symeig(matrix_p, eigenvectors=False)[0][1]

    matrix_n = torch.zeros(Ndim * len(old_xyz[0]),
                           Ndim * len(old_xyz[0]))

    for i in range(Ndim * len(old_xyz[0]) - 1):
        matrix_n[i][i] = eigenvalues[i + 1]
        matrix_n[i][Ndim * len(old_xyz[0]) - 1] = F[i + 1]
        matrix_n[Ndim * len(old_xyz[0]) - 1][i] = F[i + 1]

    lambda_n = torch.symeig(matrix_n, eigenvectors=False)[0][0]

    lambda_n = lambda_n.new_full(
        (Ndim * len(old_xyz[0]) - 1,), lambda_n.item()).to(device)

    h_p = -1.0 * F[0] * eigvecs_t[0] / (eigenvalues[0] - lambda_p)
    h_n = -1.0 * F[1:] * eigvecs_t[1:] / \
        ((eigenvalues[1:] - lambda_n).reshape(-1, 1))

    h = torch.add(h_p, torch.sum(h_n, dim=0)
                  ).reshape(-1, len(old_xyz[0]), Ndim)

    step_size = h.norm()
    if step_size <= maxstepsize:
        new_xyz = old_xyz + h
    else:
        new_xyz = old_xyz + (h / (step_size / maxstepsize))

    output = (new_xyz.detach(), grad.detach(),
              hessian.detach(), h.reshape(-1).detach())

    return output
//  [python]
def richardson_iter(grads, model, tol=1e-6):
    def hessian_max_min_eig(grads, model):
        hess = hessian(grads, model)#.to(device)
        eigvalues = torch.eig(hess)[0]
        real_eigvalues = [real for real, img in eigvalues if torch.abs(img) < 1e-6]
        return max(real_eigvalues), min(real_eigvalues)
    newton_direction = [torch.zeros_like(param.data) for param in model.parameters()]
    eigmax, eigmin = hessian_max_min_eig(grads, model)
    alpha = 2 / (eigmax + eigmin)
    num_iters = int(torch.ceil(((eigmax + eigmin) / (2 * eigmin)) * np.log(1/tol)))
    # print("alpha:", alpha)
    for it in range(1, num_iters + 1):
        model.zero_grad()
        Hv = torch.autograd.grad(grads,
                                 model.parameters(),
                                 grad_outputs=newton_direction,
                                 create_graph=False,
                                 retain_graph=True)
        for dt, hv, grad in zip(newton_direction, Hv, grads):
            dt.data = dt.data - alpha * hv - alpha * grad.data
    return newton_direction
//  [python]
def richardson_iter(grads, model, tol=1e-6):
    def hessian_max_min_eig(grads, model):
        hess = hessian(grads, model)#.to(device)
        eigvalues = torch.linalg.eig(hess)[0]
        real_eigvalues = [real for real, img in eigvalues if torch.abs(img) < 1e-6]
        return max(real_eigvalues), min(real_eigvalues)
    newton_direction = [torch.zeros_like(param.data) for param in model.parameters()]
    eigmax, eigmin = hessian_max_min_eig(grads, model)
    alpha = 2 / (eigmax + eigmin)
    num_iters = int(torch.ceil(((eigmax + eigmin) / (2 * eigmin)) * np.log(1/tol)))
    # print("alpha:", alpha)
    for it in range(1, num_iters + 1):
        model.zero_grad()
        Hv = torch.autograd.grad(grads,
                                 model.parameters(),
                                 grad_outputs=newton_direction,
                                 create_graph=False,
                                 retain_graph=True)
        for dt, hv, grad in zip(newton_direction, Hv, grads):
            dt.data = dt.data - alpha * hv - alpha * grad.data
    return newton_direction
//  [python]
def compute_beta(self,xm):

        M = torch.tensor(xm.shape[0], dtype=torch.float)

        xm = xm.detach().requires_grad_(True)

        Kxx = self.K(xm, xm.detach())

        dKxx_dx = autograd.grad(Kxx.sum(), xm)[0] #TODO: is this doing the right thing ?

        # Kxx = Kxx + eta * I
        if self.eta is not None:
            Kxx += self.eta * torch.eye(Kxx.shape[1])

        eigen_vals, eigen_vecs = zip(*[torch.eig(x,eigenvectors = True) for x in Kxx])

        eigen_vals = torch.stack(eigen_vals)
        eigen_vecs = torch.stack(eigen_vecs)

        if self.num_eigs is not None:
            eigen_vals = eigen_vals[:,self.num_eigs]
            eigen_vecs = eigen_vecs[:,:, :self.num_eigs]



        # Compute the Monte Carlo estimate of the gradient of
        # the eigenfunction at x
        dKxx_dx_avg = -dKxx_dx/xm.shape[1] # [M x D]

        #beta = - torch.sqrt(M) * eigen_vecs.t() @ dKxx_dx_avg
        beta = - torch.sqrt(M) * torch.einsum('jki,kjm->jim',(eigen_vecs,dKxx_dx_avg))
        beta *= 1/eigen_vals[:,:,0].unsqueeze(-1)

        return beta, eigen_vals, eigen_vecs
//  [python]
def compute_beta(self,xm):

        M = torch.tensor(xm.shape[0], dtype=torch.float)

        xm = xm.detach().requires_grad_(True)

        Kxx = self.K(xm, xm.detach())

        dKxx_dx = autograd.grad(Kxx.sum(), xm)[0] #TODO: is this doing the right thing ?

        # Kxx = Kxx + eta * I
        if self.eta is not None:
            Kxx += self.eta * torch.eye(Kxx.shape[1])

        eigen_vals, eigen_vecs = zip(*[torch.linalg.eig(x) for x in Kxx])

        eigen_vals = torch.stack(eigen_vals)
        eigen_vecs = torch.stack(eigen_vecs)

        if self.num_eigs is not None:
            eigen_vals = eigen_vals[:,self.num_eigs]
            eigen_vecs = eigen_vecs[:,:, :self.num_eigs]



        # Compute the Monte Carlo estimate of the gradient of
        # the eigenfunction at x
        dKxx_dx_avg = -dKxx_dx/xm.shape[1] # [M x D]

        #beta = - torch.sqrt(M) * eigen_vecs.t() @ dKxx_dx_avg
        beta = - torch.sqrt(M) * torch.einsum('jki,kjm->jim',(eigen_vecs,dKxx_dx_avg))
        beta *= 1/eigen_vals[:,:,0].unsqueeze(-1)

        return beta, eigen_vals, eigen_vecs
//  [python]
def average_quaternion(q: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    """
    calculate the average of quaternions
    :param q: (W, J, 4)
    :param w: (W)
    :return: (J, 4)
    """
    w = w.view(-1, 1, 1, 1)
    joint_num = q.shape[1]
    q = q.reshape(-1, 4, 1)
    q_t = q.reshape(-1, 1, 4)
    qq_t = torch.bmm(q, q_t)
    qq_t = qq_t.reshape(w.shape[0], joint_num, 4, 4)
    averages = []
    m = torch.sum(w * qq_t, dim=0)
    for jm in m:
        eigenvalues, eigenvectors = torch.eig(jm, eigenvectors=True)
        q_avg = eigenvectors[:, torch.argmax(eigenvalues[:, 0])]
        averages.append(q_avg)
    return torch.stack(averages)
//  [python]
def average_quaternion(q: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    """
    calculate the average of quaternions
    :param q: (W, J, 4)
    :param w: (W)
    :return: (J, 4)
    """
    w = w.view(-1, 1, 1, 1)
    joint_num = q.shape[1]
    q = q.reshape(-1, 4, 1)
    q_t = q.reshape(-1, 1, 4)
    qq_t = torch.bmm(q, q_t)
    qq_t = qq_t.reshape(w.shape[0], joint_num, 4, 4)
    averages = []
    m = torch.sum(w * qq_t, dim=0)
    for jm in m:
        eigenvalues, eigenvectors = torch.linalg.eig(jm)
        q_avg = eigenvectors[:, torch.argmax(eigenvalues[:, 0])]
        averages.append(q_avg)
    return torch.stack(averages)
//  [python]
def generate_mesh(coeff, A_6x6):
    np.random.seed(222)
    color = np.random.rand(100, 3)

    polyorder = torch.from_numpy(gen_polynomial_orders(4)).to(device)
    valid_idx = [i for i in range(A_6x6.shape[3]) if torch.min(torch.eig(A_6x6[0,:,:,i])[0][:,0])<0]

    N = 128
    x, y, z = np.mgrid[:N, :N, :N]
    x = torch.from_numpy((x / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)
    y = torch.from_numpy((y / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)
    z = torch.from_numpy((z / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)

    ### generate meshes 
    meshlist = []
    feasible_idx = []
    total_func_val = torch.ones(N, N, N).to(device)*100
    for i, idx in enumerate(valid_idx):
        func_val = (coeff[0,:,idx] * (x**polyorder[:,0] * y**polyorder[:,1] * z**polyorder[:,2])).sum(dim=3).detach()
        verts0, faces0 = marching_cubes(func_val, 0.0)
        if verts0.shape[0]>3: # if valid
            mesh0 = mesh_from_vf(verts0, faces0, color[valid_idx[i],:])
            meshlist.append(mesh0)
            feasible_idx.append(idx)

        # for whole shape
        total_func_val = torch.min(total_func_val, func_val) 

    ### generate a mesh for whole shape
    verts0, faces0 = marching_cubes(total_func_val, 0.0)
    total_mesh = mesh_from_vf(verts0, faces0, np.array([0.5, 0.5, 0.5])) # gray color

    return meshlist, total_mesh, feasible_idx
//  [python]
def generate_mesh(coeff, A_6x6):
    np.random.seed(222)
    color = np.random.rand(100, 3)

    polyorder = torch.from_numpy(gen_polynomial_orders(4)).to(device)
    valid_idx = [i for i in range(A_6x6.shape[3]) if torch.min(torch.linalg.eig(A_6x6[0,:,:,i])[0].real)<0]

    N = 128
    x, y, z = np.mgrid[:N, :N, :N]
    x = torch.from_numpy((x / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)
    y = torch.from_numpy((y / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)
    z = torch.from_numpy((z / (N-1)*2.2 -1.1).astype('float32')).cuda().unsqueeze(dim=3)#.reshape(-1,1)

    ### generate meshes 
    meshlist = []
    feasible_idx = []
    total_func_val = torch.ones(N, N, N).to(device)*100
    for i, idx in enumerate(valid_idx):
        func_val = (coeff[0,:,idx] * (x**polyorder[:,0] * y**polyorder[:,1] * z**polyorder[:,2])).sum(dim=3).detach()
        verts0, faces0 = marching_cubes(func_val, 0.0)
        if verts0.shape[0]>3: # if valid
            mesh0 = mesh_from_vf(verts0, faces0, color[valid_idx[i],:])
            meshlist.append(mesh0)
            feasible_idx.append(idx)

        # for whole shape
        total_func_val = torch.min(total_func_val, func_val) 

    ### generate a mesh for whole shape
    verts0, faces0 = marching_cubes(total_func_val, 0.0)
    total_mesh = mesh_from_vf(verts0, faces0, np.array([0.5, 0.5, 0.5])) # gray color

    return meshlist, total_mesh, feasible_idx
//  [python]
def Ising(self, **kwargs):
        self.g = 1.0
        self.name = 'Ising'
        self.description = ''
        self.dtype=torch.double
        self.g = torch.tensor(1.0, dtype=self.dtype)
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        
        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sZ = torch.tensor([[1.0, 0], [0, -1.0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) + 0.5 * self.g * (torch.einsum('xz, wy -> zyxw', sZ, torch.eye(2)) + torch.einsum('xz, wy -> zyxw', torch.eye(2), sZ))
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        integ, _ = integrate.quad(lambda x: np.sqrt((self.g - np.cos(x))**2 + np.sin(x)**2), -np.pi, np.pi) 
        self.E_exact = - integ / (2 * np.pi)
//  [python]
def Ising(self, **kwargs):
        self.g = 1.0
        self.name = 'Ising'
        self.description = ''
        self.dtype=torch.double
        self.g = torch.tensor(1.0, dtype=self.dtype)
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        
        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sZ = torch.tensor([[1.0, 0], [0, -1.0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) + 0.5 * self.g * (torch.einsum('xz, wy -> zyxw', sZ, torch.eye(2)) + torch.einsum('xz, wy -> zyxw', torch.eye(2), sZ))
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.linalg.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        integ, _ = integrate.quad(lambda x: np.sqrt((self.g - np.cos(x))**2 + np.sin(x)**2), -np.pi, np.pi) 
        self.E_exact = - integ / (2 * np.pi)
//  [python]
def is_pos_def(X):
  eig_values,_ = torch.eig(X,eigenvectors=False)
  if_psd = torch.all(eig_values[:,0]>0)
  return if_psd
//  [python]
def is_pos_def(X):
  eig_values,_ = torch.linalg.eig(X)
  if_psd = torch.all(eig_values[:,0]>0)
  return if_psd
//  [python]
def get_eig_dense(adj):
    adj = adj.float()
    rowsum = adj.sum(1)
    r_inv =rowsum.pow(-0.5)
    r_mat_inv = torch.diag(r_inv)
    nr_adj = torch.matmul(torch.matmul(r_mat_inv,adj),r_mat_inv)
    graph_laplacian = torch.eye(adj.shape[0])-nr_adj
    L,V = torch.eig(graph_laplacian,eigenvectors=True)
    return L.T[0],V
//  [python]
def get_eig_dense(adj):
    adj = adj.float()
    rowsum = adj.sum(1)
    r_inv =rowsum.pow(-0.5)
    r_mat_inv = torch.diag(r_inv)
    nr_adj = torch.matmul(torch.matmul(r_mat_inv,adj),r_mat_inv)
    graph_laplacian = torch.eye(adj.shape[0])-nr_adj
    L,V = torch.linalg.eig(graph_laplacian)
    return L.T[0],V
//  [python]
def __gt__(self, other):
        """
        Greater than (abstraction relationship
        :param other: Second operand
        :return: True/False
        """
        # Eigen values of C - D
        eigv = torch.eig(other.C - self.C, eigenvectors=False)
        return float(torch.max(eigv)) > 0.0
//  [python]
def __gt__(self, other):
        """
        Greater than (abstraction relationship
        :param other: Second operand
        :return: True/False
        """
        # Eigen values of C - D
        eigv = torch.linalg.eig(other.C - self.C)
        return float(torch.max(eigv)) > 0.0
//  [python]
def XY(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Heisenberg XY'
        self.description = ''
        self.dtype=torch.double

        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sY = torch.tensor([[0, -1.0], [1.0, 0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) - self.g * torch.einsum('xz, wy -> zyxw', sY, sY)
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        if self.g == 1.0:
            self.E_exact = -4 / math.pi
        else:
            self.E_exact = 0
//  [python]
def XY(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Heisenberg XY'
        self.description = ''
        self.dtype=torch.double

        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sY = torch.tensor([[0, -1.0], [1.0, 0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) - self.g * torch.einsum('xz, wy -> zyxw', sY, sY)
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.linalg.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        if self.g == 1.0:
            self.E_exact = -4 / math.pi
        else:
            self.E_exact = 0
//  [python]
def __ge__(self, other):
        """
        Greater or equal (abstraction relationship)
        :param other: Second operand
        :return: True/False
        """
        # Eigen values of C - D
        eigv = torch.eig(other.C - self.C, eigenvectors=False)
        return float(torch.max(eigv)) >= 0.0
//  [python]
def __ge__(self, other):
        """
        Greater or equal (abstraction relationship)
        :param other: Second operand
        :return: True/False
        """
        # Eigen values of C - D
        eigv = torch.linalg.eig(other.C - self.C)
        return float(torch.max(eigv)) >= 0.0
//  [python]
def XXZ(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Heisenberg XXZ'
        self.description = ''
        self.dtype=torch.double

        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sY = torch.tensor([[0, -1.0], [1.0, 0]], dtype=self.dtype)
        sZ = torch.tensor([[1.0, 0], [0, -1.0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) - torch.einsum('xz, wy -> zyxw', sY, sY) + self.g * torch.einsum('xz, wy -> zyxw', sZ, sZ)
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        if self.g == 1.0:
            self.E_exact = 1 - 4 * np.log(2)
        else:
            self.E_exact = 0
//  [python]
def XXZ(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Heisenberg XXZ'
        self.description = ''
        self.dtype=torch.double

        sX = torch.tensor([[0, 1.0], [1.0, 0]], dtype=self.dtype)
        sY = torch.tensor([[0, -1.0], [1.0, 0]], dtype=self.dtype)
        sZ = torch.tensor([[1.0, 0], [0, -1.0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', sX, sX) - torch.einsum('xz, wy -> zyxw', sY, sY) + self.g * torch.einsum('xz, wy -> zyxw', sZ, sZ)
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.linalg.eig(H.reshape(4, 4))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(4).reshape(2, 2, 2, 2)

        if self.g == 1.0:
            self.E_exact = 1 - 4 * np.log(2)
        else:
            self.E_exact = 0
//  [python]
def __lt__(self, other):
        """
        Less than (abstraction relationship)
        :param other: Second operand
        :return: True/False
        """
        # Eigen value of C - D
        eigv = torch.eig(other.C - self.C, eigenvectors=False)
        return float(torch.max(eigv)) < 0.0
//  [python]
def __lt__(self, other):
        """
        Less than (abstraction relationship)
        :param other: Second operand
        :return: True/False
        """
        # Eigen value of C - D
        eigv = torch.linalg.eig(other.C - self.C)
        return float(torch.max(eigv)) < 0.0
//  [python]
def __le__(self, other):
        """
        Less or equal than (abstraction relatioship)
        :param other: Second operand
        :return: True/False
        """
        eigv = torch.eig(other.C - self.C, eigenvectors=False)
        return float(torch.max(eigv)) <= 0.0
//  [python]
def __le__(self, other):
        """
        Less or equal than (abstraction relatioship)
        :param other: Second operand
        :return: True/False
        """
        eigv = torch.linalg.eig(other.C - self.C)
        return float(torch.max(eigv)) <= 0.0
//  [python]
def apply(C, func):
    e, v = torch.eig(C, eigenvectors=True)
    reg = 0
    with torch.no_grad():
        reg = 1e-5 -  e[:,0].min()
    e = func(e[:,0] + reg)
    C_ = torch.einsum('id,d,jd->ij', (v, e, v))
    return C_
//  [python]
def apply(C, func):
    e, v = torch.linalg.eig(C)
    reg = 0
    with torch.no_grad():
        reg = 1e-5 -  e[:,0].min()
    e = func(e[:,0] + reg)
    C_ = torch.einsum('id,d,jd->ij', (v, e, v))
    return C_
//  [python]
def Potts(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Potts'
        self.description = ''
        self.dtype=torch.double

        Mz = torch.tensor([[2, 0, 0], [0, -1.0, 0], [0, 0, -1.0]], dtype=self.dtype)
        Mx1 = torch.tensor([[0, 1.0, 0], [0, 0, 1.0], [1.0, 0, 0]], dtype=self.dtype)
        Mx2 = torch.tensor([[0, 0, 1.0], [1.0, 0, 0], [0, 1.0, 0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', Mx1, Mx2) + torch.einsum('xz, wy -> zyxw', Mx2, Mx1) + 0.5 * self.g * (torch.einsum('xz, wy -> zyxw', Mz, torch.eye(3)) + torch.einsum('xz, wy -> zyxw', torch.eye(3), Mz))
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.eig(H.reshape(9, 9))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(9).reshape(3, 3, 3, 3)

        # E_exact = -1.8156
        # l = (3 + 2 * math.sqrt(3) + 2) / (math.sqrt(3) + 2)
        # h = l / 3
        # E_exact = 1 - 0.5 * (l + h) - math.sqrt(0.25 * (l-h)**2 + 2*h**2)

        self.E_exact = 0
//  [python]
def Potts(self, **kwargs):
        self.g = 1.0
        if 'g' in list(kwargs.keys()):
            self.g = kwargs['g']
        self.name = 'Potts'
        self.description = ''
        self.dtype=torch.double

        Mz = torch.tensor([[2, 0, 0], [0, -1.0, 0], [0, 0, -1.0]], dtype=self.dtype)
        Mx1 = torch.tensor([[0, 1.0, 0], [0, 0, 1.0], [1.0, 0, 0]], dtype=self.dtype)
        Mx2 = torch.tensor([[0, 0, 1.0], [1.0, 0, 0], [0, 1.0, 0]], dtype=self.dtype)
        H = torch.einsum('xz, wy -> zyxw', Mx1, Mx2) + torch.einsum('xz, wy -> zyxw', Mx2, Mx1) + 0.5 * self.g * (torch.einsum('xz, wy -> zyxw', Mz, torch.eye(3)) + torch.einsum('xz, wy -> zyxw', torch.eye(3), Mz))
        H = 0.5 * (H + H.permute(1, 0, 3, 2))
        self.bias = np.amax(torch.linalg.eig(H.reshape(9, 9))[0].view(-1).numpy())
        self.ham = H - self.bias * torch.eye(9).reshape(3, 3, 3, 3)

        # E_exact = -1.8156
        # l = (3 + 2 * math.sqrt(3) + 2) / (math.sqrt(3) + 2)
        # h = l / 3
        # E_exact = 1 - 0.5 * (l + h) - math.sqrt(0.25 * (l-h)**2 + 2*h**2)

        self.E_exact = 0