interpolation.py
``````import numpy as np
import tntorch as tn
import torch
import time

def als_completion(X, y, ranks_tt, shape=None, ws=None, x0=None, niter=10, verbose=True):
"""
Complete an N-dimensional TT from P samples using alternating least squares (ALS).
We assume only low-rank structure, and no smoothness/spatial locality. Such assumption requires that there is at
least one sample for each tensor hyperslice. This is usually the case for categorical variables, since it is
meaningless to have a class or label for which no instances exist.

Note that this method may not converge (or be extremely slow to do so) if the number of available samples is below
or near a certain proportion of the overall tensor. Such proportion, unfortunately, depends on the data set and its
true rank structure
("Riemannian optimization for high-dimensional tensor completion", M. Steinlechner, 2015)

:param X: a P X N matrix of integers (tensor indices)
:param y: a vector with P elements
:param ranks_tt: an integer (or list). Ignored if x0 is given
:param shape: list of N integers. If None, the smallest shape that accommodates `X` will be chosen
:param ws: a vector with P elements, with the weight of each sample (if None, 1 is assumed)
:param x0: initial solution (a TT tensor). If None, a random tensor will be used
:param niter: number of ALS sweeps. Default is 10
:param verbose:
:return: a `tntorch.Tensor'
"""

assert not X.dtype.is_floating_point
assert X.dim() == 2
assert y.dim() == 1
if ws is None:
ws = torch.ones(len(y))
X = X.long()
if shape is None:
shape = [val.item() for val in torch.max(X, dim=0)[0] + 1]
y = y.to(torch.get_default_dtype())
P = X.shape[0]
N = X.shape[1]
if x0 is None:
x0 = tn.rand(shape, ranks_tt=ranks_tt)
# All tensor slices must contain at least one sample point
for dim in range(N):
if torch.unique(X[:, dim]).numel() != x0.shape[dim]:
raise ValueError('One groundtruth sample is needed for every tensor slice')

if verbose:
print('Completing a {}D tensor of size {} using {} samples...'.format(N, list(shape), P))

normy = torch.norm(y)
x0.orthogonalize(0)
cores = x0.cores

# Memoized product chains for all groundtruth points
# lefts will be initialized on the go
lefts = [torch.ones(1, P, x0.cores[n].shape[0]) for n in range(N)]
# rights, however, needs to be initialized now
rights = [None] * N
rights[-1] = torch.ones(1, P, 1)
for dim in range(N - 2, -1, -1):
rights[dim] = torch.einsum('ijk,kjl->ijl', (cores[dim + 1][:, X[:, dim + 1], :], rights[dim + 1]))

def optimize_core(cores, mu, direction):
sse = 0
for index in range(cores[mu].shape[1]):
idx = torch.where(X[:, mu] == index)[0]
leftside = lefts[mu][0, idx, :]
rightside = rights[mu][:, idx, 0]
lhs = rightside.t()[:, :, None]
rhs = leftside[:, None, :]
A = torch.reshape(lhs*rhs, [len(idx), -1])*ws[idx, None]
b = y[idx] * ws[idx]
sol = torch.linalg.lstsq(A, b).solution[:A.shape[1]]
residuals = torch.norm(A.matmul(sol)[0] - b) ** 2
cores[mu][:, index, :] = torch.reshape(sol, cores[mu][:, index, :].shape)#.t()
sse += residuals
# Update product chains for next core
if direction == 'right':
x0.left_orthogonalize(mu)
lefts[mu+1] = torch.einsum('ijk,kjl->ijl', (lefts[mu], cores[mu][:, X[:, mu], :]))
else:
x0.right_orthogonalize(mu)
rights[mu-1] = torch.einsum('ijk,kjl->ijl', (cores[mu][:, X[:, mu], :], rights[mu]))
return sse

start = time.time()
for swp in range(niter):

# Sweep: left-to-right
for mu in range(N-1):
optimize_core(cores, mu, direction="right")

# Sweep: right-to-left
for mu in range(N-1, 0, -1):
sse = optimize_core(cores, mu, direction="left")
eps = torch.sqrt(sse) / normy

if verbose:
print('iter: {: <{}}'.format(swp, len('{}'.format(niter)) + 1), end='')
print('| eps: {:.3e}'.format(eps), end='')
print(' | time: {:8.4f}'.format(time.time() - start))

