jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _forward(a, formatter, eye):
        if torch.any(a[..., 0] == 0):
            raise RuntimeError("Leading coefficient must be non-zero.")

        # Make companion matrix.
        a = -a[..., 1:] / a[..., :1]  # (..., M)
        E = eye.expand(a.size()[:-1] + eye.size())
        A = torch.cat((a.unsqueeze(-2), E), dim=-2)  # (..., M, M)

        # Find roots as eigenvalues.
        x, _ = torch.eig(A)
        x = formatter(x)
        return x
//  [python]
def _forward(a, formatter, eye):
        if torch.any(a[..., 0] == 0):
            raise RuntimeError("Leading coefficient must be non-zero.")

        # Make companion matrix.
        a = -a[..., 1:] / a[..., :1]  # (..., M)
        E = eye.expand(a.size()[:-1] + eye.size())
        A = torch.cat((a.unsqueeze(-2), E), dim=-2)  # (..., M, M)

        # Find roots as eigenvalues.
        x, _ = torch.linalg.eig(A)
        x = formatter(x)
        return x
//  [python]
def wasserstein2_loss(mean_reals: torch.Tensor,
                      mean_gen: torch.Tensor,
                      cov_reals: torch.Tensor,
                      cov_gen: torch.Tensor,
                      eps: float = 1e-12) -> torch.Tensor:
    """Computes 2-Wasserstein distance."""
    mean_term = torch.sum(torch.square(mean_reals - mean_gen.squeeze(0)))
    eigenvalues = torch.real(torch.symeig(torch.matmul(cov_gen, cov_reals),eigenvectors=False)[0])
    cov_term = torch.trace(cov_reals) + torch.trace(cov_gen) - 2 * torch.sum(torch.sqrt(abs(eigenvalues) + eps))
    return mean_term + cov_term
//  [python]
def wasserstein2_loss(mean_reals: torch.Tensor,
                      mean_gen: torch.Tensor,
                      cov_reals: torch.Tensor,
                      cov_gen: torch.Tensor,
                      eps: float = 1e-12) -> torch.Tensor:
    """Computes 2-Wasserstein distance."""
    mean_term = torch.sum(torch.square(mean_reals - mean_gen.squeeze(0)))
    eigenvalues = torch.real(torch.linalg.eig(torch.matmul(cov_gen, cov_reals))[0])
    cov_term = torch.trace(cov_reals) + torch.trace(cov_gen) - 2 * torch.sum(torch.sqrt(abs(eigenvalues) + eps))
    return mean_term + cov_term
//  [python]
def test_eigenvalues_property(self, factored_matrices):
        for factored_matrix in factored_matrices:
            if factored_matrix.ldim == factored_matrix.rdim:
                eigenvalues = factored_matrix.eigenvalues
                expected_eigenvalues = torch.eig(
                    factored_matrix.B @ factored_matrix.A
                ).eigenvalues
                assert torch.allclose(
                    torch.abs(eigenvalues), torch.abs(expected_eigenvalues), atol=1e-5
                )
            else:
                with pytest.raises(AssertionError):
                    _ = factored_matrix.eigenvalues
//  [python]
def test_eigenvalues_property(self, factored_matrices):
        for factored_matrix in factored_matrices:
            if factored_matrix.ldim == factored_matrix.rdim:
                eigenvalues = factored_matrix.eigenvalues
                expected_eigenvalues = torch.linalg.eig(
                    factored_matrix.B @ factored_matrix.A
                ).eigenvalues
                assert torch.allclose(
                    torch.abs(eigenvalues), torch.abs(expected_eigenvalues), atol=1e-5
                )
            else:
                with pytest.raises(AssertionError):
                    _ = factored_matrix.eigenvalues
