jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.eig(C)  # v[:,j] is j-th eigenvector
        e, v = torch.view_as_real(e), v.real

        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
//  [python]
def solve_quadratic(num_users, Vars, Hs):
    device = Hs[0][0].device
    num_cls = Hs[0].shape[0] # number of classes
    d = Hs[0].shape[1] # dimension of feature representation
    avg_weight = []
    for i in range(num_users):
        # ---------------------------------------------------------------------------
        # variance ter
        v = torch.tensor(Vars, device=device)
        # ---------------------------------------------------------------------------
        # bias term
        h_ref = Hs[i]
        dist = torch.zeros((num_users, num_users), device=device)
        for j1, j2 in pairwise(tuple(range(num_users))):
            h_j1 = Hs[j1]
            h_j2 = Hs[j2]
            h = torch.zeros((d, d), device=device)
            for k in range(num_cls):
                h += torch.mm((h_ref[k]-h_j1[k]).reshape(d,1), (h_ref[k]-h_j2[k]).reshape(1,d))
            dj12 = torch.trace(h)
            dist[j1][j2] = dj12
            dist[j2][j1] = dj12

        # QP solver
        p_matrix = torch.diag(v) + dist
        p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
        evals, evecs = torch.eig(torch.tensor(p_matrix))
        
        # for numerical stablity
        p_matrix_new = 0
        p_matrix_new = 0
        for ii in range(num_users):
            if evals[ii].real >= 0.01:
                p_matrix_new += evals[ii].real*torch.mm(evecs[:,ii].reshape(num_users,1), evecs[:,ii].reshape(1, num_users))
        p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix)>=0.0) else p_matrix
        
        # solve QP
        alpha = 0
        eps = 1e-3
        if np.all(np.linalg.eigvals(p_matrix)>=0):
            alphav = cvx.Variable(num_users)
            obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
            prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
            prob.solve()
            alpha = alphav.value
            alpha = [(i)*(i>eps) for i in alpha] # zero-out small weights (<eps)
        else:
            alpha = None # if no solution for the optimization problem, use local classifier only
        
        avg_weight.append(alpha)

    return avg_weight
//  [python]
def solve_quadratic(num_users, Vars, Hs):
    device = Hs[0][0].device
    num_cls = Hs[0].shape[0] # number of classes
    d = Hs[0].shape[1] # dimension of feature representation
    avg_weight = []
    for i in range(num_users):
        # ---------------------------------------------------------------------------
        # variance ter
        v = torch.tensor(Vars, device=device)
        # ---------------------------------------------------------------------------
        # bias term
        h_ref = Hs[i]
        dist = torch.zeros((num_users, num_users), device=device)
        for j1, j2 in pairwise(tuple(range(num_users))):
            h_j1 = Hs[j1]
            h_j2 = Hs[j2]
            h = torch.zeros((d, d), device=device)
            for k in range(num_cls):
                h += torch.mm((h_ref[k]-h_j1[k]).reshape(d,1), (h_ref[k]-h_j2[k]).reshape(1,d))
            dj12 = torch.trace(h)
            dist[j1][j2] = dj12
            dist[j2][j1] = dj12

        # QP solver
        p_matrix = torch.diag(v) + dist
        p_matrix = p_matrix.cpu().numpy()  # coefficient for QP problem
        evals, evecs = torch.linalg.eig(torch.tensor(p_matrix))
        
        # for numerical stablity
        p_matrix_new = 0
        p_matrix_new = 0
        for ii in range(num_users):
            if evals[ii].real >= 0.01:
                p_matrix_new += evals[ii].real*torch.mm(evecs[:,ii].reshape(num_users,1), evecs[:,ii].reshape(1, num_users))
        p_matrix = p_matrix_new.numpy() if not np.all(np.linalg.eigvals(p_matrix)>=0.0) else p_matrix
        
        # solve QP
        alpha = 0
        eps = 1e-3
        if np.all(np.linalg.eigvals(p_matrix)>=0):
            alphav = cvx.Variable(num_users)
            obj = cvx.Minimize(cvx.quad_form(alphav, p_matrix))
            prob = cvx.Problem(obj, [cvx.sum(alphav) == 1.0, alphav >= 0])
            prob.solve()
            alpha = alphav.value
            alpha = [(i)*(i>eps) for i in alpha] # zero-out small weights (<eps)
        else:
            alpha = None # if no solution for the optimization problem, use local classifier only
        
        avg_weight.append(alpha)

    return avg_weight
//  [python]
def decode_ori(ori, b):
    """Decode orientation predicted by the neural network.
    This code is optimized compared to Proença code: vectorization and pre-compute

    Compute the average quaternion q of a set of quaternions Q,
    based on a Linear Least Squares Solution of the form: Ax = 0

    The sum of squared dot products between quaternions:
        L(q) = sum_i w_i(Q_i^T*q)^T(Q_i^T*q)^T

    achieves its maximum when its derivative wrt q is zero, i.e.:
        Aq = 0 where A = sum_i (Q_i*Q_i^T)

    Therefore, the optimal q is simply the right null space of A.

    For more solutions check:
    F. Landis Markley et al. "Averaging quaternions." Journal of Guidance, Control, and Dynamics (2007)
    https://ntrs.nasa.gov/api/citations/20070017872/downloads/20070017872.pdf

    Arguments:
        b: The pre-computed decode ori variable based on the histogram
        ori: The weights associated to each orientation bin (prediction of the neural network)
    Returns:
        q_avg: The solution
        h_inv: The uncertainty in the maximum likelihood sense
    """

    n_bins = b.size(0)

    # Remark: referenced as small "a" in the code but big "A" in the article
    a = torch.sum(b * torch.reshape(ori, (n_bins, 1, 1)), dim=0)

    if True in torch.isnan(a):
        raise ValueError("Error during orientation decoding")

    s, v = torch.eig(a, eigenvectors=True)
    s, v = torch.real(s), torch.real(v)

    idx = torch.argsort(s)

    q_avg = v[:, idx[-1]]

    # Due to numerical errors, we need to enforce normalization (comes from Proença code)
    q_avg = q_avg / torch.linalg.norm(q_avg)

    h_inv = torch.inverse(a)

    return q_avg, h_inv
