jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.eigh(M)
        elif impl == "svd":
            _, self.evals, self.evecs = torch.svd(M, some=False)
        else:
            raise NotImplementedError
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        r"""
        Returns a randomly sampled matrix on the manifold as

        .. math::

            WW^\intercal \qquad W_{i,j} \sim \texttt{init_}

        By default ``init\_`` is a (xavier) normal distribution, so that the
        returned matrix follows a Wishart distribution.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = PSSD(layer.weight.size())
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\_ (callable): Optional.
                    A function that takes a tensor and fills it in place according
                    to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(
                *(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype
            )
            init_(X)
            X = X @ X.transpose(-2, -1)
            L, Q = torch.symeig(X, eigenvectors=True)
            L = L[..., -self.rank :]
            Q = Q[..., -self.rank :]
            if factorized:
                return L, Q
            else:
                return (Q * L.unsqueeze(-2)) @ Q.transpose(-2, -1)
//  [python]
def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
        r"""
        Returns a randomly sampled matrix on the manifold as

        .. math::

            WW^\intercal \qquad W_{i,j} \sim \texttt{init_}

        By default ``init\_`` is a (xavier) normal distribution, so that the
        returned matrix follows a Wishart distribution.

        The output of this method can be used to initialize a parametrized tensor
        that has been parametrized with this or any other manifold as::

            >>> layer = nn.Linear(20, 20)
            >>> M = PSSD(layer.weight.size())
            >>> geotorch.register_parametrization(layer, "weight", M)
            >>> layer.weight = M.sample()

        Args:
            init\_ (callable): Optional.
                    A function that takes a tensor and fills it in place according
                    to some distribution. See
                    `torch.init <https://pytorch.org/docs/stable/nn.init.html>`_.
                    Default: ``torch.nn.init.xavier_normal_``
        """
        with torch.no_grad():
            device = self[0].base.device
            dtype = self[0].base.dtype
            X = torch.empty(
                *(self.tensorial_size + (self.n, self.n)), device=device, dtype=dtype
            )
            init_(X)
            X = X @ X.transpose(-2, -1)
            L, Q = torch.linalg.eigh(X)
            L = L[..., -self.rank :]
            Q = Q[..., -self.rank :]
            if factorized:
                return L, Q
            else:
                return (Q * L.unsqueeze(-2)) @ Q.transpose(-2, -1)
//  [python]
def forward(ctx, input, eps=1e-2):
        dim = input.dim()
        ctx.dim = dim
        use_cuda = input.is_cuda
        if input.size(0) < 300:
            input = input.cpu()
        # e, v = torch.symeig(input, eigenvectors=True)
        e, v = torch.linalg.eigh(input, UPLO='U')
        if use_cuda and input.size(0) < 300:
            e = e.cuda()
            v = v.cuda()
        e = e.clamp(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.save_for_backward(e_sqrt, v)
        e_rsqrt = e_sqrt.reciprocal()

        if dim > 2:
            output = v.bmm(v.permute(0, 2, 1) * e_rsqrt.unsqueeze(-1))
        else:
            output = v.mm(v.t() * e_rsqrt.view(-1, 1))
        return output
//  [python]
def _project_impl(self, pt: Point) -> Point:
        assert pt.shape == (self.m, self.m), \
            f"Input to AngularCKA.project() must be a m by m matrix but is size {pt.shape}!"
        # 1. Ensure matrix is symmetric
        pt = (pt + pt.T) / 2
        if self._unbiased:
            # 2a. Normally we would enforce positive (semi-) definiteness, but we can't guarantee that here due to
            # the masking of the diagonal of the pt matrix
            pass
            # 3a. If the unbiased flag is set, we cannot risk calling center() a second time on the same data, since
            # this would be essentially center(mask * center(data)), and the outer center will interact with the mask
            mask = torch.ones_like(pt) - torch.eye(self.m, dtype=pt.dtype, device=pt.device)
            pt = _matrix_unit_norm(mask * pt)
        else:
            # 2b. Ensure matrix is positive (semi-) definite by clipping its eigenvalues
            e, v = torch.symeig(pt)
            e = torch.clip(e, min=0., max=None)
            pt = v @ torch.diag(e) @ v.T
            # 3b. If the unbiased flag is not set, we can be extra careful here and both center and normalize. Centering
            # twice is the same as centering once, so we don't need to worry whether pt is already centered.
            pt = _matrix_unit_norm(center(pt))
        return pt
//  [python]
def _project_impl(self, pt: Point) -> Point:
        assert pt.shape == (self.m, self.m), \
            f"Input to AngularCKA.project() must be a m by m matrix but is size {pt.shape}!"
        # 1. Ensure matrix is symmetric
        pt = (pt + pt.T) / 2
        if self._unbiased:
            # 2a. Normally we would enforce positive (semi-) definiteness, but we can't guarantee that here due to
            # the masking of the diagonal of the pt matrix
            pass
            # 3a. If the unbiased flag is set, we cannot risk calling center() a second time on the same data, since
            # this would be essentially center(mask * center(data)), and the outer center will interact with the mask
            mask = torch.ones_like(pt) - torch.eye(self.m, dtype=pt.dtype, device=pt.device)
            pt = _matrix_unit_norm(mask * pt)
        else:
            # 2b. Ensure matrix is positive (semi-) definite by clipping its eigenvalues
            e, v = torch.linalg.eigh(pt)
            e = torch.clip(e, min=0., max=None)
            pt = v @ torch.diag(e) @ v.T
            # 3b. If the unbiased flag is not set, we can be extra careful here and both center and normalize. Centering
            # twice is the same as centering once, so we don't need to worry whether pt is already centered.
            pt = _matrix_unit_norm(center(pt))
        return pt
//  [python]
def Rotation_invariance(local_cor, local_weight=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    LRF_direction = torch.symeig(cov_matrix, eigenvectors=True)[1]  # 计算特征值（和向量），这里得到他的最大的特征向量

    return LRF_direction
//  [python]
def Rotation_invariance(local_cor, local_weight=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    LRF_direction = torch.linalg.eigh(cov_matrix)[1]  # 计算特征值（和向量），这里得到他的最大的特征向量

    return LRF_direction
//  [python]
scaling_factor = (cont_sz / total_var) * args.sigma_scale
sigma *= scaling_factor 
S = torch.diag(1.0 / torch.diag(sigma))
sqrtinvL = (1.0 / scaling_factor)
eigs, _ = torch.symeig(sigma)
//  [python]
scaling_factor = (cont_sz / total_var) * args.sigma_scale
sigma *= scaling_factor 
S = torch.diag(1.0 / torch.diag(sigma))
sqrtinvL = (1.0 / scaling_factor)
eigs, _ = torch.linalg.eigh(sigma)
//  [python]
def __getitem__(self, idx):
        '''
        neighbor: n-1
        '''
        x = self.data[idx + self.x_offsets, :, :self.input_dim] #[t, node, dim]
        y = self.data[idx + self.y_offsets, :, :self.output_dim]
        
        par = self._adj2par(self.edge_index, self.n_nodes, self.edge_index.shape[1]).to_dense()
        L0 = torch.matmul(par, par.T)
        lambda0, _ = torch.symeig(L0, eigenvectors=False)
        L1 = torch.matmul(par.T, par)
        edge_index_t, edge_weight_t = dense_to_sparse(2*L0/lambda0.max())
        edge_index_s, edge_weight_s = dense_to_sparse(2*L1/lambda0.max())
        
        time_delay_sim = self._time_delay_similarity_full(x)
        x_s = torch.cat([self.edge_attr, time_delay_sim], -1) #[n_edge, 2 + l//scaler]
        x_t = x.permute(1,0,2)
        
        graph = PairData(x_s=x_s, edge_index_s=edge_index_s, edge_weight_s=edge_weight_s, # edge
                            x_t=x_t, edge_index_t=edge_index_t, edge_weight_t=edge_weight_t, # node
                            y = y)
        graph.edge_index = self.edge_index
        graph.num_nodes = self.n_nodes
            
        return [graph,  y]
//  [python]
def __getitem__(self, idx):
        '''
        neighbor: n-1
        '''
        x = self.data[idx + self.x_offsets, :, :self.input_dim] #[t, node, dim]
        y = self.data[idx + self.y_offsets, :, :self.output_dim]
        
        par = self._adj2par(self.edge_index, self.n_nodes, self.edge_index.shape[1]).to_dense()
        L0 = torch.matmul(par, par.T)
        lambda0, _ = torch.linalg.eigh(L0)
        L1 = torch.matmul(par.T, par)
        edge_index_t, edge_weight_t = dense_to_sparse(2*L0/lambda0.max())
        edge_index_s, edge_weight_s = dense_to_sparse(2*L1/lambda0.max())
        
        time_delay_sim = self._time_delay_similarity_full(x)
        x_s = torch.cat([self.edge_attr, time_delay_sim], -1) #[n_edge, 2 + l//scaler]
        x_t = x.permute(1,0,2)
        
        graph = PairData(x_s=x_s, edge_index_s=edge_index_s, edge_weight_s=edge_weight_s, # edge
                            x_t=x_t, edge_index_t=edge_index_t, edge_weight_t=edge_weight_t, # node
                            y = y)
        graph.edge_index = self.edge_index
        graph.num_nodes = self.n_nodes
            
        return [graph,  y]
//  [python]
def forward(self, z):
        assert z.shape[1] == self.dim_z

        # construct input for declarative node
        x = self.mlp(z)
        x = torch.reshape(x, (z.shape[0], self.m, self.m))

        if self.matrix_type == 'general':
            pass
        elif self.matrix_type == 'psd':
            x = torch.matmul(x, x.transpose(1, 2)) # positive definite
        elif self.matrix_type == 'nsd':
            x = -1.0 * torch.matmul(x, x.transpose(1, 2))  # negative definite
        elif self.matrix_type == '1pev':
            u = x
            x = -1.0 * torch.matmul(u, u.transpose(1, 2)) # negative definite
            x = x + 2.0 * torch.einsum("bi,bj->bij", u[:, :, 0], u[:, :, 0]) # 1 positive eigenvalue
        elif self.matrix_type == 'rank1':
            u = x
            x = torch.matmul(u[:, :, 0], u[:, :, 0].transpose(1, 2))
        elif self.matrix_type == 'rank2':
            u = x
            x = torch.matmul(u[:, :, 0:1], u[:, :, 0:1].transpose(1, 2))
        else:
            assert False, "unknown matrix_type"

        if self.method == 'pytorch':
            x = 0.5 * (x + x.transpose(1, 2))
            v, y = torch.symeig(x)
        elif self.method == 'exact':
            y = EigenDecompositionFcn().apply(x, self.top_k)
        elif self.method == 'approx':
            y = InstrumentedApproxEigenDecompositionFcn().apply(x, self.top_k)
        else:
            assert False

        return y
//  [python]
def forward(self, z):
        assert z.shape[1] == self.dim_z

        # construct input for declarative node
        x = self.mlp(z)
        x = torch.reshape(x, (z.shape[0], self.m, self.m))

        if self.matrix_type == 'general':
            pass
        elif self.matrix_type == 'psd':
            x = torch.matmul(x, x.transpose(1, 2)) # positive definite
        elif self.matrix_type == 'nsd':
            x = -1.0 * torch.matmul(x, x.transpose(1, 2))  # negative definite
        elif self.matrix_type == '1pev':
            u = x
            x = -1.0 * torch.matmul(u, u.transpose(1, 2)) # negative definite
            x = x + 2.0 * torch.einsum("bi,bj->bij", u[:, :, 0], u[:, :, 0]) # 1 positive eigenvalue
        elif self.matrix_type == 'rank1':
            u = x
            x = torch.matmul(u[:, :, 0], u[:, :, 0].transpose(1, 2))
        elif self.matrix_type == 'rank2':
            u = x
            x = torch.matmul(u[:, :, 0:1], u[:, :, 0:1].transpose(1, 2))
        else:
            assert False, "unknown matrix_type"

        if self.method == 'pytorch':
            x = 0.5 * (x + x.transpose(1, 2))
            v, y = torch.linalg.eigh(x)
        elif self.method == 'exact':
            y = EigenDecompositionFcn().apply(x, self.top_k)
        elif self.method == 'approx':
            y = InstrumentedApproxEigenDecompositionFcn().apply(x, self.top_k)
        else:
            assert False

        return y
//  [python]
def match_histogram(target_tensor, source_tensor, mode='avg'):
    if mode == 'False':
        return target_tensor
    backup = target_tensor.clone()
    try:
        if mode == 'avg':
            elementwise = True
            random_frame = False
        else:
            elementwise = False
            random_frame = True
        if not isinstance(source_tensor, list):
            source_tensor = [source_tensor]
        output_tensor = torch.zeros_like(target_tensor)
        for source in source_tensor:
            target = target_tensor.permute(0, 3, 2, 1)
            source = source.permute(0, 3, 2, 1)
            if elementwise:
                source = source.mean(0).unsqueeze(0)
            if random_frame:
                source = source[np.random.randint(0, source.shape[0])].unsqueeze(0)
            matched_tensor = torch.zeros_like(target)
            for idx in range(target.shape[0] if elementwise else 1):
                frame = target[idx].unsqueeze(0) if elementwise else target
                _, t, Ct = get_histogram(frame + 0.001 * torch.randn(size=frame.shape, device=frame.device))
                mu_s, _, Cs = get_histogram(source + 0.001 * torch.randn(size=source.shape, device=frame.device))
                eva_t, eve_t = torch.symeig(Ct, eigenvectors=True)
                Et = torch.sqrt(torch.diagflat(eva_t))
                Et[Et != Et] = 0
                Qt = torch.mm(torch.mm(eve_t, Et), eve_t.T)
                eva_s, eve_s = torch.symeig(Cs, eigenvectors=True)
                Es = torch.sqrt(torch.diagflat(eva_s))
                Es[Es != Es] = 0
                Qs = torch.mm(torch.mm(eve_s, Es), eve_s.T)
                ts = torch.mm(torch.mm(Qs, torch.inverse(Qt)), t)
                match = ts.reshape(*frame.permute(0, 3, 1, 2).shape).permute(0, 2, 3, 1)
                match += mu_s
                if elementwise:
                    matched_tensor[idx] = match
                else:
                    matched_tensor = match
            output_tensor += matched_tensor.permute(0, 3, 2, 1) / len(source_tensor)
    except RuntimeError:
        traceback.print_exc()
        None
        output_tensor = backup
    return output_tensor.clamp(min([s.min().item() for s in source_tensor]), max([s.max().item() for s in source_tensor]))
//  [python]
def match_histogram(target_tensor, source_tensor, mode='avg'):
    if mode == 'False':
        return target_tensor
    backup = target_tensor.clone()
    try:
        if mode == 'avg':
            elementwise = True
            random_frame = False
        else:
            elementwise = False
            random_frame = True
        if not isinstance(source_tensor, list):
            source_tensor = [source_tensor]
        output_tensor = torch.zeros_like(target_tensor)
        for source in source_tensor:
            target = target_tensor.permute(0, 3, 2, 1)
            source = source.permute(0, 3, 2, 1)
            if elementwise:
                source = source.mean(0).unsqueeze(0)
            if random_frame:
                source = source[np.random.randint(0, source.shape[0])].unsqueeze(0)
            matched_tensor = torch.zeros_like(target)
            for idx in range(target.shape[0] if elementwise else 1):
                frame = target[idx].unsqueeze(0) if elementwise else target
                _, t, Ct = get_histogram(frame + 0.001 * torch.randn(size=frame.shape, device=frame.device))
                mu_s, _, Cs = get_histogram(source + 0.001 * torch.randn(size=source.shape, device=frame.device))
                eva_t, eve_t = torch.linalg.eigh(Ct, UPLO='U')
                Et = torch.sqrt(torch.diagflat(eva_t))
                Et[Et != Et] = 0
                Qt = torch.mm(torch.mm(eve_t, Et), eve_t.T)
                eva_s, eve_s = torch.linalg.eigh(Cs, UPLO='U')
                Es = torch.sqrt(torch.diagflat(eva_s))
                Es[Es != Es] = 0
                Qs = torch.mm(torch.mm(eve_s, Es), eve_s.T)
                ts = torch.mm(torch.mm(Qs, torch.inverse(Qt)), t)
                match = ts.reshape(*frame.permute(0, 3, 1, 2).shape).permute(0, 2, 3, 1)
                match += mu_s
                if elementwise:
                    matched_tensor[idx] = match
                else:
                    matched_tensor = match
            output_tensor += matched_tensor.permute(0, 3, 2, 1) / len(source_tensor)
    except RuntimeError:
        traceback.print_exc()
        None
        output_tensor = backup
    return output_tensor.clamp(min([s.min().item() for s in source_tensor]), max([s.max().item() for s in source_tensor]))
//  [python]
def compute_kfe(self, mod: LayerType):
        """Computes the covariances."""
        state = self.state_storage[mod]
        x = state.x  # (N, in, xh, xw)
        gy = state.gy  # (N, out, yh, yw)

        if isinstance(mod, nn.Conv2d):
            state.num_locations = gy.size(2) * gy.size(3)  # yh * yw
            if not self.sua:
                x = self.filter_dict[mod](x)    # (N, in * kh * kw, yh, yw)
            # (in [* kh * kw], N * [yh * yw]{xh * xw})
            x = x.transpose(0, 1).flatten(1)
            gy = gy.transpose(0, 1).flatten(1)  # (out, N * yh * yw)
        else:
            state.num_locations = 1
            x = x.t()  # (in, N)
            gy = gy.t()  # (out, N)

        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones])    # (in [* kh * kw] + 1, N * yh * yw)
        # (in [* kh * kw] + 1, in [* kh * kw] + 1)
        xxt = x.mm(x.t()) / x.size(1)
        ggt = gy.mm(gy.t()) / gy.size(1)  # (out, out)

        Ex, state.kfe_x = torch.symeig(xxt, eigenvectors=True)
        Eg, state.kfe_gy = torch.symeig(ggt, eigenvectors=True)
        Ex: torch.Tensor  # (in [* kh * kw] + 1)
        Eg: torch.Tensor  # (out)
        state.m2 = state.num_locations * \
            Eg.outer(Ex)     # (out, in [* kh * kw] + 1)
        if isinstance(mod, nn.Conv2d) and self.sua:
            kh, kw = mod.kernel_size
            state.m2 = state.m2[..., None, None].repeat(
                1, 1, kh, kw)
//  [python]
def compute_kfe(self, mod: LayerType):
        """Computes the covariances."""
        state = self.state_storage[mod]
        x = state.x  # (N, in, xh, xw)
        gy = state.gy  # (N, out, yh, yw)

        if isinstance(mod, nn.Conv2d):
            state.num_locations = gy.size(2) * gy.size(3)  # yh * yw
            if not self.sua:
                x = self.filter_dict[mod](x)    # (N, in * kh * kw, yh, yw)
            # (in [* kh * kw], N * [yh * yw]{xh * xw})
            x = x.transpose(0, 1).flatten(1)
            gy = gy.transpose(0, 1).flatten(1)  # (out, N * yh * yw)
        else:
            state.num_locations = 1
            x = x.t()  # (in, N)
            gy = gy.t()  # (out, N)

        if mod.bias is not None:
            ones = torch.ones_like(x[:1])
            x = torch.cat([x, ones])    # (in [* kh * kw] + 1, N * yh * yw)
        # (in [* kh * kw] + 1, in [* kh * kw] + 1)
        xxt = x.mm(x.t()) / x.size(1)
        ggt = gy.mm(gy.t()) / gy.size(1)  # (out, out)

        Ex, state.kfe_x = torch.linalg.eigh(xxt)
        Eg, state.kfe_gy = torch.linalg.eigh(ggt)
        Ex: torch.Tensor  # (in [* kh * kw] + 1)
        Eg: torch.Tensor  # (out)
        state.m2 = state.num_locations * \
            Eg.outer(Ex)     # (out, in [* kh * kw] + 1)
        if isinstance(mod, nn.Conv2d) and self.sua:
            kh, kw = mod.kernel_size
            state.m2 = state.m2[..., None, None].repeat(
                1, 1, kh, kw)
//  [python]
def stable_symmetric_inverse(A, ret_sqrt=False):
    """
    Utility function to return stable inverse A^(-1) or sqrt of inverse A^(-0.5) of a symmetric matrix A
    Ref.: https://github.com/markovmodel/deeptime

    Proceeds by diagonalizing matrix, setting eigenvalues smaller than threshold to zero, performing operations
    on diagonlized matrix, reprojecting back to original basis

    Parameters
    ----------
    A : torch.tensor
        symmetric matrix

    ret_sqrt : bool, default = False
        indicator flag as to whether to return square root of inverse

    Return
    ------
    B : torch.tensor
        stable inverse approximation to A^(-1) or A^(-0.5) depending on ret_sqrt flag
    """
    assert torch.allclose(A, A.t(), atol=1e-6)

    w, V = torch.symeig(A)

    w = torch.where(w > torch.finfo(torch.float32).eps, w, torch.zeros_like(w))

    if ret_sqrt:
        B = torch.matmul(V, torch.matmul(torch.diag(w ** (-0.5)), V.t()))
    else:
        B = torch.matmul(V, torch.matmul(torch.diag(w ** (-1.0)), V.t()))

    return B
//  [python]
def stable_symmetric_inverse(A, ret_sqrt=False):
    """
    Utility function to return stable inverse A^(-1) or sqrt of inverse A^(-0.5) of a symmetric matrix A
    Ref.: https://github.com/markovmodel/deeptime

    Proceeds by diagonalizing matrix, setting eigenvalues smaller than threshold to zero, performing operations
    on diagonlized matrix, reprojecting back to original basis

    Parameters
    ----------
    A : torch.tensor
        symmetric matrix

    ret_sqrt : bool, default = False
        indicator flag as to whether to return square root of inverse

    Return
    ------
    B : torch.tensor
        stable inverse approximation to A^(-1) or A^(-0.5) depending on ret_sqrt flag
    """
    assert torch.allclose(A, A.t(), atol=1e-6)

    w, V = torch.linalg.eigh(A)

    w = torch.where(w > torch.finfo(torch.float32).eps, w, torch.zeros_like(w))

    if ret_sqrt:
        B = torch.matmul(V, torch.matmul(torch.diag(w ** (-0.5)), V.t()))
    else:
        B = torch.matmul(V, torch.matmul(torch.diag(w ** (-1.0)), V.t()))

    return B
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        # compute eigenvectors
        _, eigvecs = torch.symeig(cov(ODhat.T)) 
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        # compute eigenvectors
        _, eigvecs = torch.linalg.eigh(cov(ODhat.T)) 
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def identity_PCA(x, dim=3, normalize=False):
    """Reduce dimension of x based on PCA on the union of the n-simplex.
    This is a way of reducing the dimension of x while treating all
    input dimensions with the same importance, independently of the
    input distribution in x.
    """
    assert x.dim() == 2, f"Expected x.dim()=2, got x.dim()={x.dim()} instead"

    # Create z the union of the N-simplex
    input_dim = x.shape[1]
    z = torch.eye(input_dim)

    # PCA on z
    z_offset = z.mean(axis=0)
    z_centered = z - z_offset
    cov_matrix = z_centered.T.mm(z_centered) / len(z_centered)
    _, eigenvectors = torch.symeig(cov_matrix)

    # Normalize x if need be
    if normalize:
        high = x.max(dim=0).values
        low = x.min(dim=0).values
        x = (x - low) / (high - low)
        x[x.isnan() | x.isinf()] = 0

    # Apply the PCA on x
    x_reduced = (x - z_offset).mm(eigenvectors[:, -dim:])

    return x_reduced
//  [python]
def identity_PCA(x, dim=3, normalize=False):
    """Reduce dimension of x based on PCA on the union of the n-simplex.
    This is a way of reducing the dimension of x while treating all
    input dimensions with the same importance, independently of the
    input distribution in x.
    """
    assert x.dim() == 2, f"Expected x.dim()=2, got x.dim()={x.dim()} instead"

    # Create z the union of the N-simplex
    input_dim = x.shape[1]
    z = torch.eye(input_dim)

    # PCA on z
    z_offset = z.mean(axis=0)
    z_centered = z - z_offset
    cov_matrix = z_centered.T.mm(z_centered) / len(z_centered)
    _, eigenvectors = torch.linalg.eigh(cov_matrix)

    # Normalize x if need be
    if normalize:
        high = x.max(dim=0).values
        low = x.min(dim=0).values
        x = (x - low) / (high - low)
        x[x.isnan() | x.isinf()] = 0

    # Apply the PCA on x
    x_reduced = (x - z_offset).mm(eigenvectors[:, -dim:])

    return x_reduced
//  [python]
def compute_whiten_matrix(
    dataset: torch.utils.data.Dataset, take_fd: bool = False,
) -> torch.Tensor:
    """Compute whiten matrix of dataset.

    :param dataset: Time-series dataset.
    :type dataset: torch.utils.data.Dataset
    :return: whiten matrix
    :rtype: torch.Tensor
    """

    data_all = get_data_all(dataset, take_fd)
    # Demean data
    data_all -= data_all.mean(0)
    # Calc covariance
    Sigma = data_all.T @ data_all / (data_all.shape[0] - 1)
    # Calculating Eigenvalues and Eigenvectors of the covariance matrix
    Lamb, V = torch.symeig(Sigma, eigenvectors=True)
    # Deal with multiple constant features
    Lamb[Lamb < torch.finfo(torch.float32).eps] = 1
    # ZCA Whitening Matrix
    return V @ torch.diag(1./torch.sqrt(Lamb)) @ V.T
//  [python]
def compute_whiten_matrix(
    dataset: torch.utils.data.Dataset, take_fd: bool = False,
) -> torch.Tensor:
    """Compute whiten matrix of dataset.

    :param dataset: Time-series dataset.
    :type dataset: torch.utils.data.Dataset
    :return: whiten matrix
    :rtype: torch.Tensor
    """

    data_all = get_data_all(dataset, take_fd)
    # Demean data
    data_all -= data_all.mean(0)
    # Calc covariance
    Sigma = data_all.T @ data_all / (data_all.shape[0] - 1)
    # Calculating Eigenvalues and Eigenvectors of the covariance matrix
    Lamb, V = torch.linalg.eigh(Sigma)
    # Deal with multiple constant features
    Lamb[Lamb < torch.finfo(torch.float32).eps] = 1
    # ZCA Whitening Matrix
    return V @ torch.diag(1./torch.sqrt(Lamb)) @ V.T
//  [python]
def EDExperiments(dim_z=5, m=10, batch=10, iters=500, trials=10, method='exact', loss_on='all', mat_type='general'):

    learning_curves = [[] for i in range(trials)]
    cosine_sim_curves = [[] for i in range(trials)]

    for trial in range(trials):
        # prepare data and model
        torch.manual_seed(22 + trial)
        X_true = torch.rand((batch, m, m), dtype=torch.float, requires_grad=False)
        #X_true = torch.matmul(X_true, X_true.transpose(1, 2))
        V_true, Q_true = torch.symeig(X_true)
        z_init = torch.randn((batch, dim_z), dtype=torch.float, requires_grad=False)

        if method == 'approx':
            InstrumentedApproxEigenDecompositionFcn.trace = cosine_sim_curves[trial]

        model = EDNetwork(dim_z, m, method=method, top_k=1 if loss_on=='max' else None, matrix_type=mat_type)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3)

        # do optimisation
        for i in range(iters):
            optimizer.zero_grad(set_to_none=True)
            Q_pred = model(z_init)
            if loss_on == 'all':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred, Q_true, dim=1))) # all ev
            elif loss_on == 'max':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred[:, :, -1], Q_true[:, :, -1]))) # largest
            elif loss_on == 'min':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred[:, :, 0], Q_true[:, :, 0]))) # smallest
            else:
                assert False, "loss_on must be one of ('all', 'max', 'min')"
            learning_curves[trial].append(float(loss.item()))
            loss.backward()
            optimizer.step()

            if (i % 100 == 0):
                print("{: 3} {: 6}: {}".format(trial, i, loss.item()))

        InstrumentedApproxEigenDecompositionFcn.trace = None

    return learning_curves, cosine_sim_curves
//  [python]
def EDExperiments(dim_z=5, m=10, batch=10, iters=500, trials=10, method='exact', loss_on='all', mat_type='general'):

    learning_curves = [[] for i in range(trials)]
    cosine_sim_curves = [[] for i in range(trials)]

    for trial in range(trials):
        # prepare data and model
        torch.manual_seed(22 + trial)
        X_true = torch.rand((batch, m, m), dtype=torch.float, requires_grad=False)
        #X_true = torch.matmul(X_true, X_true.transpose(1, 2))
        V_true, Q_true = torch.linalg.eigh(X_true)
        z_init = torch.randn((batch, dim_z), dtype=torch.float, requires_grad=False)

        if method == 'approx':
            InstrumentedApproxEigenDecompositionFcn.trace = cosine_sim_curves[trial]

        model = EDNetwork(dim_z, m, method=method, top_k=1 if loss_on=='max' else None, matrix_type=mat_type)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3)

        # do optimisation
        for i in range(iters):
            optimizer.zero_grad(set_to_none=True)
            Q_pred = model(z_init)
            if loss_on == 'all':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred, Q_true, dim=1))) # all ev
            elif loss_on == 'max':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred[:, :, -1], Q_true[:, :, -1]))) # largest
            elif loss_on == 'min':
                loss = 1.0 - torch.mean(torch.abs(torch.nn.functional.cosine_similarity(Q_pred[:, :, 0], Q_true[:, :, 0]))) # smallest
            else:
                assert False, "loss_on must be one of ('all', 'max', 'min')"
            learning_curves[trial].append(float(loss.item()))
            loss.backward()
            optimizer.step()

            if (i % 100 == 0):
                print("{: 3} {: 6}: {}".format(trial, i, loss.item()))

        InstrumentedApproxEigenDecompositionFcn.trace = None

    return learning_curves, cosine_sim_curves
//  [python]
def _lambda_one_plus(self, T, device):

        eigen_pairs = torch.symeig(T)

        Lambda, U = eigen_pairs.eigenvalues, eigen_pairs.eigenvectors
        lambda_n, u_n = Lambda[0].to(device=device), U[:, 0].to(device=device)

        return torch.maximum(-lambda_n, torch.tensor([0], device=device)), lambda_n, u_n[:, None]
//  [python]
def _lambda_one_plus(self, T, device):

        eigen_pairs = torch.linalg.eigh(T)

        Lambda, U = eigen_pairs.eigenvalues, eigen_pairs.eigenvectors
        lambda_n, u_n = Lambda[0].to(device=device), U[:, 0].to(device=device)

        return torch.maximum(-lambda_n, torch.tensor([0], device=device)), lambda_n, u_n[:, None]
//  [python]
def build_graph(self, config):
        x = self.variable(name='x', shape=config.x_shape, dtype=config.x_dtype)
        out_w, out_v = torch.symeig(input=x, eigenvectors=True, upper=config.UPLO=='U')

        self.feed_list = [x]
        self.fetch_list = [out_w, out_v]
        if config.backward:
            self.append_gradients(out_w.sum() + torch.abs(out_v).sum(), [x])
//  [python]
def build_graph(self, config):
        x = self.variable(name='x', shape=config.x_shape, dtype=config.x_dtype)
        out_w, out_v = torch.linalg.eigh(input=x, UPLO=config.UPLO)

        self.feed_list = [x]
        self.fetch_list = [out_w, out_v]
        if config.backward:
            self.append_gradients(out_w.sum() + torch.abs(out_v).sum(), [x])
//  [python]
def safe_cholesky(A: torch.Tensor) -> torch.Tensor:
    eigenvalues, eigenvectors = torch.symeig(A)
    eigenvalues = torch.clamp(
        eigenvalues, min=1e-10
    )  # set any negative eigenvalues to a small positive value
    # L = eigenvectors @ torch.diag(eigenvalues.sqrt()) @ eigenvectors.t() #ZCA
    __, Rqr = torch.linalg.qr(torch.diag(eigenvalues.sqrt()) @ eigenvectors.t())
    Dg = torch.diag(torch.sign(torch.diag(Rqr)))
    Rqr = Dg @ Rqr
    return Rqr.T
//  [python]
def safe_cholesky(A: torch.Tensor) -> torch.Tensor:
    eigenvalues, eigenvectors = torch.linalg.eigh(A)
    eigenvalues = torch.clamp(
        eigenvalues, min=1e-10
    )  # set any negative eigenvalues to a small positive value
    # L = eigenvectors @ torch.diag(eigenvalues.sqrt()) @ eigenvectors.t() #ZCA
    __, Rqr = torch.linalg.qr(torch.diag(eigenvalues.sqrt()) @ eigenvectors.t())
    Dg = torch.diag(torch.sign(torch.diag(Rqr)))
    Rqr = Dg @ Rqr
    return Rqr.T
//  [python]
def frame_averaging_3D(pos, cell=None, fa_method="stochastic", check=False):
    """Computes new positions for the graph atoms using
    frame averaging, which itself builds on the PCA of atom positions.
    Base case for 3D inputs.

    Args:
        pos (tensor): positions of atoms in the graph
        cell (tensor): unit cell of the graph. None if no pbc.
        fa_method (str): FA method used
            (stochastic, det, all, se3-all, se3-det, se3-stochastic)
        check (bool): check if constraints are satisfied. Default: False.

    Returns:
        (tensor): updated atom positions
        (tensor): updated unit cell
        (tensor): the rotation matrix used (PCA)
    """

    # Compute centroid and covariance
    pos = pos - pos.mean(dim=0, keepdim=True)
    C = torch.matmul(pos.t(), pos)

    # Eigendecomposition
    eigenval, eigenvec = torch.symeig(C, eigenvectors=True)

    # Sort, if necessary
    idx = eigenval.argsort(descending=True)
    eigenvec = eigenvec[:, idx]
    eigenval = eigenval[idx]

    # Check if constraints are satisfied
    if check:
        check_constraints(eigenval, eigenvec, 3)

    # Compute fa_pos
    fa_pos, fa_cell, fa_rot = compute_frames(eigenvec, pos, cell, fa_method)

    # No need to update distances, they are preserved.

    return fa_pos, fa_cell, fa_rot
//  [python]
def frame_averaging_3D(pos, cell=None, fa_method="stochastic", check=False):
    """Computes new positions for the graph atoms using
    frame averaging, which itself builds on the PCA of atom positions.
    Base case for 3D inputs.

    Args:
        pos (tensor): positions of atoms in the graph
        cell (tensor): unit cell of the graph. None if no pbc.
        fa_method (str): FA method used
            (stochastic, det, all, se3-all, se3-det, se3-stochastic)
        check (bool): check if constraints are satisfied. Default: False.

    Returns:
        (tensor): updated atom positions
        (tensor): updated unit cell
        (tensor): the rotation matrix used (PCA)
    """

    # Compute centroid and covariance
    pos = pos - pos.mean(dim=0, keepdim=True)
    C = torch.matmul(pos.t(), pos)

    # Eigendecomposition
    eigenval, eigenvec = torch.linalg.eigh(C)

    # Sort, if necessary
    idx = eigenval.argsort(descending=True)
    eigenvec = eigenvec[:, idx]
    eigenval = eigenval[idx]

    # Check if constraints are satisfied
    if check:
        check_constraints(eigenval, eigenvec, 3)

    # Compute fa_pos
    fa_pos, fa_cell, fa_rot = compute_frames(eigenvec, pos, cell, fa_method)

    # No need to update distances, they are preserved.

    return fa_pos, fa_cell, fa_rot
//  [python]
def make_pos_def_torch(x):
    """
    Force a square symmetric matrix to be positive semi definite
    """
    w, v = torch.eig(torch.real(x))
    w = torch.real(w[:,0])
    w_pos = torch.clip(w, 0, None)
    nonzero_w = w_pos[w_pos > 0]
    w_new = w_pos
    if nonzero_w.shape[0] > 0:
        if nonzero_w.shape[0] < w.shape[0]:
            # min_w = max(torch.max(nonzero_w)*1e-5,torch.min(nonzero_w))
            min_w = max(torch.max(nonzero_w) * 0.1, torch.min(nonzero_w))
            w_new = w_pos + min_w
    elif nonzero_w.shape[0] == 0:
        eprint(
            "No positive eigenvalues for A {}. w={} {}, w_pos={} {}, nonzero_w={}".format(
                x, w, w.shape, w_pos, w_pos.shape, nonzero_w
            )
        )
        w_new = torch.ones_like(w)
    x_star = v @ torch.diag(w_new) @ v.T
    p = torch.sqrt(torch.sum(torch.abs(w)) / torch.sum(torch.abs(w_new)))
    return p * x_star
//  [python]
def make_pos_def_torch(x):
    """
    Force a square symmetric matrix to be positive semi definite
    """
    w, v = torch.linalg.eigh(torch.real(x))
    w_pos = torch.clip(w, 0, None)
    nonzero_w = w_pos[w_pos > 0]
    w_new = w_pos
    if nonzero_w.shape[0] > 0:
        if nonzero_w.shape[0] < w.shape[0]:
            # min_w = max(torch.max(nonzero_w)*1e-5,torch.min(nonzero_w))
            min_w = max(torch.max(nonzero_w) * 0.1, torch.min(nonzero_w))
            w_new = w_pos + min_w
    elif nonzero_w.shape[0] == 0:
        eprint(
            "No positive eigenvalues for A {}. w={} {}, w_pos={} {}, nonzero_w={}".format(
                x, w, w.shape, w_pos, w_pos.shape, nonzero_w
            )
        )
        w_new = torch.ones_like(w)
    x_star = v @ np.diag(w_new) @ v.T
    p = torch.sqrt(torch.sum(torch.abs(w)) / torch.sum(torch.abs(w_new)))
    return p * x_star
//  [python]
def frame_averaging_2D(pos, cell=None, fa_method="stochastic", check=False):
    """Computes new positions for the graph atoms using
    frame averaging, which itself builds on the PCA of atom positions.
    2D case: we project the atoms on the plane orthogonal to the z-axis.
    Motivation: sometimes, the z-axis is not the most relevant one (e.g. fixed).

    Args:
        pos (tensor): positions of atoms in the graph
        cell (tensor): unit cell of the graph. None if no pbc.
        fa_method (str): FA method used (stochastic, det, all, se3)
        check (bool): check if constraints are satisfied. Default: False.

    Returns:
        (tensor): updated atom positions
        (tensor): updated unit cell
        (tensor): the rotation matrix used (PCA)
    """

    # Compute centroid and covariance
    pos_2D = pos[:, :2] - pos[:, :2].mean(dim=0, keepdim=True)
    C = torch.matmul(pos_2D.t(), pos_2D)

    # Eigendecomposition
    eigenval, eigenvec = torch.symeig(C, eigenvectors=True)
    # Sort eigenvalues
    idx = eigenval.argsort(descending=True)
    eigenval = eigenval[idx]
    eigenvec = eigenvec[:, idx]

    # Check if constraints are satisfied
    if check:
        check_constraints(eigenval, eigenvec, 3)

    # Compute all frames
    fa_pos, fa_cell, fa_rot = compute_frames(
        eigenvec,
        pos_2D,
        cell,
        fa_method,
        pos[:, 2],
    )
    # No need to update distances, they are preserved.

    return fa_pos, fa_cell, fa_rot
//  [python]
def frame_averaging_2D(pos, cell=None, fa_method="stochastic", check=False):
    """Computes new positions for the graph atoms using
    frame averaging, which itself builds on the PCA of atom positions.
    2D case: we project the atoms on the plane orthogonal to the z-axis.
    Motivation: sometimes, the z-axis is not the most relevant one (e.g. fixed).

    Args:
        pos (tensor): positions of atoms in the graph
        cell (tensor): unit cell of the graph. None if no pbc.
        fa_method (str): FA method used (stochastic, det, all, se3)
        check (bool): check if constraints are satisfied. Default: False.

    Returns:
        (tensor): updated atom positions
        (tensor): updated unit cell
        (tensor): the rotation matrix used (PCA)
    """

    # Compute centroid and covariance
    pos_2D = pos[:, :2] - pos[:, :2].mean(dim=0, keepdim=True)
    C = torch.matmul(pos_2D.t(), pos_2D)

    # Eigendecomposition
    eigenval, eigenvec = torch.linalg.eigh(C)
    # Sort eigenvalues
    idx = eigenval.argsort(descending=True)
    eigenval = eigenval[idx]
    eigenvec = eigenvec[:, idx]

    # Check if constraints are satisfied
    if check:
        check_constraints(eigenval, eigenvec, 3)

    # Compute all frames
    fa_pos, fa_cell, fa_rot = compute_frames(
        eigenvec,
        pos_2D,
        cell,
        fa_method,
        pos[:, 2],
    )
    # No need to update distances, they are preserved.

    return fa_pos, fa_cell, fa_rot
//  [python]
def map_positional_encoding(g, pos_enc_dim, use_unique_sign=True, use_unique_basis=True, use_eig_val=True):
    """
        Graph positional encoding v/ universal spectral embeddings
    """
    A = g.adjacency_matrix(scipy_fmt="csr").astype(np.double)
    N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=np.double)
    A = torch.from_numpy((N * A * N).toarray())
    n, k = A.shape[0], pos_enc_dim
    E, U = torch.eig(A, eigenvectors=True)
    E, U = E.round(decimals=14)[:-1], U[:, :-1]
    dim = min(n - 1, k)
    _, mult = torch.unique(E[-dim:], return_counts=True)
    ind = torch.cat([torch.LongTensor([0]), torch.cumsum(mult, dim=0)]) + max(n - 1 - k, 0)
    if use_unique_sign:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                U[:, ind[i]:ind[i + 1]] = unique_sign(U[:, ind[i]:ind[i + 1]])  # eliminate sign ambiguity
    if use_unique_basis:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                continue  # single eigenvector, no basis ambiguity
            try:
                U[:, ind[i]:ind[i + 1]] = unique_basis(U[:, ind[i]:ind[i + 1]])  # eliminate basis ambiguity
            except AssertionError:
                continue  # assumption violated, skip
    if use_eig_val:
        Lambda = torch.nn.ReLU()(torch.diag(E))
        U = U @ torch.sqrt(Lambda)
    if n - 1 < k:
        zeros = torch.zeros([n, k - n + 1])
        U = torch.cat([U, zeros], dim=-1)
    g.ndata['pos_enc'] = U[:, -k:]  # last k non-trivial eigenvectors
    return g
//  [python]
def map_positional_encoding(g, pos_enc_dim, use_unique_sign=True, use_unique_basis=True, use_eig_val=True):
    """
        Graph positional encoding v/ universal spectral embeddings
    """
    A = g.adjacency_matrix(scipy_fmt="csr").astype(np.double)
    N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=np.double)
    A = torch.from_numpy((N * A * N).toarray())
    n, k = A.shape[0], pos_enc_dim
    E, U = torch.linalg.eigh(A)
    E, U = E.round(decimals=14)[:-1], U[:, :-1]
    dim = min(n - 1, k)
    _, mult = torch.unique(E[-dim:], return_counts=True)
    ind = torch.cat([torch.LongTensor([0]), torch.cumsum(mult, dim=0)]) + max(n - 1 - k, 0)
    if use_unique_sign:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                U[:, ind[i]:ind[i + 1]] = unique_sign(U[:, ind[i]:ind[i + 1]])  # eliminate sign ambiguity
    if use_unique_basis:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                continue  # single eigenvector, no basis ambiguity
            try:
                U[:, ind[i]:ind[i + 1]] = unique_basis(U[:, ind[i]:ind[i + 1]])  # eliminate basis ambiguity
            except AssertionError:
                continue  # assumption violated, skip
    if use_eig_val:
        Lambda = torch.nn.ReLU()(torch.diag(E))
        U = U @ torch.sqrt(Lambda)
    if n - 1 < k:
        zeros = torch.zeros([n, k - n + 1])
        U = torch.cat([U, zeros], dim=-1)
    g.ndata['pos_enc'] = U[:, -k:]  # last k non-trivial eigenvectors
    return g
//  [python]
def matrix_sqrt(M):
    S, V = torch.symeig(M)
    S[S<0] = 0
    S = torch.diag(S**0.5)
    return (V @ S @ V.T).cpu()
//  [python]
def matrix_sqrt(M):
    S, V = torch.linalg.eigh(M)
    S[S<0] = 0
    S = torch.diag(S**0.5)
    return (V @ S @ V.T).cpu()
//  [python]
def admm_simple(A: Tensor, n_eig: Optional[int] = None, min_eig: float = 0.0) -> Tensor:
    N = A.shape[1]
    if n_eig:  # if specify the number of eigen components
        if (0 < n_eig <= N) and (n_eig % 2 == 0):
            pass
        else:
            raise ValueError(
                f" {n_eig} not a valid number of eigen component to reserve"
            )
    else:
        n_eig = N

    # eigenvalue decomposition of A, ascending
    delta, V = torch.symeig(A, eigenvectors=True)

    # to descending order
    delta = torch.flip(delta, [0])
    V = torch.flip(V, [1])

    # compute eigenvalue of the bipartite bga of A
    lambda_ = A.new_zeros(n_eig)
    if N % 2 == 0:
        lambda_[: N // 2] = 0.5 * (delta[: N // 2] - torch.flip(delta[N // 2 :], [0]))
        lambda_[N // 2 :] = torch.flip(-lambda_[: N // 2], [0])
    elif N % 2 != 0:
        lambda_[: N // 2] = 0.5 * (
            delta[: N // 2] - torch.flip(delta[N // 2 + 1 :], [0])
        )
        lambda_[N // 2 + 1 :] = torch.flip(-lambda_[: N // 2], [0])

    if n_eig < N:  # delete the middle N-n_eig eigenvalues of B
        lambda_[n_eig / 2 : -n_eig / 2] = 0

    if min_eig > 0:
        idx = (torch.abs(lambda_) < min_eig).nonzero()[:, 0]
        lambda_[idx] = (lambda_[idx]).sign() * min_eig

    # recover B
    B = V.mm(torch.diag(lambda_)).mm(V.t())
    return (B + B.t()) / 2
//  [python]
def admm_simple(A: Tensor, n_eig: Optional[int] = None, min_eig: float = 0.0) -> Tensor:
    N = A.shape[1]
    if n_eig:  # if specify the number of eigen components
        if (0 < n_eig <= N) and (n_eig % 2 == 0):
            pass
        else:
            raise ValueError(
                f" {n_eig} not a valid number of eigen component to reserve"
            )
    else:
        n_eig = N

    # eigenvalue decomposition of A, ascending
    delta, V = torch.linalg.eigh(A)

    # to descending order
    delta = torch.flip(delta, [0])
    V = torch.flip(V, [1])

    # compute eigenvalue of the bipartite bga of A
    lambda_ = A.new_zeros(n_eig)
    if N % 2 == 0:
        lambda_[: N // 2] = 0.5 * (delta[: N // 2] - torch.flip(delta[N // 2 :], [0]))
        lambda_[N // 2 :] = torch.flip(-lambda_[: N // 2], [0])
    elif N % 2 != 0:
        lambda_[: N // 2] = 0.5 * (
            delta[: N // 2] - torch.flip(delta[N // 2 + 1 :], [0])
        )
        lambda_[N // 2 + 1 :] = torch.flip(-lambda_[: N // 2], [0])

    if n_eig < N:  # delete the middle N-n_eig eigenvalues of B
        lambda_[n_eig / 2 : -n_eig / 2] = 0

    if min_eig > 0:
        idx = (torch.abs(lambda_) < min_eig).nonzero()[:, 0]
        lambda_[idx] = (lambda_[idx]).sign() * min_eig

    # recover B
    B = V.mm(torch.diag(lambda_)).mm(V.t())
    return (B + B.t()) / 2