//  [python]
def complex_whiten(self, complex_image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Whiten complex image.

        Parameters
        ----------
        complex_image: torch.Tensor
            Complex image tensor to whiten.

        Returns
        -------
        mean, std, whitened_image: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        """
        # From: https://github.com/facebookresearch/fastMRI
        #       blob/da1528585061dfbe2e91ebbe99a5d4841a5c3f43/banding_removal/fastmri/data/transforms.py#L464  # noqa
        real = complex_image[..., 0]
        imag = complex_image[..., 1]

        # Center around mean.
        mean = complex_image.mean()
        centered_complex_image = complex_image - mean

        # Determine covariance between real and imaginary.
        n_elements = real.nelement()
        real_real = (real.mul(real).sum() - real.mean().mul(real.mean())) / n_elements
        real_imag = (real.mul(imag).sum() - real.mean().mul(imag.mean())) / n_elements
        imag_imag = (imag.mul(imag).sum() - imag.mean().mul(imag.mean())) / n_elements
        eig_input = torch.Tensor([[real_real, real_imag], [real_imag, imag_imag]])

        # Remove correlation by rotating around covariance eigenvectors.
        eig_values, eig_vecs = torch.eig(eig_input)

        # Scale by eigenvalues for unit variance.
        std = (eig_values[:, 0].real + self.epsilon).sqrt()
        whitened_image = torch.matmul(centered_complex_image, eig_vecs.real) / std

        return mean, std, whitened_image
//  [python]
def complex_whiten(self, complex_image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Whiten complex image.

        Parameters
        ----------
        complex_image: torch.Tensor
            Complex image tensor to whiten.

        Returns
        -------
        mean, std, whitened_image: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        """
        # From: https://github.com/facebookresearch/fastMRI
        #       blob/da1528585061dfbe2e91ebbe99a5d4841a5c3f43/banding_removal/fastmri/data/transforms.py#L464  # noqa
        real = complex_image[..., 0]
        imag = complex_image[..., 1]

        # Center around mean.
        mean = complex_image.mean()
        centered_complex_image = complex_image - mean

        # Determine covariance between real and imaginary.
        n_elements = real.nelement()
        real_real = (real.mul(real).sum() - real.mean().mul(real.mean())) / n_elements
        real_imag = (real.mul(imag).sum() - real.mean().mul(imag.mean())) / n_elements
        imag_imag = (imag.mul(imag).sum() - imag.mean().mul(imag.mean())) / n_elements
        eig_input = torch.Tensor([[real_real, real_imag], [real_imag, imag_imag]])

        # Remove correlation by rotating around covariance eigenvectors.
        eig_values, eig_vecs = torch.linalg.eig(eig_input)

        # Scale by eigenvalues for unit variance.
        std = (eig_values.real + self.epsilon).sqrt()
        whitened_image = torch.matmul(centered_complex_image, eig_vecs.real) / std

        return mean, std, whitened_image
//  [python]
def forward(ctx, A, eps = 1E-6):
        
        # Perform the eigendecomposition.
        eigenvalues, eigenvectors = torch.eig(A)
        
        # Sort the eigenvalues in non-descencing order, according to their real parts.
        indices = torch.real(eigenvalues).argsort()
        
        # Apply the ordering to the (complex) eigenvalues and their eigenvectors.
        eigenvalues = eigenvalues.gather(dim=-1, index=indices)
        indices = indices[:,:,:,:,None,:]
        indices = indices.repeat((1,1,1,1,indices.shape[-1],1))
        eigenvectors = eigenvectors.gather(dim=-1, index=indices)
        
        # Save the decomposition for the backwards pass.
        ctx.save_for_backward(A, eigenvalues, eigenvectors)
        ctx.eps = eps
        
        return eigenvalues, eigenvectors
//  [python]
def forward(ctx, A, eps = 1E-6):
        
        # Perform the eigendecomposition.
        eigenvalues, eigenvectors = torch.linalg.eig(A)
        
        # Sort the eigenvalues in non-descencing order, according to their real parts.
        indices = torch.real(eigenvalues).argsort()
        
        # Apply the ordering to the (complex) eigenvalues and their eigenvectors.
        eigenvalues = eigenvalues.gather(dim=-1, index=indices)
        indices = indices[:,:,:,:,None,:]
        indices = indices.repeat((1,1,1,1,indices.shape[-1],1))
        eigenvectors = eigenvectors.gather(dim=-1, index=indices)
        
        # Save the decomposition for the backwards pass.
        ctx.save_for_backward(A, eigenvalues, eigenvectors)
        ctx.eps = eps
        
        return eigenvalues, eigenvectors
//  [python]
def _eigen(cov):
        try:
            diag, rot = torch.eig(cov, eigenvectors=True)
            assert (diag[:,1].abs() < 1e-6).all(), "`cov` possesses complex valued eigenvalues"
            diag = diag[:,0] 
        except AttributeError:
            # old implementation
            diag, rot = torch.eig(cov, eigenvectors=True)
            assert (diag[:,1].abs() < 1e-6).all(), "`cov` possesses complex valued eigenvalues"
            diag = diag[:,0] 
        return diag + 1e-6, rot
//  [python]
def _eigen(cov):
        try:
            diag, rot = torch.linalg.eig(cov)
            assert (diag.imag.abs() < 1e-6).all(), "`cov` possesses complex valued eigenvalues"
            diag, rot = diag.real, rot.real
        except AttributeError:
            # old implementation
            diag, rot = torch.eig(cov, eigenvectors=True)
            assert (diag[:,1].abs() < 1e-6).all(), "`cov` possesses complex valued eigenvalues"
            diag = diag[:,0] 
        return diag + 1e-6, rot
//  [python]
def _eigen_decomposition(self):
        # H to E transformation matirx
        P_tmp = torch.matmul(torch.vstack((self.Kx_norm,self.Ky_norm)), torch.linalg.inv(self.eps_conv[-1]))
        self.P.append(torch.hstack((torch.vstack((torch.zeros_like(self.mu_conv[-1]),-self.mu_conv[-1])),
            torch.vstack((self.mu_conv[-1],torch.zeros_like(self.mu_conv[-1]))))) + torch.matmul(P_tmp, torch.hstack((self.Ky_norm,-self.Kx_norm))))
        # E to H transformation matrix
        Q_tmp = torch.matmul(torch.vstack((self.Kx_norm,self.Ky_norm)), torch.linalg.inv(self.mu_conv[-1]))
        self.Q.append(torch.hstack((torch.vstack((torch.zeros_like(self.eps_conv[-1]),self.eps_conv[-1])),
            torch.vstack((-self.eps_conv[-1],torch.zeros_like(self.eps_conv[-1]))))) + torch.matmul(Q_tmp, torch.hstack((-self.Ky_norm,self.Kx_norm))))
        
        # Eigen-decomposition
        if self.stable_eig_grad is True:
            kz_norm, E_eigvec = Eig.apply(torch.matmul(self.P[-1],self.Q[-1]))
        else:
            kz_norm, E_eigvec = torch.eig(torch.matmul(self.P[-1],self.Q[-1]))
        
        kz_norm = torch.sqrt(kz_norm)
        self.kz_norm.append(torch.where(torch.imag(kz_norm)<0,-kz_norm,kz_norm)) # Normalized kz for positive mode
        self.E_eigvec.append(E_eigvec)
//  [python]
def _eigen_decomposition(self):
        # H to E transformation matirx
        P_tmp = torch.matmul(torch.vstack((self.Kx_norm,self.Ky_norm)), torch.linalg.inv(self.eps_conv[-1]))
        self.P.append(torch.hstack((torch.vstack((torch.zeros_like(self.mu_conv[-1]),-self.mu_conv[-1])),
            torch.vstack((self.mu_conv[-1],torch.zeros_like(self.mu_conv[-1]))))) + torch.matmul(P_tmp, torch.hstack((self.Ky_norm,-self.Kx_norm))))
        # E to H transformation matrix
        Q_tmp = torch.matmul(torch.vstack((self.Kx_norm,self.Ky_norm)), torch.linalg.inv(self.mu_conv[-1]))
        self.Q.append(torch.hstack((torch.vstack((torch.zeros_like(self.eps_conv[-1]),self.eps_conv[-1])),
            torch.vstack((-self.eps_conv[-1],torch.zeros_like(self.eps_conv[-1]))))) + torch.matmul(Q_tmp, torch.hstack((-self.Ky_norm,self.Kx_norm))))
        
        # Eigen-decomposition
        if self.stable_eig_grad is True:
            kz_norm, E_eigvec = Eig.apply(torch.matmul(self.P[-1],self.Q[-1]))
        else:
            kz_norm, E_eigvec = torch.linalg.eig(torch.matmul(self.P[-1],self.Q[-1]))
        
        kz_norm = torch.sqrt(kz_norm)
        self.kz_norm.append(torch.where(torch.imag(kz_norm)<0,-kz_norm,kz_norm)) # Normalized kz for positive mode
        self.E_eigvec.append(E_eigvec)
//  [python]
def decomp(faces):
    data = []; row = []; col = []
    for _, triangle in enumerate(faces):
        for j in range(3):
            data.append(1.0); row.append(triangle[j]); col.append(triangle[(j + 1) % 3]);
            data.append(1.0); col.append(triangle[j]); row.append(triangle[(j + 1) % 3]);
    A = sparse.coo_matrix((data, (row, col)), dtype=float)
    A = A.todense().astype(np.float64) / 2
    L = A.copy()
    for i in range(L.shape[0]):
        L[i, i] = -np.sum(L[i])
    L = -L

    S = A.copy()
    for i in range(S.shape[0]):
        S[i, i] = 1
        S[i, :] /= np.sum(S[i])
    # np.set_printoptions(threshold=np.inf)
    # with open('test.txt', 'w+') as f:
    #     print(A, file=f)
    #test = L - L.T
    start = time.time()
    if 0:
        with torch.no_grad():
            L = torch.from_numpy(L).float().cuda()
            lmbd, U = torch.eig(L)
            lmbd = lmbd.detach().cpu().numpy()
            U = U.detach().cpu().numpy()
    else:
        lmbd, U = eig(L)
    end = time.time()
    lmbd = np.float64(lmbd)
    idx = np.argsort(lmbd)
    lmbd = sorted(lmbd)
    U = U[:, idx]
    return U, lmbd
//  [python]
def decomp(faces):
    data = []; row = []; col = []
    for _, triangle in enumerate(faces):
        for j in range(3):
            data.append(1.0); row.append(triangle[j]); col.append(triangle[(j + 1) % 3]);
            data.append(1.0); col.append(triangle[j]); row.append(triangle[(j + 1) % 3]);
    A = sparse.coo_matrix((data, (row, col)), dtype=float)
    A = A.todense().astype(np.float64) / 2
    L = A.copy()
    for i in range(L.shape[0]):
        L[i, i] = -np.sum(L[i])
    L = -L

    S = A.copy()
    for i in range(S.shape[0]):
        S[i, i] = 1
        S[i, :] /= np.sum(S[i])
    # np.set_printoptions(threshold=np.inf)
    # with open('test.txt', 'w+') as f:
    #     print(A, file=f)
    #test = L - L.T
    start = time.time()
    if 0:
        with torch.no_grad():
            L = torch.from_numpy(L).float().cuda()
            lmbd, U = torch.linalg.eig(L)
            lmbd = lmbd.detach().cpu().numpy()
            U = U.detach().cpu().numpy()
    else:
        lmbd, U = eig(L)
    end = time.time()
    lmbd = np.float64(lmbd)
    idx = np.argsort(lmbd)
    lmbd = sorted(lmbd)
    U = U[:, idx]
    return U, lmbd
//  [python]
def compute_layer_to_mode_matrices(self):
        """Matrices to change from layers to modes."""
        A = self.A[0] if self.n_ens > 0 else self.A
        lambd_r, R = torch.eig(A)
        lambd_l, L = torch.eig(A.T)
        self.lambd = lambd_r[:,0].real
        R, L = R.real, L.real
        self.Cl2m = torch.diag(1./torch.diag(L.T @ R)) @ L.T
        self.Cm2l = R
        if self.n_ens > 0:
            self.Cl2m.unsqueeze_(0), self.Cm2l.unsqueeze_(0)
//  [python]
def compute_layer_to_mode_matrices(self):
        """Matrices to change from layers to modes."""
        A = self.A[0] if self.n_ens > 0 else self.A
        lambd_r, R = torch.linalg.eig(A)
        lambd_l, L = torch.linalg.eig(A.T)
        self.lambd = lambd_r.real
        R, L = R.real, L.real
        self.Cl2m = torch.diag(1./torch.diag(L.T @ R)) @ L.T
        self.Cm2l = R
        if self.n_ens > 0:
            self.Cl2m.unsqueeze_(0), self.Cm2l.unsqueeze_(0)
//  [python]
def forward(matrix):
        res = torch.eig(matrix)
        # print(111, res[0][0], res[1][0,0])
        return res
//  [python]
def forward(matrix):
        res = torch.linalg.eig(matrix)
        # print(111, res[0][0], res[1][0,0])
        return res
//  [python]
def initialize(self):
        grid = plt.get_grid_uniform(100, self.grid_boundary)
        z = []
        points = grid['grid_points']
        for pnts in track(torch.split(points, 1000000, dim=0)):
            z.append(self.model.implicit_network(pnts)[:,0].detach().cpu().numpy())
        z = np.concatenate(z, axis=0).astype(np.float32)
        verts, faces, normals, values = measure.marching_cubes(
        volume=z.reshape(grid['xyz'][1].shape[0], grid['xyz'][0].shape[0],
                         grid['xyz'][2].shape[0]).transpose([1, 0, 2]),
                         level=0,
                         spacing=(grid['xyz'][0][2] - grid['xyz'][0][1],
                         grid['xyz'][0][2] - grid['xyz'][0][1],
                         grid['xyz'][0][2] - grid['xyz'][0][1]))
        verts = verts + np.array([grid['xyz'][0][0], grid['xyz'][1][0], grid['xyz'][2][0]])
        mesh_low_res = trimesh.Trimesh(verts, faces, normals)
        recon_pc = trimesh.sample.sample_surface(mesh_low_res, 10000)[0]
        recon_pc = torch.from_numpy(recon_pc).float().cuda()
        s_mean = recon_pc.mean(dim=0)
        s_cov = recon_pc - s_mean
        s_cov = torch.mm(s_cov.transpose(0, 1), s_cov)
        self.vecs = torch.view_as_real(torch.symeig(s_cov, eigenvectors=True)[1].transpose(0, 1))[:, :, 0]
        if torch.det(self.vecs) < 0:
            self.vecs = torch.mm(torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]]).cuda().float(), self.vecs)
        helper = torch.bmm(self.vecs.unsqueeze(0).repeat(recon_pc.shape[0], 1, 1),
                        (recon_pc - s_mean).unsqueeze(-1)).squeeze().cpu()
        grid_aligned = plt.get_grid(helper, self.resolution)
        grid_points = grid_aligned['grid_points']
        g = []
        for pnts in track(torch.split(grid_points, 1000000, dim=0)):
            g.append(torch.bmm(self.vecs.unsqueeze(0).repeat(pnts.shape[0], 1, 1).transpose(1, 2),
                           pnts.unsqueeze(-1)).squeeze() + s_mean)
        grid_points = torch.cat(g, dim=0)
        points = grid_points.cpu()
        self.test_dataset = dataset.GridDataset(points, grid_aligned['xyz'])
        self.grid_points = grid_points
        self.initialized = True
