"""Module for matrices in the TT format""" import tensorly as tl # Note how tt_matrix_to_tensor is implemented in tenalg to allow for more efficient implementations # (e.g. using the einsum backend) from .tenalg import _tt_matrix_to_tensor as tt_matrix_to_tensor from ._factorized_tensor import FactorizedTensor import numpy as np def validate_tt_matrix_rank(tensorized_shape, rank='same'): """Returns the rank of a TT-Matrix Decomposition Parameters ---------- tensor_shape : tupe shape of the tensorized matrix to decompose rank : {'same', float, tuple, int}, default is same way to determine the rank, by default 'same' if 'same': rank is computed to keep the number of parameters (at most) the same if float, computes a rank so as to keep rank percent of the original number of parameters if int or tuple, just returns rank constant_rank : bool, default is False * if True, the *same* rank will be chosen for each modes * if False (default), the rank of each mode will be proportional to the corresponding tensor_shape *used only if rank == 'same' or 0 < rank <= 1* rounding = {'round', 'floor', 'ceil'} Returns ------- rank : int tuple rank of the decomposition """ n_dim = len(tensorized_shape) // 2 if n_dim*2 != len(tensorized_shape): msg = (f'The order of the give tensorized shape is not a multiple of 2.' 'However, there should be as many dimensions for the left side (number of rows)' ' as of the right side (number of columns). ' ' For instance, to convert a matrix of size (8, 9) to the TT-format, ' ' it can be tensorized to (2, 4, 3, 3) but NOT to (2, 2, 2, 3, 3).') raise ValueError(msg) left_shape = tensorized_shape[:n_dim] right_shape = tensorized_shape[n_dim:] full_shape = tuple(i*o for i, o in zip(left_shape, right_shape)) return tl.tt_tensor.validate_tt_rank(full_shape, rank) def _tt_matrix_n_param(tensorized_shape, rank): """Number of parameters of a TT-Matrix decomposition for a given `rank` and full `tensor_shape`. Parameters ---------- tensorized_shape : int tuple shape of the full tensorized matrix to decompose (or approximate) rank : tuple rank of the TT-Matrix decomposition Returns ------- n_params : int Number of parameters of a TT-Matrix decomposition of rank `rank` of a full tensor of shape `tensor_shape` """ n_dim = len(tensorized_shape) // 2 if n_dim*2 != len(tensorized_shape): msg = (f'The order of the give tensorized shape is not a multiple of 2.' 'However, there should be as many dimensions for the left side (number of rows)' ' as of the right side (number of columns). ' ' For instance, to convert a matrix of size (8, 9) to the TT-format, ' ' it can be tensorized to (2, 4, 3, 3) but NOT to (2, 2, 2, 3, 3).') raise ValueError(msg) left_shape = tensorized_shape[:n_dim] right_shape = tensorized_shape[n_dim:] factor_params = [] for i, (ls, rs) in enumerate(zip(left_shape, right_shape)): factor_params.append(rank[i]*ls*rs*rank[i+1]) return np.sum(factor_params) def tt_matrix_to_matrix(tt_matrix): """Reconstruct the original matrix that was tensorized and compressed in the TT-Matrix format Re-assembles 'factors', which represent a tensor in TT-Matrix format into the corresponding matrix Parameters ---------- factors: list of 4D-arrays TT-Matrix factors (known as core) of shape (rank_k, left_dim_k, right_dim_k, rank_{k+1}) Returns ------- output_matrix: 2D-array matrix whose TT-Matrix decomposition was given by 'factors' """ in_shape = tuple(c.shape[1] for c in tt_matrix) return tl.reshape(tt_matrix_to_tensor(tt_matrix), (np.prod(in_shape), -1)) def tt_matrix_to_unfolded(tt_matrix, mode): """Returns the unfolding matrix of a tensor given in TT-Matrix format Reassembles a full tensor from 'factors' and returns its unfolding matrix with mode given by 'mode' Parameters ---------- factors : list of 3D-arrays TT-Matrix factors mode : int unfolding matrix to be computed along this mode Returns ------- 2-D array unfolding matrix at mode given by 'mode' """ return tl.unfold(tt_matrix_to_tensor(tt_matrix), mode) def tt_matrix_to_vec(tt_matrix): """Returns the tensor defined by its TT-Matrix format ('factors') into its vectorized format Parameters ---------- factors : list of 3D-arrays TT factors Returns ------- 1-D array format of tensor defined by 'factors' """ return tl.tensor_to_vec(tt_matrix_to_tensor(tt_matrix)) def _validate_tt_matrix(tt_tensor): factors = tt_tensor n_factors = len(factors) if n_factors < 1: raise ValueError('A Tensor-Train (MPS) tensor should be composed of at least one factor.' 'However, {} factor was given.'.format(n_factors)) rank = [] left_shape = [] right_shape = [] for index, factor in enumerate(factors): current_rank, current_left_shape, current_right_shape, next_rank = tl.shape(factor) # Check that factors are third order tensors if not tl.ndim(factor)==4: raise ValueError('A TTMatrix expresses a tensor as fourth order factors (tt-cores).\n' 'However, tl.ndim(factors[{}]) = {}'.format( index, tl.ndim(factor))) # Consecutive factors should have matching ranks if index and tl.shape(factors[index - 1])[-1] != current_rank: raise ValueError('Consecutive factors should have matching ranks\n' ' -- e.g. tl.shape(factors[0])[-1]) == tl.shape(factors[1])[0])\n' 'However, tl.shape(factor[{}])[-1] == {} but' ' tl.shape(factor[{}])[0] == {} '.format( index - 1, tl.shape(factors[index - 1])[-1], index, current_rank)) # Check for boundary conditions if (index == 0) and current_rank != 1: raise ValueError('Boundary conditions dictate factor[0].shape[0] == 1.' 'However, got factor[0].shape[0] = {}.'.format( current_rank)) if (index == n_factors - 1) and next_rank != 1: raise ValueError('Boundary conditions dictate factor[-1].shape[2] == 1.' 'However, got factor[{}].shape[2] = {}.'.format( n_factors, next_rank)) left_shape.append(current_left_shape) right_shape.append(current_right_shape) rank.append(current_rank) # Add last rank (boundary condition) rank.append(next_rank) return tuple(left_shape) + tuple(right_shape), tuple(rank) class TTMatrix(FactorizedTensor): def __init__(self, factors, inplace=False): super().__init__() # Will raise an error if invalid shape, rank = _validate_tt_matrix(factors) self.shape = tuple(shape) self.order = len(self.shape) // 2 self.left_shape = self.shape[:self.order] self.right_shape = self.shape[self.order:] self.rank = tuple(rank) self.factors = factors def __getitem__(self, index): return self.factors[index] def __setitem__(self, index, value): self.factors[index] = value def __iter__(self): for index in range(len(self)): yield self[index] def __len__(self): return len(self.factors) def __repr__(self): message = (f'factors list : rank-{self.rank} TT-Matrix of tensorized shape {self.shape}' f' corresponding to a matrix of size {np.prod(self.left_shape)} x {np.prod(self.right_shape)}') return message def to_tensor(self): return tt_matrix_to_tensor(self) def to_matrix(self): return tt_matrix_to_matrix(self) def to_unfolding(self, mode): return tt_matrix_to_unfolded(self, mode) def to_vec(self): return tt_matrix_to_vec(self)