Revision d38084d993f6b218862f3aa0693aacc2e3b4b1b3 authored by TUNA Caglayan on 29 April 2021, 08:44:25 UTC, committed by TUNA Caglayan on 12 May 2021, 12:26:22 UTC
1 parent e54a8ff
tucker_tensor.py
"""
Core operations on Tucker tensors.
"""

from ._factorized_tensor import FactorizedTensor
from .base import unfold, tensor_to_vec
from .tenalg import multi_mode_dot, mode_dot
from . import backend as tl
import numpy as np
from scipy.optimize import brentq
import warnings

# Author: Jean Kossaifi <jean.kossaifi+tensors@gmail.com>

# License: BSD 3 clause

def _validate_tucker_tensor(tucker_tensor):
core, factors = tucker_tensor

if len(factors) < 2:
raise ValueError('A Tucker tensor should be composed of at least two factors and a core.'
'However, {} factor was given.'.format(len(factors)))

if len(factors) != tl.ndim(core):
raise ValueError('Tucker decompositions should have one factor per more of the core tensor.'
'However, core has {} modes but {} factors have been provided'.format(
tl.ndim(core), len(factors)))

shape = []
rank = []
for i, factor in enumerate(factors):
current_shape, current_rank = tl.shape(factor)
if current_rank != tl.shape(core)[i]:
raise ValueError('Factor n of Tucker decomposition should verify:\n'
'factors[n].shape[1] = core.shape[n].'
'However, factors[{0}].shape[1]={1} but core.shape[{0}]={2}.'.format(
i, tl.shape(factor)[1], tl.shape(core)[i]))
shape.append(current_shape)
rank.append(current_rank)

return tuple(shape), tuple(rank)

def tucker_to_tensor(tucker_tensor, skip_factor=None, transpose_factors=False):
"""Converts the Tucker tensor into a full tensor

Parameters
----------
tucker_tensor : tl.TuckerTensor or (core, factors)
core tensor and list of factor matrices
skip_factor : None or int, optional, default is None
if not None, index of a matrix to skip
Note that in any case, modes, if provided, should have a lengh of tensor.ndim
transpose_factors : bool, optional, default is False
if True, the matrices or vectors in in the list are transposed

Returns
-------
2D-array
full tensor of shape (factors[0].shape[0], ..., factors[-1].shape[0])
"""
core, factors = tucker_tensor
return multi_mode_dot(core, factors, skip=skip_factor, transpose=transpose_factors)

def tucker_to_unfolded(tucker_tensor, mode=0, skip_factor=None, transpose_factors=False):
"""Converts the Tucker decomposition into an unfolded tensor (i.e. a matrix)

Parameters
----------
tucker_tensor : tl.TuckerTensor or (core, factors)
core tensor and list of factor matrices
mode : None or int list, optional, default is None
skip_factor : None or int, optional, default is None
if not None, index of a matrix to skip
Note that in any case, modes, if provided, should have a length of tensor.ndim
transpose_factors : bool, optional, default is False
if True, the matrices or vectors in in the list are transposed

Returns
-------
2D-array
unfolded tensor
"""
return unfold(tucker_to_tensor(tucker_tensor, skip_factor=skip_factor, transpose_factors=transpose_factors), mode)

def tucker_to_vec(tucker_tensor, skip_factor=None, transpose_factors=False):
"""Converts a Tucker decomposition into a vectorised tensor

Parameters
----------
tucker_tensor : tl.TuckerTensor or (core, factors)
core tensor and list of factor matrices
skip_factor : None or int, optional, default is None
if not None, index of a matrix to skip
Note that in any case, modes, if provided, should have a length of tensor.ndim
transpose_factors : bool, optional, default is False
if True, the matrices or vectors in in the list are transposed

Returns
-------
1D-array
vectorised tensor

Notes
-----
Mathematically equivalent but much slower,
you can obtain the same result using:

>>> def tucker_to_vec(core, factors):
...     return kronecker(factors).dot(tensor_to_vec(core))
"""
return tensor_to_vec(tucker_to_tensor(tucker_tensor, skip_factor=skip_factor, transpose_factors=transpose_factors))