return x0

def sparse_tt_svd(X, y, eps, shape=None, rmax=None):
"""
TT-SVD for sparse tensors.

:param X: matrix P X N of sample coordinates (integers)
:param y: P-sized vector of sample values
:param eps: prescribed accuracy (resulting relative error is guaranteed to be not larger than this)
:param shape: input tensor shape. If not specified, a tensor will be chosen such that `X` fits in
:param rmax: optionally, cap all ranks above this value
:param verbose:
:return: a TT
"""

def sparse_covariance(Xs, ys, nrows):
u, v = torch.unique(Xs[:, 1:], dim=0, return_inverse=True)
D = torch.zeros(nrows, len(u))
D[Xs[:, 0], v] = ys
return D.mm(D.t())

def full_times_sparse(F, Xs, ys):
F = torch.Tensor(F)
u, v = torch.unique(Xs[:, 1:], dim=0, return_inverse=True)
idx = np.unique(v, return_index=True)[1]

D = torch.zeros(F.shape[1], len(u))
D[Xs[:, 0], v] = ys
FD = F.mm(D)
new_row = torch.remainder(torch.arange(FD.numel()), FD.shape[0])
newcols = Xs[idx, 1:][:, None, :].repeat(1, FD.shape[0], 1)  # , FD.shape[0], axis=1)
newcols = newcols.reshape(len(idx) * FD.shape[0], -1)

def sparse_truncate_svd(Xs, ys, nrows, delta, rmax):

cov = sparse_covariance(Xs, ys, nrows)

w, v = torch.linalg.eigh(cov)
w[w < 0] = 0
w = torch.sqrt(w)
svd = [v, w]

# Sort eigenvalues and eigenvectors in decreasing importance
idx = np.argsort(svd[1])[torch.arange(len(svd[1])-1, -1, -1)]
svd[0] = svd[0][:, idx]
svd[1] = svd[1][idx]

S = svd[1]**2
where = torch.where(np.cumsum(S[torch.arange(len(S)-1, -1, -1)]) <= delta**2)[0]
if len(where) == 0:
rank = max(1, int(np.min([rmax, len(S)])))
else:
rank = max(1, int(np.min([rmax, len(S) - 1 - where[-1]])))
left = svd[0]
left = left[:, :rank]

Xs, ys = full_times_sparse(left.T, Xs, ys)
return left, Xs, ys

if isinstance(X, np.ndarray):
X = torch.Tensor(X)
if isinstance(y, np.ndarray):
y = torch.Tensor(y)
assert not X.dtype.is_floating_point
assert X.dim() == 2
assert y.dim() == 1
N = X.shape[1]
if shape is None:
shape = [val.item() for val in torch.max(X, dim=0)[0] + 1]
assert N == len(shape)
if rmax is None:
rmax = np.iinfo(np.int32).max

delta = eps / np.sqrt(N - 1) * torch.norm(y).item()

