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
def square_root_velocity(self, curve): """Compute the square root velocity representation of a curve. The velocity is computed using the log map. The case of several curves is handled through vectorization. In that case, an index selection procedure allows to get rid of the log between the end point of curve[k, :, :] and the starting point of curve[k + 1, :, :]. Parameters ---------- curve : Returns ------- srv : """ curve = gs.to_ndarray(curve, to_ndim=3) n_curves, n_sampling_points, n_coords = curve.shape srv_shape = (n_curves, n_sampling_points - 1, n_coords) curve = gs.reshape(curve, (n_curves * n_sampling_points, n_coords)) coef = gs.cast(gs.array(n_sampling_points - 1), gs.float32) velocity = coef * self.ambient_metric.log(point=curve[1:, :], base_point=curve[:-1, :]) velocity_norm = self.ambient_metric.norm(velocity, curve[:-1, :]) srv = velocity / gs.sqrt(velocity_norm) index = gs.arange(n_curves * n_sampling_points - 1) mask = ~gs.equal((index + 1) % n_sampling_points, 0) index_select = gs.gather(index, gs.squeeze(gs.where(mask))) srv = gs.reshape(gs.gather(srv, index_select), srv_shape) return srv
def assertAllEqual(self, a, b): if tf_backend(): return tf.test.TestCase().assertAllEqual(a, b) elif np_backend() or autograd_backend(): np.testing.assert_array_equal(a, b) else: self.assertTrue(gs.equal(a, b))
def _exp_translation_transform(self, rot_vec): """Compute matrix associated to rot_vec for the translation part in exp. Parameters ---------- rot_vec : array-like, shape=[..., 3] Returns ------- transform : array-like, shape=[..., 3, 3] Matrix to be applied to the translation part in exp. """ n_samples = rot_vec.shape[0] angle = gs.linalg.norm(rot_vec, axis=-1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_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. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle ** 2 + TAYLOR_COEFFS_1_AT_0[4] * angle ** 4 + TAYLOR_COEFFS_1_AT_0[6] * angle ** 6) coef_2 += mask_close_0_float * ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle ** 2 + TAYLOR_COEFFS_2_AT_0[4] * angle ** 4 + TAYLOR_COEFFS_2_AT_0[6] * angle ** 6) angle += mask_0_float * 1. coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle ** 2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle ** 3) term_1 = gs.einsum('...i,...ij->...ij', coef_1, skew_mat) term_2 = gs.einsum('...i,...ij->...ij', coef_2, sq_skew_mat) term_id = gs.array([gs.eye(3)] * n_samples) transform = term_id + term_1 + term_2 return transform
def _exponential_matrix(self, rot_vec): """Compute exponential of rotation matrix represented by rot_vec. Parameters ---------- rot_vec : array-like, shape=[..., 3] Returns ------- exponential_mat : Matrix exponential of rot_vec """ # TODO (nguigs): find usecase for this method rot_vec = self.rotations.regularize(rot_vec) n_rot_vecs = 1 if rot_vec.ndim == 1 else len(rot_vec) angle = gs.linalg.norm(rot_vec, axis=-1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_rot_vec = self.rotations.skew_matrix_from_vector(rot_vec) coef_1 = gs.empty_like(angle) coef_2 = gs.empty_like(coef_1) mask_0 = gs.equal(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_close_to_0 = gs.isclose(angle, 0) mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1) mask_else = ~mask_0 & ~mask_close_to_0 coef_1[mask_close_to_0] = (1. / 2. - angle[mask_close_to_0] ** 2 / 24.) coef_2[mask_close_to_0] = (1. / 6. - angle[mask_close_to_0] ** 3 / 120.) # TODO (nina): Check if the discontinuity at 0 is expected. coef_1[mask_0] = 0 coef_2[mask_0] = 0 coef_1[mask_else] = (angle[mask_else] ** (-2) * (1. - gs.cos(angle[mask_else]))) coef_2[mask_else] = (angle[mask_else] ** (-2) * (1. - (gs.sin(angle[mask_else]) / angle[mask_else]))) term_1 = gs.zeros((n_rot_vecs, self.n, self.n)) term_2 = gs.zeros_like(term_1) for i in range(n_rot_vecs): term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i] term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i] exponential_mat = term_1 + term_2 return exponential_mat
def belongs(self, mat): """ Check if mat belongs to GL(n). """ mat = gs.to_ndarray(mat, to_ndim=3) n_mats, _, _ = mat.shape mat_rank = gs.zeros((n_mats, 1)) mat_rank = gs.linalg.matrix_rank(mat) mat_rank = gs.to_ndarray(mat_rank, to_ndim=1) return gs.equal(mat_rank, self.n)
def exponential_matrix(self, rot_vec): """ Compute the exponential of the rotation matrix represented by rot_vec. :param rot_vec: 3D rotation vector :returns exponential_mat: 3x3 matrix """ rot_vec = self.rotations.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) skew_rot_vec = so_group.skew_matrix_from_vector(rot_vec) coef_1 = gs.empty_like(angle) coef_2 = gs.empty_like(coef_1) mask_0 = gs.equal(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_close_to_0 = gs.isclose(angle, 0) mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1) mask_else = ~mask_0 & ~mask_close_to_0 coef_1[mask_close_to_0] = (1. / 2. - angle[mask_close_to_0] ** 2 / 24.) coef_2[mask_close_to_0] = (1. / 6. - angle[mask_close_to_0] ** 3 / 120.) # TODO(nina): check if the discountinuity as 0 is expected. coef_1[mask_0] = 0 coef_2[mask_0] = 0 coef_1[mask_else] = (angle[mask_else] ** (-2) * (1. - gs.cos(angle[mask_else]))) coef_2[mask_else] = (angle[mask_else] ** (-2) * (1. - (gs.sin(angle[mask_else]) / angle[mask_else]))) term_1 = gs.zeros((n_rot_vecs, self.n, self.n)) term_2 = gs.zeros_like(term_1) for i in range(n_rot_vecs): term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i] term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i] exponential_mat = term_1 + term_2 assert exponential_mat.ndim == 3 return exponential_mat
def log(self, point, base_point): """ Compute the Riemannian logarithm at point base_point, of point wrt the metric obtained by embedding of the n-dimensional sphere in the (n+1)-dimensional euclidean space. This gives a tangent vector at point base_point. :param base_point: point on the n-dimensional sphere :param point: point on the n-dimensional sphere :return log: tangent vector at 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.0, 1.0) angle = gs.arccos(cos_angle) mask_0 = gs.isclose(angle, 0.0) mask_else = gs.equal(mask_0, False) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1[mask_0] = ( 1. + INV_SIN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2 + INV_SIN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4 + INV_SIN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6 + INV_SIN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8) coef_2[mask_0] = ( 1. + INV_TAN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2 + INV_TAN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4 + INV_TAN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6 + INV_TAN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8) coef_1[mask_else] = angle[mask_else] / gs.sin(angle[mask_else]) coef_2[mask_else] = angle[mask_else] / gs.tan(angle[mask_else]) log = (gs.einsum('ni,nj->nj', coef_1, point) - gs.einsum('ni,nj->nj', coef_2, base_point)) return log
def dist(self, point_a, point_b): """ Geodesic distance between two points. """ if gs.all(gs.equal(point_a, point_b)): return 0. sq_norm_a = self.embedding_metric.squared_norm(point_a) sq_norm_b = self.embedding_metric.squared_norm(point_b) inner_prod = self.embedding_metric.inner_product(point_a, point_b) cosh_angle = -inner_prod / gs.sqrt(sq_norm_a * sq_norm_b) cosh_angle = gs.clip(cosh_angle, 1, None) dist = gs.arccosh(cosh_angle) return dist
def dist(self, point_a, point_b): """ Compute the distance induced on the hyperbolic space, from its embedding in the Minkowski space. """ if gs.all(gs.equal(point_a, point_b)): return 0. sq_norm_a = self.embedding_metric.squared_norm(point_a) sq_norm_b = self.embedding_metric.squared_norm(point_b) inner_prod = self.embedding_metric.inner_product(point_a, point_b) cosh_angle = -inner_prod / gs.sqrt(sq_norm_a * sq_norm_b) cosh_angle = gs.clip(cosh_angle, 1, None) dist = gs.arccosh(cosh_angle) return dist
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.0, 1.0) angle = gs.arccos(cos_angle) mask_0 = gs.isclose(angle, 0.0) mask_else = gs.equal(mask_0, gs.cast(gs.array(False), gs.int8)) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1[mask_0] = ( 1. + INV_SIN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2 + INV_SIN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4 + INV_SIN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6 + INV_SIN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8) coef_2[mask_0] = ( 1. + INV_TAN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2 + INV_TAN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4 + INV_TAN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6 + INV_TAN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8) coef_1[mask_else] = angle[mask_else] / gs.sin(angle[mask_else]) coef_2[mask_else] = angle[mask_else] / gs.tan(angle[mask_else]) log = (gs.einsum('ni,nj->nj', coef_1, point) - gs.einsum('ni,nj->nj', coef_2, base_point)) return log
def log(self, point, base_point): """Compute the Riemannian logarithm of a point. Parameters ---------- point : array-like, shape=[..., dim + 1] Point on the hypersphere. base_point : array-like, shape=[..., dim + 1] Point on the hypersphere. Returns ------- log : array-like, shape=[..., dim + 1] Tangent vector at the base point equal to the Riemannian logarithm of point at the base point. """ 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('...i,...j->...j', coef_1, point) - gs.einsum('...i,...j->...j', 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
def group_exp_from_identity(self, tangent_vec): """ Compute the group exponential of the tangent vector at the identity. """ tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) mask_close_pi = gs.isclose(angle, gs.pi) mask_close_pi = gs.squeeze(mask_close_pi, axis=1) rot_vec[mask_close_pi] = rotations.regularize( rot_vec[mask_close_pi]) skew_mat = so_group.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0) mask_close_0 = gs.isclose(angle, 0) & ~mask_0 mask_0 = gs.squeeze(mask_0, axis=1) mask_close_0 = gs.squeeze(mask_close_0, axis=1) mask_else = ~mask_0 & ~mask_close_0 coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1[mask_0] = 1. / 2. coef_2[mask_0] = 1. / 6. coef_1[mask_close_0] = (1. / 2. - angle[mask_close_0] ** 2 / 24. + angle[mask_close_0] ** 4 / 720. - angle[mask_close_0] ** 6 / 40320.) coef_2[mask_close_0] = (1. / 6. - angle[mask_close_0] ** 2 / 120. + angle[mask_close_0] ** 4 / 5040. - angle[mask_close_0] ** 6 / 362880.) coef_1[mask_else] = ((1. - gs.cos(angle[mask_else])) / angle[mask_else] ** 2) coef_2[mask_else] = ((angle[mask_else] - gs.sin(angle[mask_else])) / angle[mask_else] ** 3) n_tangent_vecs, _ = tangent_vec.shape group_exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) group_exp_translation[i] = translation_i + term_1_i + term_2_i group_exp = gs.zeros_like(tangent_vec) group_exp[:, :dim_rotations] = rot_vec group_exp[:, dim_rotations:] = group_exp_translation group_exp = self.regularize(group_exp) return group_exp
def exp_from_identity(self, tangent_vec, point_type=None): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- group_exp: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the group exponential of the tangent vectors calculated at the identity """ if point_type == 'vector': rotations = self.rotations dim_rotations = rotations.dim rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec, point_type=point_type) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) exp_translation += gs.outer( mask_i_float, translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp, point_type=point_type) return group_exp if point_type == 'matrix': return GeneralLinear.exp(tangent_vec) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.')
def group_exp_from_identity(self, tangent_vec, point_type=None): """ Compute the group exponential of the tangent vector at the identity. """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec, point_type=point_type) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) mask_close_pi = gs.isclose(angle, gs.pi) mask_close_pi = gs.squeeze(mask_close_pi, axis=1) rot_vec[mask_close_pi] = rotations.regularize( rot_vec[mask_close_pi], point_type=point_type) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0) mask_close_0 = gs.isclose(angle, 0) & ~mask_0 mask_0 = gs.squeeze(mask_0, axis=1) mask_close_0 = gs.squeeze(mask_close_0, axis=1) mask_else = ~mask_0 & ~mask_close_0 coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1[mask_0] = 1. / 2. coef_2[mask_0] = 1. / 6. coef_1[mask_close_0] = ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle[mask_close_0] ** 2 + TAYLOR_COEFFS_1_AT_0[4] * angle[mask_close_0] ** 4 + TAYLOR_COEFFS_1_AT_0[6] * angle[mask_close_0] ** 6) coef_2[mask_close_0] = ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle[mask_close_0] ** 2 + TAYLOR_COEFFS_2_AT_0[4] * angle[mask_close_0] ** 4 + TAYLOR_COEFFS_2_AT_0[6] * angle[mask_close_0] ** 6) coef_1[mask_else] = ((1. - gs.cos(angle[mask_else])) / angle[mask_else] ** 2) coef_2[mask_else] = ((angle[mask_else] - gs.sin(angle[mask_else])) / angle[mask_else] ** 3) n_tangent_vecs, _ = tangent_vec.shape group_exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) group_exp_translation[i] = translation_i + term_1_i + term_2_i group_exp = gs.zeros_like(tangent_vec) group_exp[:, :dim_rotations] = rot_vec group_exp[:, dim_rotations:] = group_exp_translation group_exp = self.regularize(group_exp, point_type=point_type) return group_exp elif point_type == 'matrix': raise NotImplementedError()
def group_exp_from_identity(self, tangent_vec, point_type=None): """ Compute the group exponential of the tangent vector at the identity. """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec, point_type=point_type) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape group_exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) group_exp_translation += mask_i_float * (translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, group_exp_translation], axis=1) group_exp = self.regularize(group_exp, point_type=point_type) return group_exp elif point_type == 'matrix': raise NotImplementedError()
def get_mask_i_float(i, n): range_n = gs.arange(n) i_float = gs.cast(gs.array([i]), gs.int32)[0] mask_i = gs.equal(range_n, i_float) mask_i_float = gs.cast(mask_i, gs.float32) return mask_i_float
def exp_from_identity(self, tangent_vec): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[..., 3] Returns ------- group_exp: array-like, shape=[..., 3] The group exponential of the tangent vectors calculated at the identity. """ rotations = self.rotations dim_rotations = rotations.dim rot_vec = tangent_vec[..., :dim_rotations] rot_vec = self.rotations.regularize(rot_vec) translation = tangent_vec[..., dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=-1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * (TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * (TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) exp_translation += gs.outer(mask_i_float, translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp) return group_exp