def tucker_mode_dot(tucker_tensor, matrix_or_vector, mode, keep_dim=False, copy=False):
"""n-mode product of a Tucker tensor and a matrix or vector at the specified mode

Parameters
----------
tucker_tensor : tl.TuckerTensor or (core, factors)

matrix_or_vector : ndarray
1D or 2D array of shape (J, i_k) or (i_k, )
matrix or vectors to which to n-mode multiply the tensor
mode : int

Returns
-------
TuckerTensor = (core, factors)
mode-mode product of tensor by matrix_or_vector
* of shape :math:(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N) if matrix_or_vector is a matrix
* of shape :math:(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N) if matrix_or_vector is a vector

--------
tucker_multi_mode_dot : chaining several mode_dot in one call
"""
shape, rank = _validate_tucker_tensor(tucker_tensor)
core, factors = tucker_tensor
contract = False

if tl.ndim(matrix_or_vector) == 2:  # Tensor times matrix
# Test for the validity of the operation
if matrix_or_vector.shape[1] != shape[mode]:
raise ValueError(
'shapes {0} and {1} not aligned in mode-{2} multiplication: {3} (mode {2}) != {4} (dim 1 of matrix)'.format(
shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[1]
))

elif tl.ndim(matrix_or_vector) == 1:  # Tensor times vector
if matrix_or_vector.shape[0] != shape[mode]:
raise ValueError(
'shapes {0} and {1} not aligned for mode-{2} multiplication: {3} (mode {2}) != {4} (vector size)'.format(
shape, matrix_or_vector.shape, mode, shape[mode], matrix_or_vector.shape[0]
))
if not keep_dim:
contract = True # Contract over that mode
else:
raise ValueError('Can only take n_mode_product with a vector or a matrix.')

if copy:
factors = [tl.copy(f) for f in factors]
core = tl.copy(core)

if contract:
print('contracting mode')
f = factors.pop(mode)
core = mode_dot(core, tl.dot(matrix_or_vector, f), mode=mode)
else:
factors[mode] = tl.dot(matrix_or_vector, factors[mode])

return TuckerTensor((core, factors))

class TuckerTensor(FactorizedTensor):
def __init__(self, tucker_tensor):
super().__init__()

shape, rank = _validate_tucker_tensor(tucker_tensor)
core, factors = tucker_tensor
self.shape = tuple(shape)
self.rank = tuple(rank)
self.factors = factors
self.core = core

def __getitem__(self, index):
if index == 0:
return self.core
elif index == 1:
return self.factors
else:
raise IndexError('You tried to access index {} of a Tucker tensor.\n'
'You can only access index 0 and 1 of a Tucker tensor'
'(corresponding respectively to core and factors)'.format(index))

def __setitem__(self, index, value):
if index == 0:
self.core = value
elif index == 1:
self.factors = value
else:
raise IndexError('You tried to set index {} of a Tucker tensor.\n'
'You can only set index 0 and 1 of a Tucker tensor'
'(corresponding respectively to core and factors)'.format(index))

def __iter__(self):
yield self.core
yield self.factors

def __len__(self):
return 2

def __repr__(self):
message = 'Decomposed rank-{} TuckerTensor of shape {} '.format(self.rank, self.shape)
return message

def to_tensor(self):
return tucker_to_tensor(self)

def to_unfolded(self, mode):
return tucker_to_unfolded(self, mode)

def to_vec(self):
return tucker_to_vec(self)

def mode_dot(self, matrix_or_vector, mode, keep_dim=False, copy=False):
"""n-mode product with a matrix or vector at the specified mode

Parameters
----------
matrix_or_vector : ndarray
1D or 2D array of shape (J, i_k) or (i_k, )
matrix or vectors to which to n-mode multiply the tensor
mode : int

Returns
-------
TuckerTensor = (core, factors)
mode-mode product of tensor by matrix_or_vector
* of shape :math:(i_1, ..., i_{k-1}, J, i_{k+1}, ..., i_N) if matrix_or_vector is a matrix
* of shape :math:(i_1, ..., i_{k-1}, i_{k+1}, ..., i_N) if matrix_or_vector is a vector

--------
tucker_multi_mode_dot : chaining several mode_dot in one call
"""
return tucker_mode_dot(self, matrix_or_vector, mode, keep_dim=keep_dim, copy=copy)

