##### https://github.com/rballester/tntorch
Tip revision: 3af563a
derivatives.py
import tntorch as tn
import torch

"""
Given a tensor, compute another one that contains all partial derivatives of certain order(s) and according to some optional mask.

This function does not use padding.

:Examples:

>>> t = tn.rand([10, 10, 10])  # A 3D tensor
>>> x, y, z = tn.symbols(3)
>>> partialset(t, 1, x)  # x
>>> partialset(t, 2, x)  # xx, xy, xz
>>> partialset(t, 2, tn.only(y | z))  # yy, yz, zz

:param t: a :class:Tensor
:param order: an int or list of ints. Default is 1
:param mask: an optional mask to select only a subset of partials
:param bounds: a list of pairs [lower bound, upper bound] specifying parameter ranges (used to compute derivative steps). If None (default), all steps will be 1

:return: a :class:Tensor
"""

if t.batch:
raise ValueError('Batched tensors are not supproted.')

if bounds is None:
bounds = [[0, sh - 1] for sh in t.shape]
if not hasattr(order, '__len__'):
order = [order]

max_order = max(order)
def diff(core, n):
if core.shape == 1:
raise ValueError('Tensor size {} along dimension {} not enough to compute high-order derivative'.format(t.shape[n], n))
step = (bounds[n] - bounds[n]) / (core.shape[-2] - 1)
return (core[..., 1:, :] - core[..., :-1, :]) / step
cores = []
idxs = []
for n in range(t.dim()):
if t.Us[n] is None:
stack = [t.cores[n]]
else:
stack = [torch.einsum('ijk,aj->iak', (t.cores[n], t.Us[n]))]
idx = torch.zeros([t.shape[n]])
for o in range(1, max_order + 1):
stack.append(diff(stack[-1], n))
idx = torch.cat((idx, torch.ones(stack[-1].shape[-2]) * o))
if o == max_order:
break
cores.append(torch.cat(stack, dim=-2))
idxs.append(idx)
d = tn.Tensor(cores, idxs=idxs)
result.idxs = idxs
return result

def partial(t, dim, order=1, bounds=None, periodic=False):
"""
Compute a single partial derivative.

:param t: a :class:Tensor
:param dim: int or list of ints
:param order: how many times to derive. Default is 1
:param bounds: variable(s) range bounds (to compute the derivative step). If None (default), step 1 will be assumed
:param periodic: int or list of ints (same as dim), mark dimensions with periodicity

:return: a :class:Tensor
"""

if not hasattr(dim, '__len__'):
dim = [dim]
if bounds is None:
bounds = [[0, t.shape[n]] for n in range(t.dim())]
if not hasattr(bounds, '__len__'):
bounds = [bounds]
if not hasattr(periodic, '__len__'):
periodic = [periodic]*len(dim)

t2 = t.clone()
for i, d in enumerate(dim):
step = (bounds[i] - bounds[i]) / (t.shape[d]+1) * 2
for o in range(1, order+1):
if periodic[i]:
if t2.Us[d] is None:
t2.cores[d] = (t2.cores[d][:, list(range(1, t2.cores[d].shape))+, :] -
t2.cores[d][:, [-1]+list(range(0, t2.cores[d].shape-1)), :]) / step
else:
t2.Us[d] = (t2.Us[d][list(range(1, t2.Us[d].shape)) + , :] -
t2.Us[d][[-1] + list(range(0, t2.Us[d].shape - 1)), :]) / step
else:
if t2.Us[d] is None:
t2.cores[d] = t2.cores[d][:, +list(range(t2.shape[d]))+[t2.shape[d]-1], :]
t2.cores[d][:, 0, :] -= t2.cores[d][:, 2, :] - t2.cores[d][:, 1, :]
t2.cores[d][:, -1, :] += t2.cores[d][:, -2, :] - t2.cores[d][:, -3, :]
t2.cores[d] = (t2.cores[d][:, 2:, :] - t2.cores[d][:, :-2, :])/step
else:
t2.Us[d] = t2.Us[d][ + list(range(t2.shape[d])) + [t2.shape[d] - 1], :]
t2.Us[d][0, :] -= t2.Us[d][2, :] - t2.Us[d][1, :]
t2.Us[d][-1, :] += t2.Us[d][-2, :] - t2.Us[d][-3, :]
t2.Us[d] = (t2.Us[d][2:, :] - t2.Us[d][:-2, :]) / step
return t2

"""
Compute the gradient of a tensor.

:param t: a :class:Tensor
:param dim: an integer (or list of integers). Default is all
:param bounds: a pair (or list of pairs) of reals, or None. The bounds for each variable

:return: a :class:Tensor (or a list thereof)
"""

if t.batch:
raise ValueError('Batched tensors are not supproted.')