//  [python]
def initialize(self):
        grid = plt.get_grid_uniform(100, self.grid_boundary)
        z = []
        points = grid['grid_points']
        for pnts in track(torch.split(points, 1000000, dim=0)):
            z.append(self.model.implicit_network(pnts)[:,0].detach().cpu().numpy())
        z = np.concatenate(z, axis=0).astype(np.float32)
        verts, faces, normals, values = measure.marching_cubes(
        volume=z.reshape(grid['xyz'][1].shape[0], grid['xyz'][0].shape[0],
                         grid['xyz'][2].shape[0]).transpose([1, 0, 2]),
                         level=0,
                         spacing=(grid['xyz'][0][2] - grid['xyz'][0][1],
                         grid['xyz'][0][2] - grid['xyz'][0][1],
                         grid['xyz'][0][2] - grid['xyz'][0][1]))
        verts = verts + np.array([grid['xyz'][0][0], grid['xyz'][1][0], grid['xyz'][2][0]])
        mesh_low_res = trimesh.Trimesh(verts, faces, normals)
        recon_pc = trimesh.sample.sample_surface(mesh_low_res, 10000)[0]
        recon_pc = torch.from_numpy(recon_pc).float().cuda()
        s_mean = recon_pc.mean(dim=0)
        s_cov = recon_pc - s_mean
        s_cov = torch.mm(s_cov.transpose(0, 1), s_cov)
        self.vecs = torch.view_as_real(torch.linalg.eig(s_cov)[1].transpose(0, 1))[:, :, 0]
        if torch.det(self.vecs) < 0:
            self.vecs = torch.mm(torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]]).cuda().float(), self.vecs)
        helper = torch.bmm(self.vecs.unsqueeze(0).repeat(recon_pc.shape[0], 1, 1),
                        (recon_pc - s_mean).unsqueeze(-1)).squeeze().cpu()
        grid_aligned = plt.get_grid(helper, self.resolution)
        grid_points = grid_aligned['grid_points']
        g = []
        for pnts in track(torch.split(grid_points, 1000000, dim=0)):
            g.append(torch.bmm(self.vecs.unsqueeze(0).repeat(pnts.shape[0], 1, 1).transpose(1, 2),
                           pnts.unsqueeze(-1)).squeeze() + s_mean)
        grid_points = torch.cat(g, dim=0)
        points = grid_points.cpu()
        self.test_dataset = dataset.GridDataset(points, grid_aligned['xyz'])
        self.grid_points = grid_points
        self.initialized = True
