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 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 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 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 draw_vector(self, tangent_vec, base_point, tol=1e-03, **kwargs): """Draw one vector in the tangent space to disk at a base point.""" r_bp, th_bp = self.convert_to_polar_coordinates(base_point) bp = gs.array([ gs.cos(th_bp) * gs.sin(2 * r_bp), gs.sin(th_bp) * gs.sin(2 * r_bp), gs.cos(2 * r_bp) ]) r_exp, th_exp = self.convert_to_polar_coordinates( METRIC_S33.exp( tol * tangent_vec / METRIC_S33.norm(tangent_vec, base_point), base_point)) exp = gs.array([ gs.cos(th_exp) * gs.sin(2 * r_exp), gs.sin(th_exp) * gs.sin(2 * r_exp), gs.cos(2 * r_exp) ]) pole = gs.array([0., 0., 1.]) tv = exp - gs.dot(exp, bp) * bp u_tv = tv / gs.linalg.norm(tv) u_r = (gs.dot(pole, bp) * bp - pole) / gs.linalg.norm(gs.dot(pole, bp) * bp - pole) u_th = gs.cross(bp, u_r) x_r, x_th = gs.dot(u_tv, u_r), gs.dot(u_tv, u_th) bp = self.convert_to_planar_coordinates(base_point) u_r = bp / gs.linalg.norm(bp) u_th = gs.array([[0., -1.], [1., 0.]]) @ u_r tv = METRIC_S33.norm(tangent_vec, base_point) * (x_r * u_r + x_th * u_th) self.ax.quiver(bp[0], bp[1], tv[0], tv[1], **kwargs)
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 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 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 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 _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 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 _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_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 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 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 __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 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 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 quaternion_from_rotation_vector(self, rot_vec): """ Convert a rotation vector into a unit quaternion. """ assert self.n == 3, ('The quaternion representation does not exist' ' for rotations in %d dimensions.' % self.n) rot_vec = self.regularize(rot_vec) n_rot_vecs, _ = rot_vec.shape angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) rotation_axis = gs.zeros_like(rot_vec) mask_0 = gs.isclose(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_not_0 = ~mask_0 rotation_axis[mask_not_0] = rot_vec[mask_not_0] / angle[mask_not_0] n_quaternions, _ = rot_vec.shape quaternion = gs.zeros((n_quaternions, 4)) quaternion[:, :1] = gs.cos(angle / 2) quaternion[:, 1:] = gs.sin(angle / 2) * rotation_axis[:] return quaternion