nitorch_fastmath.lie

meanm

meanm(mats, max_iter=1024, tol=1e-20)

Compute the exponential barycentre of a set of matrices.

Parameters:

Name	Type	Description	Default
`mats`	`(N, M, M) tensor`	Set of square invertible matrices	required
`max_iter`	`int`	Maximum number of iterations	`1024`
`tol`	`float`	Tolerance for early stopping. The tolerance criterion is the sum-of-squares of the residuals in log-space, i.e., \(\lVert\frac{1}{N}\sum_n \log_M(\mathbf{A}_n)\rVert_\mathrm{F}^2\)	`1e-20`

Returns:

Name	Type	Description
`mean_mat`	`(M, M) tensor`	Mean matrix.

References

Pennec, X. and Arsigny, V., 2012. Exponential barycenters of the canonical Cartan connection and invariant means on Lie groups. In Matrix information geometry (pp. 123-166). Berlin, Heidelberg: Springer Berlin Heidelberg.

@incollection{pennec2012,
  title={Exponential barycenters of the canonical Cartan connection and invariant means on Lie groups},
  author={Pennec, Xavier and Arsigny, Vincent},
  booktitle={Matrix information geometry},
  pages={123--166},
  year={2012},
  publisher={Springer},
  url={https://hal.inria.fr/hal-00699361}
}

expm

expm(X, basis=None, max_order=10000, tol=1e-32)

Matrix exponential.

Notes

This function evaluates the matrix exponential using a Taylor approximation. A faster integration technique, based e.g. on scaling and squaring, could have been used instead.
PyTorch/NumPy broadcasting rules apply.
This function is automatically differentiable.

Parameters:

Name	Type	Description	Default
`X`		If `basis` is None: log-matrix. Else: parameters of the log-matrix in the basis set.	`(...`
`basis`	`(..., F, D, D) tensor, optional.`	Basis set.	`None`
`max_order`	`int`	Order of the Taylor expansion	`10000`
`tol`	`float`	Tolerance for early stopping The criterion is based on the Frobenius norm of the last term of the Taylor series.	`1e-32`

Returns:

Name	Type	Description
`eX`	`(..., D, D) tensor`	Matrix exponential

expm_derivatives

expm_derivatives(X, basis=None, grad_X=False, grad_basis=False, hess_X=False, max_order=10000, tol=1e-32)

Matrix exponential (and its derivatives).

Notes

This function evaluates the matrix exponential and its derivatives using a Taylor approximation. A faster integration technique, based e.g. on scaling and squaring, could have been used instead.
PyTorch/NumPy broadcasting rules apply.

The output shapes of dX and dB can be different from the shapes of X and basis, because of broadcasting. When computing the backward pass, you should take their dot product with the output gradient and then reduce across axes that have been expanded by broadcasting. E.g.:

def backward(grad, X, B):
    # X.shape     == (*batch_X, F)
    # B.shape     == (*batch_B, F, D, D)
    # grad.shape  == (*batch_XB, D, D)
    _, dX, dB = _dexpm(X, B, grad_X=True, grad_basis=True)
    # dX.shape    == (*batch_XB, F, D, D)
    # dB.shape    == (*batch_XB, F, D, D, D, D)
    dX = torch.sum(dX * grad[..., None, :, :], dim=[-1, -2])
    dX = broadcast_backward(dX, X.shape)
    dB = torch.sum(dB * grad[..., None, None, None, :, :], dim=[-1, -2])
    dB = broadcast_backward(dB, B.shape)
    return dX, dB

Parameters:

Name	Type	Description	Default
`X`		If `basis` is None: log-matrix. Else: parameters of the log-matrix in the basis set.	`(...`
`basis`	`(..., F, D, D) tensor`	Basis set. If None, basis of all DxD matrices and F = D**2.	`None`
`grad_X`	`bool`	Compute derivatives with respect to `X`.	`False`
`grad_basis`	`bool`	Compute derivatives with respect to `basis`.	`False`
`max_order`	`int`	Order of the Taylor expansion	`10000`
`tol`	`float`	Tolerance for early stopping The criterion is based on the Frobenius norm of the last term of the Taylor series.	`1e-32`

Returns:

Name	Type	Description
`eX`	`(..., D, D) tensor`	Matrix exponential
`dX`	(..., F, D, D) tensor, if `grad_X is True`	Derivative of the matrix exponential with respect to the parameters in the basis set
`dB`	(..., F, D, D, D, D) tensor, if `grad_basis is True`	Derivative of the matrix exponential with respect to the basis.
`hX`	(..., F, F, D, D) tensor, if `hess_X is True`	Second derivative of the matrix exponential with respect to the parameters in the basis set

logm

logm(mat)

Batched matrix logarithm.

This implementation actually use scipy, so the data will be transferred to cpu and transferred back to device afterwards.

Parameters:

Name	Type	Description	Default
`mat`	`(..., N, N) tensor`	Input matrix or batch of matrices	required

Returns:

Name	Type	Description
`logmat`	`(..., N, N) tensor`	Input log-matrix or batch of log-matrices