//  [python]
def _spectral_embedding(self, x: DNDarray) -> Tuple[DNDarray, DNDarray]:
        """
        Helper function for dataset x embedding.
        Returns Tupel(Eigenvalues, Eigenvectors) of the graph's Laplacian matrix.

        Parameters
        ----------
        x : DNDarray
            Sample Matrix for which the embedding should be calculated

        Notes
        -----
        This will throw out the complex side of the eigenvalues found during this.

        """
        L = self._laplacian.construct(x)
        # 3. Eigenvalue and -vector calculation via Lanczos Algorithm
        v0 = ht.full(
            (L.shape[0],),
            fill_value=1.0 / math.sqrt(L.shape[0]),
            dtype=L.dtype,
            split=0,
            device=L.device,
        )
        V, T = ht.lanczos(L, self.n_lanczos, v0)

        # if int(torch.__version__.split(".")[1]) >= 9:
        try:
            # 4. Calculate and Sort Eigenvalues and Eigenvectors of tridiagonal matrix T
            eval, evec = torch.eig(T.larray, eigenvectors=True)

            # If x is an Eigenvector of T, then y = V@x is the corresponding Eigenvector of L
            eval, idx = torch.sort(eval[:, 0], dim=0)
            eigenvalues = ht.array(eval)
            eigenvectors = ht.matmul(V, ht.array(evec))[:, idx]

            return eigenvalues, eigenvectors
        except AttributeError:  # torch version is less than 1.9.0
            # 4. Calculate and Sort Eigenvalues and Eigenvectors of tridiagonal matrix T
            eval, evec = torch.eig(T.larray, eigenvectors=True)
            # If x is an Eigenvector of T, then y = V@x is the corresponding Eigenvector of L
            eval, idx = torch.sort(eval[:, 0], dim=0)
            eigenvalues = ht.array(eval)
            eigenvectors = ht.matmul(V, ht.array(evec))[:, idx]

            return eigenvalues, eigenvectors
//  [python]
def _spectral_embedding(self, x: DNDarray) -> Tuple[DNDarray, DNDarray]:
        """
        Helper function for dataset x embedding.
        Returns Tupel(Eigenvalues, Eigenvectors) of the graph's Laplacian matrix.

        Parameters
        ----------
        x : DNDarray
            Sample Matrix for which the embedding should be calculated

        Notes
        -----
        This will throw out the complex side of the eigenvalues found during this.

        """
        L = self._laplacian.construct(x)
        # 3. Eigenvalue and -vector calculation via Lanczos Algorithm
        v0 = ht.full(
            (L.shape[0],),
            fill_value=1.0 / math.sqrt(L.shape[0]),
            dtype=L.dtype,
            split=0,
            device=L.device,
        )
        V, T = ht.lanczos(L, self.n_lanczos, v0)

        # if int(torch.__version__.split(".")[1]) >= 9:
        try:
            # 4. Calculate and Sort Eigenvalues and Eigenvectors of tridiagonal matrix T
            eval, evec = torch.linalg.eig(T.larray)

            # If x is an Eigenvector of T, then y = V@x is the corresponding Eigenvector of L
            eval, idx = torch.sort(eval.real, dim=0)
            eigenvalues = ht.array(eval)
            eigenvectors = ht.matmul(V, ht.array(evec))[:, idx]

            return eigenvalues.real, eigenvectors.real
        except AttributeError:  # torch version is less than 1.9.0
            # 4. Calculate and Sort Eigenvalues and Eigenvectors of tridiagonal matrix T
            eval, evec = torch.eig(T.larray, eigenvectors=True)
            # If x is an Eigenvector of T, then y = V@x is the corresponding Eigenvector of L
            eval, idx = torch.sort(eval[:, 0], dim=0)
            eigenvalues = ht.array(eval)
            eigenvectors = ht.matmul(V, ht.array(evec))[:, idx]

            return eigenvalues, eigenvectors