//  [python]
def decode_ori(ori, b):
    """Decode orientation predicted by the neural network.
    This code is optimized compared to Proença code: vectorization and pre-compute

    Compute the average quaternion q of a set of quaternions Q,
    based on a Linear Least Squares Solution of the form: Ax = 0

    The sum of squared dot products between quaternions:
        L(q) = sum_i w_i(Q_i^T*q)^T(Q_i^T*q)^T

    achieves its maximum when its derivative wrt q is zero, i.e.:
        Aq = 0 where A = sum_i (Q_i*Q_i^T)

    Therefore, the optimal q is simply the right null space of A.

    For more solutions check:
    F. Landis Markley et al. "Averaging quaternions." Journal of Guidance, Control, and Dynamics (2007)
    https://ntrs.nasa.gov/api/citations/20070017872/downloads/20070017872.pdf

    Arguments:
        b: The pre-computed decode ori variable based on the histogram
        ori: The weights associated to each orientation bin (prediction of the neural network)
    Returns:
        q_avg: The solution
        h_inv: The uncertainty in the maximum likelihood sense
    """

    n_bins = b.size(0)

    # Remark: referenced as small "a" in the code but big "A" in the article
    a = torch.sum(b * torch.reshape(ori, (n_bins, 1, 1)), dim=0)

    if True in torch.isnan(a):
        raise ValueError("Error during orientation decoding")

    # s, v = torch.eig(a, eigenvectors=True)
    s, v = torch.linalg.eig(a)
    s, v = torch.real(s), torch.real(v)

    idx = torch.argsort(s)

    q_avg = v[:, idx[-1]]

    # Due to numerical errors, we need to enforce normalization (comes from Proença code)
    q_avg = q_avg / torch.linalg.norm(q_avg)

    h_inv = torch.inverse(a)

    return q_avg, h_inv
//  [python]
def cast(cls, dp_optimizer: DPOptimizer, pub_mat_inv):
        assert isinstance(dp_optimizer, DPOptimizer)
        dp_optimizer.__class__ = cls
        assert isinstance(dp_optimizer, PdaPdmdLinearOptimizer)
        dp_optimizer.pub_mat_inv = pub_mat_inv
        w, v = torch.eig(pub_mat_inv)
        print(max(w[:,0].real))
        return dp_optimizer
//  [python]
def cast(cls, dp_optimizer: DPOptimizer, pub_mat_inv):
        assert isinstance(dp_optimizer, DPOptimizer)
        dp_optimizer.__class__ = cls
        assert isinstance(dp_optimizer, PdaPdmdLinearOptimizer)
        dp_optimizer.pub_mat_inv = pub_mat_inv
        w, v = torch.linalg.eig(pub_mat_inv)
        print(max(w.real))
        return dp_optimizer
//  [python]
def __init__(self, num_classes, vec_length, num_examples):
        super().__init__()
        assert vec_length >= num_classes

        random.seed(RAND_SEED)

        #Generate the class label at this point
        self.targets = [None] * num_examples
        for idx in range(num_examples):
            self.targets[idx] = random.randrange(num_classes)

        # To get R, make a random symmetric matrix and use eigenvalue
        # decomposition
        torch.manual_seed(RAND_SEED)
        R = torch.rand([vec_length, vec_length])
        R = R + R.transpose(0, 1)
        _, eigenvectors = torch.symeig(R)
        self._R = eigenvectors.to(torch.float)

        # # For now, use identity for R
        # self._R = torch.eye(vec_length, vec_length)

        self._dist = torch.distributions.normal.Normal(0, 0.2)
        self._dist = self._dist.expand([vec_length])

        self._vec_length = vec_length
//  [python]
def __init__(self, num_classes, vec_length, num_examples):
        super().__init__()
        assert vec_length >= num_classes

        random.seed(RAND_SEED)

        #Generate the class label at this point
        self.targets = [None] * num_examples
        for idx in range(num_examples):
            self.targets[idx] = random.randrange(num_classes)

        # To get R, make a random symmetric matrix and use eigenvalue
        # decomposition
        torch.manual_seed(RAND_SEED)
        R = torch.rand([vec_length, vec_length])
        R = R + R.transpose(0, 1)
        _, eigenvectors = torch.linalg.eig(R)
        self._R = eigenvectors.to(torch.float)

        # # For now, use identity for R
        # self._R = torch.eye(vec_length, vec_length)

        self._dist = torch.distributions.normal.Normal(0, 0.2)
        self._dist = self._dist.expand([vec_length])

        self._vec_length = vec_length
//  [python]
def esprit(Rz: torch.Tensor, M: int, batch_size: int):
    """Implementation of the model-based Esprit algorithm, support Pytorch, intended for
        MB-DL models. the model sets for nominal and ideal condition (Narrow-band, ULA, non-coherent)
        as it accepts the surrogate covariance matrix.
        it is equivalent to src.methods: RootMUSIC.narrowband() method.

    Args:
    -----
        Rz (torch.Tensor): Focused covariance matrix
        M (int): Number of sources
        batch_size: the number of batches

    Returns:
    --------
        doa_batches (torch.Tensor): The predicted doa, over all batches.
    """

    doa_batches = []

    Bs_Rz = Rz
    for iter in range(batch_size):
        R = Bs_Rz[iter]
        # Extract eigenvalues and eigenvectors using EVD
        eigenvalues, eigenvectors = torch.eig(R)

        # Get signal subspace
        Us = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, :M]
        # Separate the signal subspace into 2 overlapping subspaces
        Us_upper, Us_lower = (
            Us[0 : R.shape[0] - 1],
            Us[1 : R.shape[0]],
        )
        # Generate Phi matrix
        phi = torch.linalg.pinv(Us_upper) @ Us_lower
        # Find eigenvalues and eigenvectors (EVD) of Phi
        phi_eigenvalues, _ = torch.eig(phi)
        # Calculate the phase component of the roots
        eigenvalues_angels = torch.angle(phi_eigenvalues)
        # Calculate the DoA out of the phase component
        doa_predictions = -1 * torch.arcsin((1 / np.pi) * eigenvalues_angels)
        doa_batches.append(doa_predictions)

    return torch.stack(doa_batches, dim=0)
