def compute_coordinates(self, point):
"""Compute the ellipse coordinates of a 2D SPD matrix.
Parameters
----------
point : array-like, shape=[2, 2]
SPD matrix.
Returns
-------
x_coords : array-like, shape=[n_sampling_points,]
x_coords coordinates of the sampling points on the discretized ellipse.
Y: array-like, shape = [n_sampling_points,]
y coordinates of the sampling points on the discretized ellipse.
"""
eigvalues, eigvectors = gs.linalg.eigh(point)
eigvalues = gs.where(eigvalues < gs.atol, gs.atol, eigvalues)
[eigvalue1, eigvalue2] = eigvalues
rot_sin = eigvectors[1, 0]
rot_cos = eigvectors[0, 0]
thetas = gs.linspace(0.0, 2 * gs.pi, self.n_sampling_points + 1)
x_coords = eigvalue1 * gs.cos(thetas) * rot_cos
x_coords -= rot_sin * eigvalue2 * gs.sin(thetas)
y_coords = eigvalue1 * gs.cos(thetas) * rot_sin
y_coords += rot_cos * eigvalue2 * gs.sin(thetas)
return x_coords, y_coords
def test_belongs(self):
theta = gs.pi / 3
point_1 = gs.array([[gs.cos(theta), -gs.sin(theta), 2.],
[gs.sin(theta), gs.cos(theta), 3.], [0., 0., 1.]])
result = self.group.belongs(point_1)
self.assertTrue(result)
point_2 = gs.array([[gs.cos(theta), -gs.sin(theta), 2.],
[gs.sin(theta), gs.cos(theta), 3.], [0., 0., 0.]])
result = self.group.belongs(point_2)
self.assertFalse(result)
point = gs.array([point_1, point_2])
expected = gs.array([True, False])
result = self.group.belongs(point)
self.assertAllClose(result, expected)
point = point_1[0]
result = self.group.belongs(point)
self.assertFalse(result)
point = gs.zeros((2, 3))
result = self.group.belongs(point)
self.assertFalse(result)
point = gs.zeros((2, 2, 3))
result = self.group.belongs(point)
self.assertFalse(gs.all(result))
self.assertAllClose(result.shape, (2, ))
def spherical_to_extrinsic(self, point_spherical):
"""Convert point from spherical to extrensic coordinates.
Convert from the spherical coordinates in the Hypersphere
to the extrinsic coordinates in Euclidean space.
Only implemented in dimension 2.
Parameters
----------
point_spherical : array-like, shape=[n_samples, dimension]
Returns
-------
point_extrinsic : array_like, shape=[n_samples, dimension + 1]
"""
if self.dimension != 2:
raise NotImplementedError(
'The conversion from spherical coordinates'
' to extrinsic coordinates is implemented'
' only in dimension 2.')
point_spherical = gs.to_ndarray(point_spherical, to_ndim=2)
theta = point_spherical[:, 0]
phi = point_spherical[:, 1]
point_extrinsic = gs.zeros(
(point_spherical.shape[0], self.dimension + 1))
point_extrinsic[:, 0] = gs.sin(theta) * gs.cos(phi)
point_extrinsic[:, 1] = gs.sin(theta) * gs.sin(phi)
point_extrinsic[:, 2] = gs.cos(theta)
assert self.belongs(point_extrinsic).all()
return point_extrinsic
def rotation(theta, phi):
"""Rotation sending a triangle at pole to location theta, phi."""
rot_th = gs.array([[gs.cos(theta), -gs.sin(theta), 0.],
[gs.sin(theta), gs.cos(theta), 0.], [0., 0., 1.]])
rot_phi = gs.array([[gs.cos(phi), 0., gs.sin(phi)], [0., 1., 0.],
[-gs.sin(phi), 0, gs.cos(phi)]])
return rot_th @ rot_phi @ gs.transpose(rot_th)
def christoffels(self, point, point_type='spherical'):
"""Compute Christoffel symbols.
Only implemented in dimension 2 and for spherical coordinates.
Parameters
----------
point : array-like, shape=[n_samples, dimension]
Returns
-------
christoffel : array-like, shape=[n_samples,
contravariant index,
first covariant index,
second covariant index]
"""
if self.dimension != 2 or point_type != 'spherical':
raise NotImplementedError(
'The Christoffel symbols are only implemented'
' for spherical coordinates in the 2-sphere')
point = gs.to_ndarray(point, to_ndim=2)
n_samples = point.shape[0]
christoffel = gs.zeros(
(n_samples, self.dimension, self.dimension, self.dimension))
christoffel[:, 0, 1, 1] = -gs.sin(point[:, 0]) * gs.cos(point[:, 0])
christoffel[:, 1, 0, 1] = gs.cos(point[:, 0]) / gs.sin(point[:, 0])
christoffel[:, 1, 1, 0] = gs.cos(point[:, 0]) / gs.sin(point[:, 0])
return christoffel
def test_assignment_with_matrices(self):
np_array = _np.zeros((2, 3, 3))
gs_array = gs.zeros((2, 3, 3))
np_array[:, 0, 1] = 44.
gs_array = gs.assignment(gs_array, 44., (0, 1), axis=0)
self.assertAllCloseToNp(gs_array, np_array)
n_samples = 3
theta = _np.random.rand(5)
phi = _np.random.rand(5)
np_array = _np.zeros((n_samples, 5, 4))
gs_array = gs.array(np_array)
np_array[0, :, 0] = gs.cos(theta) * gs.cos(phi)
np_array[0, :, 1] = -gs.sin(theta) * gs.sin(phi)
gs_array = gs.assignment(gs_array,
gs.cos(theta) * gs.cos(phi), (0, 0),
axis=1)
gs_array = gs.assignment(gs_array,
-gs.sin(theta) * gs.sin(phi), (0, 1),
axis=1)
self.assertAllCloseToNp(gs_array, np_array)
def random_uniform(self, n_samples=1, bound=0.5):
"""
Sample in the Hypersphere with the uniform distribution.
Parameters
----------
n_samples : int, optional
bound: float, optional
Returns
-------
point : array-like, shape=[n_samples, dimension + 1]
"""
size = (n_samples, self.dimension)
if bound is None:
spherical_coord = gs.random.rand(*size) * gs.pi
spherical_coord[:, -1] *= 2
point = gs.zeros((n_samples, self.dimension + 1))
for i in range(self.dimension):
point[:,
i] = (gs.prod(gs.sin(spherical_coord[:, :i]), axis=1) *
gs.cos(spherical_coord[:, i]))
point[:, -1] = gs.prod(gs.sin(spherical_coord), axis=1)
else:
assert bound <= 0.5
point = bound * (2 * gs.random.rand(*size) - 1)
point = self.intrinsic_to_extrinsic_coords(point)
return point
def christoffels(self, point, point_type='spherical'):
"""Compute Christoffel symbols.
Only implemented in dimension 2 and for spherical coordinates.
Parameters
----------
point : array-like, shape=[n_samples, dimension]
point_type: str
Returns
-------
christoffel : array-like, shape=[n_samples,
contravariant index,
first covariant index,
second covariant index]
"""
if self.dimension != 2 or point_type != 'spherical':
raise NotImplementedError(
'The Christoffel symbols are only implemented'
' for spherical coordinates in the 2-sphere')
point = gs.to_ndarray(point, to_ndim=2)
christoffel = []
for sample in point:
gamma_0 = gs.array([[0, 0],
[0, -gs.sin(sample[0]) * gs.cos(sample[0])]])
gamma_1 = gs.array([[0, gs.cos(sample[0]) / gs.sin(sample[0])],
[gs.cos(sample[0]) / gs.sin(sample[0]), 0]])
christoffel.append(gs.stack([gamma_0, gamma_1]))
return gs.stack(christoffel)
def christoffels(self, point, point_type='spherical'):
"""Compute the Christoffel symbols at a point.
Only implemented in dimension 2 and for spherical coordinates.
Parameters
----------
point : array-like, shape=[n_samples, dim]
Point on hypersphere where the Christoffel symbols are computed.
point_type: str, {'spherical', 'intrinsic', 'extrinsic'}
Coordinates in which to express the Christoffel symbols.
Returns
-------
christoffel : array-like, shape=[n_samples, contravariant index, 1st
covariant index, 2nd covariant index]
Christoffel symbols at point.
"""
if self.dim != 2 or point_type != 'spherical':
raise NotImplementedError(
'The Christoffel symbols are only implemented'
' for spherical coordinates in the 2-sphere')
point = gs.to_ndarray(point, to_ndim=2)
christoffel = []
for sample in point:
gamma_0 = gs.array([[0, 0],
[0, -gs.sin(sample[0]) * gs.cos(sample[0])]])
gamma_1 = gs.array([[0, gs.cos(sample[0]) / gs.sin(sample[0])],
[gs.cos(sample[0]) / gs.sin(sample[0]), 0]])
christoffel.append(gs.stack([gamma_0, gamma_1]))
return gs.stack(christoffel)
def spherical_to_extrinsic(self, point_spherical):
"""Convert point from spherical to extrinsic coordinates.
Convert from the spherical coordinates in the hypersphere
to the extrinsic coordinates in Euclidean space.
Only implemented in dimension 2.
Parameters
----------
point_spherical : array-like, shape=[n_samples, dim]
Point on the sphere, in spherical coordinates.
Returns
-------
point_extrinsic : array_like, shape=[n_samples, dim + 1]
Point on the sphere, in extrinsic coordinates in Euclidean space.
"""
if self.dim != 2:
raise NotImplementedError(
'The conversion from spherical coordinates'
' to extrinsic coordinates is implemented'
' only in dimension 2.')
theta = point_spherical[:, 0]
phi = point_spherical[:, 1]
point_extrinsic = gs.zeros((point_spherical.shape[0], self.dim + 1))
point_extrinsic[:, 0] = gs.sin(theta) * gs.cos(phi)
point_extrinsic[:, 1] = gs.sin(theta) * gs.sin(phi)
point_extrinsic[:, 2] = gs.cos(theta)
if not gs.all(self.belongs(point_extrinsic)):
raise ValueError('Points do not belong to the manifold.')
return point_extrinsic
def _sphere_immersion(spherical_coords):
theta = spherical_coords[..., 0]
phi = spherical_coords[..., 1]
return gs.array([
gs.cos(phi) * gs.sin(theta),
gs.sin(phi) * gs.sin(theta),
gs.cos(theta),
])
def group_useful_matrix(theta, elem_33=1.0):
return gs.array(
[
[gs.cos(theta), -gs.sin(theta), 2.0],
[gs.sin(theta), gs.cos(theta), 3.0],
[0.0, 0.0, elem_33],
]
)
def convert_to_spherical_coordinates(self, points):
"""Convert polar coordinates to spherical one."""
coords_theta, coords_phi = self.convert_to_polar_coordinates(points)
coords_x = 0.5 * gs.cos(coords_theta) * gs.sin(coords_phi)
coords_y = 0.5 * gs.sin(coords_theta) * gs.sin(coords_phi)
coords_z = 0.5 * gs.cos(coords_phi)
spherical_coords = gs.transpose(gs.stack((coords_x, coords_y, coords_z)))
return spherical_coords
def test_rotation_vector_from_matrix(self):
angle = 0.12
rot_mat = gs.array([[gs.cos(angle), -gs.sin(angle)],
[gs.sin(angle), gs.cos(angle)]])
result = self.group.rotation_vector_from_matrix(rot_mat)
expected = gs.array([0.12])
self.assertAllClose(result, expected)
def test_Localization_rotation_matrix(self):
initial_state = gs.array([0.5, 1.0, 2.0])
angle = initial_state[0]
rotation = gs.array([[gs.cos(angle), -gs.sin(angle)],
[gs.sin(angle), gs.cos(angle)]])
expected = rotation
result = self.nonlinear_model.rotation_matrix(angle)
self.assertAllClose(expected, result)
def test_Localization_innovation(self):
initial_state = gs.array([0.5, 1.0, 2.0])
measurement = gs.array([0.7, 2.1])
angle = initial_state[0]
rotation = gs.array([[gs.cos(angle), -gs.sin(angle)],
[gs.sin(angle), gs.cos(angle)]])
expected = gs.matmul(gs.transpose(rotation), gs.array([-0.3, 0.1]))
result = self.nonlinear_model.innovation(initial_state, measurement)
self.assertAllClose(expected, result)
def _expected_jacobian_immersion(point):
theta = point[..., 0]
phi = point[..., 1]
jacobian = gs.array([
[gs.cos(phi) * gs.cos(theta), -gs.sin(phi) * gs.sin(theta)],
[gs.sin(phi) * gs.cos(theta),
gs.cos(phi) * gs.sin(theta)],
[-gs.sin(theta), 0.0],
])
return jacobian
def test_Localization_adjoint_map(self):
initial_state = gs.array([0.5, 1.0, 2.0])
angle = initial_state[0]
rotation = gs.array([[gs.cos(angle), -gs.sin(angle)],
[gs.sin(angle), gs.cos(angle)]])
first_line = gs.eye(1, 3)
last_lines = gs.hstack((gs.array([[2.0], [-1.0]]), rotation))
expected = gs.vstack((first_line, last_lines))
result = self.nonlinear_model.adjoint_map(initial_state)
self.assertAllClose(expected, result)
def matrix_from_tait_bryan_angles_extrinsic_zyx(self, tait_bryan_angles):
"""Convert Tait-Bryan angles to rot mat in extrensic coords (zyx).
Convert a rotation given in terms of the tait bryan angles,
[angle_1, angle_2, angle_3] in extrinsic (fixed) coordinate system
in order zyx, into a rotation matrix.
rot_mat = X(angle_1).Y(angle_2).Z(angle_3)
where:
- X(angle_1) is a rotation of angle angle_1 around axis x.
- Y(angle_2) is a rotation of angle angle_2 around axis y.
- Z(angle_3) is a rotation of angle angle_3 around axis z.
Parameters
----------
tait_bryan_angles : array-like, shape=[..., 3]
Returns
-------
rot_mat : array-like, shape=[..., n, n]
"""
n_tait_bryan_angles, _ = tait_bryan_angles.shape
rot_mat = gs.zeros((n_tait_bryan_angles,) + (self.n,) * 2)
angle_1 = tait_bryan_angles[:, 0]
angle_2 = tait_bryan_angles[:, 1]
angle_3 = tait_bryan_angles[:, 2]
for i in range(n_tait_bryan_angles):
cos_angle_1 = gs.cos(angle_1[i])
sin_angle_1 = gs.sin(angle_1[i])
cos_angle_2 = gs.cos(angle_2[i])
sin_angle_2 = gs.sin(angle_2[i])
cos_angle_3 = gs.cos(angle_3[i])
sin_angle_3 = gs.sin(angle_3[i])
column_1 = [[cos_angle_2 * cos_angle_3],
[(cos_angle_1 * sin_angle_3
+ cos_angle_3 * sin_angle_1 * sin_angle_2)],
[(sin_angle_1 * sin_angle_3
- cos_angle_1 * cos_angle_3 * sin_angle_2)]]
column_2 = [[- cos_angle_2 * sin_angle_3],
[(cos_angle_1 * cos_angle_3
- sin_angle_1 * sin_angle_2 * sin_angle_3)],
[(cos_angle_3 * sin_angle_1
+ cos_angle_1 * sin_angle_2 * sin_angle_3)]]
column_3 = [[sin_angle_2],
[- cos_angle_2 * sin_angle_1],
[cos_angle_1 * cos_angle_2]]
rot_mat[i] = gs.hstack((column_1, column_2, column_3))
return rot_mat
def draw(self, n_theta=25, n_phi=13, scale=0.05, elev=60.0, azim=0.0):
"""Draw the sphere regularly sampled with corresponding triangles."""
self.set_ax()
self.set_view(elev=elev, azim=azim)
self.ax.set_axis_off()
plt.tight_layout()
coords_theta = gs.linspace(0.0, 2.0 * gs.pi, n_theta)
coords_phi = gs.linspace(0.0, gs.pi, n_phi)
coords_x = gs.to_numpy(0.5 * gs.outer(gs.sin(coords_phi), gs.cos(coords_theta)))
coords_y = gs.to_numpy(0.5 * gs.outer(gs.sin(coords_phi), gs.sin(coords_theta)))
coords_z = gs.to_numpy(
0.5 * gs.outer(gs.cos(coords_phi), gs.ones_like(coords_theta))
)
self.ax.plot_surface(
coords_x,
coords_y,
coords_z,
rstride=1,
cstride=1,
color="grey",
linewidth=0,
alpha=0.1,
zorder=-1,
)
self.ax.plot_wireframe(
coords_x,
coords_y,
coords_z,
linewidths=0.6,
color="grey",
alpha=0.6,
zorder=-1,
)
def lim(theta):
return (
gs.pi
- self.elev
+ (2.0 * self.elev - gs.pi) / gs.pi * abs(self.azim - theta)
)
for theta in gs.linspace(0.0, 2.0 * gs.pi, n_theta // 2 + 1):
for phi in gs.linspace(0.0, gs.pi, n_phi):
if theta <= self.azim + gs.pi and phi <= lim(theta):
self.draw_triangle(theta, phi, scale)
if theta > self.azim + gs.pi and phi < lim(
2.0 * self.azim + 2.0 * gs.pi - theta
):
self.draw_triangle(theta, phi, scale)
def tangent_spherical_to_extrinsic(self, tangent_vec_spherical,
base_point_spherical):
"""Convert tangent vector from spherical to extrinsic coordinates.
Convert from the spherical coordinates in the hypersphere
to the extrinsic coordinates in Euclidean space for a tangent
vector. Only implemented in dimension 2.
Parameters
----------
tangent_vec_spherical : array-like, shape=[..., dim]
Tangent vector to the sphere, in spherical coordinates.
base_point_spherical : array-like, shape=[..., dim]
Point on the sphere, in spherical coordinates.
Returns
-------
tangent_vec_extrinsic : array-like, shape=[..., dim + 1]
Tangent vector to the sphere, at base point,
in extrinsic coordinates in Euclidean space.
"""
if self.dim != 2:
raise NotImplementedError(
'The conversion from spherical coordinates'
' to extrinsic coordinates is implemented'
' only in dimension 2.')
n_samples = base_point_spherical.shape[0]
theta = base_point_spherical[:, 0]
phi = base_point_spherical[:, 1]
jac = gs.zeros((n_samples, self.dim + 1, self.dim))
zeros = gs.zeros(n_samples)
jac = gs.concatenate([
gs.array([[[
gs.cos(theta[i]) * gs.cos(phi[i]),
-gs.sin(theta[i]) * gs.sin(phi[i])
],
[
gs.cos(theta[i]) * gs.sin(phi[i]),
gs.sin(theta[i]) * gs.cos(phi[i])
], [-gs.sin(theta[i]), zeros[i]]]])
for i in range(n_samples)
],
axis=0)
tangent_vec_extrinsic = gs.einsum('...ij,...j->...i', jac,
tangent_vec_spherical)
return tangent_vec_extrinsic
def tangent_spherical_to_extrinsic(
self, tangent_vec_spherical, base_point_spherical
):
"""Convert tangent vector from spherical to extrinsic coordinates.
Convert from the spherical coordinates in the hypersphere
to the extrinsic coordinates in Euclidean space for a tangent
vector. Only implemented in dimension 2.
Parameters
----------
tangent_vec_spherical : array-like, shape=[..., dim]
Tangent vector to the sphere, in spherical coordinates.
base_point_spherical : array-like, shape=[..., dim]
Point on the sphere, in spherical coordinates.
Returns
-------
tangent_vec_extrinsic : array-like, shape=[..., dim + 1]
Tangent vector to the sphere, at base point,
in extrinsic coordinates in Euclidean space.
"""
if self.dim != 2:
raise NotImplementedError(
"The conversion from spherical coordinates"
" to extrinsic coordinates is implemented"
" only in dimension 2."
)
axes = (2, 0, 1) if base_point_spherical.ndim == 2 else (0, 1)
theta = base_point_spherical[..., 0]
phi = base_point_spherical[..., 1]
phi = gs.where(theta == 0.0, 0.0, phi)
zeros = gs.zeros_like(theta)
jac = gs.array(
[
[gs.cos(theta) * gs.cos(phi), -gs.sin(theta) * gs.sin(phi)],
[gs.cos(theta) * gs.sin(phi), gs.sin(theta) * gs.cos(phi)],
[-gs.sin(theta), zeros],
]
)
jac = gs.transpose(jac, axes)
tangent_vec_extrinsic = gs.einsum(
"...ij,...j->...i", jac, tangent_vec_spherical
)
return tangent_vec_extrinsic
def draw_triangle(self, r, theta, scale):
"""Draw the corresponding triangle on the disk at r, theta."""
u_theta = gs.cos(theta) * self.ua + gs.sin(theta) * self.na
triangle = gs.cos(r) * self.pole + gs.sin(r) * u_theta
triangle = scale * triangle
x = list(r * gs.cos(theta) + triangle[:, 0])
x = x + [x[0]]
y = list(r * gs.sin(theta) + triangle[:, 1])
y = y + [y[0]]
self.ax.plot(x, y, 'grey', zorder=1)
c = ['red', 'green', 'blue']
for i in range(3):
self.ax.scatter(x[i], y[i], color=c[i], s=10, alpha=1, zorder=1)
def parallel_transport(tangent_vec_a, tangent_vec_b, base_point):
r"""Compute the parallel transport of a tangent vector.
Closed-form solution for the parallel transport of a tangent vector a
along the geodesic defined by :math: `t \mapsto exp_(base_point)(t*
tangent_vec_b)`.
Parameters
----------
tangent_vec_a : array-like, shape=[..., dim + 1]
Tangent vector at base point to be transported.
tangent_vec_b : array-like, shape=[..., dim + 1]
Tangent vector at base point, along which the parallel transport
is computed.
base_point : array-like, shape=[..., dim + 1]
Point on the hypersphere.
Returns
-------
transported_tangent_vec: array-like, shape=[..., dim + 1]
Transported tangent vector at `exp_(base_point)(tangent_vec_b)`.
"""
theta = gs.linalg.norm(tangent_vec_b, axis=-1)
normalized_b = gs.einsum('...,...i->...i', 1 / theta, tangent_vec_b)
pb = gs.einsum('...i,...i->...', tangent_vec_a, normalized_b)
p_orth = tangent_vec_a - gs.einsum('...,...i->...i', pb, normalized_b)
transported = \
- gs.einsum('...,...i->...i', gs.sin(theta) * pb, base_point)\
+ gs.einsum('...,...i->...i', gs.cos(theta) * pb, normalized_b)\
+ p_orth
return transported
def __init__(self, n_angles=100, points=None):
angles = gs.linspace(0, 2 * gs.pi, n_angles)
self.circle_x = gs.cos(angles)
self.circle_y = gs.sin(angles)
self.points = []
if points is not None:
self.add_points(points)
def fibonnaci_points(self, n_points=16000):
"""Spherical Fibonacci point sets yield nearly uniform point
distributions on the unit sphere."""
x_vals = []
y_vals = []
z_vals = []
offset = 2. / n_points
increment = gs.pi * (3. - gs.sqrt(5.))
for i in range(n_points):
y = ((i * offset) - 1) + (offset / 2)
r = gs.sqrt(1 - pow(y, 2))
phi = ((i + 1) % n_points) * increment
x = gs.cos(phi) * r
z = gs.sin(phi) * r
x_vals.append(x)
y_vals.append(y)
z_vals.append(z)
x_vals = [(self.radius * i) for i in x_vals]
y_vals = [(self.radius * i) for i in y_vals]
z_vals = [(self.radius * i) for i in z_vals]
return gs.array([x_vals, y_vals, z_vals])
def __init__(self, n_meridians=40, n_circles_latitude=None, points=None):
if n_circles_latitude is None:
n_circles_latitude = max(n_meridians / 2, 4)
u, v = gs.meshgrid(gs.arange(0, 2 * gs.pi, 2 * gs.pi / n_meridians),
gs.arange(0, gs.pi, gs.pi / n_circles_latitude))
self.center = gs.zeros(3)
self.radius = 1
self.sphere_x = self.center[0] + self.radius * gs.cos(u) * gs.sin(v)
self.sphere_y = self.center[1] + self.radius * gs.sin(u) * gs.sin(v)
self.sphere_z = self.center[2] + self.radius * gs.cos(v)
self.points = []
if points is not None:
self.add_points(points)
def test_Localization_propagate(self):
initial_state = gs.array([0.5, 1.0, 2.0])
time_step = gs.array([0.5])
linear_vel = gs.array([1.0, 0.5])
angular_vel = gs.array([0.0])
increment = gs.concatenate((time_step, linear_vel, angular_vel),
axis=0)
angle = initial_state[0]
rotation = gs.array([[gs.cos(angle), -gs.sin(angle)],
[gs.sin(angle), gs.cos(angle)]])
next_position = initial_state[1:] + time_step * gs.matmul(
rotation, linear_vel)
expected = gs.concatenate((gs.array([angle]), next_position), axis=0)
result = self.nonlinear_model.propagate(initial_state, increment)
self.assertAllClose(expected, result)
def convert_to_planar_coordinates(self, points):
"""Convert polar coordinates to spherical one."""
coords_r, coords_theta = self.convert_to_polar_coordinates(points)
coords_x = coords_r * gs.cos(coords_theta)
coords_y = coords_r * gs.sin(coords_theta)
planar_coords = gs.transpose(gs.stack((coords_x, coords_y)))
return planar_coords
def log(self, point, base_point):
"""
Riemannian logarithm of a point wrt a base point.
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
norm_base_point = self.embedding_metric.norm(base_point)
norm_point = self.embedding_metric.norm(point)
inner_prod = self.embedding_metric.inner_product(base_point, point)
cos_angle = inner_prod / (norm_base_point * norm_point)
cos_angle = gs.clip(cos_angle, -1., 1.)
angle = gs.arccos(cos_angle)
angle = gs.to_ndarray(angle, to_ndim=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_0 = gs.isclose(angle, 0.)
mask_else = gs.equal(mask_0, gs.array(False))
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1 += mask_0_float * (1. + INV_SIN_TAYLOR_COEFFS[1] * angle**2 +
INV_SIN_TAYLOR_COEFFS[3] * angle**4 +
INV_SIN_TAYLOR_COEFFS[5] * angle**6 +
INV_SIN_TAYLOR_COEFFS[7] * angle**8)
coef_2 += mask_0_float * (1. + INV_TAN_TAYLOR_COEFFS[1] * angle**2 +
INV_TAN_TAYLOR_COEFFS[3] * angle**4 +
INV_TAN_TAYLOR_COEFFS[5] * angle**6 +
INV_TAN_TAYLOR_COEFFS[7] * angle**8)
# This avoids division by 0.
angle += mask_0_float * 1.
coef_1 += mask_else_float * angle / gs.sin(angle)
coef_2 += mask_else_float * angle / gs.tan(angle)
log = (gs.einsum('ni,nj->nj', coef_1, point) -
gs.einsum('ni,nj->nj', coef_2, base_point))
mask_same_values = gs.isclose(point, base_point)
mask_else = gs.equal(mask_same_values, gs.array(False))
mask_else_float = gs.cast(mask_else, gs.float32)
mask_else_float = gs.to_ndarray(mask_else_float, to_ndim=1)
mask_else_float = gs.to_ndarray(mask_else_float, to_ndim=2)
mask_not_same_points = gs.sum(mask_else_float, axis=1)
mask_same_points = gs.isclose(mask_not_same_points, 0.)
mask_same_points = gs.cast(mask_same_points, gs.float32)
mask_same_points = gs.to_ndarray(mask_same_points, to_ndim=2, axis=1)
mask_same_points_float = gs.cast(mask_same_points, gs.float32)
log -= mask_same_points_float * log
return log