//  [python]
def torus_representative(self, x):
        if self.n == 3:
            # In SO(3) the torus representative is determined by the non-trivial pair of eigenvalues,
            # which can be calculated from the trace
            trace = torch.einsum('...ii->...', x)
            real = (trace - 1) / 2
            imag = torch.sqrt(torch.max(1-torch.square(real), torch.zeros_like(real)))
            return torch.view_as_complex(torch.cat((real.unsqueeze(-1), imag.unsqueeze(-1)), -1)).unsqueeze(-1)
        elif self.n % 2 == 1:
            # In SO(2n+1) the torus representative is determined by the (unordered) non-trivial eigenvalues
            eigvals = torch.eig(x, eigenvectors=False).eigenvalues
            sorted_ind = torch.sort(torch.view_as_real(eigvals), dim=-2).indices[..., 0]
            eigvals = torch.gather(eigvals, dim=-1, index=sorted_ind)
            gamma = eigvals[..., 0:-1:2]
            return gamma
        else:
            # In SO(2n) each unordered set of eigenvalues determines two conjugacy classes
            eigvals, eigvecs = torch.eig(x, eigenvectors=True).eigenvalues, torch.eig(x, eigenvectors=True).eigenvectors
            sorted_ind = torch.sort(torch.view_as_real(eigvals), dim=-2).indices[..., 0]
            eigvals = torch.gather(eigvals, dim=-1, index=sorted_ind)
            eigvecs = torch.gather(eigvecs, dim=-1, index=sorted_ind.unsqueeze(-2).broadcast_to(eigvecs.shape))
            # c is a matrix transforming x into its canonical form (with 2x2 blocks)
            c = torch.zeros_like(eigvecs)
            c[..., ::2] = eigvecs[..., ::2] + eigvecs[..., 1::2]
            c[..., 1::2] = (eigvecs[..., ::2] - eigvecs[..., 1::2])
            # eigenvectors calculated by LAPACK are either real or purely imaginary, make everything real
            # WARNING: might depend on the implementation of the eigendecomposition!
            c = c.real + c.imag
            # normalize s.t. det(c)≈±1, probably unnecessary
            c /= math.sqrt(2)
            torch.pow(eigvals[..., 0], torch.det(c).sgn(), out=eigvals[..., 0])
            gamma = eigvals[..., ::2]
            return gamma
//  [python]
def torus_representative(self, x):
        if self.n == 3:
            # In SO(3) the torus representative is determined by the non-trivial pair of eigenvalues,
            # which can be calculated from the trace
            trace = torch.einsum('...ii->...', x)
            real = (trace - 1) / 2
            imag = torch.sqrt(torch.max(1-torch.square(real), torch.zeros_like(real)))
            return torch.view_as_complex(torch.cat((real.unsqueeze(-1), imag.unsqueeze(-1)), -1)).unsqueeze(-1)
        elif self.n % 2 == 1:
            # In SO(2n+1) the torus representative is determined by the (unordered) non-trivial eigenvalues
            eigvals = torch.linalg.eigvals(x)
            sorted_ind = torch.sort(torch.view_as_real(eigvals), dim=-2).indices[..., 0]
            eigvals = torch.gather(eigvals, dim=-1, index=sorted_ind)
            gamma = eigvals[..., 0:-1:2]
            return gamma
        else:
            # In SO(2n) each unordered set of eigenvalues determines two conjugacy classes
            eigvals, eigvecs = torch.linalg.eig(x)
            sorted_ind = torch.sort(torch.view_as_real(eigvals), dim=-2).indices[..., 0]
            eigvals = torch.gather(eigvals, dim=-1, index=sorted_ind)
            eigvecs = torch.gather(eigvecs, dim=-1, index=sorted_ind.unsqueeze(-2).broadcast_to(eigvecs.shape))
            # c is a matrix transforming x into its canonical form (with 2x2 blocks)
            c = torch.zeros_like(eigvecs)
            c[..., ::2] = eigvecs[..., ::2] + eigvecs[..., 1::2]
            c[..., 1::2] = (eigvecs[..., ::2] - eigvecs[..., 1::2])
            # eigenvectors calculated by LAPACK are either real or purely imaginary, make everything real
            # WARNING: might depend on the implementation of the eigendecomposition!
            c = c.real + c.imag
            # normalize s.t. det(c)≈±1, probably unnecessary
            c /= math.sqrt(2)
            torch.pow(eigvals[..., 0], torch.det(c).sgn(), out=eigvals[..., 0])
            gamma = eigvals[..., ::2]
            return gamma
//  [python]
def eig_vector(data, K):
    b, n, _ = data.shape
    _, idx, _ = pytorch3d.ops.knn_points(data, data, K=K)  # idx (b,n,K)

    idx0 = torch.arange(0,b,device=data.device).reshape((b,1)).expand(-1,n*K).reshape((1,b*n*K))
    idx1 = torch.arange(0,n,device=data.device).reshape((1,n,1)).expand(b,n,K).reshape((1,b*n*K))
    idx = idx.reshape((1,b*n*K))
    idx = torch.cat([idx0, idx1, idx], dim=0) # (3, b*n*K)
    # print(b, n, K, idx.shape)
    ones = torch.ones(idx.shape[1], dtype=bool, device=data.device)
    A = torch.sparse_coo_tensor(idx, ones, (b, n, n)).to_dense() # (b,n,n)
    A = A | A.transpose(1, 2)
    A = A.float()
    deg = torch.diag_embed(torch.sum(A, dim=2))
    laplacian = deg - A
    u, v = torch.linalg.eig(laplacian) # (b,n,n)
    return v.real, laplacian, u.real
//  [python]
def find_nearest_orthogonal_set(fset):
    cset = (fset ** 2).sum(dim=0).sqrt()
    F = fset / cset
    A, V = torch.eig((F.T @ F).inverse())
    Fhat = F @ (V @ A.diag().sqrt() @ V.inverse()).real
    return cset * Fhat
//  [python]
def find_nearest_orthogonal_set(fset):
    cset = (fset ** 2).sum(dim=0).sqrt()
    F = fset / cset
    A, V = torch.linalg.eig((F.T @ F).inverse())
    Fhat = F @ (V @ A.diag().sqrt() @ V.inverse()).real
    return cset * Fhat
//  [python]
def pca(X, device, th=True):
    if th == True:
        mean = torch.mean(X, dim=0)
        X = X - mean
        X = X / torch.std(X, dim=0)
        cov = 1 / X.shape[0] * (X.t() @ X)
        Lambda, V = torch.symeig(cov, eigenvectors=True)
        max_index = torch.argmax(Lambda)
        max_eigval = torch.max(Lambda)
        # print("Torch")
        # mean :no grad
        return V[:, max_index], mean
    else:
        X_ = X.cpu().detach().numpy()
        mean_ = np.mean(X_, axis=0)
        X_ -= mean_
        X_ /= np.std(X_, axis=0)
        cov = 1 / X_.shape[0] * (X_.transpose() @ X_)
        Lambda, V = np.linalg.eig(cov)
        max_index = np.argmax(Lambda)
        max_eigval = np.max(Lambda)
        # mean :no grad
        return torch.from_numpy(V[:, max_index]).to(device), torch.from_numpy(mean_).to(device)
