jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def mol2conf_loss_one_graph(pos_y, pos_x, pos_y_mean, pos_x_mean, num_nodes, total_iso):
        with torch.no_grad():
            total_nodes = pos_y.shape[0]
            y = pos_y - pos_y_mean
            x = pos_x - pos_x_mean
            a = y + x
            b = y - x
            a = a.view(-1, 1, 3)
            b = b.view(-1, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(-1, num_nodes, 16)
            tmpb = torch.mean(tmpb, dim=1).view(-1, 4, 4)
            w, v = torch.symeig(tmpb, eigenvectors=True)
            min_q = v[:, :, 0]
            rotation = GNNet.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean.expand(total_iso, -1), rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
            pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
            loss = (pos_y - pos_x).norm(dim=-1, keepdim=True).view(-1, num_nodes,).mean(-1)
            return torch.argmin(loss)
//  [python]
def mol2conf_loss_one_graph(pos_y, pos_x, pos_y_mean, pos_x_mean, num_nodes, total_iso):
        with torch.no_grad():
            total_nodes = pos_y.shape[0]
            y = pos_y - pos_y_mean
            x = pos_x - pos_x_mean
            a = y + x
            b = y - x
            a = a.view(-1, 1, 3)
            b = b.view(-1, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(-1, num_nodes, 16)
            tmpb = torch.mean(tmpb, dim=1).view(-1, 4, 4)
            w, v = torch.linalg.eigh(tmpb)
            min_q = v[:, :, 0]
            rotation = GNNet.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean.expand(total_iso, -1), rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
            pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
            loss = (pos_y - pos_x).norm(dim=-1, keepdim=True).view(-1, num_nodes,).mean(-1)
            return torch.argmin(loss)
//  [python]
def compute_a_inv(self, damping: float = 0.001) -> None:
        """Compute A inverse on assigned rank.

        update_a_factor() must be called at least once before this function.

        Args:
            damping (float, optional): damping value to condition inverse
                (default: 0.001).
        """
        if not isinstance(self.a_factor, torch.Tensor):
            raise RuntimeError(
                'Cannot eigendecompose A before A has been computed',
            )

        if self.symmetric_factors:
            self.da, self.qa = torch.symeig(
                self.a_factor.to(torch.float32),
            )
        else:
            self.da, self.qa = torch.eig(
                self.a_factor.to(torch.float32),
            )
        self.qa = cast(torch.Tensor, self.qa).to(self.inv_dtype)
        self.da = cast(torch.Tensor, self.da).to(self.inv_dtype)
        self.da = torch.clamp(self.da, min=0.0)
//  [python]
def compute_a_inv(self, damping: float = 0.001) -> None:
        """Compute A inverse on assigned rank.

        update_a_factor() must be called at least once before this function.

        Args:
            damping (float, optional): damping value to condition inverse
                (default: 0.001).
        """
        if not isinstance(self.a_factor, torch.Tensor):
            raise RuntimeError(
                'Cannot eigendecompose A before A has been computed',
            )

        if self.symmetric_factors:
            self.da, self.qa = torch.linalg.eigh(
                self.a_factor.to(torch.float32),
            )
        else:
            self.da, self.qa = torch.linalg.eig(
                self.a_factor.to(torch.float32),
            )
        self.qa = cast(torch.Tensor, self.qa).to(self.inv_dtype)
        self.da = cast(torch.Tensor, self.da).to(self.inv_dtype)
        self.da = torch.clamp(self.da, min=0.0)
//  [python]
def compute_g_inv(self, damping: float = 0.001) -> None:
        """See `compute_g_inv`."""
        if not isinstance(self.g_factor, torch.Tensor):
            raise RuntimeError(
                'Cannot eigendecompose G before G has been computed',
            )

        if self.symmetric_factors:
            self.dg, self.qg = torch.symeig(
                self.g_factor.to(torch.float32),
            )
        else:
            self.dg, self.qg = torch.eig(
                self.g_factor.to(torch.float32),
            )
        self.qg = cast(torch.Tensor, self.qg).to(self.inv_dtype)
        self.dg = self.dg.to(self.inv_dtype)
        self.dg = torch.clamp(self.dg, min=0.0)
        if self.prediv_eigenvalues:
            self.dgda = 1 / (torch.outer(self.dg, self.da) + damping)
            self.dg = None
            self.da = None
//  [python]
def compute_g_inv(self, damping: float = 0.001) -> None:
        """See `compute_g_inv`."""
        if not isinstance(self.g_factor, torch.Tensor):
            raise RuntimeError(
                'Cannot eigendecompose G before G has been computed',
            )

        if self.symmetric_factors:
            self.dg, self.qg = torch.linalg.eigh(
                self.g_factor.to(torch.float32),
            )
        else:
            self.dg, self.qg = torch.linalg.eig(
                self.g_factor.to(torch.float32),
            )
        self.qg = cast(torch.Tensor, self.qg).to(self.inv_dtype)
        self.dg = self.dg.to(self.inv_dtype)
        self.dg = torch.clamp(self.dg, min=0.0)
        if self.prediv_eigenvalues:
            self.dgda = 1 / (torch.outer(self.dg, self.da) + damping)
            self.dg = None
            self.da = None
//  [python]
def alignment_loss(
        pos_y, pos_x, batch,
    ):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.symeig(tmpb, eigenvectors=True)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNNet.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch).mean()
        return loss
//  [python]
def alignment_loss(
        pos_y, pos_x, batch,
    ):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.linalg.eigh(tmpb)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNNet.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch).mean()
        return loss