//  [python]
def esprit(Rz: torch.Tensor, M: int, batch_size: int):
    """Implementation of the model-based Esprit algorithm, support Pytorch, intended for
        MB-DL models. the model sets for nominal and ideal condition (Narrow-band, ULA, non-coherent)
        as it accepts the surrogate covariance matrix.
        it is equivalent to src.methods: RootMUSIC.narrowband() method.

    Args:
    -----
        Rz (torch.Tensor): Focused covariance matrix
        M (int): Number of sources
        batch_size: the number of batches

    Returns:
    --------
        doa_batches (torch.Tensor): The predicted doa, over all batches.
    """

    doa_batches = []

    Bs_Rz = Rz
    for iter in range(batch_size):
        R = Bs_Rz[iter]
        # Extract eigenvalues and eigenvectors using EVD
        eigenvalues, eigenvectors = torch.linalg.eig(R)

        # Get signal subspace
        Us = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, :M]
        # Separate the signal subspace into 2 overlapping subspaces
        Us_upper, Us_lower = (
            Us[0 : R.shape[0] - 1],
            Us[1 : R.shape[0]],
        )
        # Generate Phi matrix
        phi = torch.linalg.pinv(Us_upper) @ Us_lower
        # Find eigenvalues and eigenvectors (EVD) of Phi
        phi_eigenvalues, _ = torch.linalg.eig(phi)
        # Calculate the phase component of the roots
        eigenvalues_angels = torch.angle(phi_eigenvalues)
        # Calculate the DoA out of the phase component
        doa_predictions = -1 * torch.arcsin((1 / np.pi) * eigenvalues_angels)
        doa_batches.append(doa_predictions)

    return torch.stack(doa_batches, dim=0)
//  [python]
def rotation_vector(R: torch.Tensor) -> torch.Tensor:
    """Returns the Rodrigues parametrization for the given rotation matrices.

    Params:
        R: (N,...,3,3) rotation matrices

    Returns:
        u: (N,...,3) unit length rotation axes
        theta (N,...) angles of rotation

    References:
     - See last page of https://courses.cs.duke.edu/fall13/compsci527/notes/rodrigues.pdf
     - Gohlke https://github.com/malcolmreynolds/transformations/blob/af825900e981ec232ee0fc4955e592cf3ffe3cd2/transformations/transformations.py
    """
    # Most robust approach is to look for eigenvector with eval=1 as the direction.
    # There is only one real eigenvector close to one.
    R_dtype = R.dtype
    R = R.to(torch.float64)
    bshape = R.shape[:-2]
    e, E = torch.eig(R.transpose(-2, -1), eigenvectors=True)  # (N,...,3) and (N,...,3,3)
    th = (abs(e[..., 0] - 1) < 1e-5) & (abs(e[..., 1]) < 1e-5)
    assert th.any(-1).all(), "No unit eigenvector corresponding to eigenvalue 1"
    i = torch.argmax(th.float(), -1)
    i = i.unsqueeze(-1).unsqueeze(-1).expand((-1,) * len(bshape) + (3, -1))
    u = torch.take_along_dim(E, i, dim=-1).squeeze(-1)

    # trace(R) = 1 + 2*cos(theta)
    cosa = (torch.diagonal(R, dim1=-2, dim2=-1).sum(-1) - 1.0) / 2.0
    # theta only up to sign
    sina = torch.empty_like(cosa)
    mask = abs(u[..., 2]) > 1e-8
    sina[mask] = (
        R[mask][:, 1, 0] + (cosa[mask] - 1.0) * u[mask][:, 0] * u[mask][:, 1]
    ) / u[mask][:, 2]
    mask = abs(u[..., 1]) > 1e-8
    sina[mask] = (
        R[mask][:, 0, 2] + (cosa[mask] - 1.0) * u[mask][:, 0] * u[mask][:, 2]
    ) / u[mask][:, 1]
    mask = abs(u[..., 0]) > 1e-8
    sina[mask] = (
        R[mask][:, 2, 1] + (cosa[mask] - 1.0) * u[mask][:, 1] * u[mask][:, 2]
    ) / u[mask][:, 0]
    theta = torch.atan2(sina, cosa)

    return u.to(R_dtype), theta.to(R_dtype)