//  [python]
def pca(X, device, th=True):
    if th == True:
        mean = torch.mean(X, dim=0)
        X = X - mean
        X = X / torch.std(X, dim=0)
        cov = 1 / X.shape[0] * (X.t() @ X)
        Lambda, V = [torch.real(i) for i in torch.linalg.eig(cov)]
        max_index = torch.argmax(Lambda)
        max_eigval = torch.max(Lambda)
        # print("Torch")
        # mean :no grad
        return V[:, max_index], mean
    else:
        X_ = X.cpu().detach().numpy()
        mean_ = np.mean(X_, axis=0)
        X_ -= mean_
        X_ /= np.std(X_, axis=0)
        cov = 1 / X_.shape[0] * (X_.transpose() @ X_)
        Lambda, V = np.linalg.eig(cov)
        max_index = np.argmax(Lambda)
        max_eigval = np.max(Lambda)
        # mean :no grad
        return torch.from_numpy(V[:, max_index]).to(device), torch.from_numpy(mean_).to(device)
//  [python]
def find_grain_aspect(ii):
    
    tmp = torch.from_numpy(np.cov(ii.cpu().numpy())).to(ii.device)
    
    a0, a1 = torch.eig(tmp)
    a0, a1 = a0[:,0], a1.real
    i_maj = torch.argmax(a0)
    i_min = torch.argmin(a0)
    v_maj = a1[:,i_maj] #major axis
    v_min = a1[:,i_min] #minor axis
    tmp = torch.sqrt(a0)
    r_maj = tmp[i_maj] #radius
    r_min = tmp[i_min] #radius
    return r_maj, r_min, v_maj, v_min
//  [python]
def find_grain_aspect(ii):
    
    tmp = torch.from_numpy(np.cov(ii.cpu().numpy())).to(ii.device)
    
    a0, a1 = torch.linalg.eig(tmp)
    a0, a1 = a0.real, a1.real
    i_maj = torch.argmax(a0)
    i_min = torch.argmin(a0)
    v_maj = a1[:,i_maj] #major axis
    v_min = a1[:,i_min] #minor axis
    tmp = torch.sqrt(a0)
    r_maj = tmp[i_maj] #radius
    r_min = tmp[i_min] #radius
    return r_maj, r_min, v_maj, v_min
//  [python]
def get_log_like_gradient_and_hessian(self, X, eps=1e-2, compute_hessian=False):

        log_pi = self.log_prob(X)
        if compute_hessian:
            self.hessian_counter += 1
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.autograd.functional.hessian(self.log_prob, X, create_graph=True)
            lamda, Q = torch.eig(H)
            lamda, Q = torch.real(lamda), torch.real(Q)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / lamda_prime
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def get_log_like_gradient_and_hessian(self, X, eps=1e-2, compute_hessian=False):

        log_pi = self.log_prob(X)
        if compute_hessian:
            self.hessian_counter += 1
            g = torch.autograd.grad(log_pi, X, retain_graph=True, create_graph=True)[0]
            H = torch.autograd.functional.hessian(self.log_prob, X, create_graph=True)
            lamda, Q = torch.linalg.eig(H)
            lamda, Q = torch.real(lamda), torch.real(Q)
            lamda_prime = torch.sqrt(lamda**2 + eps)
            lamda_prime_inv = 1.0 / lamda_prime
            H = Q @ torch.diag(lamda_prime) @ Q.T
            Hinv = Q @ torch.diag(lamda_prime_inv) @ Q.T
            log_det_Hinv = torch.sum(torch.log(lamda_prime_inv))
            return log_pi, g, H, Hinv, log_det_Hinv
        else:
            return log_pi
//  [python]
def make_random_matrix(num_samples: int, dim: int, mu: float = 0., ell: Optional[float] = None) -> torch.Tensor:
    matrix = torch.randn(num_samples, dim, dim)
    eigs: torch.Tensor
    eigs, V = torch.linalg.eig(matrix)
    eigs.real = abs(eigs.real)
    if ell is not None:
        R_0 = ((1/eigs).real).min(-1, keepdim=True)[0]
        eigs = eigs*R_0*ell
    matrix = (V @ torch.diag_embed(eigs) @ torch.linalg.inv(V)).real
    s = torch.linalg.eigvals(matrix)
    #print(s.real.max(dim=-1)[0].min(), s.real.max(dim=-1)[0].max(), s.real.min(), s.imag.max(), abs(s.imag).min())
    return matrix
//  [python]
def track_grad_flow(model, dataset, data, opt, L_decomp=False, rep=None, pos_encoding=None):
    if L_decomp:
        i = dataset.data.edge_index
        n = dataset.data.num_nodes
        adj = model.odeblock.odefunc.symm_norm_adj
        coo_adj = torch.sparse_coo_tensor(i, adj, (n, n))
        torch_eye = torch.eye(n)
        coo_lap = torch_eye - coo_adj
        L_evals, L_evecs = torch.eig(coo_lap)

        mag_evals = torch.abs(L_evals)
        high_eval = torch.max(mag_evals)
        low_eval = torch.min(mag_evals)
        # assuming eigenvalues of graph laplacian are real
        high_evec = L_evecs[:, torch.argmax(mag_evals)].real.unsqueeze(1)
        low_evec = L_evecs[:, torch.argmin(mag_evals)].real.unsqueeze(1)

        model.odeblock.odefunc.high_evec = high_evec
        model.odeblock.odefunc.low_evec = low_evec

    else:
        model.odeblock.odefunc.high_evec = torch.zeros((dataset.data.num_nodes,1))
        model.odeblock.odefunc.low_evec = torch.zeros((dataset.data.num_nodes,1))

    wandb.config.update({'track_grad_flow_switch': True}, allow_val_change=True)
    model.eval()
    feat = data.x
    if model.opt['use_labels']:
        feat = add_labels(feat, data.y, data.train_mask, model.num_classes, model.device)
    logits, accs = model(feat, pos_encoding), []
    wandb.config.update({'track_grad_flow_switch': False}, allow_val_change=True)

    grad_flow_DE = model.odeblock.odefunc.grad_flow_DE
    grad_flow_WDE = model.odeblock.odefunc.grad_flow_WDE
    grad_flow_RQ = model.odeblock.odefunc.grad_flow_RQ
    grad_flow_cos_high = model.odeblock.odefunc.grad_flow_cos_high
    grad_flow_cos_low = model.odeblock.odefunc.grad_flow_cos_low
    grad_flow_train_acc = model.odeblock.odefunc.grad_flow_train_acc
    grad_flow_val_acc = model.odeblock.odefunc.grad_flow_val_acc
    grad_flow_test_acc = model.odeblock.odefunc.grad_flow_test_acc

    if opt['wandb']:
        W = model.odeblock.odefunc.W
        W_evals, W_evecs = torch.eig(W)
        if rep is not None:
            wandb.log({f"W_evals_t_{opt['time']}_rep_{rep}": W_evals, f"W_evecs_t_{opt['time']}_rep_{rep}": W_evecs})
        else:
            wandb.log({f"W_evals": W_evals, f"W_evecs": W_evecs})

        # https://docs.wandb.ai/guides/track/log/log-tables
        times = np.arange(0, opt['time'] + opt['step_size'], opt['step_size'])
        data = [[time, DE, WDE, RQ, CH, CL, train, val, test] for (time, DE, WDE, RQ, CH, CL, train, val, test) in
                zip(times, grad_flow_DE, grad_flow_WDE, grad_flow_RQ, grad_flow_cos_high, grad_flow_cos_low,
                    grad_flow_train_acc, grad_flow_val_acc, grad_flow_test_acc)]
        wand_table = wandb.Table(data=data,
                                 columns=["time", "DE", "WDE", "RQ", "CH", "CL", "train_acc", "val_acc", "test_acc"])
        if rep is not None:
            wandb.log({f"gf_table_t_{opt['time']}_rep_{rep}": wand_table})
        else:
            wandb.log({f"gf_table_t_{opt['time']}": wand_table})