//  [python]
def rot_to_quat_pth(rot, unstack_inputs:bool=False):
  if unstack_inputs:
    rot = torch.stack([torch.moveaxis(x, -1, 0) for x in torch.moveaxis(rot, -2, 0)])
  # [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
  xx, xy, xz, yx, yy, yz, zx, zy, zz = rot[0][0],rot[0][1],rot[0][2],rot[1][0],rot[1][1],rot[1][2],rot[2][0],rot[2][1],rot[2][2]
  # pylint: disable=bad-whitespace
  k = [[ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
       [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
       [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
       [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]]
  # pylint: enable=bad-whitespace
  k = (1./3.) * torch.stack([torch.stack(x, dim=-1) for x in k], dim=-2)
  # Get eigenvalues in non-decreasing order and associated.
  _, qs = torch.symeig(k, eigenvectors=True)
  return qs[..., -1]
//  [python]
def rot_to_quat_pth(rot, unstack_inputs:bool=False):
  if unstack_inputs:
    rot = torch.stack([torch.moveaxis(x, -1, 0) for x in torch.moveaxis(rot, -2, 0)])
  # [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
  xx, xy, xz, yx, yy, yz, zx, zy, zz = rot[0][0],rot[0][1],rot[0][2],rot[1][0],rot[1][1],rot[1][2],rot[2][0],rot[2][1],rot[2][2]
  # pylint: disable=bad-whitespace
  k = [[ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
       [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
       [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
       [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]]
  # pylint: enable=bad-whitespace
  k = (1./3.) * torch.stack([torch.stack(x, dim=-1) for x in k], dim=-2)
  # Get eigenvalues in non-decreasing order and associated.
  _, qs = torch.linalg.eigh(k)
  return qs[..., -1]
//  [python]
def eigen(self, Q: torch.Tensor) -> torch.Tensor:
        return torch.symeig(Q)
//  [python]
def eigen(self, Q: torch.Tensor) -> torch.Tensor:
        return torch.linalg.eigh(Q)
//  [python]
def test_eigh3():
    
    def rand_C3():
        x = torch.randn(3, 3)
        return x @ x.t()
    
    n = 2
    C = torch.stack([rand_C3() for _ in range(n)])

    eigvals_torch, eigvecs_torch = torch.symeig(C,eigenvectors=True)
    # print('eigvals_torch:\n', eigvals_torch)
    # print('eigvecs_torch:\n', eigvecs_torch)
        
    eigvals, eigvecs = eigh3(C)
    # print('eigvals:\n', eigvals)
    # print('eigvecs:\n', eigvecs)

    assert torch.allclose(eigvals, eigvals_torch, atol=1e-6), \
            (eigvals, eigvals_torch, eigvals - eigvals_torch)
    assert torch.all(torch.isclose(eigvecs, eigvecs_torch, atol=1e-5)
                     | torch.isclose(-eigvecs, eigvecs_torch, atol=1e-5)), \
            (eigvecs, eigvecs_torch, eigvecs - eigvecs_torch)
//  [python]
def test_eigh3():
    
    def rand_C3():
        x = torch.randn(3, 3)
        return x @ x.t()
    
    n = 2
    C = torch.stack([rand_C3() for _ in range(n)])

    eigvals_torch, eigvecs_torch = torch.linalg.eigh(C)
    # print('eigvals_torch:\n', eigvals_torch)
    # print('eigvecs_torch:\n', eigvecs_torch)
        
    eigvals, eigvecs = eigh3(C)
    # print('eigvals:\n', eigvals)
    # print('eigvecs:\n', eigvecs)

    assert torch.allclose(eigvals, eigvals_torch, atol=1e-6), \
            (eigvals, eigvals_torch, eigvals - eigvals_torch)
    assert torch.all(torch.isclose(eigvecs, eigvecs_torch, atol=1e-5)
                     | torch.isclose(-eigvecs, eigvecs_torch, atol=1e-5)), \
            (eigvecs, eigvecs_torch, eigvecs - eigvecs_torch)
//  [python]
def EOSP_fast_eps(self, num_f=-1, oee_eps=1e-6, eps=1e-14):
        # !!!! 未调试 !!!!
        self.eps = eps  # avoid NAN in OEE
        pos_eosp = list()
        if num_f == -1:
            num_f = len(self.tensors)
        pos_rec = list(range(len(self.tensors)))
        pos_large_oee = pos_rec
        qs_num = 0.0
        for n in range(num_f):
            bf.print_progress_bar(n, num_f)
            if len(pos_rec) == 1:
                pos_eosp.append(pos_rec[0])
            else:
                self.center_orthogonalization(
                    c=0, way='qr', normalize=True)
                OEE = self.entanglement_ent_onsite(which=pos_large_oee, print_info=False)
                if len(pos_large_oee) == 0:
                    order1 = tc.argsort(OEE, descending=True)
                    pos1 = [pos_rec[x] for x in order1[:(num_f-len(pos_eosp))]]
                    pos_eosp = pos_eosp + pos1
                    return pos_eosp, qs_num
                qs_num += OEE.sum().item() / OEE.numel() / np.log(
                    self.tensors[0].shape[1])
                p_max = OEE.argmax()
                pos_eosp.append(pos_rec.pop(p_max))
                rho = self.one_body_RDM(p_max)
                lm, u = tc.eig(rho)
                u = u[:, lm.argmax()]
                self.project_multi_qubits(
                    [p_max], u.reshape(-1, 1))
                pos_large_oee = [x for x in range(OEE.numel()) if OEE[x] > oee_eps]
                pos_large_oee.pop(p_max)
        pos_eosp = tc.tensor(
            pos_eosp, dtype=tc.int64, device=self.device)
        return pos_eosp, qs_num
//  [python]
def EOSP_fast_eps(self, num_f=-1, oee_eps=1e-6, eps=1e-14):
        # !!!! 未调试 !!!!
        self.eps = eps  # avoid NAN in OEE
        pos_eosp = list()
        if num_f == -1:
            num_f = len(self.tensors)
        pos_rec = list(range(len(self.tensors)))
        pos_large_oee = pos_rec
        qs_num = 0.0
        for n in range(num_f):
            bf.print_progress_bar(n, num_f)
            if len(pos_rec) == 1:
                pos_eosp.append(pos_rec[0])
            else:
                self.center_orthogonalization(
                    c=0, way='qr', normalize=True)
                OEE = self.entanglement_ent_onsite(which=pos_large_oee, print_info=False)
                if len(pos_large_oee) == 0:
                    order1 = tc.argsort(OEE, descending=True)
                    pos1 = [pos_rec[x] for x in order1[:(num_f-len(pos_eosp))]]
                    pos_eosp = pos_eosp + pos1
                    return pos_eosp, qs_num
                qs_num += OEE.sum().item() / OEE.numel() / np.log(
                    self.tensors[0].shape[1])
                p_max = OEE.argmax()
                pos_eosp.append(pos_rec.pop(p_max))
                rho = self.one_body_RDM(p_max)
                lm, u = tc.linalg.eigh(rho)
                u = u[:, lm.argmax()]
                self.project_multi_qubits(
                    [p_max], u.reshape(-1, 1))
                pos_large_oee = [x for x in range(OEE.numel()) if OEE[x] > oee_eps]
                pos_large_oee.pop(p_max)
        pos_eosp = tc.tensor(
            pos_eosp, dtype=tc.int64, device=self.device)
        return pos_eosp, qs_num
//  [python]
def EOSP_self(self, num_f=-1, recalculate=False):
        pos_eosp = list()
        length = len(self.tensors)
        if num_f == -1:
            num_f = len(self.tensors)
        if recalculate or (self.eosp_ordering is None) \
                or self.eosp_ordering.numel() != len(self.tensors):
            pos_rec = list(range(len(self.tensors)))
            qs_num = 0.0
            eosp_oee_av = list()
            for n in range(num_f):
                bf.print_progress_bar(n, num_f)
                if len(pos_rec) <= 1:
                    pos_eosp.append(pos_rec[0])
                else:
                    self.center_orthogonalization(
                        c=0, way='qr', normalize=True)
                    OEE = self.entanglement_ent_onsite(
                        print_info=False)
                    eosp_oee_av.append(
                        OEE.sum().item() / OEE.numel())
                    qs_num += eosp_oee_av[-1] / np.log(
                        self.tensors[0].shape[1])
                    p_max = OEE.argmax()
                    pos_eosp.append(pos_rec.pop(p_max))
                    rho = self.one_body_RDM(p_max)
                    lm, u = tc.eig(rho)
                    u = u[:, lm.argmax()]
                    self.project_multi_qubits(
                        [p_max], u.reshape(-1, 1))
            pos_eosp = tc.tensor(
                pos_eosp, dtype=tc.int64, device=self.device)
            eosp_oee_av = tc.tensor(
                eosp_oee_av, dtype=self.dtype, device=self.device)
            if pos_eosp.numel() == length:
                self.eosp_ordering = pos_eosp
                self.eosp_oee_av = eosp_oee_av
                self.qs_number = qs_num
        else:
            print('Use existing EOSP ordering...')
            pos_eosp = self.eosp_ordering[:num_f]
        return pos_eosp
//  [python]
def EOSP_self(self, num_f=-1, recalculate=False):
        pos_eosp = list()
        length = len(self.tensors)
        if num_f == -1:
            num_f = len(self.tensors)
        if recalculate or (self.eosp_ordering is None) \
                or self.eosp_ordering.numel() != len(self.tensors):
            pos_rec = list(range(len(self.tensors)))
            qs_num = 0.0
            eosp_oee_av = list()
            for n in range(num_f):
                bf.print_progress_bar(n, num_f)
                if len(pos_rec) <= 1:
                    pos_eosp.append(pos_rec[0])
                else:
                    self.center_orthogonalization(
                        c=0, way='qr', normalize=True)
                    OEE = self.entanglement_ent_onsite(
                        print_info=False)
                    eosp_oee_av.append(
                        OEE.sum().item() / OEE.numel())
                    qs_num += eosp_oee_av[-1] / np.log(
                        self.tensors[0].shape[1])
                    p_max = OEE.argmax()
                    pos_eosp.append(pos_rec.pop(p_max))
                    rho = self.one_body_RDM(p_max)
                    lm, u = tc.linalg.eigh(rho)
                    u = u[:, lm.argmax()]
                    self.project_multi_qubits(
                        [p_max], u.reshape(-1, 1))
            pos_eosp = tc.tensor(
                pos_eosp, dtype=tc.int64, device=self.device)
            eosp_oee_av = tc.tensor(
                eosp_oee_av, dtype=self.dtype, device=self.device)
            if pos_eosp.numel() == length:
                self.eosp_ordering = pos_eosp
                self.eosp_oee_av = eosp_oee_av
                self.qs_number = qs_num
        else:
            print('Use existing EOSP ordering...')
            pos_eosp = self.eosp_ordering[:num_f]
        return pos_eosp
//  [python]
def __expm__(self, matrix: Tensor, symmetric: bool) -> Tensor:
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.symeig(matrix, eigenvectors=True)
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat = torch.from_numpy(expm(matrix.cpu().numpy()))
            diff_mat = diff_mat.to(matrix.device, matrix.dtype)
        return diff_mat
//  [python]
def __expm__(self, matrix: Tensor, symmetric: bool) -> Tensor:
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.linalg.eigh(matrix, UPLO='U')
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat = torch.from_numpy(expm(matrix.cpu().numpy()))
            diff_mat = diff_mat.to(matrix.device, matrix.dtype)
        return diff_mat
//  [python]
def compute_gradients(self, samples, x=None):
        if x is None:
            kernel_width = self.heuristic_kernel_width(samples, samples)
            x = samples
        else:
            _samples = torch.cat([samples, x], dim=-2)
            kernel_width = self.heuristic_kernel_width(_samples, _samples)

        M = samples.size()[-2]
        Kq, grad_K1, grad_K2 = self.grad_gram(samples, samples, kernel_width)
        if self._eta is not None:
            Kq += self._eta * torch.eye(M)

        # eigen_values, eigen_vectors = torch.symeig(Kq, eigenvectors=True, upper=True) torch==1.3
        eigen_values, eigen_vectors = torch.symeig(Kq, eigenvectors=True)
        if (self._n_eigen is None) and (self._n_eigen_threshold is not None):
            eigen_arr = torch.mean(torch.reshape(eigen_values, [-1, M]), dim=0)

            eigen_arr = torch.flip(eigen_arr, [-1])
            eigen_arr /= torch.sum(eigen_arr)
            eigen_cum = torch.cumsum(eigen_arr, dim=-1)
            eigen_lt = torch.lt(eigen_cum, self._n_eigen_threshold)
            self._n_eigen = torch.sum(eigen_lt)
        if self._n_eigen is not None:
            eigen_values = eigen_values[..., -self._n_eigen :]
            eigen_vectors = eigen_vectors[..., -self._n_eigen :]
        eigen_ext = self.nystrom_ext(samples, x, eigen_vectors, eigen_values, kernel_width)
        grad_K1_avg = torch.mean(grad_K1, dim=-3)
        M = torch.tensor(M).to(samples.device)
        beta = (
            -torch.sqrt(M.float())
            * torch.matmul(torch.transpose(eigen_vectors, -1, -2), grad_K1_avg)
            / torch.unsqueeze(eigen_values, -1)
        )
        grads = torch.matmul(eigen_ext, beta)
        self._n_eigen = None
        return grads
//  [python]
def compute_gradients(self, samples, x=None):
        if x is None:
            kernel_width = self.heuristic_kernel_width(samples, samples)
            x = samples
        else:
            _samples = torch.cat([samples, x], dim=-2)
            kernel_width = self.heuristic_kernel_width(_samples, _samples)

        M = samples.size()[-2]
        Kq, grad_K1, grad_K2 = self.grad_gram(samples, samples, kernel_width)
        if self._eta is not None:
            Kq += self._eta * torch.eye(M)

        # eigen_values, eigen_vectors = torch.symeig(Kq, eigenvectors=True, upper=True) torch==1.3
        eigen_values, eigen_vectors = torch.linalg.eigh(Kq, UPLO="U")
        if (self._n_eigen is None) and (self._n_eigen_threshold is not None):
            eigen_arr = torch.mean(torch.reshape(eigen_values, [-1, M]), dim=0)

            eigen_arr = torch.flip(eigen_arr, [-1])
            eigen_arr /= torch.sum(eigen_arr)
            eigen_cum = torch.cumsum(eigen_arr, dim=-1)
            eigen_lt = torch.lt(eigen_cum, self._n_eigen_threshold)
            self._n_eigen = torch.sum(eigen_lt)
        if self._n_eigen is not None:
            eigen_values = eigen_values[..., -self._n_eigen :]
            eigen_vectors = eigen_vectors[..., -self._n_eigen :]
        eigen_ext = self.nystrom_ext(samples, x, eigen_vectors, eigen_values, kernel_width)
        grad_K1_avg = torch.mean(grad_K1, dim=-3)
        M = torch.tensor(M).to(samples.device)
        beta = (
            -torch.sqrt(M.float())
            * torch.matmul(torch.transpose(eigen_vectors, -1, -2), grad_K1_avg)
            / torch.unsqueeze(eigen_values, -1)
        )
        grads = torch.matmul(eigen_ext, beta)
        self._n_eigen = None
        return grads
//  [python]
def __init__(self,config,device='cuda'):
        self.grid = grid = config.diffusion.grid
        self.device= device
        self.metric =metric =config.diffusion.metric 
        self.initial_point = config.diffusion.initial_point
        self.end_point = config.diffusion.end_point

        self.x =  torch.linspace(config.diffusion.initial_point, config.diffusion.end_point, grid).to(self.device)
        self.y = self.x
        
        self.lens = lens= config.diffusion.lens
        self.gain = gain = config.diffusion.gain
        
        K = kernel(self.x,self.y,lens=lens,gain=gain, metric=metric, device=device)  # (2grid,2grid)
        K = rearrange(K, 'i j k l -> (i j) (k l)')
        eig_val, eig_vec = torch.symeig(K+1e-6*torch.eye(K.shape[0], K.shape[0]).to(self.device))
        self.lens =lens
        self.eig_val =  eig_val.to(self.device) 
        self.eig_vec = eig_vec.to(torch.float32).to(self.device) 
        print('eig_val', eig_val.min(), eig_val.max())
        self.D = torch.diag(self.eig_val).to(torch.float32).to(self.device) 
        self.M = torch.matmul(self.eig_vec, self.D).to(self.device) 
        self.M = rearrange(self.M, '(i j) (k l) -> i j k l ', j=grid, k=grid ).to(self.device)#(grid**2, grid**2)

        self.gain=gain
//  [python]
def __init__(self,config,device='cuda'):
        self.grid = grid = config.diffusion.grid
        self.device= device
        self.metric =metric =config.diffusion.metric 
        self.initial_point = config.diffusion.initial_point
        self.end_point = config.diffusion.end_point

        self.x =  torch.linspace(config.diffusion.initial_point, config.diffusion.end_point, grid).to(self.device)
        self.y = self.x
        
        self.lens = lens= config.diffusion.lens
        self.gain = gain = config.diffusion.gain
        
        K = kernel(self.x,self.y,lens=lens,gain=gain, metric=metric, device=device)  # (2grid,2grid)
        K = rearrange(K, 'i j k l -> (i j) (k l)')
        eig_val, eig_vec = torch.linalg.eigh(K+1e-6*torch.eye(K.shape[0], K.shape[0]).to(self.device))
        self.lens =lens
        self.eig_val =  eig_val.to(self.device) 
        self.eig_vec = eig_vec.to(torch.float32).to(self.device) 
        print('eig_val', eig_val.min(), eig_val.max())
        self.D = torch.diag(self.eig_val).to(torch.float32).to(self.device) 
        self.M = torch.matmul(self.eig_vec, self.D).to(self.device) 
        self.M = rearrange(self.M, '(i j) (k l) -> i j k l ', j=grid, k=grid ).to(self.device)#(grid**2, grid**2)

        self.gain=gain
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if True:
        # In old autograd notebook translations gradients look good with this.
        if(rot.shape[-2:] != (3, 3)):
            raise ValueError("Input rotation is incorrectly shaped")

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

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

        k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

        _, vectors = torch.symeig(k, eigenvectors=True)
        return vectors[..., -1]
    else:
        # Rot scores look better with this flag being used.
        return so3_utils.quaternion_from_matrix(rot)
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if True:
        # In old autograd notebook translations gradients look good with this.
        if(rot.shape[-2:] != (3, 3)):
            raise ValueError("Input rotation is incorrectly shaped")

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

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

        k = (1./3.) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

        _, vectors = torch.linalg.eigh(k)
        return vectors[..., -1]
    else:
        # Rot scores look better with this flag being used.
        return so3_utils.quaternion_from_matrix(rot)
//  [python]
def _fit(self, model, train_loader):
        model.eval()
        self.u = - torch.pinverse(model.fc.weight).matmul(model.fc.bias)
        features = []
        logits = []
        for i, (data, _) in enumerate(tqdm(train_loader, f'Fitting covariance matrix at {self.layer_name}', disable=os.getenv('TQDM_DISABLE'))):
            data = lib.to_device(data)
            if self.features_dataset:
                z = data['avg'].float()
            else:
                z = lib.get_features(model, data, layer_names=[self.layer_name])['avg'][-1]

            logits.append(model.fc.forward(z).detach())
            features.append(z-self.u)

            if i > self.max_fit_iter:
                break

        Z = torch.cat(features)
        logits = torch.cat(logits)
        cov = np.cov(Z.T.cpu().numpy(), bias=True)
        # solved singular values pb in cholesky (see https://scicomp.stackexchange.com/questions/30631/how-to-find-the-nearest-a-near-positive-definite-from-a-given-matrix)
        cov += 1e-12 * np.eye(cov.shape[0])
        cov = torch.from_numpy(cov).float().to('cuda', non_blocking=True)

        eigenvalues, eigenvectors = torch.symeig(cov, eigenvectors=True)
        eig_sorted = torch.argsort(eigenvalues, descending=True)[self.principal_dim:]
        self.R = eigenvectors.T[eig_sorted]
        self.R = self.R.T.contiguous()

        vlogits = torch.norm(Z.matmul(self.R), dim=-1)
        self.alpha = logits.max(dim=-1)[0].mean() / vlogits.mean()
        del features, logits, vlogits
        torch.cuda.empty_cache()
//  [python]
def _fit(self, model, train_loader):
        model.eval()
        self.u = - torch.pinverse(model.fc.weight).matmul(model.fc.bias)
        features = []
        logits = []
        for i, (data, _) in enumerate(tqdm(train_loader, f'Fitting covariance matrix at {self.layer_name}', disable=os.getenv('TQDM_DISABLE'))):
            data = lib.to_device(data)
            if self.features_dataset:
                z = data['avg'].float()
            else:
                z = lib.get_features(model, data, layer_names=[self.layer_name])['avg'][-1]

            logits.append(model.fc.forward(z).detach())
            features.append(z-self.u)

            if i > self.max_fit_iter:
                break

        Z = torch.cat(features)
        logits = torch.cat(logits)
        cov = np.cov(Z.T.cpu().numpy(), bias=True)
        # solved singular values pb in cholesky (see https://scicomp.stackexchange.com/questions/30631/how-to-find-the-nearest-a-near-positive-definite-from-a-given-matrix)
        cov += 1e-12 * np.eye(cov.shape[0])
        cov = torch.from_numpy(cov).float().to('cuda', non_blocking=True)

        eigenvalues, eigenvectors = torch.linalg.eigh(cov)
        eig_sorted = torch.argsort(eigenvalues, descending=True)[self.principal_dim:]
        self.R = eigenvectors.T[eig_sorted]
        self.R = self.R.T.contiguous()

        vlogits = torch.norm(Z.matmul(self.R), dim=-1)
        self.alpha = logits.max(dim=-1)[0].mean() / vlogits.mean()
        del features, logits, vlogits
        torch.cuda.empty_cache()
//  [python]
def eigh(a: ArrayLike, UPLO="L"):
    a = _atleast_float_1(a)
    return t.linalg.eigh(a, UPLO=UPLO)
//  [python]
def eigh(a: ArrayLike, UPLO="L"):
    a = _atleast_float_1(a)
    return torch.linalg.eigh(a, UPLO=UPLO)
//  [python]
def plot_graph(G, x=None, r=.3, c0='#52b', c1='#Cab', size=(10, 10)):
    """ Plot graph using laplacian eigenmodes 1 and 2. """
    G = G.scalars()
    fig = plt.figure(figsize=size)
    plt.axis('equal')
    plt.axis('off')
    if isinstance(c0, str): c0 = [c0] * G.Nvtx
    if isinstance(c1, str): c1 = [c1] * len(G[1].keys)
    #-- eigenmodes --
    if type(x) == type(None):
        L = G.codiff(1) @ G.diff(0)
        L = L.data.to_dense()
        eigval, eigvec = torch.symeig(L)
        x = eigvec[:,1:3] * G.Nvtx
    #-- arrows --
    for p, ep in enumerate(G[1].keys):
        i, j = ep
        xi, xj = x[i], x[j]
        add_arrow(xi, xj, r, .1, c1[p])
    #-- vertices --
    for i, xi in enumerate(x):
        add_vertex(xi, r, c0[i], label=str(i))           
    return fig
//  [python]
def plot_graph(G, x=None, r=.3, c0='#52b', c1='#Cab', size=(10, 10)):
    """ Plot graph using laplacian eigenmodes 1 and 2. """
    G = G.scalars()
    fig = plt.figure(figsize=size)
    plt.axis('equal')
    plt.axis('off')
    if isinstance(c0, str): c0 = [c0] * G.Nvtx
    if isinstance(c1, str): c1 = [c1] * len(G[1].keys)
    #-- eigenmodes --
    if type(x) == type(None):
        L = G.codiff(1) @ G.diff(0)
        L = L.data.to_dense()
        eigval, eigvec = torch.linalg.eigh(L)
        x = eigvec[:,1:3] * G.Nvtx
    #-- arrows --
    for p, ep in enumerate(G[1].keys):
        i, j = ep
        xi, xj = x[i], x[j]
        add_arrow(xi, xj, r, .1, c1[p])
    #-- vertices --
    for i, xi in enumerate(x):
        add_vertex(xi, r, c0[i], label=str(i))           
    return fig
//  [python]
def step(self, step_size = 0.5):
        z = torch.randn(self.mean.size(0), self.samples).to(device)
        ss = self.mean.view(-1, 1) + step_size * self.b.matmul(self.d.matmul(z))
        f = calculate_reward(sparsemax(ss, dim=1).T, 
            valid_data[:-test_size], index[:-test_size])
        results = [{'parameters': ss.T[i], 'z': z.T[i], 
        'fitness': f.item()} for i, f in enumerate(f)]
        ranked_results = sorted(results, key=lambda x: x['fitness'])
        selected_results = ranked_results[0:self.samples//2]
        z = torch.stack([g['z'] for g in selected_results])
        g = torch.stack([g['parameters'] for g in selected_results])

        self.mean = (g * self.weights.unsqueeze(1)).sum(0)
        zmean = (z * self.weights.unsqueeze(1)).sum(0)
        self.pc *= (1 - self.cc)
        pc_cov = self.pc.unsqueeze(1) * self.pc.unsqueeze(1).T
        pc_cov = pc_cov + self.cc * (2 - self.cc) * self.c

        bdz = self.b.matmul(self.d).matmul(z.T)
        cmu_cov = bdz.matmul(self.weights.diag_embed())
        cmu_cov = cmu_cov.matmul(bdz.T)

        self.c *= (1 - self.c1 - self.cmu)
        self.c += (self.c1 * pc_cov) + (self.cmu * cmu_cov)
        self.d, self.b = torch.symeig(self.c, eigenvectors=True)
        self.d = self.d.sqrt().diag_embed()
        return ranked_results
//  [python]
def step(self, step_size = 0.5):
        z = torch.randn(self.mean.size(0), self.samples).to(device)
        ss = self.mean.view(-1, 1) + step_size * self.b.matmul(self.d.matmul(z))
        f = calculate_reward(sparsemax(ss, dim=1).T, 
            valid_data[:-test_size], index[:-test_size])
        results = [{'parameters': ss.T[i], 'z': z.T[i], 
        'fitness': f.item()} for i, f in enumerate(f)]
        ranked_results = sorted(results, key=lambda x: x['fitness'])
        selected_results = ranked_results[0:self.samples//2]
        z = torch.stack([g['z'] for g in selected_results])
        g = torch.stack([g['parameters'] for g in selected_results])

        self.mean = (g * self.weights.unsqueeze(1)).sum(0)
        zmean = (z * self.weights.unsqueeze(1)).sum(0)
        self.pc *= (1 - self.cc)
        pc_cov = self.pc.unsqueeze(1) * self.pc.unsqueeze(1).T
        pc_cov = pc_cov + self.cc * (2 - self.cc) * self.c

        bdz = self.b.matmul(self.d).matmul(z.T)
        cmu_cov = bdz.matmul(self.weights.diag_embed())
        cmu_cov = cmu_cov.matmul(bdz.T)

        self.c *= (1 - self.c1 - self.cmu)
        self.c += (self.c1 * pc_cov) + (self.cmu * cmu_cov)
        self.d, self.b = torch.linalg.eigh(self.c, UPLO='U')
        self.d = self.d.sqrt().diag_embed()
        return ranked_results
//  [python]
def eigh(self, a):
        return torch.eig(a)
//  [python]
def eigh(self, a):
        return torch.linalg.eigh(a)
//  [python]
def sqrtm(self, a):
        L, V = torch.symeig(a)
        return (V * torch.sqrt(L)[None, :]) @ V.T
//  [python]
def sqrtm(self, a):
        L, V = torch.linalg.eigh(a)
        return (V * torch.sqrt(L)[None, :]) @ V.T
//  [python]
def get_gaussian_ellipsiod(
    mean: TensorType[3],
    cov: TensorType[3, 3],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard devation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.symeig(cov, eigenvectors=True)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def get_gaussian_ellipsiod(
    mean: TensorType[3],
    cov: TensorType[3, 3],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard devation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.linalg.eigh(cov)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def matrix_decomposition(matrix=None, epsilon=1e-6, method='regularize'):

    """
    Eigen-decomposition for the input hermetian matrix.

    Parameters
    ----------
    matrix: torch.Tensor, default=None
        The hermetian matrix used to perform eigen-decomposition.
    epsilon: float, default=1e-6
        The numerical modification put on the matrix to ensure its positive rank and may used as cutoff for eigenvalues.
    method: string, default='regularize
        Use three different methods to do eigen-decomposition: 'regularize', 'cutoff', 'clamp'
        regularize: Modify the input matrix by adding epsilon driven identify matrix and take absolute value for eigenvalues;
        cutoff: Set epsilon as the cutoff for the eigenvalues of the matrix;
        clamp: Set all eigenvalues smaller than epsilon to be epsilon;
    
    Returns
    ----------
    eigenvalues, eigenvectors (torch.Tensor) of the input hermetian matrix.
    """

    methods = ['regularize', 'cutoff', 'clamp']
    assert method in methods, "Invalid method "+str(method)+", should use methods from "+str(methods)

    _matrix = matrix + torch.eye(matrix.shape[0]) * epsilon
    eigenval, eigenvec = torch.symeig(_matrix,eigenvectors=True)
    eigenvec = eigenvec.t()

    if method == "regularize":
        eigenval = torch.abs(eigenval)
#         eigenvec = torch.abs(eigenvec)
    elif method == "cutoff":
        _mask = eigenval > epsilon
        eigenval = eigenval[_mask]
        eigenvec = eigenvec[_mask]
    elif method == "clamp":
        eigval = torch.clamp_min(eigval, min=epsilon)
    return eigenval, eigenvec
//  [python]
def matrix_decomposition(matrix=None, epsilon=1e-6, method='regularize'):

    """
    Eigen-decomposition for the input hermetian matrix.

    Parameters
    ----------
    matrix: torch.Tensor, default=None
        The hermetian matrix used to perform eigen-decomposition.
    epsilon: float, default=1e-6
        The numerical modification put on the matrix to ensure its positive rank and may used as cutoff for eigenvalues.
    method: string, default='regularize
        Use three different methods to do eigen-decomposition: 'regularize', 'cutoff', 'clamp'
        regularize: Modify the input matrix by adding epsilon driven identify matrix and take absolute value for eigenvalues;
        cutoff: Set epsilon as the cutoff for the eigenvalues of the matrix;
        clamp: Set all eigenvalues smaller than epsilon to be epsilon;
    
    Returns
    ----------
    eigenvalues, eigenvectors (torch.Tensor) of the input hermetian matrix.
    """

    methods = ['regularize', 'cutoff', 'clamp']
    assert method in methods, "Invalid method "+str(method)+", should use methods from "+str(methods)

    _matrix = matrix + torch.eye(matrix.shape[0]) * epsilon
    eigenval, eigenvec = torch.linalg.eigh(_matrix)
    eigenvec = eigenvec.t()

    if method == "regularize":
        eigenval = torch.abs(eigenval)
#         eigenvec = torch.abs(eigenvec)
    elif method == "cutoff":
        _mask = eigenval > epsilon
        eigenval = eigenval[_mask]
        eigenvec = eigenvec[_mask]
    elif method == "clamp":
        eigval = torch.clamp_min(eigval, min=epsilon)
    return eigenval, eigenvec
//  [python]
def validate(self, validation_past, validation_future):

        """
        Validate the model using validation dataset to get validation score.
        """

        self._lobe.eval()
        if not self._network_type == 'SRVNet':
            self._lobe_timelagged.eval()

        if isinstance(validation_past, np.ndarray):
            validation_past = torch.from_numpy(validation_past.astype(np.float32))
        if isinstance(validation_future, np.ndarray):
            validation_future = torch.from_numpy(validation_future.astype(np.float32))
        validation_past = validation_past.to(device=self._device)
        validation_future = validation_future.to(device=self._device)

        with torch.no_grad():
            if self._network_type == "SRVNet":
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe(validation_future)
            else:
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe_timelagged(validation_future)
            _val_score = loss_function(pastdata=_val_pastproject, futuredata=_val_futureproject, criterion=self._network_type, 
                              epsilon=self.__epsilon, method=self.__covariance_method)
            if self._network_type == 'SRVNet':
                rev_propogator = reverse_propogator(pastdata=_val_pastproject, futuredata=_val_futureproject, 
                                                    epsilon=self.__epsilon, method=self.__covariance_method)
                _eigenval, _ = torch.symeig(rev_propogator, eigenvectors=False)
                _eigenval = torch.flip(_eigenval[-int(_val_pastproject.shape[1]):], dims=[0]).detach()
                return -_val_score.item(), _eigenval.cpu().numpy()
            else:
                return -_val_score.item()
//  [python]
def validate(self, validation_past, validation_future):

        """
        Validate the model using validation dataset to get validation score.
        """

        self._lobe.eval()
        if not self._network_type == 'SRVNet':
            self._lobe_timelagged.eval()

        if isinstance(validation_past, np.ndarray):
            validation_past = torch.from_numpy(validation_past.astype(np.float32))
        if isinstance(validation_future, np.ndarray):
            validation_future = torch.from_numpy(validation_future.astype(np.float32))
        validation_past = validation_past.to(device=self._device)
        validation_future = validation_future.to(device=self._device)

        with torch.no_grad():
            if self._network_type == "SRVNet":
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe(validation_future)
            else:
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe_timelagged(validation_future)
            _val_score = loss_function(pastdata=_val_pastproject, futuredata=_val_futureproject, criterion=self._network_type, 
                              epsilon=self.__epsilon, method=self.__covariance_method)
            if self._network_type == 'SRVNet':
                rev_propogator = reverse_propogator(pastdata=_val_pastproject, futuredata=_val_futureproject, 
                                                    epsilon=self.__epsilon, method=self.__covariance_method)
                _eigenval, _ = torch.linalg.eigh(rev_propogator)
                _eigenval = torch.flip(_eigenval[-int(_val_pastproject.shape[1]):], dims=[0]).detach()
                return -_val_score.item(), _eigenval.cpu().numpy()
            else:
                return -_val_score.item()