score.py

import numpy as np
import tensorly as tl
import torch
import logging


EPS = 1e-12

def tensor_train(input_tensor, rank, svd="truncated_svd", verbose=False):
    """TT decomposition via recursive SVD

        Decomposes `input_tensor` into a sequence of order-3 tensors (factors)
        -- also known as Tensor-Train decomposition [1]_.

    Parameters
    ----------
    input_tensor : tensorly.tensor
    rank : {int, int list}
            maximum allowable TT rank of the factors
            if int, then this is the same for all the factors
            if int list, then rank[k] is the rank of the kth factor
    svd : str, default is 'truncated_svd'
        function to use to compute the SVD, acceptable values in tensorly.SVD_FUNS
    verbose : boolean, optional
            level of verbosity

    Returns
    -------
    factors : TT factors
              order-3 tensors of the TT decomposition

    References
    ----------
    .. [1] Ivan V. Oseledets. "Tensor-train decomposition", SIAM J. Scientific Computing, 33(5):2295–2317, 2011.
    """
    rank = tl.validate_tt_rank(tl.shape(input_tensor), rank=rank)
    tensor_size = input_tensor.shape
    n_dim = len(tensor_size)

    unfolding = input_tensor
    factors = [None] * n_dim
    s_list = []

    # Getting the TT factors up to n_dim - 1
    for k in range(n_dim - 1):
        # Reshape the unfolding matrix of the remaining factors
        n_row = int(rank[k] * tensor_size[k])
        unfolding = tl.reshape(unfolding, (n_row, -1))

        # SVD of unfolding matrix
        (n_row, n_column) = unfolding.shape
        current_rank = min(n_row, n_column, rank[k + 1])
        U, S, V = tl.svd_interface(unfolding, n_eigenvecs=current_rank, method=svd)

        rank[k + 1] = current_rank

        # Get kth TT factor
        factors[k] = tl.reshape(U, (rank[k], tensor_size[k], rank[k + 1]))

        if verbose is True:
            print(
                "TT factor " + str(k) + " computed with shape " + str(factors[k].shape)
            )

        # Get new unfolding matrix for the remaining factors
        unfolding = tl.reshape(S, (-1, 1)) * V
        s_list.append(S)

    # Getting the last factor
    (prev_rank, last_dim) = unfolding.shape
    factors[-1] = tl.reshape(unfolding, (prev_rank, last_dim, 1))

    if verbose is True:
        print(
            "TT factor "
            + str(n_dim - 1)
            + " computed with shape "
            + str(factors[n_dim - 1].shape)
        )

    return factors, s_list

def tt_reconstruction_loss(X, X_hat, return_recon=False):
    """
    X: original tensor (tensorly tensor or numpy)
    factors: TT cores from tensor_train
    """
    X = tl.tensor(X)

    diff = X_hat - X
    # Frobenius norms
    diff_norm = tl.norm(diff, 2)      # ||X_hat - X||_F
    x_norm = tl.norm(X, 2) + 1e-12    # ||X||_F

    rel_err = (diff_norm / x_norm).item() if hasattr(diff_norm, "item") else float(diff_norm / x_norm)

    # MSE
    # number of elements
    numel = np.prod(tl.shape(X))
    mse = (diff_norm**2 / numel).item() if hasattr(diff_norm, "item") else float(diff_norm**2 / numel)

    if return_recon:
        return rel_err, mse, X_hat
    return rel_err, mse

def to_matrix(factors):
    full_shape = [f.shape[1] for f in factors]
    full_tensor = tl.reshape(factors[0], (full_shape[0], -1))

    for factor in factors[1:]:
        rank_prev, _, rank_next = factor.shape
        factor = tl.reshape(factor, (rank_prev, -1))
        full_tensor = tl.dot(full_tensor, factor)
        full_tensor = tl.reshape(full_tensor, (-1, rank_next))
    return full_tensor

