def coefficients(ind_k):
"""Christoffel symbols for contravariant index ind_k."""
param_k = base_point[..., ind_k]
param_sum = gs.sum(base_point, -1)
c1 = (
1
/ gs.polygamma(1, param_k)
/ (
1 / gs.polygamma(1, param_sum)
- gs.sum(1 / gs.polygamma(1, base_point), -1)
)
)
c2 = -c1 * gs.polygamma(2, param_sum) / gs.polygamma(1, param_sum)
mat_ones = gs.ones((n_points, self.dim, self.dim))
mat_diag = from_vector_to_diagonal_matrix(
-gs.polygamma(2, base_point) / gs.polygamma(1, base_point)
)
arrays = [
gs.zeros((1, ind_k)),
gs.ones((1, 1)),
gs.zeros((1, self.dim - ind_k - 1)),
]
vec_k = gs.tile(gs.hstack(arrays), (n_points, 1))
val_k = gs.polygamma(2, param_k) / gs.polygamma(1, param_k)
vec_k = gs.einsum("i,ij->ij", val_k, vec_k)
mat_k = from_vector_to_diagonal_matrix(vec_k)
mat = (
gs.einsum("i,ijk->ijk", c2, mat_ones)
- gs.einsum("i,ijk->ijk", c1, mat_diag)
+ mat_k
)
return 1 / 2 * mat
def coefficients(ind_k):
param_k = base_point[..., ind_k]
param_sum = gs.sum(base_point, -1)
c1 = 1 / gs.polygamma(
1, param_k) / (1 / gs.polygamma(1, param_sum) -
gs.sum(1 / gs.polygamma(1, base_point), -1))
c2 = -c1 * gs.polygamma(2, param_sum) / gs.polygamma(1, param_sum)
mat_ones = gs.ones((n_points, self.dim, self.dim))
mat_diag = from_vector_to_diagonal_matrix(
-gs.polygamma(2, base_point) / gs.polygamma(1, base_point))
arrays = [
gs.zeros((1, ind_k)),
gs.ones((1, 1)),
gs.zeros((1, self.dim - ind_k - 1))
]
vec_k = gs.tile(gs.hstack(arrays), (n_points, 1))
val_k = gs.polygamma(2, param_k) / gs.polygamma(1, param_k)
vec_k = gs.einsum('i,ij->ij', val_k, vec_k)
mat_k = from_vector_to_diagonal_matrix(vec_k)
mat = gs.einsum('i,ijk->ijk', c2, mat_ones)\
- gs.einsum('i,ijk->ijk', c1, mat_diag) + mat_k
return 1 / 2 * mat
def test_from_vector_to_diagonal_matrix(self):
vec = gs.array([1.0, 2.0, 3.0])
mat_diag = utils.from_vector_to_diagonal_matrix(vec, -1)
result = mat_diag.shape
expected = (4, 4)
self.assertAllClose(result, expected)
vec = gs.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
mat_diag = utils.from_vector_to_diagonal_matrix(vec, 0)
expected = gs.array([
[[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]],
[[4.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 6.0]],
])
self.assertAllClose(mat_diag, expected)
mat_plus = utils.from_vector_to_diagonal_matrix(vec, 1)
expected = gs.array([
[
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 2.0, 0.0],
[0.0, 0.0, 0.0, 3.0],
[0.0, 0.0, 0.0, 0.0],
],
[
[0.0, 4.0, 0.0, 0.0],
[0.0, 0.0, 5.0, 0.0],
[0.0, 0.0, 0.0, 6.0],
[0.0, 0.0, 0.0, 0.0],
],
])
self.assertAllClose(mat_plus, expected)
mat_minus = utils.from_vector_to_diagonal_matrix(vec, -1)
expected = gs.array([
[
[0.0, 0.0, 0.0, 0.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 2.0, 0.0, 0.0],
[0.0, 0.0, 3.0, 0.0],
],
[
[0.0, 0.0, 0.0, 0.0],
[4.0, 0.0, 0.0, 0.0],
[0.0, 5.0, 0.0, 0.0],
[0.0, 0.0, 6.0, 0.0],
],
])
self.assertAllClose(mat_minus, expected)
def retraction(tangent_vec, base_point):
"""Compute the retraction of a tangent vector.
This computation is based on the QR-decomposition.
e.g. :math:`P_x(V) = qf(X + V)`.
Parameters
----------
tangent_vec : array-like, shape=[..., n, p]
Tangent vector at a base point.
base_point : array-like, shape=[..., n, p]
Point in the Stiefel manifold.
Returns
-------
exp : array-like, shape=[..., n, p]
Point in the Stiefel manifold equal to the retraction
of tangent_vec at the base point.
"""
n_tangent_vecs, _, _ = tangent_vec.shape
n_base_points, _, _ = base_point.shape
if not (n_tangent_vecs == n_base_points or n_tangent_vecs == 1
or n_base_points == 1):
raise NotImplementedError
matrix_q, matrix_r = gs.linalg.qr(base_point + tangent_vec)
diagonal = gs.diagonal(matrix_r, axis1=1, axis2=2)
sign = gs.sign(gs.sign(diagonal) + 0.5)
diag = algebra_utils.from_vector_to_diagonal_matrix(sign)
result = gs.einsum('nij,njk->nik', matrix_q, diag)
return result
def retraction(tangent_vec, base_point):
"""Compute the retraction of a tangent vector.
This computation is based on the QR-decomposition.
e.g. :math:`P_x(V) = qf(X + V)`.
Parameters
----------
tangent_vec : array-like, shape=[..., n, p]
Tangent vector at a base point.
base_point : array-like, shape=[..., n, p]
Point in the Stiefel manifold.
Returns
-------
exp : array-like, shape=[..., n, p]
Point in the Stiefel manifold equal to the retraction
of tangent_vec at the base point.
"""
matrix_q, matrix_r = gs.linalg.qr(base_point + tangent_vec)
diagonal = gs.diagonal(matrix_r, axis1=-2, axis2=-1)
sign = gs.sign(gs.sign(diagonal) + 0.5)
diag = algebra_utils.from_vector_to_diagonal_matrix(sign)
result = Matrices.mul(matrix_q, diag)
return result
def apply_func_to_eigvals(x, function, check_positive=False):
"""
Apply function to eigenvalues and reconstruct the matrix.
Parameters
----------
x : array_like, shape=[..., n, n]
Symmetric matrix.
function : callable
Function to apply to eigenvalues.
Returns
-------
x : array_like, shape=[..., n, n]
Symmetric matrix.
"""
eigvals, eigvecs = gs.linalg.eigh(x)
if check_positive:
if gs.any(gs.cast(eigvals, gs.float32) < 0.):
logging.warning('Negative eigenvalue encountered in'
' {}'.format(function.__name__))
eigvals = function(eigvals)
eigvals = algebra_utils.from_vector_to_diagonal_matrix(eigvals)
transp_eigvecs = Matrices.transpose(eigvecs)
reconstuction = gs.matmul(eigvecs, eigvals)
reconstuction = gs.matmul(reconstuction, transp_eigvecs)
return reconstuction
def _procrustes_preprocessing(p, matrix_v, matrix_m, matrix_n):
"""Procrustes preprocessing.
Parameters
----------
matrix_v : array-like
matrix_m : array-like
matrix_n : array-like
Returns
-------
matrix_v : array-like
"""
[matrix_d, _, matrix_r] = gs.linalg.svd(matrix_v[..., p:, p:])
j_matrix = gs.eye(p)
for i in range(1, p):
matrix_rd = Matrices.mul(
matrix_r, j_matrix, Matrices.transpose(matrix_d))
sub_matrix_v = gs.matmul(matrix_v[..., :, p:], matrix_rd)
matrix_v = gs.concatenate([
gs.concatenate([matrix_m, matrix_n], axis=-2),
sub_matrix_v], axis=-1)
det = gs.linalg.det(matrix_v)
if gs.all(det > 0):
break
ones = gs.ones(p)
reflection_vec = gs.concatenate(
[ones[:-i], gs.array([-1.] * i)], axis=0)
mask = gs.cast(det < 0, gs.float32)
sign = (mask[..., None] * reflection_vec
+ (1. - mask)[..., None] * ones)
j_matrix = algebra_utils.from_vector_to_diagonal_matrix(sign)
return matrix_v
def metric_matrix(self, base_point=None):
"""Compute the inner-product matrix.
Compute the inner-product matrix of the Fisher information metric
at the tangent space at base point.
Parameters
----------
base_point : array-like, shape=[..., dim]
Base point.
Returns
-------
mat : array-like, shape=[..., dim, dim]
Inner-product matrix.
"""
if base_point is None:
raise ValueError(
"A base point must be given to compute the " "metric matrix"
)
base_point = gs.to_ndarray(base_point, to_ndim=2)
n_points = base_point.shape[0]
mat_ones = gs.ones((n_points, self.dim, self.dim))
poly_sum = gs.polygamma(1, gs.sum(base_point, -1))
mat_diag = from_vector_to_diagonal_matrix(gs.polygamma(1, base_point))
mat = mat_diag - gs.einsum("i,ijk->ijk", poly_sum, mat_ones)
return gs.squeeze(mat)
def metric_matrix(self, base_point=None):
"""Compute the inner-product matrix.
Compute the inner-product matrix of the Fisher information metric
at the tangent space at base point.
Parameters
----------
base_point : array-like, shape=[..., 2]
Base point.
Returns
-------
mat : array-like, shape=[..., 2, 2]
Inner-product matrix.
"""
if base_point is None:
raise ValueError("A base point must be given to compute the "
"metric matrix")
base_point = gs.to_ndarray(base_point, to_ndim=2)
kappa, gamma = base_point[:, 0], base_point[:, 1]
mat_diag = gs.transpose(
gs.array([gs.polygamma(1, kappa) - 1 / kappa, kappa / gamma**2]))
mat = from_vector_to_diagonal_matrix(mat_diag)
return gs.squeeze(mat)
def test_skew_matrix_from_vector(self):
rot_vec = gs.array([0.9])
skew_matrix = self.group.skew_matrix_from_vector(rot_vec)
result = gs.matmul(skew_matrix, skew_matrix)
diag = gs.array([-0.81, -0.81])
expected = algebra_utils.from_vector_to_diagonal_matrix(diag)
self.assertAllClose(result, expected)
def orthonormal_basis_test_data(self):
smoke_data = [
dict(group=SpecialOrthogonal(3), metric_mat_at_identity=None),
dict(
group=SpecialOrthogonal(3),
metric_mat_at_identity=from_vector_to_diagonal_matrix(
gs.array([1.0, 2.0, 3.0])),
),
]
return self.generate_tests(smoke_data)
def test_update(self):
self.kalman.state = gs.zeros(2)
self.kalman.initialize_covariances(self.prior_cov, self.process_cov,
self.obs_cov)
measurement = gs.array([0.6])
expected_cov = from_vector_to_diagonal_matrix(gs.array([2. / 3., 1.]))
expected_state = gs.array([0.2, 0.])
self.kalman.update(measurement)
self.assertAllClose(expected_state, self.kalman.state)
self.assertAllClose(expected_cov, self.kalman.covariance)
def orthonormal_basis_se3_test_data(self):
smoke_data = [
dict(group=SpecialEuclidean(3), metric_mat_at_identity=None),
dict(
group=SpecialEuclidean(3),
metric_mat_at_identity=from_vector_to_diagonal_matrix(
gs.cast(gs.arange(1,
SpecialEuclidean(3).dim + 1),
gs.float32)),
),
]
return self.generate_tests(smoke_data)
def metric_matrix(self, base_point=None):
"""Compute the inner product matrix, independent of the base point.
Parameters
----------
base_point : array-like, shape=[..., dim]
Base point.
Returns
-------
inner_prod_mat : array-like, shape=[..., dim, dim]
Inner-product matrix.
"""
q, p = self.signature
diagonal = gs.array([-1.0] * p + [1.0] * q)
return from_vector_to_diagonal_matrix(diagonal)
def projection(self, point):
"""Project a matrix on SO(n) using the Frobenius norm.
Parameters
----------
mat : array-like, shape=[..., n, n]
Returns
-------
rot_mat : array-like, shape=[..., n, n]
"""
mat = point
n_mats, n, _ = mat.shape
if n == 3:
mat_unitary_u, _, mat_unitary_v = gs.linalg.svd(mat)
rot_mat = gs.einsum('nij,njk->nik', mat_unitary_u, mat_unitary_v)
mask = gs.less(gs.linalg.det(rot_mat), 0.)
mask_float = gs.cast(mask, gs.float32) + self.epsilon
diag = gs.array([[1., 1., -1.]])
diag = gs.to_ndarray(
algebra_utils.from_vector_to_diagonal_matrix(diag),
to_ndim=3) + self.epsilon
new_mat_diag_s = gs.tile(diag, [n_mats, 1, 1])
aux_mat = gs.einsum(
'nij,njk->nik',
mat_unitary_u,
new_mat_diag_s)
rot_mat += gs.einsum(
'n,njk->njk',
mask_float,
gs.einsum(
'nij,njk->nik',
aux_mat,
mat_unitary_v))
else:
aux_mat = gs.matmul(gs.transpose(mat, axes=(0, 2, 1)), mat)
inv_sqrt_mat = gs.linalg.inv(
gs.linalg.sqrtm(aux_mat))
rot_mat = gs.matmul(mat, inv_sqrt_mat)
return rot_mat
def align(self, point, base_point, **kwargs):
"""Align point to base_point.
Find the optimal rotation R in SO(m) such that the base point and
R.point are well positioned.
Parameters
----------
point : array-like, shape=[..., k_landmarks, m_ambient]
Point on the manifold.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
Point on the manifold.
Returns
-------
aligned : array-like, shape=[..., k_landmarks, m_ambient]
R.point.
"""
mat = gs.matmul(Matrices.transpose(point), base_point)
left, singular_values, right = gs.linalg.svd(mat)
det = gs.linalg.det(mat)
conditioning = (
(singular_values[..., -2]
+ gs.sign(det) * singular_values[..., -1]) /
singular_values[..., 0])
if gs.any(conditioning < 5e-4):
logging.warning(f'Singularity close, ill-conditioned matrix '
f'encountered: {conditioning}')
if gs.any(gs.isclose(conditioning, 0.)):
logging.warning("Alignment matrix is not unique.")
if gs.any(det < 0):
ones = gs.ones(self.m_ambient)
reflection_vec = gs.concatenate(
[ones[:-1], gs.array([-1.])], axis=0)
mask = gs.cast(det < 0, gs.float32)
sign = (mask[..., None] * reflection_vec
+ (1. - mask)[..., None] * ones)
j_matrix = from_vector_to_diagonal_matrix(sign)
rotation = Matrices.mul(
Matrices.transpose(right), j_matrix, Matrices.transpose(left))
else:
rotation = gs.matmul(
Matrices.transpose(right), Matrices.transpose(left))
return gs.matmul(point, Matrices.transpose(rotation))
def test_orthonormal_basis(self):
group = SpecialOrthogonal(3)
lie_algebra = SkewSymmetricMatrices(3)
metric = InvariantMetric(group=group)
basis = metric.normal_basis(lie_algebra.basis)
result = metric.inner_product_at_identity(basis[0], basis[1])
self.assertAllClose(result, 0.0)
result = metric.inner_product_at_identity(basis[1], basis[1])
self.assertAllClose(result, 1.0)
metric_mat = from_vector_to_diagonal_matrix(gs.array([1.0, 2.0, 3.0]))
metric = InvariantMetric(group=group, metric_mat_at_identity=metric_mat)
basis = metric.normal_basis(lie_algebra.basis)
result = metric.inner_product_at_identity(basis[0], basis[1])
self.assertAllClose(result, 0.0)
result = metric.inner_product_at_identity(basis[1], basis[1])
self.assertAllClose(result, 1.0)
def apply_func_to_eigvals(mat, function, check_positive=False):
"""
Apply function to eigenvalues and reconstruct the matrix.
Parameters
----------
mat : array_like, shape=[..., n, n]
Symmetric matrix.
function : callable, list of callables
Function to apply to eigenvalues. If a list of functions is passed,
a list of results will be returned.
check_positive : bool
Whether to check positivity of the eigenvalues.
Optional. Default: False.
Returns
-------
mat : array_like, shape=[..., n, n]
Symmetric matrix.
"""
eigvals, eigvecs = gs.linalg.eigh(mat)
if check_positive and gs.any(gs.cast(eigvals, gs.float32) < 0.0):
try:
name = function.__name__
except AttributeError:
name = function[0].__name__
logging.warning("Negative eigenvalue encountered in"
" {}".format(name))
return_list = True
if not isinstance(function, list):
function = [function]
return_list = False
reconstruction = []
transp_eigvecs = Matrices.transpose(eigvecs)
for fun in function:
eigvals_f = fun(eigvals)
eigvals_f = algebra_utils.from_vector_to_diagonal_matrix(eigvals_f)
reconstruction.append(
Matrices.mul(eigvecs, eigvals_f, transp_eigvecs))
return reconstruction if return_list else reconstruction[0]
def test_orthonormal_basis_se3(self):
group = SpecialEuclidean(3)
lie_algebra = group.lie_algebra
metric = InvariantMetric(group=group)
basis = metric.normal_basis(lie_algebra.basis)
for i, x in enumerate(basis):
for y in basis[i:]:
result = metric.inner_product_at_identity(x, y)
expected = 0.0 if gs.any(x != y) else 1.0
self.assertAllClose(result, expected)
metric_mat = from_vector_to_diagonal_matrix(
gs.cast(gs.arange(1, group.dim + 1), gs.float32)
)
metric = InvariantMetric(group=group, metric_mat_at_identity=metric_mat)
basis = metric.normal_basis(lie_algebra.basis)
for i, x in enumerate(basis):
for y in basis[i:]:
result = metric.inner_product_at_identity(x, y)
expected = 0.0 if gs.any(x != y) else 1.0
self.assertAllClose(result, expected)
def metric_matrix(self, base_point=None):
"""Compute the inner-product matrix.
Compute the inner-product matrix of the Fisher information metric
at the tangent space at base point.
Parameters
----------
base_point : array-like, shape=[..., dim + 1]
Base point.
Returns
-------
mat : array-like, shape=[..., dim, dim]
Inner-product matrix.
"""
if base_point is None:
raise ValueError("A base point must be given to compute the "
"metric matrix")
base_point = gs.to_ndarray(base_point, to_ndim=2)
mat = from_vector_to_diagonal_matrix(1 / base_point)
return gs.squeeze(mat)
def is_diagonal(cls, mat, atol=gs.atol):
"""Check if a matrix is square and diagonal.
Parameters
----------
mat : array-like, shape=[..., n, n]
Matrix.
atol : float
Absolute tolerance.
Optional, default: backend atol.
Returns
-------
is_diagonal : array-like, shape=[...,]
Boolean evaluating if the matrix is square and diagonal.
"""
is_square = cls.is_square(mat)
if not gs.all(is_square):
return False
diagonal_mat = from_vector_to_diagonal_matrix(cls.diagonal(mat))
is_diagonal = gs.all(gs.isclose(mat, diagonal_mat, atol=atol), axis=(-2, -1))
return is_diagonal
def jacobian_christoffels(self, base_point):
"""Compute the Jacobian of the Christoffel symbols.
Compute the Jacobian of the Christoffel symbols of the
Fisher information metric.
Parameters
----------
base_point : array-like, shape=[..., dim]
Base point.
Returns
-------
jac : array-like, shape=[..., dim, dim, dim, dim]
Jacobian of the Christoffel symbols.
:math: 'jac[..., i, j, k, l] = dGamma^i_{jk} / dx_l'
"""
n_dim = base_point.ndim
param = gs.transpose(base_point)
sum_param = gs.sum(param, 0)
term_1 = 1 / gs.polygamma(1, param)
term_2 = 1 / gs.polygamma(1, sum_param)
term_3 = -gs.polygamma(2, param) / gs.polygamma(1, param) ** 2
term_4 = -gs.polygamma(2, sum_param) / gs.polygamma(1, sum_param) ** 2
term_5 = term_3 / term_1
term_6 = term_4 / term_2
term_7 = (
gs.polygamma(2, param) ** 2
- gs.polygamma(1, param) * gs.polygamma(3, param)
) / gs.polygamma(1, param) ** 2
term_8 = (
gs.polygamma(2, sum_param) ** 2
- gs.polygamma(1, sum_param) * gs.polygamma(3, sum_param)
) / gs.polygamma(1, sum_param) ** 2
term_9 = term_2 - gs.sum(term_1, 0)
jac_1 = term_1 * term_8 / term_9
jac_1_mat = gs.squeeze(gs.tile(jac_1, (self.dim, self.dim, self.dim, 1, 1)))
jac_2 = (
-term_6
/ term_9**2
* gs.einsum("j...,i...->ji...", term_4 - term_3, term_1)
)
jac_2_mat = gs.squeeze(gs.tile(jac_2, (self.dim, self.dim, 1, 1, 1)))
jac_3 = term_3 * term_6 / term_9
jac_3_mat = gs.transpose(from_vector_to_diagonal_matrix(gs.transpose(jac_3)))
jac_3_mat = gs.squeeze(gs.tile(jac_3_mat, (self.dim, self.dim, 1, 1, 1)))
jac_4 = (
1
/ term_9**2
* gs.einsum("k...,j...,i...->kji...", term_5, term_4 - term_3, term_1)
)
jac_4_mat = gs.transpose(from_vector_to_diagonal_matrix(gs.transpose(jac_4)))
jac_5 = -gs.einsum("j...,i...->ji...", term_7, term_1) / term_9
jac_5_mat = from_vector_to_diagonal_matrix(gs.transpose(jac_5))
jac_5_mat = gs.transpose(from_vector_to_diagonal_matrix(jac_5_mat))
jac_6 = -gs.einsum("k...,j...->kj...", term_5, term_3) / term_9
jac_6_mat = gs.transpose(from_vector_to_diagonal_matrix(gs.transpose(jac_6)))
jac_6_mat = (
gs.transpose(
from_vector_to_diagonal_matrix(gs.transpose(jac_6_mat, [0, 1, 3, 2])),
[0, 1, 3, 4, 2],
)
if n_dim > 1
else from_vector_to_diagonal_matrix(jac_6_mat)
)
jac_7 = -from_vector_to_diagonal_matrix(gs.transpose(term_7))
jac_7_mat = from_vector_to_diagonal_matrix(jac_7)
jac_7_mat = gs.transpose(from_vector_to_diagonal_matrix(jac_7_mat))
jac = (
1
/ 2
* (
jac_1_mat
+ jac_2_mat
+ jac_3_mat
+ jac_4_mat
+ jac_5_mat
+ jac_6_mat
+ jac_7_mat
)
)
return (
gs.transpose(jac, [3, 1, 0, 2])
if n_dim == 1
else gs.transpose(jac, [4, 3, 1, 0, 2])
)
def christoffels(self, base_point):
"""Compute the Christoffel symbols.
Compute the Christoffel symbols of the Fisher information metric.
For computation purposes, we replace the value of
(gs.polygamma(1, x) - 1/x) by an equivalent (close lower-bound) when it becomes
too difficult to compute, as per in the second reference.
References
----------
.. [AD2008] Arwini, K. A., & Dodson, C. T. (2008).
Information geometry (pp. 31-54). Springer Berlin Heidelberg.
.. [GQ2015] Guo, B. N., Qi, F., Zhao, J. L., & Luo, Q. M. (2015).
Sharp inequalities for polygamma functions.
Mathematica Slovaca, 65(1), 103-120.
Parameters
----------
base_point : array-like, shape=[..., 2]
Base point.
Returns
-------
christoffels : array-like, shape=[..., 2, 2, 2]
Christoffel symbols, with the contravariant index on
the first dimension.
:math: 'christoffels[..., i, j, k] = Gamma^i_{jk}'
"""
base_point = gs.to_ndarray(base_point, to_ndim=2)
kappa, gamma = base_point[:, 0], base_point[:, 1]
if gs.any(kappa > 4e15):
raise ValueError(
"Christoffels computation overflows with values of kappa. "
"All values of kappa < 4e15 work.")
shape = kappa.shape
c111 = gs.where(
gs.polygamma(1, kappa) - 1 / kappa > gs.atol,
(gs.polygamma(2, kappa) + gs.array(kappa, dtype=gs.float32)**-2) /
(2 * (gs.polygamma(1, kappa) - 1 / kappa)),
0.25 * (kappa**2 * gs.polygamma(2, kappa) + 1),
)
c122 = gs.where(
gs.polygamma(1, kappa) - 1 / kappa > gs.atol,
-1 / (2 * gamma**2 * (gs.polygamma(1, kappa) - 1 / kappa)),
-(kappa**2) / (4 * gamma**2),
)
c1 = gs.squeeze(
from_vector_to_diagonal_matrix(gs.transpose(gs.array([c111,
c122]))))
c2 = gs.squeeze(
gs.transpose(
gs.array([[gs.zeros(shape), 1 / (2 * kappa)],
[1 / (2 * kappa), -1 / gamma]])))
christoffels = gs.array([c1, c2])
if len(christoffels.shape) == 4:
christoffels = gs.transpose(christoffels, [1, 0, 2, 3])
return gs.squeeze(christoffels)
def main():
"""Carry out two examples of state estimation on groups.
Both examples are localization problems, where only a part of the system
is observed. The first one is a linear system, while the second one is
non-linear.
"""
np.random.seed(12345)
model = LocalizationLinear()
kalman = KalmanFilter(model)
n_traj = 1000
obs_freq = 50
dt = 0.1
init_cov = gs.array([10.0, 1.0])
init_cov = algebra_utils.from_vector_to_diagonal_matrix(init_cov)
prop_cov = 0.001 * gs.eye(model.dim_noise)
obs_cov = 10 * gs.eye(model.dim_obs)
initial_covs = (init_cov, prop_cov, obs_cov)
kalman.initialize_covariances(*initial_covs)
true_state = gs.array([0.0, 0.0])
true_acc = gs.random.uniform(-1, 1, (n_traj, 1))
dt_vectorized = dt * gs.ones((n_traj, 1))
true_inputs = gs.hstack((dt_vectorized, true_acc))
true_traj, inputs, observations = create_data(
kalman, true_state, true_inputs, obs_freq
)
initial_state = np.random.multivariate_normal(true_state, init_cov)
estimate, uncertainty = estimation(
kalman, initial_state, inputs, observations, obs_freq
)
plt.figure()
plt.plot(true_traj[:, 0], label="Ground Truth")
plt.plot(estimate[:, 0], label="Kalman")
plt.plot(estimate[:, 0] + uncertainty[:, 0], color="k", linestyle=":")
plt.plot(
estimate[:, 0] - uncertainty[:, 0],
color="k",
linestyle=":",
label="3_sigma envelope",
)
plt.plot(
range(obs_freq, n_traj + 1, obs_freq),
observations,
marker="*",
linestyle="",
label="Observation",
)
plt.legend()
plt.title("1D Localization - Position")
plt.figure()
plt.plot(true_traj[:, 1], label="Ground Truth")
plt.plot(estimate[:, 1], label="Kalman")
plt.plot(estimate[:, 1] + uncertainty[:, 1], color="k", linestyle=":")
plt.plot(
estimate[:, 1] - uncertainty[:, 1],
color="k",
linestyle=":",
label="3_sigma envelope",
)
plt.legend()
plt.title("1D Localization - Speed")
model = Localization()
kalman = KalmanFilter(model)
init_cov = gs.array([1.0, 10.0, 10.0])
init_cov = algebra_utils.from_vector_to_diagonal_matrix(init_cov)
prop_cov = 0.001 * gs.eye(model.dim_noise)
obs_cov = 0.1 * gs.eye(model.dim_obs)
initial_covs = (init_cov, prop_cov, obs_cov)
kalman.initialize_covariances(*initial_covs)
true_state = gs.zeros(model.dim)
true_inputs = [gs.array([dt, 0.5, 0.0, 0.05]) for _ in range(n_traj)]
true_traj, inputs, observations = create_data(
kalman, true_state, true_inputs, obs_freq
)
initial_state = gs.array(np.random.multivariate_normal(true_state, init_cov))
initial_state = gs.cast(initial_state, true_state.dtype)
estimate, uncertainty = estimation(
kalman, initial_state, inputs, observations, obs_freq
)
plt.figure()
plt.plot(true_traj[:, 1], true_traj[:, 2], label="Ground Truth")
plt.plot(estimate[:, 1], estimate[:, 2], label="Kalman")
plt.scatter(observations[:, 0], observations[:, 1], s=2, c="k", label="Observation")
plt.legend()
plt.axis("equal")
plt.title("2D Localization")
plt.show()
