jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering."""
    if largest is None:
        largest = False
    E, Z = torch.linalg.eigh(A, UPLO="U")
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
d = 3
H = tc.randn((d, d), dtype=tc.complex128)
H = H + H.conj().t()
print('随机厄密矩阵 H = \n', H)

lm, u = tc.linalg.eigh(H)
//  [python]
def compute_loss(output, labels, infos, eikr_m1, eikr_0, eikr_1):
    
    label_list  = [i for i in labels.numpy()]
    num_catoms  = [infos[i]['num_catom'] for i in label_list]
    num_c       = [infos[i]['num_c'] for i in label_list]
    num_v       = [infos[i]['num_v'] for i in label_list]
    references  = [infos[i]['bands'] for i in label_list]
    num_matelements = [2*i**2 for i in num_catoms] 
    
    loss = []
    o, h = output
    o_s = torch.split(o, num_catoms)
    h_s = torch.split(h, num_matelements)
    for onsite, hopping, catoms, cbands, vbands, rbands in zip(o_s, h_s, num_catoms, num_c, num_v, references):
        h0,h1 = torch.chunk(hopping,2,0)
        h0 = torch.reshape(h0,[catoms,catoms])
        on = torch.reshape(onsite,[-1])
        h0 = h0 - torch.diag(torch.diag(h0)) + torch.diag(on)
        h1 = torch.reshape(h1,[catoms,catoms])
        hm1 = h1.t()
        h0 = h0.expand(26, catoms,catoms)
        h1 = h1.expand(26, catoms,catoms)
        hm1 = hm1.expand(26, catoms,catoms)
        hk = eikr_m1*hm1 + eikr_0*h0 + eikr_1*h1
        w, v = torch.symeig(hk,eigenvectors=True)
        e_r = w[:,ceil(catoms/2) - vbands:ceil(catoms/2) + cbands]
        loss.append(F.smooth_l1_loss(e_r,rbands.to(device)).view(-1,1))
    
    return torch.mean(torch.cat(loss))
//  [python]
def compute_loss(output, labels, infos, eikr_m1, eikr_0, eikr_1):
    
    label_list  = [i for i in labels.numpy()]
    num_catoms  = [infos[i]['num_catom'] for i in label_list]
    num_c       = [infos[i]['num_c'] for i in label_list]
    num_v       = [infos[i]['num_v'] for i in label_list]
    references  = [infos[i]['bands'] for i in label_list]
    num_matelements = [2*i**2 for i in num_catoms] 
    
    loss = []
    o, h = output
    o_s = torch.split(o, num_catoms)
    h_s = torch.split(h, num_matelements)
    for onsite, hopping, catoms, cbands, vbands, rbands in zip(o_s, h_s, num_catoms, num_c, num_v, references):
        h0,h1 = torch.chunk(hopping,2,0)
        h0 = torch.reshape(h0,[catoms,catoms])
        on = torch.reshape(onsite,[-1])
        h0 = h0 - torch.diag(torch.diag(h0)) + torch.diag(on)
        h1 = torch.reshape(h1,[catoms,catoms])
        hm1 = h1.t()
        h0 = h0.expand(26, catoms,catoms)
        h1 = h1.expand(26, catoms,catoms)
        hm1 = hm1.expand(26, catoms,catoms)
        hk = eikr_m1*hm1 + eikr_0*h0 + eikr_1*h1
        w, v = torch.linalg.eigh(hk)
        e_r = w[:,ceil(catoms/2) - vbands:ceil(catoms/2) + cbands]
        loss.append(F.smooth_l1_loss(e_r,rbands.to(device)).view(-1,1))
    
    return torch.mean(torch.cat(loss))
//  [python]
def compute_spectra(W: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the spectra of graph Laplacian from its adjacency matrix.

    Performs an eigendecomposition (w/o assuming additional structure)
    using `torch.linalg.eigh` (previously used `torch.symeig`) on a normalized
    graph laplacian. Converts from the adjacency matrix to the laplacian
    internally.

    Parameters
    ----------
    W: torch.Tensor
        Batch of graph adjacency matricies.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        Batch of eigenvalues and eigenvectors of the graph Laplacian.
    """
    # compute laplacian
    L = adjacency_to_laplacian(W)
    # normalize laplacian
    L_norm = normalize_laplacian(L)
    # perform eigen decomp
    # come out in ascending order,
    E, V = torch.symeig(L_norm, eigenvectors=True)

    return E, V