//  [python]
def rotation_vector(R: torch.Tensor) -> torch.Tensor:
    """Returns the Rodrigues parametrization for the given rotation matrices.

    Params:
        R: (N,...,3,3) rotation matrices

    Returns:
        u: (N,...,3) unit length rotation axes
        theta (N,...) angles of rotation

    References:
     - See last page of https://courses.cs.duke.edu/fall13/compsci527/notes/rodrigues.pdf
     - Gohlke https://github.com/malcolmreynolds/transformations/blob/af825900e981ec232ee0fc4955e592cf3ffe3cd2/transformations/transformations.py
    """
    # Most robust approach is to look for eigenvector with eval=1 as the direction.
    # There is only one real eigenvector close to one.
    R_dtype = R.dtype
    R = R.to(torch.float64)
    bshape = R.shape[:-2]
    e, E = torch.linalg.eig(R.transpose(-2, -1))  # (N,...,3) and (N,...,3,3)
    th = (abs(torch.real(e) - 1) < 1e-5) & (abs(torch.imag(e)) < 1e-5)
    assert th.any(-1).all(), "No unit eigenvector corresponding to eigenvalue 1"
    i = torch.argmax(th.float(), -1)
    i = i.unsqueeze(-1).unsqueeze(-1).expand((-1,) * len(bshape) + (3, -1))
    u = torch.take_along_dim(torch.real(E), i, dim=-1).squeeze(-1)

    # trace(R) = 1 + 2*cos(theta)
    cosa = (torch.diagonal(R, dim1=-2, dim2=-1).sum(-1) - 1.0) / 2.0
    # theta only up to sign
    sina = torch.empty_like(cosa)
    mask = abs(u[..., 2]) > 1e-8
    sina[mask] = (
        R[mask][:, 1, 0] + (cosa[mask] - 1.0) * u[mask][:, 0] * u[mask][:, 1]
    ) / u[mask][:, 2]
    mask = abs(u[..., 1]) > 1e-8
    sina[mask] = (
        R[mask][:, 0, 2] + (cosa[mask] - 1.0) * u[mask][:, 0] * u[mask][:, 2]
    ) / u[mask][:, 1]
    mask = abs(u[..., 0]) > 1e-8
    sina[mask] = (
        R[mask][:, 2, 1] + (cosa[mask] - 1.0) * u[mask][:, 1] * u[mask][:, 2]
    ) / u[mask][:, 0]
    theta = torch.atan2(sina, cosa)

    return u.to(R_dtype), theta.to(R_dtype)
//  [python]
def root_music(Rz: torch.Tensor, M: int, batch_size: int):
    """Implementation of the model-based Root-MUSIC algorithm, support Pytorch, intended for
        MB-DL models. the model sets for nominal and ideal condition (Narrow-band, ULA, non-coherent)
        as it accepts the surrogate covariance matrix.
        it is equivalent tosrc.methods: RootMUSIC.narrowband() method.

    Args:
    -----
        Rz (torch.Tensor): Focused covariance matrix
        M (int): Number of sources
        batch_size: the number of batches

    Returns:
    --------
        doa_batches (torch.Tensor): The predicted doa, over all batches.
        doa_all_batches (torch.Tensor): All doa predicted, given all roots, over all batches.
        roots_to_return (torch.Tensor): The unsorted roots.
    """

    dist = 0.5
    f = 1
    doa_batches = []
    doa_all_batches = []
    Bs_Rz = Rz
    for iter in range(batch_size):
        R = Bs_Rz[iter]
        # Extract eigenvalues and eigenvectors using EVD
        eigenvalues, eigenvectors = torch.eig(R, eigenvectors=True)
        # Assign noise subspace as the eigenvectors associated with M greatest eigenvalues
        Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues[:,0])).flip(0)][:, M:]
        # Generate hermitian noise subspace matrix
        F = torch.matmul(Un, torch.t(torch.conj(Un)))
        # Calculates the sum of F matrix diagonals
        diag_sum = sum_of_diags_torch(F)
        # Calculates the roots of the polynomial defined by F matrix diagonals
        roots = find_roots_torch(diag_sum)
        # Calculate the phase component of the roots
        roots_angels_all = torch.angle(roots)
        # Calculate doa
        doa_pred_all = torch.arcsin((1 / (2 * np.pi * dist * f)) * roots_angels_all)
        doa_all_batches.append(doa_pred_all)
        roots_to_return = roots
        # Take only roots which inside the unit circle
        roots = roots[
            sorted(range(roots.shape[0]), key=lambda k: abs(abs(roots[k]) - 1))
        ]
        mask = (torch.abs(roots) - 1) < 0
        roots = roots[mask][:M]
        # Calculate the phase component of the roots
        roots_angels = torch.angle(roots)
        # Calculate doa
        doa_pred = torch.arcsin((1 / (2 * np.pi * dist * f)) * roots_angels)
        doa_batches.append(doa_pred)

    return (
        torch.stack(doa_batches, dim=0),
        torch.stack(doa_all_batches, dim=0),
        roots_to_return,
    )
//  [python]
def root_music(Rz: torch.Tensor, M: int, batch_size: int):
    """Implementation of the model-based Root-MUSIC algorithm, support Pytorch, intended for
        MB-DL models. the model sets for nominal and ideal condition (Narrow-band, ULA, non-coherent)
        as it accepts the surrogate covariance matrix.
        it is equivalent tosrc.methods: RootMUSIC.narrowband() method.

    Args:
    -----
        Rz (torch.Tensor): Focused covariance matrix
        M (int): Number of sources
        batch_size: the number of batches

    Returns:
    --------
        doa_batches (torch.Tensor): The predicted doa, over all batches.
        doa_all_batches (torch.Tensor): All doa predicted, given all roots, over all batches.
        roots_to_return (torch.Tensor): The unsorted roots.
    """

    dist = 0.5
    f = 1
    doa_batches = []
    doa_all_batches = []
    Bs_Rz = Rz
    for iter in range(batch_size):
        R = Bs_Rz[iter]
        # Extract eigenvalues and eigenvectors using EVD
        eigenvalues, eigenvectors = torch.linalg.eig(R)
        # Assign noise subspace as the eigenvectors associated with M greatest eigenvalues
        Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, M:]
        # Generate hermitian noise subspace matrix
        F = torch.matmul(Un, torch.t(torch.conj(Un)))
        # Calculates the sum of F matrix diagonals
        diag_sum = sum_of_diags_torch(F)
        # Calculates the roots of the polynomial defined by F matrix diagonals
        roots = find_roots_torch(diag_sum)
        # Calculate the phase component of the roots
        roots_angels_all = torch.angle(roots)
        # Calculate doa
        doa_pred_all = torch.arcsin((1 / (2 * np.pi * dist * f)) * roots_angels_all)
        doa_all_batches.append(doa_pred_all)
        roots_to_return = roots
        # Take only roots which inside the unit circle
        roots = roots[
            sorted(range(roots.shape[0]), key=lambda k: abs(abs(roots[k]) - 1))
        ]
        mask = (torch.abs(roots) - 1) < 0
        roots = roots[mask][:M]
        # Calculate the phase component of the roots
        roots_angels = torch.angle(roots)
        # Calculate doa
        doa_pred = torch.arcsin((1 / (2 * np.pi * dist * f)) * roots_angels)
        doa_batches.append(doa_pred)

    return (
        torch.stack(doa_batches, dim=0),
        torch.stack(doa_all_batches, dim=0),
        roots_to_return,
    )