if dim == 'all':
dim = range(t.dim())
if bounds is None:
bounds = [[0, t.shape[d]] for d in dim]
if not hasattr(bounds, '__len__'):
bounds = [bounds]*len(dim)

if not hasattr(dim, '__len__'):
return tn.partial(t, dim, bounds)
else:
return [tn.partial(t, d, order=1, bounds=b) for d, b in zip(dim, bounds)]

def active_subspace(t, bounds, marginals=None):
"""
Compute the main variational directions of a tensor.

Reference: P. Constantine et al. "Discovering an Active Subspace in a Single-Diode Solar Cell Model" (2017) <https://arxiv.org/pdf/1406.7607.pdf>_

See also P. Constantine's data set repository <https://github.com/paulcon/as-data-sets/blob/master/README.md>_.

:param t: input tensor
:param bounds: a pair (or list of pairs) of reals, or None. The bounds for each variable
:param marginals: a list of vectors. If None (default), uniform marginals will be used
:return: (eigvals, eigvecs): an array and a matrix, encoding the eigenpairs in descending order
"""

if t.batch:
raise ValueError('Batched tensors are not supproted.')

if marginals is None:
marginals = [torch.ones(sh)/sh for sh in t.shape]
assert all([len(marginals[n]) == t.shape[n] for n in range(t.dim())])
cores = []
for n in range(t.dim()):
marg = (marginals[n][:-1] + marginals[n][1:]) / 2
marg /= marg.sum()
marg = torch.cat([marg, torch.zeros(1)])
cores.append(marg[None, :, None])
pdf = tn.Tensor(cores)

M = torch.zeros(t.dim(), t.dim())
for i in range(t.dim()):
for j in range(i, t.dim()):
M[j, i] = M[i, j]

w, v = torch.linalg.eigh(M)
idx = range(t.dim() - 1, -1, -1)
w = w[idx]
v = v[:, idx]
return w, v

def dgsm(t, bounds, marginals):
"""
Compute the derivative-based global sensitivity measure \nu from , defined for each i-th variable as:

$\nu_i := \int_{\Omega} \left(\frac{\partial f}{\partial x_i}\right) \, d\pmb{x}$

 "Derivative-Based Global Sensitivity Measures", by Kucherenko and Iooss (2016)

:param t: input tensor
:param bounds: a pair (or list of pairs) of reals, or None. The bounds for each variable
:param marginals: a list of vectors. If None (default), uniform marginals will be used
:return: a vector of size N
"""

if marginals is None:
marginals = [torch.ones(sh)/sh for sh in t.shape]
assert all([len(marginals[n]) == t.shape[n] for n in range(t.dim())])
cores = []
for n in range(t.dim()):
# marg = (marginals[n][:-1] + marginals[n][1:]) / 2
marg = marginals[n]
marg /= marg.sum()
# marg = torch.cat([marg, torch.zeros(1)])
cores.append(marg[None, :, None])
pdf = tn.Tensor(cores)

result = torch.zeros(t.dim())
for n in range(t.dim()):
return result

def divergence(ts, bounds=None):
"""
Computes the divergence (scalar field) out of a vector field encoded in a tensor.

:param ts: an ND vector field, encoded as a list of N ND tensors
:param bounds:

:return: a scalar field
"""

assert ts.dim() == len(ts)
assert all([t.shape == ts.shape for t in ts[1:]])
if bounds is None:
bounds = [None]*len(ts)
elif not hasattr(bounds, '__len__'):
bounds = [bounds for n in range(len(ts))]
assert len(bounds) == len(ts)

return sum([tn.partial(ts[n], n, order=1, bounds=bounds[n]) for n in range(len(ts))])

def curl(ts, bounds=None):
"""
Compute the curl of a 3D vector field.

:param ts: three 3D tensors encoding the :math:x, y, z vector coordinates respectively
:param bounds:

:return: three tensors of the same shape
"""

assert [t.dim() == 3 for t in ts]
assert len(ts) == 3
if bounds is None:
bounds = [None for n in range(3)]
elif not hasattr(bounds, '__len__'):
bounds = [bounds for n in range(3)]
assert len(bounds) == 3

return [tn.partial(ts, 1, bounds=bounds) - tn.partial(ts, 2, bounds=bounds),
tn.partial(ts, 2, bounds=bounds) - tn.partial(ts, 0, bounds=bounds),
tn.partial(ts, 0, bounds=bounds) - tn.partial(ts, 1, bounds=bounds)]

def laplacian(t, bounds=None):
"""
Computes the Laplacian of a scalar field.

:param t: a :class:Tensor
:param bounds:

:return: a :class:Tensor
"""

if bounds is None:
bounds = [None]*t.dim()
elif not hasattr(bounds, '__len__'):
bounds = [bounds for n in range(t.dim())]
assert len(bounds) == t.dim()

return sum([tn.partial(t, n, order=2, bounds=bounds[n]) for n in range(t.dim())])