//  [python]
def compute_spectra(W: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the spectra of graph Laplacian from its adjacency matrix.

    Performs an eigendecomposition (w/o assuming additional structure)
    using `torch.linalg.eigh` (previously used `torch.symeig`) on a normalized
    graph laplacian. Converts from the adjacency matrix to the laplacian
    internally.

    Parameters
    ----------
    W: torch.Tensor
        Batch of graph adjacency matricies.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        Batch of eigenvalues and eigenvectors of the graph Laplacian.
    """
    # compute laplacian
    L = adjacency_to_laplacian(W)
    # normalize laplacian
    L_norm = normalize_laplacian(L)
    # perform eigen decomp
    # come out in ascending order,
    E, V = torch.linalg.eigh(L_norm, UPLO="L")

    return E, V
//  [python]
def test_project_changes_torch():
    n, d, seed = 20, 2, None
    X = torch.randn(n, d)

    G = gram_utils.gram_matrix(X.numpy())
    lamduhs, vecs = LA.eigh(G)
    lamduhs[0] = -1
    G = vecs.T @ np.diag(lamduhs) @ vecs
    G = torch.from_numpy(G)
    e, v = torch.symeig(G, eigenvectors=True)
    assert e.min().item() < -0.5

    before = G.numpy().copy()
    after = gram_utils.onto_psd(G.numpy(), out=G.numpy())

    assert not np.allclose(before, after)
    assert not torch.allclose(torch.from_numpy(before), G)
    e, v = torch.symeig(G, eigenvectors=True)
    assert e.min() > -0.25
//  [python]
def test_project_changes_torch():
    n, d, seed = 20, 2, None
    X = torch.randn(n, d)

    G = gram_utils.gram_matrix(X.numpy())
    lamduhs, vecs = LA.eigh(G)
    lamduhs[0] = -1
    G = vecs.T @ np.diag(lamduhs) @ vecs
    G = torch.from_numpy(G)
    e, v = torch.linalg.eigh(G, UPLO="U")
    assert e.min().item() < -0.5

    before = G.numpy().copy()
    after = gram_utils.onto_psd(G.numpy(), out=G.numpy())

    assert not np.allclose(before, after)
    assert not torch.allclose(torch.from_numpy(before), G)
    e, v = torch.linalg.eigh(G, UPLO="U")
    assert e.min() > -0.25
//  [python]
def measure_largest_eigval(W):
    assert(len(W.shape) == 2)
    ## W: embedding matrix, shape=[n_data, dim]
    eig_vals, _ = torch.symeig(torch.matmul(W.T, W))
    largest_eigval_ratio = torch.max(eig_vals) / torch.sum(eig_vals)
    return float(largest_eigval_ratio)
//  [python]
def measure_largest_eigval(W):
    assert(len(W.shape) == 2)
    ## W: embedding matrix, shape=[n_data, dim]
    eig_vals, _ = torch.linalg.eigh(torch.matmul(W.T, W))
    largest_eigval_ratio = torch.max(eig_vals) / torch.sum(eig_vals)
    return float(largest_eigval_ratio)
//  [python]
def rotmat2qvec_tensor(R):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R[0, 0].item(), R[0, 1].item(), R[0, 2].item(), \
        R[1, 0].item(), R[1, 1].item(), R[1, 2].item(), R[2, 0].item(), R[2, 1].item(), R[2, 2].item()
    K = torch.tensor([
        [Rxx - Ryy - Rzz, 0, 0, 0],
        [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
        [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
        [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
    eigvals, eigvecs = torch.eig(K, eigenvectors=True)
    qvec = eigvecs[:, torch.argmax(eigvals.real)][[3, 0, 1, 2]]
    if qvec[0] < 0:
        qvec *= -1
    return qvec
//  [python]
def rotmat2qvec_tensor(R):
    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R[0, 0].item(), R[0, 1].item(), R[0, 2].item(), \
        R[1, 0].item(), R[1, 1].item(), R[1, 2].item(), R[2, 0].item(), R[2, 1].item(), R[2, 2].item()
    K = torch.tensor([
        [Rxx - Ryy - Rzz, 0, 0, 0],
        [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
        [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
        [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
    eigvals, eigvecs = torch.linalg.eigh(K)
    qvec = eigvecs[[3, 0, 1, 2], torch.argmax(eigvals)]
    if qvec[0] < 0:
        qvec *= -1
    return qvec
//  [python]
def kl(mu_x,cov_x, mu_y, cov_y, epsilon, device):
    
    epsilon=epsilon.to(device)
    
    eigvals_x = torch.linalg.eigvals(cov_x).real
    eigvals_y = torch.linalg.eigvals(cov_y).real
    for tens in [eigvals_x, eigvals_y]:
        aux = torch.maximum(tens, torch.zeros_like(tens))
        tens = aux+epsilon # add small constant to avoid infinite gradients from ln(x)
    
    first = torch.log(eigvals_y).sum()-torch.log(eigvals_x).sum()-mu_x.shape[0]
    
    evals_y, Q = torch.symeig(cov_y,eigenvectors=True) 
    Q = Q.real
    # much more stable eigvals than torch.linalg.eigh
    second = torch.dot(1/evals_y.flatten(),torch.diag(Q@cov_x@Q.t()).flatten())
    mu_diff = mu_x.flatten()-mu_y.flatten()
    aux = torch.matmul(cov_y,mu_diff)
    third = torch.dot(mu_diff,aux)
    
    return 0.5*(first+second+third)
//  [python]
def kl(mu_x,cov_x, mu_y, cov_y, epsilon, device):
    
    epsilon=epsilon.to(device)
    
    eigvals_x = torch.linalg.eigvals(cov_x).real
    eigvals_y = torch.linalg.eigvals(cov_y).real
    for tens in [eigvals_x, eigvals_y]:
        aux = torch.maximum(tens, torch.zeros_like(tens))
        tens = aux+epsilon # add small constant to avoid infinite gradients from ln(x)
    
    first = torch.log(eigvals_y).sum()-torch.log(eigvals_x).sum()-mu_x.shape[0]
    
    _, Q = torch.linalg.eigh(cov_y) 
    Q = Q.real
    # much more stable eigvals than torch.linalg.eigh
    second = torch.dot(1/eigvals_y.flatten(),torch.diag(Q@cov_x@Q.t()).flatten())
    mu_diff = mu_x.flatten()-mu_y.flatten()
    aux = torch.matmul(cov_y,mu_diff)
    third = torch.dot(mu_diff,aux)
    
    return 0.5*(first+second+third)
//  [python]
def cal_orbital_and_energies(overlap_matrix, full_hamiltonian):
    eigvals, eigvecs = torch.symeig(overlap_matrix, eigenvectors=True)
    eps = 1e-8 * torch.ones_like(eigvals)
    eigvals = torch.where(eigvals > 1e-8, eigvals, eps)
    frac_overlap = eigvecs / torch.sqrt(eigvals).unsqueeze(-2)

    Fs = torch.bmm(torch.bmm(frac_overlap.transpose(-1, -2), full_hamiltonian), frac_overlap).to('cpu')
    orbital_energies, orbital_coefficients = torch.symeig(Fs, eigenvectors=True)
    orbital_coefficients = torch.bmm(frac_overlap, orbital_coefficients)
    return orbital_energies, orbital_coefficients
//  [python]
def cal_orbital_and_energies(overlap_matrix, full_hamiltonian):
    eigvals, eigvecs = torch.linalg.eigh(overlap_matrix)
    eps = 1e-8 * torch.ones_like(eigvals)
    eigvals = torch.where(eigvals > 1e-8, eigvals, eps)
    frac_overlap = eigvecs / torch.sqrt(eigvals).unsqueeze(-2)

    Fs = torch.bmm(torch.bmm(frac_overlap.transpose(-1, -2), full_hamiltonian), frac_overlap).to('cpu')
    orbital_energies, orbital_coefficients = torch.linalg.eigh(Fs)
    orbital_coefficients = torch.bmm(frac_overlap, orbital_coefficients)
    return orbital_energies, orbital_coefficients
//  [python]
def _sym_pos_def_matrix(rdm, sym_pos_def=False, verbosity=0, who="unknown", **kwargs):
    rdm_asym = 0.5 * (rdm - rdm.conj().t())
    rdm = 0.5 * (rdm + rdm.conj().t())
    if verbosity > 0:
        log.info(f"{who} norm(rdm_sym) {rdm.norm()} norm(rdm_asym) {rdm_asym.norm()}")
    if sym_pos_def:
        with torch.no_grad():
            D, U=  torch.symeig(rdm, eigenvectors=True)
            if D.min() < 0:
                log.info(f"{who} max(diag(rdm)) {D.max()} min(diag(rdm)) {D.min()}")
                D = torch.clamp(D, min=0)
                rdm_posdef = U @ torch.diag(D) @ U.conj().t()
                rdm.copy_(rdm_posdef)
    norm= _cast_to_real(rdm.diagonal().sum(),who=who,**kwargs)
    rdm = rdm / norm
    return rdm
//  [python]
def _sym_pos_def_matrix(rdm, sym_pos_def=False, verbosity=0, who="unknown", **kwargs):
    rdm_asym = 0.5 * (rdm - rdm.conj().t())
    rdm = 0.5 * (rdm + rdm.conj().t())
    if verbosity > 0:
        log.info(f"{who} norm(rdm_sym) {rdm.norm()} norm(rdm_asym) {rdm_asym.norm()}")
    if sym_pos_def:
        with torch.no_grad():
            if _torch_version_check("1.8.1"):
                D, U=  torch.linalg.eigh(rdm)
            else:
                D, U= torch.symeig(rdm, eigenvectors=True)
            if D.min() < 0:
                log.info(f"{who} max(diag(rdm)) {D.max()} min(diag(rdm)) {D.min()}")
                D = torch.clamp(D, min=0)
                rdm_posdef = U @ torch.diag(D) @ U.conj().t()
                rdm.copy_(rdm_posdef)
    norm= _cast_to_real(rdm.diagonal().sum(),who=who,**kwargs)
    rdm = rdm / norm
    return rdm
//  [python]
def _sym_funcm_impl(func, **kwargs):
    func = partial(func, **kwargs)

    def _impl(x):
        e, v = torch.symeig(x, eigenvectors=True)
        return v @ torch.diag_embed(func(e)) @ v.transpose(-1, -2)

    return torch.jit.script(_impl)
//  [python]
def _sym_funcm_impl(func, **kwargs):
    func = partial(func, **kwargs)

    def _impl(x):
        e, v = torch.linalg.eigh(x, "U")
        return v @ torch.diag_embed(func(e)) @ v.transpose(-1, -2)

    return torch.jit.script(_impl)
//  [python]
def compute_eigendecomposition(self, impl="eigh"):
        if impl == "eigh":
            self.evals, self.evecs = torch.symeig(self.data)
        elif impl == "svd":
            _, self.evals, self.evecs = torch.svd(self.data, some=False)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl="eigh"):
        if impl == "eigh":
            self.evals, self.evecs = torch.linalg.eigh(self.data)
        elif impl == "svd":
            _, self.evals, self.evecs = torch.svd(self.data, some=False)
        else:
            raise NotImplementedError
//  [python]
def Dataset_train(data_num, name, M, N, K, ty_state, device='cpu'):
    """
    The data needed to generate the network pre-training, including probability and 
    density matrix reverse mapping data, see paper ``Ultrafast quantum state tomography 
    with feed-forward neural networks``.

    Args:
        data_num (int): The number of dataset.
        name (str): The name of saved file.
        M (tensor): The POVM, size (K, 2, 2).
        N (int): The number of qubits.
        K (int): The number of POVM elements.
        ty_state (str): The type of state, include 'mixed' and 'pure'.
        device (torch.device): GPU or CPU.

    Save:
        The dataset is saved using h5py at ``datasets/``.
    """
    num_group = num_to_groups(data_num, 5000)  # 5000 is the size of each packet
    for idx, num in enumerate(num_group):
        X = np.zeros((num, K**N))
        Y = np.zeros((num, 4**N))

        for i in range(num):
            _, rho = State().Get_real_random_state(N)
            rho = torch.from_numpy(rho).to(torch.complex64).to(device)

            if ty_state == 'pure':
                P = qmt_torch_pure(rho, [M] * N)
            else:
                P = qmt_torch(rho, [M] * N)

            # reverse state-mapping precess
            # chol
            #L, _ = torch.linalg.cholesky_ex(rho)
            #Y_i = (L.real + L.imag.T).reshape(-1)

            # chol_h
            eigenvalues, eigenvecs = torch.eig(rho, eigenvectors=True)
            eigenvalues_sqrt = torch.sqrt(torch.abs(eigenvalues[:,0]))
            A = eigenvecs @ torch.diag(eigenvalues_sqrt)
            H = torch.matmul(A, eigenvecs.T.conj())
            L = torch.tril(H)
            Y_i = (L.real + L.imag.T).reshape(-1)
            #print(torch.dist(L @ L.mT, rho))

            # data
            X[i] = P.cpu().numpy()
            Y[i] = Y_i.cpu().numpy()

            if i % 1000 == 0:
                print('---Generating data {:d}th:{:>5d} / {:d}'.format(idx + 1, idx * 5000 + i + 1000, data_num))

        # save dataset
        h5f = h5py.File('datasets/'+name+str(idx)+'.h5', 'w')
        h5f['data'] = X
        h5f['target'] = Y
        h5f.close()

    return
//  [python]
def Dataset_train(data_num, name, M, N, K, ty_state, device='cpu'):
    """
    The data needed to generate the network pre-training, including probability and 
    density matrix reverse mapping data, see paper ``Ultrafast quantum state tomography 
    with feed-forward neural networks``.

    Args:
        data_num (int): The number of dataset.
        name (str): The name of saved file.
        M (tensor): The POVM, size (K, 2, 2).
        N (int): The number of qubits.
        K (int): The number of POVM elements.
        ty_state (str): The type of state, include 'mixed' and 'pure'.
        device (torch.device): GPU or CPU.

    Save:
        The dataset is saved using h5py at ``datasets/``.
    """
    num_group = num_to_groups(data_num, 5000)  # 5000 is the size of each packet
    for idx, num in enumerate(num_group):
        X = np.zeros((num, K**N))
        Y = np.zeros((num, 4**N))

        for i in range(num):
            _, rho = State().Get_real_random_state(N)
            rho = torch.from_numpy(rho).to(torch.complex64).to(device)

            if ty_state == 'pure':
                P = qmt_torch_pure(rho, [M] * N)
            else:
                P = qmt_torch(rho, [M] * N)

            # reverse state-mapping precess
            # chol
            #L, _ = torch.linalg.cholesky_ex(rho)
            #Y_i = (L.real + L.imag.T).reshape(-1)

            # chol_h
            eigenvalues, eigenvecs = torch.linalg.eigh(rho)
            eigenvalues_sqrt = torch.sqrt(torch.abs(eigenvalues))
            A = eigenvecs * eigenvalues_sqrt
            H = torch.matmul(A, eigenvecs.T.conj())
            L = torch.tril(H)
            Y_i = (L.real + L.imag.T).reshape(-1)
            #print(torch.dist(L @ L.mT, rho))

            # data
            X[i] = P.cpu().numpy()
            Y[i] = Y_i.cpu().numpy()

            if i % 1000 == 0:
                print('---Generating data {:d}th:{:>5d} / {:d}'.format(idx + 1, idx * 5000 + i + 1000, data_num))

        # save dataset
        h5f = h5py.File('datasets/'+name+str(idx)+'.h5', 'w')
        h5f['data'] = X
        h5f['target'] = Y
        h5f.close()

    return
//  [python]
def sym_inv_sqrtm2(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Symmetric matrix inverse square root, with square root return also.

    Parameters
    ----------
    x : torch.Tensor
        symmetric matrix

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        :math:`x^{-1/2}`, :math:`x^{1/2}`
    """
    e, v = torch.symeig(x, eigenvectors=True)
    sqrt_e = torch.sqrt(e)
    inv_sqrt_e = torch.reciprocal(sqrt_e)
    return (
        v @ torch.diag_embed(inv_sqrt_e) @ v.transpose(-1, -2),
        v @ torch.diag_embed(sqrt_e) @ v.transpose(-1, -2),
    )
//  [python]
def sym_inv_sqrtm2(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Symmetric matrix inverse square root, with square root return also.

    Parameters
    ----------
    x : torch.Tensor
        symmetric matrix

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        :math:`x^{-1/2}`, :math:`x^{1/2}`
    """
    e, v = torch.linalg.eigh(x, "U")
    sqrt_e = torch.sqrt(e)
    inv_sqrt_e = torch.reciprocal(sqrt_e)
    return (
        v @ torch.diag_embed(inv_sqrt_e) @ v.transpose(-1, -2),
        v @ torch.diag_embed(sqrt_e) @ v.transpose(-1, -2),
    )
//  [python]
def compute_eigendecomposition(self, impl="eigh"):
        self.evals = dict()
        self.evecs = dict()
        if impl == "eigh":
            for layer_id in self.generator.layer_collection.layers.keys():
                a, g = self.data[layer_id]
                evals_a, evecs_a = torch.symeig(a)
                evals_g, evecs_g = torch.symeig(g)
                self.evals[layer_id] = (evals_a, evals_g)
                self.evecs[layer_id] = (evecs_a, evecs_g)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl="eigh"):
        self.evals = dict()
        self.evecs = dict()
        if impl == "eigh":
            for layer_id in self.generator.layer_collection.layers.keys():
                a, g = self.data[layer_id]
                evals_a, evecs_a = torch.linalg.eigh(a)
                evals_g, evecs_g = torch.linalg.eigh(g)
                self.evals[layer_id] = (evals_a, evals_g)
                self.evecs[layer_id] = (evecs_a, evecs_g)
        else:
            raise NotImplementedError
//  [python]
def __call__(self, g):
        # Step1: heat kernel diffusion
        # t A
        device = g.device
        eweight = self.t * g.edata.get(
            self.eweight_name, F.ones((g.num_edges(),), F.float32, device)
        )
        num_nodes = g.num_nodes()
        mat = F.zeros((num_nodes, num_nodes), F.float32, device)
        src, dst = g.edges()
        src, dst = F.astype(src, F.int64), F.astype(dst, F.int64)
        mat[dst, src] = eweight
        # t (A - I_n)
        nids = F.astype(g.nodes(), F.int64)
        mat[nids, nids] = mat[nids, nids] - self.t

        if is_bidirected(g):
            e, V = torch.eig(mat, eigenvectors=True)
            e = e[:,0]
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(mat.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(device)

        # Step2: sparsification
        num_nodes = g.num_nodes()
        eps = self.get_eps(num_nodes, diff_mat)
        dst, src = (diff_mat >= eps).nonzero(as_tuple=False).t()
        data_dict = {g.canonical_etypes[0]: (src, dst)}
        new_g = update_graph_structure(g, data_dict, copy_edata=False)
        new_g.edata[self.eweight_name] = diff_mat[dst, src]

        return new_g
//  [python]
def __call__(self, g):
        # Step1: heat kernel diffusion
        # t A
        device = g.device
        eweight = self.t * g.edata.get(
            self.eweight_name, F.ones((g.num_edges(),), F.float32, device)
        )
        num_nodes = g.num_nodes()
        mat = F.zeros((num_nodes, num_nodes), F.float32, device)
        src, dst = g.edges()
        src, dst = F.astype(src, F.int64), F.astype(dst, F.int64)
        mat[dst, src] = eweight
        # t (A - I_n)
        nids = F.astype(g.nodes(), F.int64)
        mat[nids, nids] = mat[nids, nids] - self.t

        if is_bidirected(g):
            e, V = torch.linalg.eigh(mat, UPLO="U")
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(mat.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(device)

        # Step2: sparsification
        num_nodes = g.num_nodes()
        eps = self.get_eps(num_nodes, diff_mat)
        dst, src = (diff_mat >= eps).nonzero(as_tuple=False).t()
        data_dict = {g.canonical_etypes[0]: (src, dst)}
        new_g = update_graph_structure(g, data_dict, copy_edata=False)
        new_g.edata[self.eweight_name] = diff_mat[dst, src]

        return new_g
//  [python]
def compute_posenc_stats(data, is_undirected=True, laplacian_norm='sym', max_freqs=10, eigvec_norm='L2'):
    """Precompute positional encodings for the given graph.

    Args:
        data: PyG graph
        pe_types: Positional encoding types to precompute statistics for.
            This can also be a combination, e.g. 'eigen+rw_landing'
        is_undirected: True if the graph is expected to be undirected
        cfg: Main configuration node

    Returns:
        Extended PyG Data object.
    """

    # Basic preprocessing of the input graph.
    if hasattr(data, 'num_nodes'):
        N = data.num_nodes  # Explicitly given number of nodes, e.g. ogbg-ppa
    else:
        N = data["region"].x.shape[0]  # Number of nodes, including disconnected nodes.
    laplacian_norm_type = laplacian_norm.lower()
    if laplacian_norm_type == 'none':
        laplacian_norm_type = None
    edge_index = torch.cat((data["region", "spatial", "region"].edge_index, data["region", "temporal", "region"].edge_index), dim=1).long()
    if is_undirected:
        undir_edge_index = edge_index
    else:
        undir_edge_index = to_undirected(edge_index)

    # Eigen values and vectors.
    evals, evects = None, None
    # Eigen-decomposition with numpy, can be reused for Heat kernels.
    L = to_scipy_sparse_matrix(
        *get_laplacian(undir_edge_index, normalization=laplacian_norm_type,
                        num_nodes=N)
    )
    evals, evects = torch.symeig(torch.tensor(L.toarray()))
    evals = evals.numpy()
    evects = evects.numpy()
        
    data.EigVals, data.EigVecs = get_lap_decomp_stats(
        evals=evals, evects=evects,
        max_freqs=max_freqs,
        eigvec_norm=eigvec_norm)

    return data
//  [python]
def compute_posenc_stats(data, is_undirected=True, laplacian_norm='sym', max_freqs=10, eigvec_norm='L2'):
    """Precompute positional encodings for the given graph.

    Args:
        data: PyG graph
        pe_types: Positional encoding types to precompute statistics for.
            This can also be a combination, e.g. 'eigen+rw_landing'
        is_undirected: True if the graph is expected to be undirected
        cfg: Main configuration node

    Returns:
        Extended PyG Data object.
    """

    # Basic preprocessing of the input graph.
    if hasattr(data, 'num_nodes'):
        N = data.num_nodes  # Explicitly given number of nodes, e.g. ogbg-ppa
    else:
        N = data["region"].x.shape[0]  # Number of nodes, including disconnected nodes.
    laplacian_norm_type = laplacian_norm.lower()
    if laplacian_norm_type == 'none':
        laplacian_norm_type = None
    edge_index = torch.cat((data["region", "spatial", "region"].edge_index, data["region", "temporal", "region"].edge_index), dim=1).long()
    if is_undirected:
        undir_edge_index = edge_index
    else:
        undir_edge_index = to_undirected(edge_index)

    # Eigen values and vectors.
    evals, evects = None, None
    # Eigen-decomposition with numpy, can be reused for Heat kernels.
    L = to_scipy_sparse_matrix(
        *get_laplacian(undir_edge_index, normalization=laplacian_norm_type,
                        num_nodes=N)
    )
    evals, evects = torch.linalg.eigh(torch.tensor(L.toarray()))
    evals = evals.numpy()
    evects = evects.numpy()
        
    data.EigVals, data.EigVecs = get_lap_decomp_stats(
        evals=evals, evects=evects,
        max_freqs=max_freqs,
        eigvec_norm=eigvec_norm)

    return data
//  [python]
def eig_coord_from_dist(D):
    M = (D[:, :1, :] + D[:, :, :1] - D) / 2
    L, V = torch.symeig(M,eigenvectors=True)
    L = torch.diag_embed(L)
    X = torch.matmul(V, L.clamp(min=0).sqrt())
    return X[:, :, -3:].detach()
//  [python]
def eig_coord_from_dist(D):
    M = (D[:, :1, :] + D[:, :, :1] - D) / 2
    L, V = torch.linalg.eigh(M)
    L = torch.diag_embed(L)
    X = torch.matmul(V, L.clamp(min=0).sqrt())
    return X[:, :, -3:].detach()
//  [python]
def __init__(self, generator, data=None, examples=None):
        self._check_data_examples(data, examples)

        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks(examples)
            for layer_id, layer in self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.symeig(a, eigenvectors=True)
                evals_g, evecs_g = torch.symeig(g, eigenvectors=True)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1), evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def __init__(self, generator, data=None, examples=None):
        self._check_data_examples(data, examples)

        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks(examples)
            for layer_id, layer in self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.linalg.eigh(a)
                evals_g, evecs_g = torch.linalg.eigh(g)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1), evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def matrix_to_quaternion(rot):
    if(rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot 

    k = [
        [ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
        [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
        [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
        [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]
    ]

    k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def matrix_to_quaternion(rot):
    if(rot.shape[-2:] != (3, 3)):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot 

    k = [
        [ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
        [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
        [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
        [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]
    ]

    k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def forward(self, x):
        """Perform PCA.

        Parameters
        ----------
        x : Tensor [shape=(..., M+1)]
            Input vectors.

        Returns
        -------
        e : Tensor [shape=(N,)]
            Eigenvalues.

        v : Tensor [shape=(M+1, N)]
            Eigenvectors.

        m : Tensor [shape=(M+1,)]
            Mean vector.

        Examples
        --------
        >>> x = diffsptk.nrand(10, 3)
        >>> x.size()
        torch.Size([10, 4])
        >>> pca = diffsptk.PCA(3, 3)
        >>> e, _, _ = pca(x)
        >>> e
        tensor([1.3465, 0.7497, 0.4447])
        >>> y = pca.transform(x)
        >>> y.size()
        torch.Size([10, 3])

        """
        check_size(x.size(-1), self.order + 1, "dimension of input")

        x = x.reshape(-1, x.size(-1)).T
        assert self.n_comp + 1 <= x.size(1), "Number of data samples is too small"

        e, v = torch.symeig(self.cov(x), eigenvectors=True)
        e = e[-self.n_comp :]
        v = v[:, -self.n_comp :]
        if self.sort == "descending":
            e = e.flip(-1)
            v = v.flip(-1)
        self.v[:] = v
        self.m[:] = x.mean(1)
        return e, self.v, self.m
//  [python]
def forward(self, x):
        """Perform PCA.

        Parameters
        ----------
        x : Tensor [shape=(..., M+1)]
            Input vectors.

        Returns
        -------
        e : Tensor [shape=(N,)]
            Eigenvalues.

        v : Tensor [shape=(M+1, N)]
            Eigenvectors.

        m : Tensor [shape=(M+1,)]
            Mean vector.

        Examples
        --------
        >>> x = diffsptk.nrand(10, 3)
        >>> x.size()
        torch.Size([10, 4])
        >>> pca = diffsptk.PCA(3, 3)
        >>> e, _, _ = pca(x)
        >>> e
        tensor([1.3465, 0.7497, 0.4447])
        >>> y = pca.transform(x)
        >>> y.size()
        torch.Size([10, 3])

        """
        check_size(x.size(-1), self.order + 1, "dimension of input")

        x = x.reshape(-1, x.size(-1)).T
        assert self.n_comp + 1 <= x.size(1), "Number of data samples is too small"

        e, v = torch.linalg.eigh(self.cov(x))
        e = e[-self.n_comp :]
        v = v[:, -self.n_comp :]
        if self.sort == "descending":
            e = e.flip(-1)
            v = v.flip(-1)
        self.v[:] = v
        self.m[:] = x.mean(1)
        return e, self.v, self.m
//  [python]
def projx(self, x: torch.Tensor, threshold: float = 1e-5) -> torch.Tensor:
        """Clamps the eigenvalues to be non-negative."""
        A = self.devectorize(x)
        dtype = A.dtype
        A = A.double()
        L, Q = torch.symeig(A)
        L = torch.clamp(L, min=threshold)
        P = (Q * L.unsqueeze(-2)) @ Q.mH
        return self.vectorize(P).to(dtype)
//  [python]
def projx(self, x: torch.Tensor, threshold: float = 1e-5) -> torch.Tensor:
        """Clamps the eigenvalues to be non-negative."""
        A = self.devectorize(x)
        dtype = A.dtype
        A = A.double()
        L, Q = torch.linalg.eigh(A)
        L = torch.clamp(L, min=threshold)
        P = (Q * L.unsqueeze(-2)) @ Q.mH
        return self.vectorize(P).to(dtype)
//  [python]
def make_pos_emb(h, w, dim):
    g = [[] for _ in range(h * w)]
    for i in range(h):
        for j in range(w):
            idx = i * h + j
            if i - 1 >= 0:
                g[idx].append(idx - h)
            if i + 1 < h:
                g[idx].append(idx + h)
            if j - 1 >= 0:
                g[idx].append(idx - 1)
            if j + 1 < w:
                g[idx].append(idx + 1)
    a = torch.zeros(h * w, h * w)
    d_isqrt = torch.zeros(h * w, h * w)
    for i, node in enumerate(g):
        for edge in node:
            a[i, edge] = 1
            d_isqrt[i, i] = 1 / math.sqrt(len(node))
    eye = torch.eye(h * w)
    lap = eye - d_isqrt @ a @ d_isqrt
    vals, vecs = torch.symeig(lap, eigenvectors=True)
    # proj = torch.nn.init.orthogonal_(torch.empty(dim, h * w), math.sqrt(h * w / dim))
    return vecs[:dim, :].T * math.sqrt(h * w)
//  [python]
def make_pos_emb(h, w, dim):
    g = [[] for _ in range(h * w)]
    for i in range(h):
        for j in range(w):
            idx = i * h + j
            if i - 1 >= 0:
                g[idx].append(idx - h)
            if i + 1 < h:
                g[idx].append(idx + h)
            if j - 1 >= 0:
                g[idx].append(idx - 1)
            if j + 1 < w:
                g[idx].append(idx + 1)
    a = torch.zeros(h * w, h * w)
    d_isqrt = torch.zeros(h * w, h * w)
    for i, node in enumerate(g):
        for edge in node:
            a[i, edge] = 1
            d_isqrt[i, i] = 1 / math.sqrt(len(node))
    eye = torch.eye(h * w)
    lap = eye - d_isqrt @ a @ d_isqrt
    vals, vecs = torch.linalg.eigh(lap)
    # proj = torch.nn.init.orthogonal_(torch.empty(dim, h * w), math.sqrt(h * w / dim))
    return vecs[:dim, :].T * math.sqrt(h * w)
//  [python]
def sqrtmh(A):
    """Compute the square root of a symmetric positive definite matrix."""
    L, Q = torch.symeig(A)
    zero = torch.zeros((), device=L.device, dtype=L.dtype)
    threshold = L.max(-1).values * L.size(-1) * torch.finfo(L.dtype).eps
    L = L.where(L > threshold.unsqueeze(-1), zero)  # zero out small components
    return (Q * L.sqrt().unsqueeze(-2)) @ Q.mH
//  [python]
def sqrtmh(A):
    """Compute the square root of a symmetric positive definite matrix."""
    L, Q = torch.linalg.eigh(A)
    zero = torch.zeros((), device=L.device, dtype=L.dtype)
    threshold = L.max(-1).values * L.size(-1) * torch.finfo(L.dtype).eps
    L = L.where(L > threshold.unsqueeze(-1), zero)  # zero out small components
    return (Q * L.sqrt().unsqueeze(-2)) @ Q.mH
//  [python]
def eigensolver(M):
    """ Eigensolver
        manually symmetrize M before the eigen decomposition
    """
    return torch.symeig(0.5*(M+M.t()))
//  [python]
def eigensolver(M):
    """ Eigensolver
        manually symmetrize M before the eigen decomposition
    """
    return torch.linalg.eigh(0.5*(M+M.t()))
//  [python]
def get_whitening_parameters(patches):
    n, c, h, w = patches.shape
    est_covariance = torch.cov(patches.view(n, c * h * w).t())
    eigenvalues, eigenvectors = torch.symeig(est_covariance, eigenvectors=True)
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.t().reshape(
        c * h * w, c, h, w
    ).flip(0)
//  [python]
def get_whitening_parameters(patches):
    # As a high-level summary, we're basically finding the high-dimensional oval that best fits the data here.
    # We can then later use this information to map the input information to a nicely distributed sphere, where also
    # the most significant features of the inputs each have their own axis. This significantly cleans things up for the
    # rest of the neural network and speeds up training.
    n, c, h, w = patches.shape
    est_covariance = torch.cov(patches.view(n, c * h * w).t())
    eigenvalues, eigenvectors = torch.linalg.eigh(
        est_covariance, UPLO="U"
    )  # this is the same as saying we want our eigenvectors, with the specification that the matrix be an upper triangular matrix (instead of a lower-triangular matrix)
    return eigenvalues.flip(0).view(-1, 1, 1, 1), eigenvectors.t().reshape(
        c * h * w, c, h, w
    ).flip(0)
//  [python]
def extract_rot_scale(cov: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # we add eps to the diagonal to make sure we never get negative eigenvalues
    S, R = torch.symeig(cov + torch.eye(3, device=cov.device) * 1e-8, eigenvectors=True)
    scaling = S.sqrt().nan_to_num(nan=1e-6)
    # we multiply with the determinant to make sure the
    # rotation matrix has determinant +1
    # the matrix_to_quaternion function assumes this
    rot = normalize(matrix_to_quaternion(R * R.det()[..., None, None]))
    return rot, scaling
//  [python]
def extract_rot_scale(cov: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # we add eps to the diagonal to make sure we never get negative eigenvalues
    S, R = torch.linalg.eigh(cov + torch.eye(3, device=cov.device) * 1e-8, UPLO="U")
    scaling = S.sqrt().nan_to_num(nan=1e-6)
    # we multiply with the determinant to make sure the
    # rotation matrix has determinant +1
    # the matrix_to_quaternion function assumes this
    rot = normalize(matrix_to_quaternion(R * R.det()[..., None, None]))
    return rot, scaling
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = (r_mat - r_mat.transpose(-1, -2)) #costruisco la amtrice skew
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2 
    c_angle = (torch.einsum('...ii', r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle) # trovo l'angolo di rotazione [B]
    scale = (angle / (2 * s_angle)) #[B]
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat # [B x 3 x 3]

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[...,0,0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.symeig(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[...,-1,:]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[...,None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def log_rmat(r_mat: torch.Tensor) -> torch.Tensor:
    skew_mat = (r_mat - r_mat.transpose(-1, -2)) #costruisco la amtrice skew
    sk_vec = skew2vec(skew_mat)
    s_angle = (sk_vec).norm(p=2, dim=-1) / 2 
    c_angle = (torch.einsum('...ii', r_mat) - 1) / 2
    angle = torch.atan2(s_angle, c_angle) # trovo l'angolo di rotazione [B]
    scale = (angle / (2 * s_angle)) #[B]
    # if s_angle = 0, i.e. rotation by 0 or pi (180), we get NaNs
    # by definition, scale values are 0 if rotating by 0.
    # This also breaks down if rotating by pi, fix further down
    scale[angle == 0.0] = 0.0
    log_r_mat = scale[..., None, None] * skew_mat # [B x 3 x 3]

    # Check for NaNs caused by 180deg rotations.
    nanlocs = log_r_mat[...,0,0].isnan()
    nanmats = r_mat[nanlocs]
    # We need to use an alternative way of finding the logarithm for nanmats,
    # Use eigendecomposition to discover axis of rotation.
    # By definition, these are symmetric, so use eigh.
    # NOTE: linalg.eig() isn't in torch 1.8,
    #       and torch.eig() doesn't do batched matrices
    eigval, eigvec = torch.linalg.eigh(nanmats)
    # Final eigenvalue == 1, might be slightly off because floats, but other two are -ve.
    # this *should* just be the last column if the docs for eigh are true.
    nan_axes = eigvec[...,-1,:]
    nan_angle = angle[nanlocs]
    nan_skew = vec2skew(nan_angle[...,None] * nan_axes)
    log_r_mat[nanlocs] = nan_skew
    return log_r_mat
//  [python]
def submersion_inv(self, X, check_in_manifold=True):
        with torch.no_grad():
            L, Q = torch.symeig(X)
        if check_in_manifold and not self.in_manifold_eigen(L):
            raise InManifoldError(X, self)
        if self.inv is None:
            raise InverseError(self)
        with torch.no_grad():
            Q = Q[..., -self.rank :]
            L = L[..., -self.rank :]
            L = self.inv(L)
        return L, Q
//  [python]
def submersion_inv(self, X, check_in_manifold=True):
        with torch.no_grad():
            L, Q = torch.linalg.eigh(X)
        if check_in_manifold and not self.in_manifold_eigen(L):
            raise InManifoldError(X, self)
        if self.inv is None:
            raise InverseError(self)
        with torch.no_grad():
            Q = Q[..., -self.rank :]
            L = L[..., -self.rank :]
            L = self.inv(L)
        return L, Q
//  [python]
def display_moments(*, scale=1, **kwargs):
    """Display moments of a shape (not a test)."""
    shape = create_shape(**kwargs)

    local_average = shape.point_moments(order=1, scale=scale)
    local_cov = shape.point_moments(order=2, scale=scale, central=True)

    local_QL = torch.eig(local_cov, eigenvectors=True)
    local_nuv = local_QL.eigenvectors  # (N, 3, 3)

    local_frame = (
        local_QL.eigenvectors * local_QL.eigenvalues.view(-1, 1, 3).sqrt()
    )
    local_frame = local_frame / scale

    # Pick 10 points at random:
    if shape.n_points > 100:
        mask = torch.randperm(shape.n_points)[:100]
    else:
        mask = torch.arange(shape.n_points)

    # To color our points, let's visualize the window around point 0:
    density = torch.zeros_like(shape.points[:, 0])
    density[0] = 1
    Conv = shape.point_convolution(scale=scale, normalize=True)
    density = Conv.T @ density
    assert density.shape == (len(shape.points),)

    # Let's also use the ratio of eigenvalues (as a proxy for curvature):
    anisotropy = local_QL.eigenvalues[:, 0] / local_QL.eigenvalues.sum(-1)

    point_frame = shape.point_frames(scale=scale).transpose(1, 2)
    curvedness = shape.point_curvedness(scale=scale)
    shape_index = shape.point_shape_indices(scale=scale).abs()

    # Our surface points:
    if shape.triangles is None:
        spheres = vd.Points(
            shape.points, c=(0.5, 0.5, 0.5), r=60 / (shape.n_points ** (1 / 3))
        )
    else:
        spheres = shape.to_vedo()

    spheres_1 = (
        spheres.clone()
        .alpha(0.5)
        .cmap("viridis", density, vmin=0)
        .add_scalarbar()
    )
    spheres_2 = (
        spheres.clone()
        .cmap("viridis", (3 * anisotropy) ** (1 / 3), vmin=0)
        .add_scalarbar()
    )
    spheres_3 = (
        spheres.clone().cmap("viridis", curvedness, vmin=0).add_scalarbar()
    )
    spheres_4 = (
        spheres.clone()
        .cmap("RdBu_r", shape_index, vmin=-1, vmax=1)
        .add_scalarbar()
    )

    # Vectors to the local average:
    quiver = vd.Arrows(
        shape.points[mask], local_average[mask], c="green", alpha=0.9
    )

    # Local tangent frames:
    s = 1.5 * shape.standard_deviation[0] / len(mask) ** (1 / 2)

    plt = vd.Plotter(shape=(2, 2), axes=1)
    plt.at(0).show(
        spheres_1,
        quiver,
        vd.Text2D("Density around point 0", pos="top-middle"),
    )
    plt.at(1).show(
        spheres_2,
        *vedo_frames(shape.points[mask], s * local_nuv[mask]),
        vd.Text2D("Anisotropy", pos="top-middle"),
    )
    plt.at(2).show(
        spheres_3,
        *vedo_frames(shape.points[mask], s * local_frame[mask]),
        vd.Text2D("Curvedness", pos="top-middle"),
    )
    plt.at(3).show(
        spheres_4,
        *vedo_frames(shape.points[mask], s * point_frame[mask]),
        vd.Text2D("Shape index", pos="top-middle"),
    )
    plt.interactive()
//  [python]
def display_moments(*, scale=1, **kwargs):
    """Display moments of a shape (not a test)."""
    shape = create_shape(**kwargs)

    local_average = shape.point_moments(order=1, scale=scale)
    local_cov = shape.point_moments(order=2, scale=scale, central=True)

    local_QL = torch.linalg.eigh(local_cov)
    local_nuv = local_QL.eigenvectors  # (N, 3, 3)

    local_frame = (
        local_QL.eigenvectors * local_QL.eigenvalues.view(-1, 1, 3).sqrt()
    )
    local_frame = local_frame / scale

    # Pick 10 points at random:
    if shape.n_points > 100:
        mask = torch.randperm(shape.n_points)[:100]
    else:
        mask = torch.arange(shape.n_points)

    # To color our points, let's visualize the window around point 0:
    density = torch.zeros_like(shape.points[:, 0])
    density[0] = 1
    Conv = shape.point_convolution(scale=scale, normalize=True)
    density = Conv.T @ density
    assert density.shape == (len(shape.points),)

    # Let's also use the ratio of eigenvalues (as a proxy for curvature):
    anisotropy = local_QL.eigenvalues[:, 0] / local_QL.eigenvalues.sum(-1)

    point_frame = shape.point_frames(scale=scale).transpose(1, 2)
    curvedness = shape.point_curvedness(scale=scale)
    shape_index = shape.point_shape_indices(scale=scale).abs()

    # Our surface points:
    if shape.triangles is None:
        spheres = vd.Points(
            shape.points, c=(0.5, 0.5, 0.5), r=60 / (shape.n_points ** (1 / 3))
        )
    else:
        spheres = shape.to_vedo()

    spheres_1 = (
        spheres.clone()
        .alpha(0.5)
        .cmap("viridis", density, vmin=0)
        .add_scalarbar()
    )
    spheres_2 = (
        spheres.clone()
        .cmap("viridis", (3 * anisotropy) ** (1 / 3), vmin=0)
        .add_scalarbar()
    )
    spheres_3 = (
        spheres.clone().cmap("viridis", curvedness, vmin=0).add_scalarbar()
    )
    spheres_4 = (
        spheres.clone()
        .cmap("RdBu_r", shape_index, vmin=-1, vmax=1)
        .add_scalarbar()
    )

    # Vectors to the local average:
    quiver = vd.Arrows(
        shape.points[mask], local_average[mask], c="green", alpha=0.9
    )

    # Local tangent frames:
    s = 1.5 * shape.standard_deviation[0] / len(mask) ** (1 / 2)

    plt = vd.Plotter(shape=(2, 2), axes=1)
    plt.at(0).show(
        spheres_1,
        quiver,
        vd.Text2D("Density around point 0", pos="top-middle"),
    )
    plt.at(1).show(
        spheres_2,
        *vedo_frames(shape.points[mask], s * local_nuv[mask]),
        vd.Text2D("Anisotropy", pos="top-middle"),
    )
    plt.at(2).show(
        spheres_3,
        *vedo_frames(shape.points[mask], s * local_frame[mask]),
        vd.Text2D("Curvedness", pos="top-middle"),
    )
    plt.at(3).show(
        spheres_4,
        *vedo_frames(shape.points[mask], s * point_frame[mask]),
        vd.Text2D("Shape index", pos="top-middle"),
    )
    plt.interactive()
//  [python]
def network_deconvolution(x, alpha=1.0):
    assert x.ndim == 3
    batch_size, N, N1 = x.shape
    assert N == N1

    vals, vecs = torch.eig(x, eigenvectors=True)
    #assert torch.max(vals) <= 1.0, f'network deconv: max eigenvalue should be normalized'

    nd_vals = vals[:,0] / (1 + alpha*vals[:,0])  # network deco
    nd = torch.matmul(torch.matmul(vecs, torch.diag_embed(nd_vals)), torch.transpose(vecs, 1, 2) )
    return nd
//  [python]
def network_deconvolution(x, alpha=1.0):
    assert x.ndim == 3
    batch_size, N, N1 = x.shape
    assert N == N1

    vals, vecs = torch.linalg.eigh(x)
    #assert torch.max(vals) <= 1.0, f'network deconv: max eigenvalue should be normalized'

    nd_vals = vals / (1 + alpha*vals)  # network deco
    nd = torch.matmul(torch.matmul(vecs, torch.diag_embed(nd_vals)), torch.transpose(vecs, 1, 2) )
    return nd
//  [python]
def compute_eigendecomposition(self, impl="eigh"):
        s = self.data.size()
        M = self.data.view(s[0] * s[1], s[2] * s[3])
        if impl == "eigh":
            self.evals, self.evecs = torch.symeig(M,eigenvectors=True)
        elif impl == "svd":
            _, self.evals, self.evecs = torch.svd(M, some=False)
        else:
            raise NotImplementedError