//  [python]
def track_grad_flow(model, dataset, data, opt, L_decomp=False, rep=None, pos_encoding=None):
    if L_decomp:
        i = dataset.data.edge_index
        n = dataset.data.num_nodes
        adj = model.odeblock.odefunc.symm_norm_adj
        coo_adj = torch.sparse_coo_tensor(i, adj, (n, n))
        torch_eye = torch.eye(n)
        coo_lap = torch_eye - coo_adj
        L_evals, L_evecs = torch.linalg.eig(coo_lap)

        mag_evals = torch.abs(L_evals)
        high_eval = torch.max(mag_evals)
        low_eval = torch.min(mag_evals)
        # assuming eigenvalues of graph laplacian are real
        high_evec = L_evecs[:, torch.argmax(mag_evals)].real.unsqueeze(1)
        low_evec = L_evecs[:, torch.argmin(mag_evals)].real.unsqueeze(1)

        model.odeblock.odefunc.high_evec = high_evec
        model.odeblock.odefunc.low_evec = low_evec

    else:
        model.odeblock.odefunc.high_evec = torch.zeros((dataset.data.num_nodes,1))
        model.odeblock.odefunc.low_evec = torch.zeros((dataset.data.num_nodes,1))

    wandb.config.update({'track_grad_flow_switch': True}, allow_val_change=True)
    model.eval()
    feat = data.x
    if model.opt['use_labels']:
        feat = add_labels(feat, data.y, data.train_mask, model.num_classes, model.device)
    logits, accs = model(feat, pos_encoding), []
    wandb.config.update({'track_grad_flow_switch': False}, allow_val_change=True)

    grad_flow_DE = model.odeblock.odefunc.grad_flow_DE
    grad_flow_WDE = model.odeblock.odefunc.grad_flow_WDE
    grad_flow_RQ = model.odeblock.odefunc.grad_flow_RQ
    grad_flow_cos_high = model.odeblock.odefunc.grad_flow_cos_high
    grad_flow_cos_low = model.odeblock.odefunc.grad_flow_cos_low
    grad_flow_train_acc = model.odeblock.odefunc.grad_flow_train_acc
    grad_flow_val_acc = model.odeblock.odefunc.grad_flow_val_acc
    grad_flow_test_acc = model.odeblock.odefunc.grad_flow_test_acc

    if opt['wandb']:
        W = model.odeblock.odefunc.W
        W_evals, W_evecs = torch.linalg.eigh(W)
        if rep is not None:
            wandb.log({f"W_evals_t_{opt['time']}_rep_{rep}": W_evals, f"W_evecs_t_{opt['time']}_rep_{rep}": W_evecs})
        else:
            wandb.log({f"W_evals": W_evals, f"W_evecs": W_evecs})

        # https://docs.wandb.ai/guides/track/log/log-tables
        times = np.arange(0, opt['time'] + opt['step_size'], opt['step_size'])
        data = [[time, DE, WDE, RQ, CH, CL, train, val, test] for (time, DE, WDE, RQ, CH, CL, train, val, test) in
                zip(times, grad_flow_DE, grad_flow_WDE, grad_flow_RQ, grad_flow_cos_high, grad_flow_cos_low,
                    grad_flow_train_acc, grad_flow_val_acc, grad_flow_test_acc)]
        wand_table = wandb.Table(data=data,
                                 columns=["time", "DE", "WDE", "RQ", "CH", "CL", "train_acc", "val_acc", "test_acc"])
        if rep is not None:
            wandb.log({f"gf_table_t_{opt['time']}_rep_{rep}": wand_table})
        else:
            wandb.log({f"gf_table_t_{opt['time']}": wand_table})
//  [python]
def pca_torch(X, no_dims=50):
    set_seed(42)
    (n, d) = X.shape
    X = X - torch.mean(X, 0)
    (l, M) = torch.eig(torch.mm(X.t(), X))

    # split M real
    # this part may be some difference for complex eigenvalue
    # but complex eignevalue is meanless here, so they are replaced by their real part
    i = 0
    while i < d:
        if l[i].imag != 0:
            M[:, i+1] = M[:, i]
            i += 2
        else:
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims].real)
    return Y
//  [python]
def pca_torch(X, no_dims=50):
    set_seed(42)
    (n, d) = X.shape
    X = X - torch.mean(X, 0)
    (l, M) = torch.linalg.eig(torch.mm(X.t(), X))

    # split M real
    # this part may be some difference for complex eigenvalue
    # but complex eignevalue is meanless here, so they are replaced by their real part
    i = 0
    while i < d:
        if l[i].imag != 0:
            M[:, i+1] = M[:, i]
            i += 2
        else:
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims].real)
    return Y
//  [python]
def unnormalized_laplacian_eigenvectors(L): # 
  el, ev = torch.eig(L, eigenvectors=True)
  el = torch.real(el[:,0])
  ev = torch.real(ev)
  idx = torch.argsort(el)
  el = el[idx] 
  ev = ev[:,idx]
  #print("L", L, el[1])
  return el, ev
//  [python]
def unnormalized_laplacian_eigenvectors(L): # 
  el, ev = torch.linalg.eig(L)
  el = torch.real(el)
  ev = torch.real(ev)
  idx = torch.argsort(el)
  el = el[idx] 
  ev = ev[:,idx]
  #print("L", L, el[1])
  return el, ev