# TT-SVD iteration
cores = []
curshape = shape.copy()
for n in range(1, N):
left, X, y = sparse_truncate_svd(X, y, curshape[0], delta=delta, rmax=rmax)
cores.append(left.reshape(left.shape[0]//shape[n-1], shape[n-1], left.shape[1]))
curshape[0] = left.shape[1]  # Rank of this unfolding

if n < N-1:  # Merge the two first indices (sparse reshape)
X = torch.cat([X[:, 0:1] * curshape[1] + X[:, 1:2], X[:, 2:]], dim=1)
curshape[1] *= curshape[0]
curshape = curshape[1:]

lastcore = torch.zeros(list(curshape))
lastcore[list(X.t())] = y
cores.append(lastcore[:, :, None])

return tn.Tensor(cores)

def get_bounding_box(X):
"""
Compute the bounding box of a set of points.

:param X: a tensor of shape ... X N, where N is the number of features
:return: a list of N pairs [bottom, top]
"""

return [(torch.min(X[..., n]).item(), torch.max(X[..., n]).item()) for n in range(X.shape[-1])]

def features2indices(X, bbox=None, I=512, domain=None):
"""
Convert floating point features into dicrete tensor indices.

:param X: a float tensor of shape ... X N, where N is the number of features
:param bbox: a list of N pairs [bottom, top]. If None (default), X's own bounding box will be used.
Values falling outside the box will be clamped.
:param I: grid resolution. Default is 512
:param domain: optional list of N vectors to specify the grid. Overrides `bbox` and `I`
:return: an integer tensor of shape ... X N
"""

assert X.dtype.is_floating_point
X = X.clone()
if domain is not None:
for n in range(X.shape[1]):
X[:, n] = torch.Tensor(np.interp(X[:, n].numpy(), domain[n].numpy(), np.arange(len(domain[n]))))
if bbox is None:
bbox = tn.get_bounding_box(X)
assert len(bbox) == X.shape[-1]
bbox = torch.Tensor(bbox)
X -= bbox[:, 0][[None] * (X.dim() - 1) + [slice(None)]]
X /= (bbox[:, 1] - bbox[:, 0])[[None] * (X.dim() - 1) + [slice(None)]]
X = torch.round(X * (I - 1)).long()
X[X < 0] = 0
X[X > I - 1] = I - 1
return X

def indices2features(X, bbox=None, I=512, domain=None):

assert not X.dtype.is_floating_point
assert X.dim() == 2

result = torch.zeros_like(X).to(torch.get_default_dtype())
if domain is None:
domain = [torch.linspace(
b[0] + (b[1]-b[0])/(2*I),
b[1] - (b[1]-b[0])/(2*I),
I) for b in bbox]
for n in range(X.shape[1]):
result[:, n] = domain[n][X[:, n]]
return result

def empirical_marginals(X, domain):
"""
Given a matrix of sample points, get its discrete marginal distribution over a
specified domain.

:param X: a P x N matrix
:param domain: a list of N vectors
:return: a list of N vectors
"""

assert X.dim() == 2
assert X.shape[1] == len(domain)
P = X.shape[0]
N = X.shape[1]

X_discrete = tn.discretize(X, domain=domain)
result = [torch.zeros(len(domain[n])) for n in range(N)]
for n in range(N):
unique, counts = torch.unique(X_discrete[:, n], return_counts=True)
result[n][unique.numpy()] = counts.double() / P
return result

def gram_schmidt(x, S):
"""
Create a truncated polynomial basis with S elements that is orthogonal
with respect to a given measure.

The elements are computed using a modified [1] version of Gram-Schmidt's orthonormalization
process, and lead to a PCE family that can be used for any probability distribution of
the inputs [2].

[1] https://en.wikipedia.org/wiki/Gram%E2%80%93Schmidt_process#Numerical_stability
[2] "Modeling Arbitrary Uncertainties Using Gram-Schmidt Polynomial Chaos", Witteveen and Bijl, 2012

:param x: list of observed inputs
:param S: an integer: how many basis elements to form.
:return: a matrix of shape S X S (one column per basis element)
"""

assert x.dim() == 1

xpowers = x[:, None] ** torch.arange(S, device=x.device)[None, :]

def proj(u, v):
xu = xpowers.matmul(u)
xv = xpowers.matmul(v)

def norm(u):
xu = xpowers.matmul(u)

Psi = torch.eye(S, S, device=x.device)
for s in range(1, S):
u = Psi[:, s]
for k in range(s):
u = u - proj(Psi[:, k], u)
Psi[:, s] = u / norm(u)

return Psi

class PCEInterpolator:
"""
Polynomial chaos expansion (PCE) interpolator. This class requires scikit-learn to run, and
therefore is CPU-only. Parameter distributions are fully empirical. Orthogonal polynomial families
are learned using Gram-Schmidt [1].

Coefficient selection follows [2]. A sparse regularization is enforced: only entries X whose coordinates satisfy
||X||_q < p are considered, where the Lq norm is used and `p` is a threshold. Additionally, among that space of
candidates, at most `nnz` can be non-zero (they are selected using Least Angle Regression, LARS).

We assume parameters are independently distributed; however, the regressor will often work well
even if this doesn't hold [2].

[1] "Modeling Arbitrary Uncertainties Using Gram-Schmidt Polynomial Chaos", Witteveen and Bijl, 2012
[2] "Data-driven Polynomial Chaos Expansion for Machine Learning Regression", Torre et al. 2020
"""

def __init__(self):
pass

def _design_matrix(self, x):
N = len(self.Psis)
S = self.Psis[0].shape[0]
M = torch.cat([(x[:, n:n+1]**torch.arange(S)[None, :]).matmul(self.Psis[n])[:, None, :] for n in range(N)], dim=1)
idx = torch.arange(N)[None, :].repeat(len(self.coords), 1)
M = M[:, idx.flatten(), self.coords.flatten()]
M = M.reshape(-1, self.coords.shape[0], self.coords.shape[1])
M = torch.prod(M, dim=2)
return M

def fit(self, X, y, p=5, q=0.75, val_split=0.1, seed=0, matrix_size_limit=5e7, retrain=True, verbose=True):
"""
Fit the model to a training dataset (X, y) using LARS. The optimal number of non-zero
coefficients to select is automatically found via a validation set.

:param X: a matrix of shape P X N floats (input features)
:param y: a vector of size P
:param p: threshold for the hyperbolic truncation norm. Default is 5
:param q: the truncation will use norm Lq. Default is 0.75
the optimal will be selected using a validation set
:param val_split: the fraction of elements to use for validation (default is 0.1)
:param seed: integer, the seed used to select the train/validation split. Default is 0
:param matrix_size_limit: abort if the design matrix exceeds this number of elements. Default is 5e7
:param retrain: Boolean. If True (default), retrain on the entire input data after selecting the best nnz
This improves interpolation quality, but setting retrain=False is faster and good for hyperparameter tuning.
:param verbose: Boolean; default is True
"""

import sklearn.linear_model
assert X.dim() == 2
assert X.dtype.is_floating_point
P = X.shape[0]
N = X.shape[1]
assert y.shape[0] == P
assert y.dim() == 1
assert 0 <= q <= 1

# Save features' bounding box
self.bbox = tn.get_bounding_box(X)

# Center each feature to improve stability of polynomial orthogonalization
self.X_mean = torch.mean(X, dim=0)
self.X_std = torch.std(X, dim=0)
X = (X.clone() - self.X_mean[None, :])/self.X_std[None, :]

# Split into train and validation sets
n_val = int(P*val_split)
rng = np.random.default_rng(seed=seed)
idx_val = rng.choice(P, n_val)
idx_train = np.delete(np.arange(P), idx_val)
y_train = y[idx_train]
y_val = y[idx_val]

if verbose:
start = time.time()
print('PCE interpolation (p={}, q={}) of {} points ({} train + {} val) in {}D'.format(p, q, P, P-n_val, n_val, N))

if verbose:
print('{:.3f}s | '.format(time.time() - start), end='')
print('Hyperbolic truncation...', end='')

# Find the coordinates of all coefficient candidates (i.e. all that satisfy ||X||_q < p)
idx = np.zeros(N, dtype=np.int)

def find_candidates(p, q):
# Traverse the whole hypercube of possible coefficients (size S**N)
# Since hyperbolic truncation selects a contiguous region, this is efficient
S = int(np.ceil(p))
coords = []
while True:
pos = N-1
while pos >= 0 and (max(idx) >= S or np.sum(idx ** q) >= p ** q):
idx[pos] = 0
idx[pos - 1] += 1
pos -= 1
if pos < 0:  # Traversed the entire hypercube
break
coords.append(idx.copy())
idx[-1] += 1
if len(coords)*P > matrix_size_limit:
raise ValueError('Design matrix exceeds matrix_size_limit ({:g} elements). Decrease p or q, or increase matrix_size_limit'.format(matrix_size_limit))

self.coords = find_candidates(p, q)
S = int(np.ceil(p))

if verbose:
ncandidates = len(self.coords)
print(' done, we kept {} / {} candidates'.format(ncandidates, S**N))
print('{:.3f}s | '.format(time.time() - start), end='')
print('Assembling a {} X {} design matrix...'.format(P, len(self.coords)), end='', flush=True)

# Build orthogonal polynomial bases based on the
# empirical marginals of the input features
self.Psis = [gram_schmidt(X[:, n], S) for n in range(N)]

# Assemble the design matrix from the training features
M = self._design_matrix(X)
M_train = M[idx_train, ...]
M_val = M[idx_val, ...]

if verbose:
print(' done')
print('{:.3f}s | '.format(time.time() - start), end='')
print('Finding best nnz in LARS...', end='', flush=True)

# Solve the sparse regression problem using LARS
lars = sklearn.linear_model.Lars(n_nonzero_coefs=M_train.shape[1], fit_intercept=False, fit_path=True)
lars.fit(M_train, y_train)

# Find the validation eps for every choice of nnz (LARS' solution path)
reco_path = torch.einsum('pc,cd->pd', M_val, torch.Tensor(lars.coef_path_))
error_path = torch.sqrt(torch.sum((reco_path - y_val[:, None])**2, dim=0))/torch.norm(y_val)
argmin = torch.argmin(error_path)  # Pick the best nnz
nnz = len(np.where(lars.coef_path_[:, argmin])[0])

if verbose:
print(' done, val eps={:.5g}'.format(error_path[argmin]))
print('{:.3f}s | '.format(time.time() - start), end='')

if retrain:

if verbose:
print('Retraining at nnz={}...'.format(nnz), end='', flush=True)

# Retrain now on the entire input dataset and optimal nnz
lars = sklearn.linear_model.Lars(n_nonzero_coefs=nnz, fit_intercept=False, fit_path=False)
lars.fit(M, y)

# Remove features not selected by LARS
lars.coef_ = lars.coef_[0, :]
nonzeros = np.where(lars.coef_)[0]
self.allcoords = self.coords
self.allcoef = torch.Tensor(lars.coef_)
self.coef = torch.Tensor(lars.coef_[nonzeros])
self.coords = self.coords[nonzeros, :]

if verbose:
reco = M[:, nonzeros].matmul(self.coef)
print(' done, training eps={:.5g}'.format(torch.norm(y-reco)/torch.norm(y)))
print('{:.3f}s'.format(time.time() - start), flush=True)
print()

else:
nonzeros = np.where(lars.coef_path_[:, argmin])[0]
self.coef = torch.Tensor(lars.coef_path_[nonzeros, argmin])
self.coords = self.coords[nonzeros, :]
print()

def predict(self, X):
"""
Predict regression values for an input.

:param X: a matrix of shape P X N floats (input features)
:return: a vector of size P
"""
return self._design_matrix((X - self.X_mean[None, :])/self.X_std[None, :]).matmul(self.coef)

def to_tensor(self, domain=512, rmax=200, eps=1e-3, verbose=True):
"""
Convert the interpolator to a TT-Tucker tensor.

:param domain: one of the following:
- list of N vectors to specify the tensor grid
- integer I: the training data set's bounding box will be used, at resolution I. Default is 512
:param rmax: the TT cores will be capped at this rank. Default is 500
:param eps: rounding error to cast PCE into TT. Default is 1e-3
:param verbose: Boolean; default is True
:return: a `tntorch.Tensor` in the TT-Tucker format
"""

N = len(self.Psis)
S = self.Psis[0].shape[0]
if not isinstance(domain, (list, tuple)):
domain = [torch.linspace(
self.bbox[n][0]+(self.bbox[n][1]-self.bbox[n][0])/(2*domain),
self.bbox[n][1]-(self.bbox[n][1]-self.bbox[n][0])/(2*domain),
domain) for n in range(N)]
assert len(domain) == N
domain_centered = [(domain[n] - self.X_mean[n])/self.X_std[n] for n in range(N)]

if verbose:
start = time.time()
print('Conversion to TT-Tucker format (rmax={}, eps={:.5g})'.format(rmax, eps))
print('{:.3f}s | '.format(time.time() - start), end='')
print('Sparse TT-SVD...', end='', flush=True)

# Assemble a TT-Tucker tensor:
# The core is retrieved from the set of PCE coefficients found,
# and is formed in the TT format using sparse TT-SVD
t = tn.sparse_tt_svd(self.coords, self.coef, rmax=rmax, eps=eps)

if verbose:
eps = torch.norm(t[self.coords].torch()-self.coef)/torch.norm(self.coef)
print(' done, rmax={}, eps={:.5g}'.format(max(t.ranks_tt), eps.item()))

# The factors are the polynomial bases evaluated on the grid vectors
Us = []
for n in range(N):
Us.append(
(domain_centered[n][:, None]**torch.arange(S)).to(torch.get_default_dtype()).matmul(
self.Psis[n][:, :t.shape[n]])
)
t.Us = Us

if verbose:
print('{:.3f}s'.format(time.time() - start), flush=True)
print()

return t
``````