//  [python]
def reduced_rank_eig(
    input_covariance: Tensor,
    lagged_covariance: Tensor,  # C_{0t}
    tikhonov_reg: float,
    rank: Optional[int] = None,
) -> tuple:
    """Reduced rank regression algorithm, as described in [1]_.

    Parameters
    ----------
    input_covariance : Tensor

    lagged_covariance : Tensor

    rank : Optional[int], optional
        Rank of the final estimator, by default None

    Returns
    -------
    tuple
        A tuple containing the eigenvalues and eigenvectors of the Koopman operator.

    References
    ----------
    .. [1] V. Kostic, P. Novelli, A. Maurer, C. Ciliberto, L. Rosasco, and M. Pontil, "Learning Dynamical Systems via Koopman Operator Regression in Reproducing Kernel Hilbert Spaces" (2022).
    """
    n = input_covariance.shape[0]
    reg_input_covariance = input_covariance + tikhonov_reg * torch.eye(
        n, dtype=input_covariance.dtype, device=input_covariance.device
    )

    _crcov = torch.mm(lagged_covariance, lagged_covariance.T)
    _, _vectors = generalized_eigh(_crcov, reg_input_covariance)

    _norms = spd_norm(_vectors, reg_input_covariance)
    vectors = _vectors * (1 / _norms)

    if rank is not None:
        _, idxs = torch.topk(vectors.values, rank)
        U = vectors[:, idxs]
    else:
        U = vectors

    # U@(U.T)@Tw = v w -> (U.T)@T@Uq = vq and w = Uq
    values, Q = torch.eig((U.T) @ (lagged_covariance @ U))
    return values, U @ Q
//  [python]
def reduced_rank_eig(
    input_covariance: Tensor,
    lagged_covariance: Tensor,  # C_{0t}
    tikhonov_reg: float,
    rank: Optional[int] = None,
) -> tuple:
    """Reduced rank regression algorithm, as described in [1]_.

    Parameters
    ----------
    input_covariance : Tensor

    lagged_covariance : Tensor

    rank : Optional[int], optional
        Rank of the final estimator, by default None

    Returns
    -------
    tuple
        A tuple containing the eigenvalues and eigenvectors of the Koopman operator.

    References
    ----------
    .. [1] V. Kostic, P. Novelli, A. Maurer, C. Ciliberto, L. Rosasco, and M. Pontil, "Learning Dynamical Systems via Koopman Operator Regression in Reproducing Kernel Hilbert Spaces" (2022).
    """
    n = input_covariance.shape[0]
    reg_input_covariance = input_covariance + tikhonov_reg * torch.eye(
        n, dtype=input_covariance.dtype, device=input_covariance.device
    )

    _crcov = torch.mm(lagged_covariance, lagged_covariance.T)
    _, _vectors = generalized_eigh(_crcov, reg_input_covariance)

    _norms = spd_norm(_vectors, reg_input_covariance)
    vectors = _vectors * (1 / _norms)

    if rank is not None:
        _, idxs = torch.topk(vectors.values, rank)
        U = vectors[:, idxs]
    else:
        U = vectors

    # U@(U.T)@Tw = v w -> (U.T)@T@Uq = vq and w = Uq
    values, Q = torch.linalg.eig((U.T) @ (lagged_covariance @ U))
    return values, U @ Q
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2]  # Only used in testing
        return w, P, B, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype)  # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0]  # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.eig(AP)  # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    B = contract('ij, j -> i', V_inv, B.to(V))  # V^* B
    P = contract('ij, ...j -> ...i', V_inv, P.to(V))  # V^* P

    return w, P, B, V