//  [python]
def _calculate_log_det(self, var: Tensor) -> Tensor:
        """Calculate the log determinant of a matrix to prevent overflow.

        Args:
            var (Tensor): A tensor of shape (1, k, d, d).

        Returns:
            Tensor: A tensor of shape (1, k, 1).
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            evals, evecs = torch.eig(var[0, k])  # type: ignore
            log_det[k] = torch.log(evals[:,0]).sum()

        return log_det.unsqueeze(-1)
//  [python]
def _calculate_log_det(self, var: Tensor) -> Tensor:
        """Calculate the log determinant of a matrix to prevent overflow.

        Args:
            var (Tensor): A tensor of shape (1, k, d, d).

        Returns:
            Tensor: A tensor of shape (1, k, 1).
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            evals, evecs = torch.linalg.eig(var[0, k])  # type: ignore
            log_det[k] = torch.log(evals.real).sum()

        return log_det.unsqueeze(-1)
//  [python]
def _dmd_eigen(self, Ktilde, interm) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Get the eigendecomposition of the Koopman matrix.

        Inputs:
            Ktilde: From _dmd_linearize(). Shape = (rank, rank).
            interm: From _dmd_linearize(). Shape = (encoded_size, rank).

        Returns:
            Omega: **Continuous index** eigenvalues of Ktilde. Shape = (rank, rank).
            eigvecs: First 'rank' exact eigenvectors of the full (i.e. not tilde) system. Shape = (encoded_size, rank).
        """
        Lambda, eigvecstilde = torch.eig(Ktilde) #shapes: Lambda = (rank,), eigvecstilde is (rank, rank)
        
        # Eigenvalues
        with open(self.log_file, 'a') as lf:
            lf.write(f"Largest magnitude among eigenvalues = {torch.max(torch.abs(Lambda))}\n")
        Omega = torch.diag(torch.log(Lambda)) #shape = (rank, rank) #NOTE: No need to divide by self.dh.tscale here because spacings are normalized to 1
        #NOTE: We can do a check if any pair (or more) of consecutive eigenvalues are so close that their difference is less than some threshold, since this will lead to very large / nan values during backprop. But doing this check requires looping over the Omega tensor (time complexity = O(len(Omega))) and is not worth it.
        
        # Eigenvectors
        eigvecs = interm.to(self.cfg.CTYPE) @ eigvecstilde #shape = (encoded_size, rank)
        S,s = torch.linalg.norm(eigvecs,2), torch.linalg.norm(eigvecs,-2)
        cond = S/s
        if cond > self.cond_threshold:
            with open(self.log_file, 'a') as lf:
                lf.write(f"Condition number = {cond} is greater than threshold = {self.cond_threshold}. This may lead to numerical instability in evaluating linearity loss. In an attempt to mitigate this, singular values smaller than 1/{self.cond_threshold} times the largest will be discarded from the pseudo-inverse computation.\n")

        return Omega, eigvecs
//  [python]
def _dmd_eigen(self, Ktilde, interm) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Get the eigendecomposition of the Koopman matrix.

        Inputs:
            Ktilde: From _dmd_linearize(). Shape = (rank, rank).
            interm: From _dmd_linearize(). Shape = (encoded_size, rank).

        Returns:
            Omega: **Continuous index** eigenvalues of Ktilde. Shape = (rank, rank).
            eigvecs: First 'rank' exact eigenvectors of the full (i.e. not tilde) system. Shape = (encoded_size, rank).
        """
        Lambda, eigvecstilde = torch.linalg.eig(Ktilde) #shapes: Lambda = (rank,), eigvecstilde is (rank, rank)
        
        # Eigenvalues
        with open(self.log_file, 'a') as lf:
            lf.write(f"Largest magnitude among eigenvalues = {torch.max(torch.abs(Lambda))}\n")
        Omega = torch.diag(torch.log(Lambda)) #shape = (rank, rank) #NOTE: No need to divide by self.dh.tscale here because spacings are normalized to 1
        #NOTE: We can do a check if any pair (or more) of consecutive eigenvalues are so close that their difference is less than some threshold, since this will lead to very large / nan values during backprop. But doing this check requires looping over the Omega tensor (time complexity = O(len(Omega))) and is not worth it.
        
        # Eigenvectors
        eigvecs = interm.to(self.cfg.CTYPE) @ eigvecstilde #shape = (encoded_size, rank)
        S,s = torch.linalg.norm(eigvecs,2), torch.linalg.norm(eigvecs,-2)
        cond = S/s
        if cond > self.cond_threshold:
            with open(self.log_file, 'a') as lf:
                lf.write(f"Condition number = {cond} is greater than threshold = {self.cond_threshold}. This may lead to numerical instability in evaluating linearity loss. In an attempt to mitigate this, singular values smaller than 1/{self.cond_threshold} times the largest will be discarded from the pseudo-inverse computation.\n")

        return Omega, eigvecs
//  [python]
def pre_MUSIC(self, Rz: torch.Tensor):
        """Applies the MUSIC operration for generating spectrum

        Args:
            Rz (torch.Tensor): Generated covariance matrix

        Returns:
            torch.Tensor: The generated MUSIC spectrum
        """
        spectrum = []
        bs_Rz = Rz
        for iter in range(self.batch_size):
            R = bs_Rz[iter]
            # Extract eigenvalues and eigenvectors using EVD
            eigenvalues, eigenvectors = torch.symeig(R, eigenvectors=True)
            # Noise subspace as the eigenvectors which associated with the M first eigenvalues
            Un = eigenvectors[:, self.M :]
            # Calculate MUSIC spectrum
            spectrum.append(self.spectrum_calculation(Un)[0])
        return torch.stack(spectrum, dim=0)
//  [python]
def pre_MUSIC(self, Rz: torch.Tensor):
        """Applies the MUSIC operration for generating spectrum

        Args:
            Rz (torch.Tensor): Generated covariance matrix

        Returns:
            torch.Tensor: The generated MUSIC spectrum
        """
        spectrum = []
        bs_Rz = Rz
        for iter in range(self.batch_size):
            R = bs_Rz[iter]
            # Extract eigenvalues and eigenvectors using EVD
            _, eigenvectors = torch.linalg.eig(R)
            # Noise subspace as the eigenvectors which associated with the M first eigenvalues
            Un = eigenvectors[:, self.M :]
            # Calculate MUSIC spectrum
            spectrum.append(self.spectrum_calculation(Un)[0])
        return torch.stack(spectrum, dim=0)
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        pos = data.pos

        if self.max_points > 0 and pos.size(0) > self.max_points:
            perm = torch.randperm(pos.size(0))
            pos = pos[perm[:self.max_points]]

        pos = pos - pos.mean(dim=0, keepdim=True)
        C = torch.matmul(pos.t(), pos)
        e, v = torch.symeig(C, eigenvectors=True)  # v[:,j] is j-th eigenvector

        if self.sort:
            indices = e[:, 0].argsort(descending=True)
            v = v.t()[indices].t()

        data.pos = torch.matmul(data.pos, v)

        if 'normal' in data:
            data.normal = F.normalize(torch.matmul(data.normal, v))

        return data