def _choose_rank_from_singular_values(S, delta2, max_rank=None):
    """Pick smallest r such that tail energy <= delta2."""
    # S is (r,) sorted descending by SVD
    s2 = tl.to_numpy(S) ** 2
    total_r = s2.shape[0]
    if total_r == 0:
        return 0

    # tail_energy[r] = sum_{i>=r} s2[i]
    tail = np.cumsum(s2[::-1])[::-1]

    # Find smallest r such that tail[r] <= delta2
    # Note: if r == total_r, tail is "empty" (=0), always satisfies.
    r = total_r
    for rr in range(total_r + 1):
        tail_energy = 0.0 if rr == total_r else tail[rr]
        if tail_energy <= delta2:
            r = rr
            break

    # rank must be at least 1 in TT-SVD (unless you want degenerate cores)
    r = max(1, r)

    if max_rank is not None:
        r = min(r, int(max_rank))
    return r

def tensor_train_auto_rank(input_tensor, eps=0.1, svd="truncated_svd", max_rank=None, verbose=False):
    """TT-SVD with automatic rank selection for relative Frobenius error <= eps."""
    tensor_size = tl.shape(input_tensor)
    n_dim = len(tensor_size)

    if n_dim < 2:
        raise ValueError("TT decomposition requires tensor order >= 2")

    # Frobenius norm squared of the original tensor
    X_fro2 = float(tl.norm(input_tensor, order=2) ** 2)

    # Per-step squared error budget
    delta2 = (eps ** 2) * X_fro2 / (n_dim - 1)

    factors = [None] * n_dim
    s_list = []
    rank = [1] + [None] * (n_dim - 1) + [1]  # will fill rank[1..n_dim-1]

    unfolding = input_tensor

    for k in range(n_dim - 1):
        n_row = int(rank[k] * tensor_size[k])
        unfolding = tl.reshape(unfolding, (n_row, -1))
        n_row, n_col = unfolding.shape

        # Compute SVD with as many singular values as possible (full or capped)
        # If your backend can't do full SVD, set n_eigenvecs=min(n_row,n_col)
        full_r = min(n_row, n_col)
        U, S, V = tl.svd_interface(unfolding, n_eigenvecs=full_r, method=svd)

        # Pick rank by tail energy criterion
        r_next = _choose_rank_from_singular_values(S, delta2=delta2, max_rank=max_rank)
        rank[k + 1] = r_next

        # Truncate U,S,V
        U = U[:, :r_next]
        S = S[:r_next]
        V = V[:r_next, :]

        factors[k] = tl.reshape(U, (rank[k], tensor_size[k], rank[k + 1]))
        if verbose:
            print(f"[TT] k={k} unfolding={n_row}x{n_col}  r_next={r_next}  delta2={delta2:.3e}")

        unfolding = tl.reshape(S, (-1, 1)) * V
        s_list.append(S)

    prev_rank, last_dim = unfolding.shape
    factors[-1] = tl.reshape(unfolding, (prev_rank, last_dim, 1))
    return factors, s_list, rank

def find_smallest_tt_rank(X, max_rank=256, target_rel_err=0.1, svd="truncated_svd"):
    """
    Binary search to find the smallest TT rank that achieves rel_err < target_rel_err.
    
    Parameters
    ----------
    X : array [N, a, b, c] or [N, d]
        Input tensor (can be reshaped if needed)
    max_rank : int
        Maximum rank to try
    target_rel_err : float
        Target relative error threshold (default 0.1)
    svd : str
        SVD method for TT decomposition
        
    Returns
    -------
    best_rank : int
        Smallest rank achieving target_rel_err
    best_rel_err : float
        Actual relative error achieved
    best_factors : list
        TT factors for the best rank
    """
    X = tl.tensor(X)

    # original_shape = tl.shape(X)
    
    # Attempt TT decomposition with current rank
    # if len(original_shape) == 4:
    #     rank = [1, original_shape[0], max_rank, X.shape[-1], 1]
    # if len(original_shape) == 5:
    #     rank = [1, original_shape[0], max_rank, max_rank, X.shape[-1], 1]
    # factors, s_list = tensor_train(X, rank=rank, svd=svd, verbose=False)
    factors, s_list, rank = tensor_train_auto_rank(X, eps=target_rel_err, svd=svd, verbose=False)
    
    # explained_var = ([s ** 2 for s in s_list[1]])
    # total_var = sum(explained_var)
    # if total_var <= 0:
    #     return 0, explained_var

    # explained_ratio = torch.tensor(explained_var / total_var)
    # cumulative = torch.cumsum(explained_ratio, dim=0)
    # rank = int(torch.searchsorted(cumulative, torch.tensor(1-target_rel_err, device=cumulative.device)).item() + 1)
    return rank