def _tucker_n_param(tensor_shape, rank):
"""Number of parameters of a Tucker decomposition for a given rank and full tensor_shape.

Parameters
----------
tensor_shape : int tuple
shape of the full tensor to decompose (or approximate)

rank : tuple
rank of the Tucker decomposition

Returns
-------
n_params : int
Number of parameters of a Tucker decomposition of rank rank of a full tensor of shape tensor_shape
"""
core_params = np.prod(rank)
factors_params = np.sum([r*s for (r, s) in zip(rank, tensor_shape)])
return core_params + factors_params

def validate_tucker_rank(tensor_shape, rank='same', rounding='round', fixed_modes=None):
r"""Returns the rank of a Tucker Decomposition

Parameters
----------
tensor_shape : tupe
shape of the tensor 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
rounding = {'round', 'floor', 'ceil'}
fixed_modes : int list or None, default is None
if not None, a list of modes for which the rank will be the same as the original shape
e.g. if i in fixed_modes, then rank[i] = tensor_shape[i]

Returns
-------
rank : int tuple
rank of the decomposition

Notes
-----
For a fractional input rank, I want to find a Tucker rank such that:
n_param_decomposition = rank*n_param_tensor

In particular, for an input of size I_1, ..., I_N:
I find a value c such that the rank will be (c I_1, ..., c I_N)

We have sn_param_tensor = I_1 x ... x I_N

We look for a Tucker decomposition of rank (c I_1, ..., c I_N )
This decomposition will have the following n_params:
For the core : \prod_k c I_k = c^N \prod I_k = c^N n_param_tensor
For the factors : \sum_k c I_k^2

In other words we want to solve:
c^N n_param_tensor + \sum_k c I_k^2 = rank*n_param_tensor
"""
if rounding == 'ceil':
rounding_fun = np.ceil
elif rounding == 'floor':
rounding_fun = np.floor
elif rounding == 'round':
rounding_fun = np.round
else:
raise ValueError(f'Rounding should be round, floor or ceil, but got {rounding}')

# rank is 'same' or float: choose rank so as to preserve a fraction of the original #parameters
if rank == 'same':
rank = float(1)

if isinstance(rank, float):
n_modes_compressed = len(tensor_shape)
n_param_tensor = np.prod(tensor_shape)

if fixed_modes is not None:
tensor_shape = list(tensor_shape)

# sorted to be careful with the order when popping and reinserting to not remove/add at wrong index.
# list (mode, shape) that we removed as they will be kept the same, rank[i] =
fixed_modes = [(mode, tensor_shape.pop(mode)) for mode in sorted(fixed_modes, reverse=True)][::-1]

# number of parameters coming from the fixed modes (these don't have a variable size as a fun of fraction_param)
n_fixed_params = np.sum([s**2 for _, s in fixed_modes]) # size of the factors
n_modes_compressed -= len(fixed_modes)
else:
n_fixed_params = 0

# Doesn't contain fixed_modes, those factors are accounted for in fixed_params
squared_dims = np.sum([s**2 for s in tensor_shape])

fun = lambda x : n_param_tensor*x**n_modes_compressed + squared_dims*x + n_fixed_params - rank*n_param_tensor
fraction_param = brentq(fun, 0.0, max(rank, 1.0))
rank = [max(int(rounding_fun(s*fraction_param)), 1) for s in tensor_shape]

if fixed_modes is not None:
for mode, size in fixed_modes:
rank.insert(mode, size)

elif isinstance(rank, int):
n_modes = len(tensor_shape)
message = "Given only one int for 'rank' for decomposition a tensor of order {}. Using this rank for all modes.".format(n_modes)
warnings.warn(message, RuntimeWarning)
if fixed_modes is None:
rank = [rank]*n_modes
else:
rank = [rank  if i not in fixed_modes else s for (i, s) in enumerate(tensor_shape)]# *n_mode

return rank



Computing file changes ...