def _exp_translation_transform(self, rot_vec): n_samples = rot_vec.shape[0] base_1 = gs.array([gs.eye(2)] * n_samples) base_2 = -gs.array( [self.rotations.skew_matrix_from_vector(gs.ones(1))] * n_samples) mask_close_0 = gs.isclose(rot_vec, 0.) mask_else = ~mask_close_0 mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) cos_coef = gs.zeros_like(rot_vec) sin_coef = gs.zeros_like(rot_vec) cos_coef += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] * rot_vec + TAYLOR_COEFFS_1_AT_0[1] * rot_vec ** 3 + TAYLOR_COEFFS_1_AT_0[2] * rot_vec ** 5) sin_coef += mask_close_0_float * ( 1 - TAYLOR_COEFFS_2_AT_0[0] * rot_vec ** 2 - TAYLOR_COEFFS_2_AT_0[2] * rot_vec ** 4) rot_vec = rot_vec + mask_close_0_float * 1e-6 cos_coef += mask_else_float * ((1. - gs.cos(rot_vec)) / rot_vec) sin_coef += mask_else_float * (gs.sin(rot_vec) / rot_vec) sin_term = gs.einsum('...i,...jk->...jk', sin_coef, base_1) cos_term = gs.einsum('...i,...jk->...jk', cos_coef, base_2) transform = sin_term + cos_term return transform
def test_integrability_tensor_derivative_is_alternate(self): r"""Integrability tensor derivatives is alternate in pre-shape. For two horizontal vector fields :math:`X,Y` the integrability tensor (hence its derivatives) is alternate: :math:`\nabla_X ( A_Y Z + A_Z Y ) = 0`. """ nabla_x_a_y_z, a_y_z = self.space.integrability_tensor_derivative( self.hor_x, self.hor_y, self.nabla_x_y, self.hor_z, self.nabla_x_z, self.base_point, ) nabla_x_a_z_y, a_z_y = self.space.integrability_tensor_derivative( self.hor_x, self.hor_z, self.nabla_x_z, self.hor_y, self.nabla_x_y, self.base_point, ) result = nabla_x_a_y_z + nabla_x_a_z_y self.assertAllClose(a_y_z + a_z_y, gs.zeros_like(result)) self.assertAllClose(result, gs.zeros_like(result))
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 test_assignment_with_booleans_single_index(self): np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result)
def submersion(point, k): r"""Submersion that defines the Grassmann manifold. The Grassmann manifold is defined here as embedded in the set of symmetric matrices, as the pre-image of the function defined around the projector on the space spanned by the first k columns of the identity matrix by (see Exercise E.25 in [Pau07]_). .. math: \begin{pmatrix} I_k + A & B^T \\ B & D \end{pmatrix} \mapsto (D - B(I_k + A)^{-1}B^T, A + A^2 + B^TB This map is a submersion and its zero space is the set of orthogonal rank-k projectors. References ---------- .. [Pau07] Paulin, Frédéric. “Géométrie différentielle élémentaire,” 2007. https://www.imo.universite-paris-saclay.fr/~paulin /notescours/cours_geodiff.pdf. """ _, eigvecs = gs.linalg.eigh(point) eigvecs = gs.flip(eigvecs, -1) flipped_point = Matrices.mul(Matrices.transpose(eigvecs), point, eigvecs) b = flipped_point[..., k:, :k] d = flipped_point[..., k:, k:] a = flipped_point[..., :k, :k] - gs.eye(k) first = d - Matrices.mul(b, GeneralLinear.inverse(a + gs.eye(k)), Matrices.transpose(b)) second = a + Matrices.mul(a, a) + Matrices.mul(Matrices.transpose(b), b) row_1 = gs.concatenate([first, gs.zeros_like(b)], axis=-1) row_2 = gs.concatenate([Matrices.transpose(gs.zeros_like(b)), second], axis=-1) return gs.concatenate([row_1, row_2], axis=-2)
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) angle = self.dist(base_point, point) angle = gs.to_ndarray(angle, to_ndim=1) angle = gs.to_ndarray(angle, to_ndim=2) mask_0 = gs.isclose(angle, 0) mask_else = ~mask_0 coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1[mask_0] = (1. + INV_SINH_TAYLOR_COEFFS[1] * angle[mask_0]**2 + INV_SINH_TAYLOR_COEFFS[3] * angle[mask_0]**4 + INV_SINH_TAYLOR_COEFFS[5] * angle[mask_0]**6 + INV_SINH_TAYLOR_COEFFS[7] * angle[mask_0]**8) coef_2[mask_0] = (1. + INV_TANH_TAYLOR_COEFFS[1] * angle[mask_0]**2 + INV_TANH_TAYLOR_COEFFS[3] * angle[mask_0]**4 + INV_TANH_TAYLOR_COEFFS[5] * angle[mask_0]**6 + INV_TANH_TAYLOR_COEFFS[7] * angle[mask_0]**8) coef_1[mask_else] = angle[mask_else] / gs.sinh(angle[mask_else]) coef_2[mask_else] = angle[mask_else] / gs.tanh(angle[mask_else]) log = (gs.einsum('ni,nj->nj', coef_1, point) - gs.einsum('ni,nj->nj', coef_2, base_point)) return log
def test_integrability_tensor_derivative_reverses_hor_ver(self): r"""Integrability tensor derivatives exchanges hor & ver in pre-shape. For :math:`X,Y,Z` horizontal and :math:`V,W` vertical, the integrability tensor (and thus its derivative) reverses horizontal and vertical subspaces: :math:`\nabla_X < A_Y Z, H > = 0` and :math:`nabla_X < A_Y V, W > = 0`. """ scal = self.space.ambient_metric.inner_product nabla_x_a_y_z, a_y_z = self.space.integrability_tensor_derivative( self.hor_x, self.hor_y, self.nabla_x_y, self.hor_z, self.nabla_x_z, self.base_point, ) result = scal(nabla_x_a_y_z, self.hor_h) + scal(a_y_z, self.nabla_x_h) self.assertAllClose(result, gs.zeros_like(result)) nabla_x_a_y_v, a_y_v = self.space.integrability_tensor_derivative( self.hor_x, self.hor_y, self.nabla_x_y, self.ver_v, self.nabla_x_v, self.base_point, ) result = scal(nabla_x_a_y_v, self.ver_w) + scal(a_y_v, self.nabla_x_w) self.assertAllClose(result, gs.zeros_like(result))
def regularize_tangent_vec_at_identity( self, tangent_vec, metric=None): """Regularize a tangent vector at the identify. In 3D, regularize a tangent_vector by getting its norm at the identity, determined by the metric, to be less than pi. Parameters ---------- tangent_vec : array-like, shape=[..., 3]] metric : RiemannianMetric, optional default: self.left_canonical_metric Returns ------- regularized_vec : array-like, shape=[..., 3]] """ if metric is None: metric = self.left_canonical_metric tangent_vec_metric_norm = metric.norm(tangent_vec) tangent_vec_canonical_norm = gs.linalg.norm(tangent_vec, axis=-1) mask_norm_0 = gs.isclose(tangent_vec_metric_norm, 0.) mask_canonical_norm_0 = gs.isclose(tangent_vec_canonical_norm, 0.) mask_0 = mask_norm_0 | mask_canonical_norm_0 mask_else = ~mask_0 # This avoids division by 0. mask_0_float = gs.cast(mask_0, gs.float32) + self.epsilon mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon regularized_vec = gs.zeros_like(tangent_vec) regularized_vec += gs.einsum( '...,...i->...i', mask_0_float, tangent_vec) tangent_vec_canonical_norm += mask_0_float coef = gs.zeros_like(tangent_vec_metric_norm) coef += mask_else_float * ( tangent_vec_metric_norm / tangent_vec_canonical_norm) coef_tangent_vec = gs.einsum( '...,...i->...i', coef, tangent_vec) regularized_vec += gs.einsum( '...,...i->...i', mask_else_float, self.regularize(coef_tangent_vec)) coef += mask_0_float regularized_vec = gs.einsum( '...,...i->...i', 1. / coef, regularized_vec) regularized_vec = gs.einsum( '...,...i->...i', mask_else_float, regularized_vec) return regularized_vec
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 exp(self, tangent_vec, base_point): """ Riemannian exponential of a tangent vector wrt to a base point. Parameters ---------- tangent_vec : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] base_point : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] Returns ------- exp : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] """ tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) base_point = gs.to_ndarray(base_point, to_ndim=2) sq_norm_tangent_vec = self.embedding_metric.squared_norm( tangent_vec) norm_tangent_vec = gs.sqrt(sq_norm_tangent_vec) mask_0 = gs.isclose(sq_norm_tangent_vec, 0.) mask_0 = gs.to_ndarray(mask_0, to_ndim=1) mask_else = ~mask_0 mask_else = gs.to_ndarray(mask_else, to_ndim=1) mask_0_float = gs.cast(mask_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) coef_1 = gs.zeros_like(norm_tangent_vec) coef_2 = gs.zeros_like(norm_tangent_vec) coef_1 += mask_0_float * ( 1. + COSH_TAYLOR_COEFFS[2] * norm_tangent_vec ** 2 + COSH_TAYLOR_COEFFS[4] * norm_tangent_vec ** 4 + COSH_TAYLOR_COEFFS[6] * norm_tangent_vec ** 6 + COSH_TAYLOR_COEFFS[8] * norm_tangent_vec ** 8) coef_2 += mask_0_float * ( 1. + SINH_TAYLOR_COEFFS[3] * norm_tangent_vec ** 2 + SINH_TAYLOR_COEFFS[5] * norm_tangent_vec ** 4 + SINH_TAYLOR_COEFFS[7] * norm_tangent_vec ** 6 + SINH_TAYLOR_COEFFS[9] * norm_tangent_vec ** 8) # This avoids dividing by 0. norm_tangent_vec += mask_0_float * 1.0 coef_1 += mask_else_float * (gs.cosh(norm_tangent_vec)) coef_2 += mask_else_float * ( (gs.sinh(norm_tangent_vec) / (norm_tangent_vec))) exp = (gs.einsum('ni,nj->nj', coef_1, base_point) + gs.einsum('ni,nj->nj', coef_2, tangent_vec)) hyperbolic_space = HyperbolicSpace(dimension=self.dimension) exp = hyperbolic_space.regularize(exp) return exp
def exp(self, tangent_vec, base_point): """Compute the Riemannian exponential of a tangent vector. Parameters ---------- tangent_vec : array-like, shape=[..., dim + 1] Tangent vector at a base point. base_point : array-like, shape=[..., dim + 1] Point in hyperbolic space. Returns ------- exp : array-like, shape=[..., dim + 1] Point in hyperbolic space equal to the Riemannian exponential of tangent_vec at the base point. """ sq_norm_tangent_vec = self.embedding_metric.squared_norm( tangent_vec) sq_norm_tangent_vec = gs.clip(sq_norm_tangent_vec, 0, math.inf) norm_tangent_vec = gs.sqrt(sq_norm_tangent_vec) mask_0 = gs.isclose(sq_norm_tangent_vec, 0.) mask_0 = gs.to_ndarray(mask_0, to_ndim=1) mask_else = ~mask_0 mask_else = gs.to_ndarray(mask_else, to_ndim=1) mask_0_float = gs.cast(mask_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) coef_1 = gs.zeros_like(norm_tangent_vec) coef_2 = gs.zeros_like(norm_tangent_vec) coef_1 += mask_0_float * ( 1. + COSH_TAYLOR_COEFFS[2] * norm_tangent_vec ** 2 + COSH_TAYLOR_COEFFS[4] * norm_tangent_vec ** 4 + COSH_TAYLOR_COEFFS[6] * norm_tangent_vec ** 6 + COSH_TAYLOR_COEFFS[8] * norm_tangent_vec ** 8) coef_2 += mask_0_float * ( 1. + SINH_TAYLOR_COEFFS[3] * norm_tangent_vec ** 2 + SINH_TAYLOR_COEFFS[5] * norm_tangent_vec ** 4 + SINH_TAYLOR_COEFFS[7] * norm_tangent_vec ** 6 + SINH_TAYLOR_COEFFS[9] * norm_tangent_vec ** 8) # This avoids dividing by 0. norm_tangent_vec += mask_0_float * 1.0 coef_1 += mask_else_float * (gs.cosh(norm_tangent_vec)) coef_2 += mask_else_float * ( (gs.sinh(norm_tangent_vec) / (norm_tangent_vec))) exp = ( gs.einsum('...,...j->...j', coef_1, base_point) + gs.einsum('...,...j->...j', coef_2, tangent_vec)) hyperbolic_space = Hyperboloid(dim=self.dim) exp = hyperbolic_space.regularize(exp) return exp
def log(self, point, base_point): """Compute Riemannian logarithm of a point wrt a base point. If point_type = 'poincare' then base_point belongs to the Poincare ball and point is a vector in the Euclidean space of the same dimension as the ball. Parameters ---------- point : array-like, shape=[..., dim + 1] Point in hyperbolic space. base_point : array-like, shape=[..., dim + 1] Point in hyperbolic space. Returns ------- log : array-like, shape=[..., dim + 1] Tangent vector at the base point equal to the Riemannian logarithm of point at the base point. """ angle = self.dist(base_point, point) / self.scale angle = gs.to_ndarray(angle, to_ndim=1) mask_0 = gs.isclose(angle, 0.) mask_else = ~mask_0 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_SINH_TAYLOR_COEFFS[1] * angle ** 2 + INV_SINH_TAYLOR_COEFFS[3] * angle ** 4 + INV_SINH_TAYLOR_COEFFS[5] * angle ** 6 + INV_SINH_TAYLOR_COEFFS[7] * angle ** 8) coef_2 += mask_0_float * ( 1. + INV_TANH_TAYLOR_COEFFS[1] * angle ** 2 + INV_TANH_TAYLOR_COEFFS[3] * angle ** 4 + INV_TANH_TAYLOR_COEFFS[5] * angle ** 6 + INV_TANH_TAYLOR_COEFFS[7] * angle ** 8) # This avoids dividing by 0. angle += mask_0_float * 1. coef_1 += mask_else_float * (angle / gs.sinh(angle)) coef_2 += mask_else_float * (angle / gs.tanh(angle)) log_term_1 = gs.einsum('...,...j->...j', coef_1, point) log_term_2 = - gs.einsum('...,...j->...j', coef_2, base_point) log = log_term_1 + log_term_2 return log
def log(self, point, base_point): """ Riemannian logarithm of a point wrt a base point. Parameters ---------- point : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] base_point : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] Returns ------- log : array-like, shape=[n_samples, dimension + 1] or shape=[1, dimension + 1] """ point = gs.to_ndarray(point, to_ndim=2) base_point = gs.to_ndarray(base_point, to_ndim=2) angle = self.dist(base_point, point) angle = gs.to_ndarray(angle, to_ndim=1) angle = gs.to_ndarray(angle, to_ndim=2) mask_0 = gs.isclose(angle, 0.) mask_else = ~mask_0 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_SINH_TAYLOR_COEFFS[1] * angle ** 2 + INV_SINH_TAYLOR_COEFFS[3] * angle ** 4 + INV_SINH_TAYLOR_COEFFS[5] * angle ** 6 + INV_SINH_TAYLOR_COEFFS[7] * angle ** 8) coef_2 += mask_0_float * ( 1. + INV_TANH_TAYLOR_COEFFS[1] * angle ** 2 + INV_TANH_TAYLOR_COEFFS[3] * angle ** 4 + INV_TANH_TAYLOR_COEFFS[5] * angle ** 6 + INV_TANH_TAYLOR_COEFFS[7] * angle ** 8) # This avoids dividing by 0. angle += mask_0_float * 1. coef_1 += mask_else_float * (angle / gs.sinh(angle)) coef_2 += mask_else_float * (angle / gs.tanh(angle)) log = (gs.einsum('ni,nj->nj', coef_1, point) - gs.einsum('ni,nj->nj', coef_2, base_point)) return log
def test_assignment(self): gs_array_1 = gs.ones(3) self.assertRaises(ValueError, gs.assignment, gs_array_1, [.1, 2., 1.], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] = 1.5 gs_result = gs.assignment(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] = 1.5 gs_result = gs.assignment(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] = 1 gs_result = gs.assignment(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] = 1 gs_result = gs.assignment(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] = 1 gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) gs_array_4_arr = gs.zeros_like(gs.array(np_array_4)) gs_result = gs.assignment(gs_array_4_arr, 1, gs.array((0, 1)), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] = 1 gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list)
def matrix_from_rotation_vector(self, rot_vec): """Convert rotation vector to rotation matrix. Parameters ---------- rot_vec: array-like, shape=[..., 3] Returns ------- rot_mat: array-like, shape=[..., 3] """ rot_vec = self.regularize(rot_vec) angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_rot_vec = self.skew_matrix_from_vector(rot_vec) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) # This avoids dividing by 0. mask_0 = gs.isclose(angle, 0.) mask_0_float = gs.cast(mask_0, gs.float32) + self.epsilon coef_1 += mask_0_float * (1. - (angle ** 2) / 6.) coef_2 += mask_0_float * (1. / 2. - angle ** 2) # This avoids dividing by 0. mask_else = ~mask_0 mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon angle += mask_0_float coef_1 += mask_else_float * (gs.sin(angle) / angle) coef_2 += mask_else_float * ( (1. - gs.cos(angle)) / (angle ** 2)) coef_1 = gs.squeeze(coef_1, axis=1) coef_2 = gs.squeeze(coef_2, axis=1) term_1 = (gs.eye(self.dim) + gs.einsum('n,njk->njk', coef_1, skew_rot_vec)) squared_skew_rot_vec = gs.einsum( 'nij,njk->nik', skew_rot_vec, skew_rot_vec) term_2 = gs.einsum('n,njk->njk', coef_2, squared_skew_rot_vec) return term_1 + term_2
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 test_fit_transform_hyperbolic(self): point = gs.array([2., 1., 1., 1.]) points = gs.array([point, point]) transformer = ToTangentSpace(geometry=self.hyperbolic.metric) result = transformer.fit_transform(X=points) expected = gs.zeros_like(points) self.assertAllClose(expected, result)
def curvature_test_data(self): group = self.matrix_so3 metric = InvariantMetric(group) x, y, z = metric.normal_basis(group.lie_algebra.basis) smoke_data = [ dict( group=group, tangent_vec_a=x, tangent_vec_b=y, tangent_vec_c=x, expected=1.0 / 8 * y, ), dict( group=group, tangent_vec_a=gs.stack([x, x]), tangent_vec_b=gs.stack([y] * 2), tangent_vec_c=gs.stack([x, x]), expected=gs.array([1.0 / 8 * y] * 2), ), dict( group=group, tangent_vec_a=y, tangent_vec_b=y, tangent_vec_c=z, expected=gs.zeros_like(z), ), ] return self.generate_tests(smoke_data)
def rotation_vector_from_quaternion(self, quaternion): """ Convert a unit quaternion into a rotation vector. """ assert self.n == 3, ('The quaternion representation does not exist' ' for rotations in %d dimensions.' % self.n) quaternion = gs.to_ndarray(quaternion, to_ndim=2) n_quaternions, _ = quaternion.shape cos_half_angle = quaternion[:, 0] cos_half_angle = gs.clip(cos_half_angle, -1, 1) half_angle = gs.arccos(cos_half_angle) half_angle = gs.to_ndarray(half_angle, to_ndim=2, axis=1) assert half_angle.shape == (n_quaternions, 1) rot_vec = gs.zeros_like(quaternion[:, 1:]) mask_0 = gs.isclose(half_angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_not_0 = ~mask_0 rotation_axis = (quaternion[mask_not_0, 1:] / gs.sin(half_angle[mask_not_0])) rot_vec[mask_not_0] = (2 * half_angle[mask_not_0] * rotation_axis) rot_vec = self.regularize(rot_vec) return rot_vec
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
def closest_rotation_matrix(mat): """ Compute the closest rotation matrix of a given matrix mat, in terms of the Frobenius norm. """ mat = gs.to_ndarray(mat, to_ndim=3) n_mats, mat_dim_1, mat_dim_2 = mat.shape assert mat_dim_1 == mat_dim_2 if mat_dim_1 == 3: mat_unitary_u, diag_s, mat_unitary_v = gs.linalg.svd(mat) rot_mat = gs.matmul(mat_unitary_u, mat_unitary_v) mask = gs.where(gs.linalg.det(rot_mat) < 0) new_mat_diag_s = gs.tile(gs.diag([1, 1, -1]), len(mask)) rot_mat[mask] = gs.matmul( gs.matmul(mat_unitary_u[mask], new_mat_diag_s), mat_unitary_v[mask]) else: aux_mat = gs.matmul(gs.transpose(mat, axes=(0, 2, 1)), mat) inv_sqrt_mat = gs.zeros_like(mat) for i in range(n_mats): sym_mat = aux_mat[i] assert spd_matrices_space.is_symmetric(sym_mat) inv_sqrt_mat[i] = gs.linalg.inv(spd_matrices_space.sqrtm(sym_mat)) rot_mat = gs.matmul(mat, inv_sqrt_mat) assert rot_mat.ndim == 3 return rot_mat
def regularize(self, point): """ In 3D, regularize the norm of the rotation vector, to be between 0 and pi, following the axis-angle representation's convention. If the angle angle is between pi and 2pi, the function computes its complementary in 2pi and inverts the direction of the rotation axis. """ point = gs.to_ndarray(point, to_ndim=2) assert self.belongs(point) n_points, vec_dim = point.shape if vec_dim == 3: angle = gs.linalg.norm(point, axis=1) regularized_point = point.astype('float64') mask_not_0 = angle != 0 k = gs.floor(angle / (2 * gs.pi) + .5) norms_ratio = gs.zeros_like(angle).astype('float64') norms_ratio[mask_not_0] = ( 1. - 2. * gs.pi * k[mask_not_0] / angle[mask_not_0]) norms_ratio[angle == 0] = 1 for i in range(n_points): regularized_point[i, :] = norms_ratio[i] * point[i] else: # TODO(nina): regularization needed in nD? regularized_point = point assert regularized_point.ndim == 2 return regularized_point
def test_curvature(self): group = self.matrix_so3 metric = InvariantMetric(group=group) x, y, z = metric.normal_basis(group.lie_algebra.basis) result = metric.curvature_at_identity(x, y, x) expected = 1.0 / 8 * y self.assertAllClose(result, expected) tan_a = gs.stack([x, x]) tan_b = gs.stack([y] * 2) result = metric.curvature(tan_a, tan_b, tan_a) self.assertAllClose(result, gs.array([expected] * 2)) point = group.random_uniform() translation_map = group.tangent_translation_map(point) tan_a = translation_map(x) tan_b = translation_map(y) result = metric.curvature(tan_a, tan_b, tan_a, point) expected = translation_map(expected) self.assertAllClose(result, expected) result = metric.curvature(y, y, z) expected = gs.zeros_like(z) self.assertAllClose(result, expected)
def test_curvature_derivative(self): group = self.matrix_so3 metric = InvariantMetric(group=group) x, y, z = metric.normal_basis(group.lie_algebra.basis) result = metric.curvature_derivative(x, y, z, x) expected = gs.zeros_like(x) self.assertAllClose(result, expected) point = group.random_uniform() translation_map = group.tangent_translation_map(point) tan_a = translation_map(x) tan_b = translation_map(y) tan_c = translation_map(z) result = metric.curvature_derivative(tan_a, tan_b, tan_c, tan_a, point) expected = gs.zeros_like(x) self.assertAllClose(result, expected)
def test_fit_riemannian_se2(self): init = (self.y_se2[0], gs.zeros_like(self.y_se2[0])) gr = GeodesicRegression( self.se2, metric=self.metric_se2, center_X=False, method="riemannian", max_iter=50, init_step_size=0.1, verbose=True, initialization=init, ) gr.fit(self.X_se2, self.y_se2, compute_training_score=True) intercept_hat, coef_hat = gr.intercept_, gr.coef_ training_score = gr.training_score_ self.assertAllClose(intercept_hat.shape, self.shape_se2) self.assertAllClose(coef_hat.shape, self.shape_se2) self.assertAllClose(training_score, 1.0, atol=1e-4) self.assertAllClose(intercept_hat, self.intercept_se2_true, atol=1e-4) tangent_vec_of_transport = self.se2.metric.log( self.intercept_se2_true, base_point=intercept_hat) transported_coef_hat = self.se2.metric.parallel_transport( tangent_vec=coef_hat, base_point=intercept_hat, direction=tangent_vec_of_transport, ) self.assertAllClose(transported_coef_hat, self.coef_se2_true, atol=0.6)
def inverse(self, point): """ Compute the group inverse in SE(n). Formula: (R, t)^{-1} = (R^{-1}, R^{-1}.(-t)) """ rotations = self.rotations dim_rotations = rotations.dimension point = self.regularize(point) n_points, _ = point.shape rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] inverse_point = gs.zeros_like(point) inverse_rotation = -rot_vec inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation) inverse_translation = gs.zeros((n_points, self.n)) for i in range(n_points): inverse_translation[i] = gs.dot(-translation[i], gs.transpose(inv_rot_mat[i])) inverse_point[:, :dim_rotations] = inverse_rotation inverse_point[:, dim_rotations:] = inverse_translation inverse_point = self.regularize(inverse_point) return inverse_point
def regularize(self, point, point_type=None): """ Regularize a point to the canonical representation chosen for SE(n). """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': point = gs.to_ndarray(point, to_ndim=2) assert self.belongs(point, point_type=point_type) rotations = self.rotations dim_rotations = rotations.dimension regularized_point = gs.zeros_like(point) rot_vec = point[:, :dim_rotations] regularized_point[:, :dim_rotations] = rotations.regularize( rot_vec, point_type=point_type) regularized_point[:, dim_rotations:] = point[:, dim_rotations:] elif point_type == 'matrix': point = gs.to_ndarray(point, to_ndim=3) # TODO(nina): regularization for matrices? regularized_point = gs.copy(point) return regularized_point
def compose(self, point_a, point_b): """Compute the group product of elements `point_a` and `point_b`. Parameters ---------- point_a : array-like, shape=[..., 3] Left factor in the product. point_b : array-like, shape=[..., 3] Right factor in the product. Returns ------- point_ab : array-like, shape=[..., 3] Product of point_a and point_b along the first dimension. """ point_ab = point_a + point_b point_ab_additional_term = gs.array( 1 / 2 * (point_a[..., 0] * point_b[..., 1] - point_a[..., 1] * point_b[..., 0])) point_ab = point_ab + gs.concatenate( [ gs.zeros_like(point_ab[..., :2]), point_ab_additional_term[..., None] ], axis=-1, ) return point_ab
def regularize_tangent_vec(self, tangent_vec, base_point, metric=None): if metric is None: metric = self.left_canonical_metric tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) base_point = gs.to_ndarray(base_point, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_tangent_vec = tangent_vec[:, :dim_rotations] rot_base_point = base_point[:, :dim_rotations] metric_mat = metric.inner_product_mat_at_identity rot_metric_mat = metric_mat[:, :dim_rotations, :dim_rotations] rot_metric = InvariantMetric( group=rotations, inner_product_mat_at_identity=rot_metric_mat, left_or_right=metric.left_or_right) regularized_vec = gs.zeros_like(tangent_vec) regularized_vec[:, :dim_rotations] = rotations.regularize_tangent_vec( tangent_vec=rot_tangent_vec, base_point=rot_base_point, metric=rot_metric) regularized_vec[:, dim_rotations:] = tangent_vec[:, dim_rotations:] return regularized_vec
def _log_translation_transform(self, rot_vec): rot_vec = gs.to_ndarray(rot_vec, to_ndim=2, axis=1) exp_transform = self._exp_translation_transform(rot_vec) if rot_vec.dtype == gs.float32: mask_close_0 = gs.isclose(rot_vec, 0., atol=1e-6) else: mask_close_0 = gs.isclose(rot_vec, 0.) mask_else = ~mask_close_0 mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) inv_determinant = gs.zeros_like(rot_vec) inv_determinant += 0.5 * mask_close_0_float / ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[1] * rot_vec ** 2 + TAYLOR_COEFFS_1_AT_0[2] * rot_vec ** 6) rot_vec = rot_vec + 1e-3 * mask_close_0_float inv_determinant += mask_else_float * ( rot_vec ** 2 / (2 * (1 - gs.cos(rot_vec)))) transform = gs.einsum( 'il, ijk -> ijk', inv_determinant, gs.transpose(exp_transform, axes=[0, 2, 1])) return transform