//  [python]
def nplr(measure, N, rank=1, dtype=torch.float):
    """ Return w, p, q, V, B such that
    (w - p q^*, B) is unitarily equivalent to the original HiPPO A, B by the matrix V
    i.e. A = V[w - p q^*]V^*, B = V B
    """
    assert dtype == torch.float or torch.cfloat
    if measure == 'random':
        dtype = torch.cfloat if dtype == torch.float else torch.cdouble
        # w = torch.randn(N//2, dtype=dtype)
        w = -torch.exp(torch.randn(N//2)) + 1j*torch.randn(N//2)
        P = torch.randn(rank, N//2, dtype=dtype)
        B = torch.randn(N//2, dtype=dtype)
        V = torch.eye(N, dtype=dtype)[..., :N//2]  # Only used in testing
        return w, P, B, V

    A, B = transition(measure, N)
    A = torch.as_tensor(A, dtype=dtype)  # (N, N)
    B = torch.as_tensor(B, dtype=dtype)[:, 0]  # (N,)

    P = rank_correction(measure, N, rank=rank, dtype=dtype)
    AP = A + torch.sum(P.unsqueeze(-2)*P.unsqueeze(-1), dim=-3)
    w, V = torch.linalg.eig(AP)  # (..., N) (..., N, N)
    # V w V^{-1} = A

    # Only keep one of the conjugate pairs
    w = w[..., 0::2].contiguous()
    V = V[..., 0::2].contiguous()

    V_inv = V.conj().transpose(-1, -2)

    B = contract('ij, j -> i', V_inv, B.to(V))  # V^* B
    P = contract('ij, ...j -> ...i', V_inv, P.to(V))  # V^* P

    return w, P, B, V
//  [python]
def dis_s(self, N, app_ord):  
        app_ord = int(app_ord / 2) 
        s = torch.cat((torch.tensor([0, 1]), torch.zeros(N-1-2*app_ord), torch.tensor([1])))
        S = self.cconvm(N,s) + torch.diag((torch.fft.fft(s)).real);

        p = N
        r = math.floor(N/2)
        P = torch.zeros((p,p))
        P[0,0] = 1
        even = 1 - (p%2)
        
        for i in range(1,r-even+1):
            P[i,i] = 1/(2**(1/2))
            P[i,p-i] = 1/(2**(1/2))
            
        if even:
            P[r,r] = 1
            
        for i in range(r+1,p):
            P[i,i] = -1/(2**(1/2))
            P[i,p-i] = 1/(2**(1/2))

        CS = torch.einsum("ij,jk,ni->nk", S, P.T, P)
        C2 = CS[0:math.floor(N/2+1), 0:math.floor(N/2+1)]
        S2 = CS[math.floor(N/2+1):N, math.floor(N/2+1):N]
        ec, vc = torch.eig(C2)
        es, vs = torch.eig(S2)
        ec = ec.real
        vc = vc.real
        es = es.real
        vs = vs.real
        qvc = torch.vstack((vc, torch.zeros([math.ceil(N/2-1), math.floor(N/2+1)])))
        SC2 = P@qvc # Even Eigenvector of S
        qvs = torch.vstack((torch.zeros([math.floor(N/2+1), math.ceil(N/2-1)]),vs))
        SS2 = P@qvs # Odd Eigenvector of S
        idx = torch.argsort(-ec)
        SC2 = SC2[:,idx]
        idx = torch.argsort(-es)
        SS2 = SS2[:,idx]
        
        if N%2 == 0:
            S2C2 = torch.zeros([N,N+1])
            SS2 = torch.hstack([SS2, torch.zeros((SS2.shape[0],1))])
            S2C2[:,range(0,N+1,2)] = SC2;
            S2C2[:,range(1,N,2)] = SS2
            S2C2 = S2C2[:, torch.arange(S2C2.size(1)) != N-1]
        else:
            S2C2 = torch.zeros([N,N])
            S2C2[:,range(0,N+1,2)] = SC2;
            S2C2[:,range(1,N,2)] = SS2
        
        return S2C2
//  [python]
def dis_s(self, N, app_ord):  
        app_ord = int(app_ord / 2) 
        s = torch.cat((torch.tensor([0, 1]), torch.zeros(N-1-2*app_ord), torch.tensor([1])))
        S = self.cconvm(N,s) + torch.diag((torch.fft.fft(s)).real);

        p = N
        r = math.floor(N/2)
        P = torch.zeros((p,p))
        P[0,0] = 1
        even = 1 - (p%2)
        
        for i in range(1,r-even+1):
            P[i,i] = 1/(2**(1/2))
            P[i,p-i] = 1/(2**(1/2))
            
        if even:
            P[r,r] = 1
            
        for i in range(r+1,p):
            P[i,i] = -1/(2**(1/2))
            P[i,p-i] = 1/(2**(1/2))

        CS = torch.einsum("ij,jk,ni->nk", S, P.T, P)
        C2 = CS[0:math.floor(N/2+1), 0:math.floor(N/2+1)]
        S2 = CS[math.floor(N/2+1):N, math.floor(N/2+1):N]
        ec, vc = torch.linalg.eig(C2)
        es, vs = torch.linalg.eig(S2)
        ec = ec.real
        vc = vc.real
        es = es.real
        vs = vs.real
        qvc = torch.vstack((vc, torch.zeros([math.ceil(N/2-1), math.floor(N/2+1)])))
        SC2 = P@qvc # Even Eigenvector of S
        qvs = torch.vstack((torch.zeros([math.floor(N/2+1), math.ceil(N/2-1)]),vs))
        SS2 = P@qvs # Odd Eigenvector of S
        idx = torch.argsort(-ec)
        SC2 = SC2[:,idx]
        idx = torch.argsort(-es)
        SS2 = SS2[:,idx]
        
        if N%2 == 0:
            S2C2 = torch.zeros([N,N+1])
            SS2 = torch.hstack([SS2, torch.zeros((SS2.shape[0],1))])
            S2C2[:,range(0,N+1,2)] = SC2;
            S2C2[:,range(1,N,2)] = SS2
            S2C2 = S2C2[:, torch.arange(S2C2.size(1)) != N-1]
        else:
            S2C2 = torch.zeros([N,N])
            S2C2[:,range(0,N+1,2)] = SC2;
            S2C2[:,range(1,N,2)] = SS2
        
        return S2C2
//  [python]
def __init__(self, 
        xyzs,           # (N, 3)
        opacities,      # (N, 1)
        covariances,    # (N, 6)
        covariances_inv # (N, 6)
    ) -> None:
        super().__init__()
        coefficients = torch.nan_to_num(torch.rsqrt(torch.linalg.det(recover_symmetric(covariances)).clamp_min(0.) + 1E-8).reshape(-1, 1)) # (N, 1)
        radius = 3.0 * torch.sqrt(torch.real(torch.eig(recover_symmetric(covariances), eigenvectors=False)[0]).max(dim=-1).values).reshape(-1)
        front_coordinate = xyzs[:, 0] - radius
        back_coordinate = xyzs[:, 0] + radius

        indices = torch.sort(back_coordinate).indices.squeeze(-1)

        self.xyzs = xyzs[indices]
        self.opacities = opacities[indices]
        self.covariances = covariances[indices]
        self.covariances_inv = covariances_inv[indices]
        self.front_coordinate = front_coordinate[indices]
        self.back_coordinate = back_coordinate[indices]
        self.coefficients = coefficients[indices]
//  [python]
def __init__(self, 
        xyzs,           # (N, 3)
        opacities,      # (N, 1)
        covariances,    # (N, 6)
        covariances_inv # (N, 6)
    ) -> None:
        super().__init__()
        coefficients = torch.nan_to_num(torch.rsqrt(torch.linalg.det(recover_symmetric(covariances)).clamp_min(0.) + 1E-8).reshape(-1, 1)) # (N, 1)
        radius = 3.0 * torch.sqrt(torch.real(torch.linalg.eig(recover_symmetric(covariances)).eigenvalues).max(dim=-1).values).reshape(-1)
        front_coordinate = xyzs[:, 0] - radius
        back_coordinate = xyzs[:, 0] + radius

        indices = torch.sort(back_coordinate).indices.squeeze(-1)

        self.xyzs = xyzs[indices]
        self.opacities = opacities[indices]
        self.covariances = covariances[indices]
        self.covariances_inv = covariances_inv[indices]
        self.front_coordinate = front_coordinate[indices]
        self.back_coordinate = back_coordinate[indices]
        self.coefficients = coefficients[indices]
//  [python]
def Root_MUSIC(self, Rz, M):
        dist = 0.5
        f = 1
        DOA_list = []
        DOA_all_list = []
        Bs_Rz = Rz
        for iter in range(self.BATCH_SIZE):
            R = Bs_Rz[iter]
            eigenvalues, eigenvectors = torch.eig(R)                                         # Find the eigenvalues and eigenvectors using EVD
            Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, M:] 
            # Un = eigenvectors[:, M:]
            F = torch.matmul(Un, torch.t(torch.conj(Un)))                                           # Set F as the matrix conatains Information, 
            coeff = self.sum_of_diags(F)                                                            # Calculate the sum of the diagonals of F
            roots = self.find_roots(coeff)                                                          # Calculate its roots
            
            roots_angels_all = torch.angle(roots)                                                   # Calculate the phase component of the roots 
            DOA_pred_all = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels_all)              # Calculate the DOA our of the phase component
            DOA_all_list.append(DOA_pred_all)
            roots_to_return = roots
            
            roots = roots[sorted(range(roots.shape[0]), key = lambda k : abs(abs(roots[k]) - 1))]   # Take only roots which are outside unit circle
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred_test = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                 # Calculate the DOA our of the phase component
            mask = (torch.abs(roots) - 1) < 0
            
            roots = roots[mask][:M]
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                      # Calculate the DOA our of the phase component
            DOA_list.append(DOA_pred)                                                               # Convert from radians to Degrees
            
        return torch.stack(DOA_list, dim = 0), torch.stack(DOA_all_list, dim = 0), roots_to_return
//  [python]
def Root_MUSIC(self, Rz, M):
        dist = 0.5
        f = 1
        DOA_list = []
        DOA_all_list = []
        Bs_Rz = Rz
        for iter in range(self.BATCH_SIZE):
            R = Bs_Rz[iter]
            eigenvalues, eigenvectors = torch.linalg.eig(R)                                         # Find the eigenvalues and eigenvectors using EVD
            Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, M:] 
            # Un = eigenvectors[:, M:]
            F = torch.matmul(Un, torch.t(torch.conj(Un)))                                           # Set F as the matrix conatains Information, 
            coeff = self.sum_of_diags(F)                                                            # Calculate the sum of the diagonals of F
            roots = self.find_roots(coeff)                                                          # Calculate its roots
            
            roots_angels_all = torch.angle(roots)                                                   # Calculate the phase component of the roots 
            DOA_pred_all = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels_all)              # Calculate the DOA our of the phase component
            DOA_all_list.append(DOA_pred_all)
            roots_to_return = roots
            
            roots = roots[sorted(range(roots.shape[0]), key = lambda k : abs(abs(roots[k]) - 1))]   # Take only roots which are outside unit circle
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred_test = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                 # Calculate the DOA our of the phase component
            mask = (torch.abs(roots) - 1) < 0
            
            roots = roots[mask][:M]
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                      # Calculate the DOA our of the phase component
            DOA_list.append(DOA_pred)                                                               # Convert from radians to Degrees
            
        return torch.stack(DOA_list, dim = 0), torch.stack(DOA_all_list, dim = 0), roots_to_return
//  [python]
def matrix_operations(A):
        """Calculate the -1/2 power of matrix A"""

        V, Q = torch.eig(A)
        V_inverse = torch.diag(V[:,0] ** (-0.5))
        A_inverse = torch.mm(torch.mm(Q, V_inverse), torch.inverse(Q))

        return A_inverse.double()
//  [python]
def matrix_operations(A):
        """Calculate the -1/2 power of matrix A"""

        V, Q = torch.linalg.eig(A)
        V_inverse = torch.diag(V ** (-0.5))
        A_inverse = torch.mm(torch.mm(Q, V_inverse), torch.linalg.inv(Q))

        return A_inverse.double()
//  [python]
def logm(A):
        """Calculate the matrix logarithm of matrix A"""

        V, Q = torch.eig(A)  # V为特征值,Q为特征向量
        V_log = torch.diag(torch.log(V[:,0]))
        A_logm = torch.mm(torch.mm(Q, V_log), Q.inverse())

        return A_logm.double()
//  [python]
def logm(A):
        """Calculate the matrix logarithm of matrix A"""

        V, Q = torch.linalg.eig(A)  # V为特征值,Q为特征向量
        V_log = torch.diag(torch.log(V))
        A_logm = torch.mm(torch.mm(Q, V_log), torch.linalg.inv(Q))

        return A_logm.double()
//  [python]
def Root_MUSIC(self, Rz, M):
        dist = 0.5 
        f = 1
        DOA_list = []
        DOA_all_list = []
        Bs_Rz = Rz
        for iter in range(self.BATCH_SIZE):
            R = Bs_Rz[iter]
            eigenvalues, eigenvectors = torch.eig(R, eigenvectors=True)                                         # Find the eigenvalues and eigenvectors using EVD
            Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues[:,0])).flip(0)][:, M:] 
            # Un = eigenvectors[:, M:]
            # Un = eigenvectors[:, M:]
            F = torch.matmul(Un, torch.t(torch.conj(Un)))                                           # Set F as the matrix conatains Information, 
            coeff = self.sum_of_diags(F)                                                            # Calculate the sum of the diagonals of F
            roots = self.find_roots(coeff)                                                          # Calculate its roots
            
            roots_angels_all = torch.angle(roots)                                                   # Calculate the phase component of the roots 
            DOA_pred_all = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels_all)              # Calculate the DOA our of the phase component
            DOA_all_list.append(DOA_pred_all)
            roots_to_return = roots
            
            roots = roots[sorted(range(roots.shape[0]), key = lambda k : abs(abs(roots[k]) - 1))]   # Take only roots which are outside unit circle
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred_test = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                 # Calculate the DOA our of the phase component
            
            # print("abs(roots)", abs(roots)-1)
            # print("DOA_pred_test", DOA_pred_test * 180 / np.pi)
            
            mask = (torch.abs(roots) - 1) < 0
            roots = roots[mask][:M]
            # indices = torch.nonzero(mask
            # print("abs(roots)", abs(roots)-1)

            
            
            # roots = [root for root in roots if (abs(root) - 1) < 0][:M]
            # roots = roots[:2 * M:2]                                                                 # Take the M most closest to the unit circle roots
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                      # Calculate the DOA our of the phase component
            # print("DOA Pred", DOA_pred * 180 / np.pi)
            DOA_list.append(DOA_pred)                                                               # Convert from radians to Deegres
        return torch.stack(DOA_list, dim = 0), torch.stack(DOA_all_list, dim = 0), roots_to_return
