def matrix_from_vector(self, vec): """Convert point in vector point-type to matrix. Parameters ---------- vec : array-like, shape=[..., dimension] Vector. Returns ------- mat : array-like, shape=[..., n+1, n+1] Matrix. """ vec = self.regularize(vec) n_vecs, _ = vec.shape rot_vec = vec[:, :self.rotations.dim] trans_vec = vec[:, self.rotations.dim:] rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) trans_vec = gs.reshape(trans_vec, (n_vecs, self.n, 1)) mat = gs.concatenate((rot_mat, trans_vec), axis=2) last_lines = gs.array(gs.get_mask_i_float(self.n, self.n + 1)) last_lines = gs.to_ndarray(last_lines, to_ndim=2) last_lines = gs.to_ndarray(last_lines, to_ndim=3) mat = gs.concatenate((mat, last_lines), axis=1) return mat
def matrix_from_quaternion(self, quaternion): """Convert a unit quaternion into a rotation vector. Parameters ---------- quaternion : array-like, shape=[..., 4] Returns ------- rot_mat : array-like, shape=[..., 3] """ n_quaternions, _ = quaternion.shape w, x, y, z = gs.hsplit(quaternion, 4) rot_mat = gs.zeros((n_quaternions, ) + (self.n, ) * 2) for i in range(n_quaternions): # TODO(nina): Vectorize by applying the composition of # quaternions to the identity matrix column_1 = [ w[i]**2 + x[i]**2 - y[i]**2 - z[i]**2, 2 * x[i] * y[i] - 2 * w[i] * z[i], 2 * x[i] * z[i] + 2 * w[i] * y[i] ] column_2 = [ 2 * x[i] * y[i] + 2 * w[i] * z[i], w[i]**2 - x[i]**2 + y[i]**2 - z[i]**2, 2 * y[i] * z[i] - 2 * w[i] * x[i] ] column_3 = [ 2 * x[i] * z[i] - 2 * w[i] * y[i], 2 * y[i] * z[i] + 2 * w[i] * x[i], w[i]**2 - x[i]**2 - y[i]**2 + z[i]**2 ] mask_i = gs.get_mask_i_float(i, n_quaternions) rot_mat_i = gs.transpose(gs.hstack([column_1, column_2, column_3])) rot_mat_i = gs.to_ndarray(rot_mat_i, to_ndim=3) rot_mat += gs.einsum('...,...ij->...ij', mask_i, rot_mat_i) return rot_mat
def group_log_from_identity(self, point, point_type=None): """ Compute the group logarithm of the point at the identity. """ if point_type is None: point_type = self.default_point_type point = self.regularize(point, point_type=point_type) rotations = self.rotations dim_rotations = rotations.dimension if point_type == 'vector': rot_vec = point[:, :dim_rotations] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) translation = point[:, dim_rotations:] skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec) sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec) mask_close_0 = gs.isclose(angle, 0.) mask_close_pi = gs.isclose(angle, gs.pi) mask_else = ~mask_close_0 & ~mask_close_pi mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_close_pi_float = gs.cast(mask_close_pi, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) mask_0 = gs.isclose(angle, 0., atol=1e-6) mask_0_float = gs.cast(mask_0, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = -0.5 * gs.ones_like(angle) coef_2 = gs.zeros_like(angle) coef_2 += mask_close_0_float * (1. / 12. + angle**2 / 720. + angle**4 / 30240. + angle**6 / 1209600.) delta_angle = angle - gs.pi coef_2 += mask_close_pi_float * ( 1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) - ((PI2 - 12.) * delta_angle**2 / (4. * PI4)) + ((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) - ((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) + ((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 / (480. * PI7)) - ((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 / (480. * PI8))) psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle)) coef_2 += mask_else_float * (1 - psi) / (angle**2) n_points, _ = point.shape group_log_translation = gs.zeros((n_points, self.n)) for i in range(n_points): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_rot_vec[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_rot_vec[i])) mask_i_float = gs.get_mask_i_float(i, n_points) group_log_translation += mask_i_float * (translation_i + term_1_i + term_2_i) group_log = gs.concatenate([rot_vec, group_log_translation], axis=1) assert gs.ndim(group_log) == 2 elif point_type == 'matrix': raise NotImplementedError() return group_log
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 log_from_identity(self, point): """Compute the group logarithm of the point at the identity. Parameters ---------- point: array-like, shape=[..., 3] Returns ------- group_log: array-like, shape=[..., 3] the group logarithm in the Lie algbra """ point = self.regularize(point) rotations = self.rotations dim_rotations = rotations.dim rot_vec = point[:, :dim_rotations] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) translation = point[:, dim_rotations:] skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec) sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec) mask_close_0 = gs.isclose(angle, 0.) mask_close_pi = gs.isclose(angle, gs.pi) mask_else = ~mask_close_0 & ~mask_close_pi mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_close_pi_float = gs.cast(mask_close_pi, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) mask_0 = gs.isclose(angle, 0., atol=1e-6) mask_0_float = gs.cast(mask_0, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = -0.5 * gs.ones_like(angle) coef_2 = gs.zeros_like(angle) coef_2 += mask_close_0_float * (1. / 12. + angle**2 / 720. + angle**4 / 30240. + angle**6 / 1209600.) delta_angle = angle - gs.pi coef_2 += mask_close_pi_float * ( 1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) - ((PI2 - 12.) * delta_angle**2 / (4. * PI4)) + ((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) - ((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) + ((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 / (480. * PI7)) - ((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 / (480. * PI8))) psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle)) coef_2 += mask_else_float * (1 - psi) / (angle**2) n_points, _ = point.shape log_translation = gs.zeros((n_points, self.n)) for i in range(n_points): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_rot_vec[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_rot_vec[i])) mask_i_float = gs.get_mask_i_float(i, n_points) log_translation += gs.outer(mask_i_float, translation_i + term_1_i + term_2_i) return gs.concatenate([rot_vec, log_translation], axis=1)
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
def jacobian_translation(self, point, left_or_right='left'): """Compute the jacobian matrix corresponding to translation. Compute the jacobian matrix of the differential of the left/right translations from the identity to point in SO(3). Parameters ---------- point : array-like, shape=[..., 3] left_or_right : str, {'left', 'right'}, optional default: 'left' point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- jacobian : array-like, shape=[..., 3, 3] """ geomstats.error.check_parameter_accepted_values( left_or_right, 'left_or_right', ['left', 'right']) point = self.regularize(point) n_points, _ = point.shape angle = gs.linalg.norm(point, axis=-1) angle = gs.expand_dims(angle, axis=-1) coef_1 = gs.zeros([n_points, 1]) coef_2 = gs.zeros([n_points, 1]) # 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 * (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_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) # This avoids dividing by 0. mask_pi = gs.isclose(angle, gs.pi) mask_pi_float = gs.cast(mask_pi, gs.float32) + self.epsilon delta_angle = angle - gs.pi coef_1 += mask_pi_float * (TAYLOR_COEFFS_1_AT_PI[1] * delta_angle + TAYLOR_COEFFS_1_AT_PI[2] * delta_angle**2 + TAYLOR_COEFFS_1_AT_PI[3] * delta_angle**3 + TAYLOR_COEFFS_1_AT_PI[4] * delta_angle**4 + TAYLOR_COEFFS_1_AT_PI[5] * delta_angle**5 + TAYLOR_COEFFS_1_AT_PI[6] * delta_angle**6) angle += mask_0_float coef_2 += mask_pi_float * ((1 - coef_1) / angle**2) # This avoids dividing by 0. mask_else = ~mask_0 & ~mask_pi mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon # This avoids division by 0. angle += mask_pi_float coef_1 += mask_else_float * ((angle / 2) / gs.tan(angle / 2)) coef_2 += mask_else_float * ((1 - coef_1) / angle**2) jacobian = gs.zeros((n_points, self.dim, self.dim)) n_points_tensor = gs.array(n_points) for i in range(n_points): # This avoids dividing by 0. mask_i_float = (gs.get_mask_i_float(i, n_points_tensor) + self.epsilon) sign = -1. if left_or_right == 'left': sign = +1. jacobian_i = (coef_1[i] * gs.eye(self.dim) + coef_2[i] * gs.outer(point[i], point[i]) + sign * self.skew_matrix_from_vector(point[i]) / 2.) jacobian += gs.einsum('n,ij->nij', mask_i_float, jacobian_i) return jacobian
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\'.')