"""Miscellaneous math and algorithm code"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/utils/math.ipynb.

# %% auto #0
__all__ = ['matrix_power', 'gram_schmidt']

# %% ../../src/utils/math.ipynb #983ccb27-8049-48ad-97e1-c1763fbb71e6
from ..imports import *

# %% ../../src/utils/math.ipynb #fbf98023-d46f-46da-8342-c73e8b08d592
def matrix_power(x: torch.Tensor, p: float) -> torch.Tensor:
    """
    Power of a matrix using Eigenspace Decomposition. Assuming decomposition of `x` exists.
    """
    
    vals, vecs = torch.linalg.eig(x)
    vals_pow   = torch.pow(vals, p)
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs)))
    
    return matrix_pow

# %% ../../src/utils/math.ipynb #35135abe-a5d2-476d-92fc-e9e1bbe300de
def gram_schmidt(X: torch.Tensor):
    """
    Perform Gram–Schmidt orthonormalization on the vectors given by the rows of matrix X.
    """
    assert X.dim() == 2, "Only 2-dim tensor supported."
    
    X_type = X.dtype
    X      = X.to(torch.float64)    
    Q = []
    for q in X:
        # Take the current row vector 
        # Subtract projec+tions onto existing basis vectors
        for v in Q:
            q = q - torch.dot(q, v) * v
        # Normalize the vector
        q = q / torch.norm(q)
        Q.append(q)
    return torch.stack(Q).to(X_type)