//  [python]
def Root_MUSIC(self, Rz, M):
        dist = 0.5 
        f = 1
        DOA_list = []
        DOA_all_list = []
        Bs_Rz = Rz
        for iter in range(self.BATCH_SIZE):
            R = Bs_Rz[iter]
            eigenvalues, eigenvectors = torch.linalg.eig(R)                                         # Find the eigenvalues and eigenvectors using EVD
            Un = eigenvectors[:, torch.argsort(torch.abs(eigenvalues)).flip(0)][:, M:] 
            # Un = eigenvectors[:, M:]
            # Un = eigenvectors[:, M:]
            F = torch.matmul(Un, torch.t(torch.conj(Un)))                                           # Set F as the matrix conatains Information, 
            coeff = self.sum_of_diags(F)                                                            # Calculate the sum of the diagonals of F
            roots = self.find_roots(coeff)                                                          # Calculate its roots
            
            roots_angels_all = torch.angle(roots)                                                   # Calculate the phase component of the roots 
            DOA_pred_all = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels_all)              # Calculate the DOA our of the phase component
            DOA_all_list.append(DOA_pred_all)
            roots_to_return = roots
            
            roots = roots[sorted(range(roots.shape[0]), key = lambda k : abs(abs(roots[k]) - 1))]   # Take only roots which are outside unit circle
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred_test = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                 # Calculate the DOA our of the phase component
            
            # print("abs(roots)", abs(roots)-1)
            # print("DOA_pred_test", DOA_pred_test * 180 / np.pi)
            
            mask = (torch.abs(roots) - 1) < 0
            roots = roots[mask][:M]
            # indices = torch.nonzero(mask
            # print("abs(roots)", abs(roots)-1)

            
            
            # roots = [root for root in roots if (abs(root) - 1) < 0][:M]
            # roots = roots[:2 * M:2]                                                                 # Take the M most closest to the unit circle roots
            roots_angels = torch.angle(roots)                                                       # Calculate the phase component of the roots 
            DOA_pred = torch.arcsin((1/(2 * np.pi * dist * f)) * roots_angels)                      # Calculate the DOA our of the phase component
            # print("DOA Pred", DOA_pred * 180 / np.pi)
            DOA_list.append(DOA_pred)                                                               # Convert from radians to Deegres
        return torch.stack(DOA_list, dim = 0), torch.stack(DOA_all_list, dim = 0), roots_to_return
//  [python]
def __init__(self, kernel, support, device='cpu'):
        assert support.shape[1] == 1, 'Only supported for d_x = 1'
        self.support = support  # (n_x, 1)
        self.device = device

        self.cov_matr = kernel(support, support).to(self.device).double()
        self.cov_matr_ddx = kernel(support, support, differentiate=True).double().to(self.device)
        dx = support[1:] - support[:-1]
        self.ddx_matrix = make_ddx_operator(support.shape[0], dx=dx).to(self.device).double()
        self.discrete_op = self.cov_matr + torch.matmul(self.cov_matr_ddx, self.ddx_matrix)
        self.inv_op = torch.linalg.inv(self.discrete_op).to(self.device)  # (n_x, n_x)
        self.sobolev_form = self.ddx_matrix.T @ self.ddx_matrix @ self.inv_op  # D^T D C^{-1}
        self.quadratic_form = self.inv_op + self.sobolev_form

        # The quadratic form is not symmetric, PSD in general. So let's project it.
        # https://nhigham.com/2021/01/26/what-is-the-nearest-positive-semidefinite-matrix/
        # https://scicomp.stackexchange.com/questions/30631/how-to-find-the-nearest-a-near-positive-definite-from-a-given-matrix
        self.symmetrized_quadratic_form = 0.5 * (self.quadratic_form + self.quadratic_form.T)
        L, V = torch.eig(self.symmetrized_quadratic_form)
        assert torch.all(torch.isreal(V)), 'Got complex eigenvectors'
        assert torch.all(torch.isreal(L)), 'Got complex eigenvalues'
        V = V.double()
        L = L.double()
        eps = 1e-5
        L[L < eps] = eps
        self.quadratic_form_proj = V @ torch.diag_embed(L) @ V.T