Пример #1
0
class TestIntegrator(geomstats.tests.TestCase):
    def setUp(self):
        self.dimension = 4
        self.dt = 0.1
        self.euclidean = Euclidean(self.dimension)
        self.matrices = Matrices(self.dimension, self.dimension)
        self.intercept = self.euclidean.random_uniform(1)
        self.slope = Matrices.to_symmetric(self.matrices.random_uniform(1))

    def function_linear(self, _, vector):
        return -gs.dot(self.slope, vector)

    @geomstats.tests.np_and_pytorch_only
    def test_symplectic_euler_step(self):
        state = (self.intercept, self.slope)
        result = len(
            integrator._symplectic_euler_step(state, self.function_linear,
                                              self.dt))
        expected = len(state)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_rk4_step(self):
        state = (self.intercept, self.slope)
        result = len(integrator.rk4_step(state, self.function_linear, self.dt))
        expected = len(state)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_integrator(self):
        initial_state = self.euclidean.random_uniform(2)

        def function(_, velocity):
            return gs.zeros_like(velocity)

        for step in ['euler', 'rk4']:
            flow, _ = integrator.integrate(function, initial_state, step=step)
            result = flow[-1]
            expected = initial_state[0] + initial_state[1]

            self.assertAllClose(result, expected)
Пример #2
0
    def test_linear_mean(self):
        euclidean = Euclidean(3)
        point = euclidean.random_uniform(self.n_samples)

        estimator = ExponentialBarycenter(euclidean)

        estimator.fit(point)
        result = estimator.estimate_

        expected = gs.mean(point, axis=0)

        self.assertAllClose(result, expected)
Пример #3
0
class _SpecialEuclideanMatrices(GeneralLinear, LieGroup):
    """Class for special orthogonal groups.

    Parameters
    ----------
    n : int
        Integer representing the shape of the matrices: n x n.
    """
    def __init__(self, n):
        super(_SpecialEuclideanMatrices,
              self).__init__(default_point_type='matrix', n=n + 1)
        self.rotations = SpecialOrthogonal(n=n)
        self.translations = Euclidean(dim=n)
        self.n = n
        self.dim = int((n * (n + 1)) / 2)

    def get_identity(self):
        """Return the identity matrix."""
        return gs.eye(self.n + 1, self.n + 1)

    identity = property(get_identity)

    def belongs(self, point):
        """Check whether point is of the form rotation, translation.

        Parameters
        ----------
        point : array-like, shape=[..., n, n].
            Point to be checked.

        Returns
        -------
        belongs : array-like, shape=[...,]
            Boolean denoting if point belongs to the group.
        """
        point_dim1, point_dim2 = point.shape[-2:]
        belongs = (point_dim1 == point_dim2 == self.n + 1)

        rotation = point[..., :self.n, :self.n]
        rot_belongs = self.rotations.belongs(rotation)

        belongs = gs.logical_and(belongs, rot_belongs)

        last_line_except_last_term = point[..., self.n:, :-1]
        all_but_last_zeros = ~gs.any(last_line_except_last_term, axis=(-2, -1))

        belongs = gs.logical_and(belongs, all_but_last_zeros)

        last_term = point[..., self.n:, self.n:]
        belongs = gs.logical_and(belongs, gs.all(last_term == 1,
                                                 axis=(-2, -1)))

        if point.ndim == 2:
            return gs.squeeze(belongs)
        return gs.flatten(belongs)

    def _is_in_lie_algebra(self, tangent_vec, atol=TOLERANCE):
        """Project vector rotation part onto skew-symmetric matrices."""
        point_dim1, point_dim2 = tangent_vec.shape[-2:]
        belongs = (point_dim1 == point_dim2 == self.n + 1)

        rotation = tangent_vec[..., :self.n, :self.n]
        rot_belongs = self.is_skew_symmetric(rotation, atol=atol)

        belongs = gs.logical_and(belongs, rot_belongs)

        last_line = tangent_vec[..., -1, :]
        all_zeros = ~gs.any(last_line, axis=-1)

        belongs = gs.logical_and(belongs, all_zeros)
        return belongs

    def _to_lie_algebra(self, tangent_vec):
        """Project vector rotation part onto skew-symmetric matrices."""
        translation_mask = gs.hstack(
            [gs.ones((self.n, ) * 2), 2 * gs.ones((self.n, 1))])
        translation_mask = gs.concatenate(
            [translation_mask, gs.zeros((1, self.n + 1))], axis=0)
        tangent_vec = tangent_vec * gs.where(translation_mask != 0.,
                                             gs.array(1.), gs.array(0.))
        tangent_vec = (tangent_vec - GeneralLinear.transpose(tangent_vec)) / 2.
        return tangent_vec * translation_mask

    def random_uniform(self, n_samples=1, tol=1e-6):
        """Sample in SE(n) from the uniform distribution.

        Parameters
        ----------
        n_samples : int
            Number of samples.
            Optional, default: 1.
        tol : unused

        Returns
        -------
        samples : array-like, shape=[..., n + 1, n + 1]
            Sample in SE(n).
        """
        random_translation = self.translations.random_uniform(n_samples)
        random_rotation = self.rotations.random_uniform(n_samples)
        random_rotation = gs.to_ndarray(random_rotation, to_ndim=3)

        random_translation = gs.to_ndarray(random_translation, to_ndim=2)
        random_translation = gs.transpose(
            gs.to_ndarray(random_translation, to_ndim=3, axis=1), (0, 2, 1))

        random_point = gs.concatenate((random_rotation, random_translation),
                                      axis=2)
        last_line = gs.zeros((n_samples, 1, self.n + 1))
        random_point = gs.concatenate((random_point, last_line), axis=1)
        random_point = gs.assignment(random_point, 1, (-1, -1), axis=0)
        if gs.shape(random_point)[0] == 1:
            random_point = gs.squeeze(random_point, axis=0)
        return random_point
Пример #4
0
class _SpecialEuclideanVectors(LieGroup):
    """Base Class for the special euclidean groups in 2d and 3d in vector form.

    i.e. the Lie group of rigid transformations. Elements of SE(2), SE(3) can
    either be represented as vectors (in 2d or 3d) or as matrices in general.
    The matrix representation corresponds to homogeneous coordinates. This
    class is specific to the vector representation of rotations. For the matrix
    representation use the SpecialEuclidean class and set `n=2` or `n=3`.

    Parameter
    ---------
    epsilon : float
        Precision to use for calculations involving potential
        division by 0 in rotations.
        Optional, default: 0.
    """
    def __init__(self, n, epsilon=0.):
        dim = n * (n + 1) // 2
        LieGroup.__init__(self, dim=dim, default_point_type='vector')

        self.n = n
        self.epsilon = epsilon
        self.rotations = SpecialOrthogonal(n=n,
                                           point_type='vector',
                                           epsilon=epsilon)
        self.translations = Euclidean(dim=n)

    def get_identity(self, point_type=None):
        """Get the identity of the group.

        Parameters
        ----------
        point_type : str, {'vector', 'matrix'}
            The point_type of the returned value.
            Optional, default: self.default_point_type

        Returns
        -------
        identity : array-like, shape={[dim], [n + 1, n + 1]}
        """
        if point_type is None:
            point_type = self.default_point_type
        identity = gs.zeros(self.dim)
        return identity

    identity = property(get_identity)

    def get_point_type_shape(self, point_type=None):
        """Get the shape of the instance given the default_point_style."""
        return self.get_identity(point_type).shape

    def belongs(self, point):
        """Evaluate if a point belongs to SE(2) or SE(3).

        Parameters
        ----------
        point : array-like, shape=[..., dimension]
            Point to check.

        Returns
        -------
        belongs : array-like, shape=[...,]
            Boolean indicating whether point belongs to SE(2) or SE(3).
        """
        point_dim = point.shape[-1]
        point_ndim = point.ndim
        belongs = gs.logical_and(point_dim == self.dim, point_ndim < 3)
        belongs = gs.logical_and(
            belongs, self.rotations.belongs(point[..., :self.rotations.dim]))
        return belongs

    def regularize(self, point):
        """Regularize a point to the default representation for SE(n).

        Parameters
        ----------
        point : array-like, shape=[..., 3]
            Point to regularize.

        Returns
        -------
        point : array-like, shape=[..., 3]
            Regularized point.
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        regularized_point = point
        rot_vec = regularized_point[..., :dim_rotations]
        regularized_rot_vec = rotations.regularize(rot_vec)

        translation = regularized_point[..., dim_rotations:]

        return gs.concatenate([regularized_rot_vec, translation], axis=-1)

    @geomstats.vectorization.decorator(['else', 'vector', 'else'])
    def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None):
        """Regularize a tangent vector at the identity.

        Parameters
        ----------
        tangent_vec: array-like, shape=[..., 3]
            Tangent vector at base point.
        metric : RiemannianMetric
            Metric.
            Optional, default: None.

        Returns
        -------
        regularized_vec : array-like, shape=[..., 3]
            Regularized vector.
        """
        return self.regularize_tangent_vec(tangent_vec, self.identity, metric)

    @geomstats.vectorization.decorator(['else', 'vector'])
    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

    @geomstats.vectorization.decorator(['else', 'vector', 'vector'])
    def compose(self, point_a, point_b):
        r"""Compose two elements of SE(2) or SE(3).

        Parameters
        ----------
        point_a : array-like, shape=[..., dimension]
            Point of the group.
        point_b : array-like, shape=[..., dimension]
            Point of the group.

        Equation
        --------
        (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`)

        Returns
        -------
        composition : array-like, shape=[..., dimension]
            Composition of point_a and point_b.
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point_a = self.regularize(point_a)
        point_b = self.regularize(point_b)

        rot_vec_a = point_a[..., :dim_rotations]
        rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a)

        rot_vec_b = point_b[..., :dim_rotations]
        rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b)

        translation_a = point_a[..., dim_rotations:]
        translation_b = point_b[..., dim_rotations:]

        composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b)
        composition_rot_vec = rotations.rotation_vector_from_matrix(
            composition_rot_mat)

        composition_translation = gs.einsum('...j,...kj->...k', translation_b,
                                            rot_mat_a) + translation_a

        composition = gs.concatenate(
            (composition_rot_vec, composition_translation), axis=-1)
        return self.regularize(composition)

    @geomstats.vectorization.decorator(['else', 'vector'])
    def inverse(self, point):
        r"""Compute the group inverse in SE(n).

        Parameters
        ----------
        point: array-like, shape=[..., dimension]
            Point.

        Returns
        -------
        inverse_point : array-like, shape=[..., dimension]
            Inverted point.

        Notes
        -----
        :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))`
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point = self.regularize(point)

        rot_vec = point[:, :dim_rotations]
        translation = point[:, dim_rotations:]

        inverse_rotation = -rot_vec

        inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation)

        inverse_translation = gs.einsum(
            'ni,nij->nj', -translation,
            gs.transpose(inv_rot_mat, axes=(0, 2, 1)))

        inverse_point = gs.concatenate([inverse_rotation, inverse_translation],
                                       axis=-1)
        return self.regularize(inverse_point)

    @geomstats.vectorization.decorator(['else', 'vector'])
    def exp_from_identity(self, tangent_vec):
        """Compute group exponential of the tangent vector at the identity.

        Parameters
        ----------
        tangent_vec: array-like, shape=[..., 3]
            Tangent vector at base point.

        Returns
        -------
        group_exp: array-like, shape=[..., 3]
            Group exponential of the tangent vectors computed
            at the identity.
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        rot_vec = tangent_vec[..., :dim_rotations]
        rot_vec_regul = self.rotations.regularize(rot_vec)
        rot_vec_regul = gs.to_ndarray(rot_vec_regul, to_ndim=2, axis=1)

        transform = self._exp_translation_transform(rot_vec_regul)

        translation = tangent_vec[..., dim_rotations:]
        exp_translation = gs.einsum('ijk, ik -> ij', transform, translation)

        group_exp = gs.concatenate([rot_vec, exp_translation], axis=1)

        group_exp = self.regularize(group_exp)
        return group_exp

    @geomstats.vectorization.decorator(['else', 'vector'])
    def log_from_identity(self, point):
        """Compute the group logarithm of the point at the identity.

        Parameters
        ----------
        point: array-like, shape=[..., 3]
            Point.

        Returns
        -------
        group_log: array-like, shape=[..., 3]
            Group logarithm in the Lie algebra.
        """
        point = self.regularize(point)

        rotations = self.rotations
        dim_rotations = rotations.dim

        rot_vec = point[:, :dim_rotations]

        transform = self._log_translation_transform(rot_vec)

        translation = point[:, dim_rotations:]

        log_translation = gs.einsum('ijk, ik -> ij', transform, translation)

        return gs.concatenate([rot_vec, log_translation], axis=1)

    def random_uniform(self, n_samples=1):
        """Sample in SE(n) with the uniform distribution.

        Parameters
        ----------
        n_samples : int
            Number of samples.
            Optional, default: 1

        Returns
        -------
        random_point : array-like, shape=[..., dimension]
            Sample.
        """
        random_translation = self.translations.random_uniform(n_samples)
        random_rot_vec = self.rotations.random_uniform(n_samples)
        return gs.concatenate([random_rot_vec, random_translation], axis=-1)
Пример #5
0
class TestEuclideanMethods(geomstats.tests.TestCase):
    def setUp(self):
        gs.random.seed(1234)

        self.dimension = 2
        self.space = Euclidean(self.dimension)
        self.metric = self.space.metric

        self.n_samples = 3

        self.one_point_a = gs.array([0., 1.])
        self.one_point_b = gs.array([2., 10.])
        self.n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        self.n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

    def test_random_uniform_and_belongs(self):
        point = self.space.random_uniform()
        result = self.space.belongs(point)
        expected = True

        self.assertAllClose(result, expected)

    def test_squared_norm_vectorization(self):
        n_samples = self.n_samples
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        result = self.metric.squared_norm(n_points)

        expected = gs.array([5., 20., 26.])

        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

    def test_norm_vectorization_single_sample(self):
        one_point = gs.array([[0., 1.]])

        result = self.metric.norm(one_point)
        expected = gs.array([1.])
        self.assertAllClose(gs.shape(result), (1,))
        self.assertAllClose(result, expected)

        one_point = gs.array([0., 1.])

        result = self.metric.norm(one_point)
        expected = 1.
        self.assertAllClose(gs.shape(result), ())
        self.assertAllClose(result, expected)

    def test_norm_vectorization_n_samples(self):
        n_samples = self.n_samples
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])

        result = self.metric.norm(n_points)

        expected = gs.array([2.2360679775, 4.472135955, 5.09901951359])

        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

    def test_exp_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension

        one_tangent_vec = gs.array([0., 1.])
        one_base_point = gs.array([2., 10.])
        n_tangent_vecs = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_base_points = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.exp(one_tangent_vec, one_base_point)
        expected = one_tangent_vec + one_base_point

        self.assertAllClose(result, expected)

        result = self.metric.exp(n_tangent_vecs, one_base_point)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.exp(one_tangent_vec, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.exp(n_tangent_vecs, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

    def test_log_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension

        one_point = gs.array([0., 1.])
        one_base_point = gs.array([2., 10.])
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_base_points = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.log(one_point, one_base_point)
        expected = one_point - one_base_point
        self.assertAllClose(result, expected)

        result = self.metric.log(n_points, one_base_point)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.log(one_point, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.log(n_points, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

    def test_squared_dist_vectorization(self):
        n_samples = self.n_samples

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.squared_dist(one_point_a, one_point_b)
        vec = one_point_a - one_point_b
        expected = gs.dot(vec, gs.transpose(vec))
        self.assertAllClose(result, expected)

        result = self.metric.squared_dist(n_points_a, one_point_b)
        self.assertAllClose(gs.shape(result), (n_samples,))

        result = self.metric.squared_dist(one_point_a, n_points_b)
        self.assertAllClose(gs.shape(result), (n_samples,))

        result = self.metric.squared_dist(n_points_a, n_points_b)
        expected = gs.array([81., 109., 29.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

    def test_dist_vectorization(self):
        n_samples = self.n_samples

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.dist(one_point_a, one_point_b)
        vec = one_point_a - one_point_b
        expected = gs.sqrt(gs.dot(vec, gs.transpose(vec)))
        self.assertAllClose(result, expected)

        result = self.metric.dist(n_points_a, one_point_b)
        self.assertAllClose(gs.shape(result), (n_samples,))

        result = self.metric.dist(one_point_a, n_points_b)
        self.assertAllClose(gs.shape(result), (n_samples,))

        result = self.metric.dist(n_points_a, n_points_b)
        expected = gs.array([9., gs.sqrt(109.), gs.sqrt(29.)])

        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

    def test_belongs(self):
        point = gs.array([0., 1.])

        result = self.space.belongs(point)
        expected = True

        self.assertAllClose(result, expected)

    def test_random_uniform(self):
        result = self.space.random_uniform()

        self.assertAllClose(gs.shape(result), (self.dimension,))

    def test_inner_product_matrix(self):
        result = self.metric.inner_product_matrix()

        expected = gs.eye(self.dimension)

        self.assertAllClose(result, expected)

    def test_inner_product(self):
        point_a = gs.array([0., 1.])
        point_b = gs.array([2., 10.])

        result = self.metric.inner_product(point_a, point_b)
        expected = 10.

        self.assertAllClose(result, expected)

    def test_inner_product_vectorization_single_sample(self):
        one_point_a = gs.array([[0., 1.]])
        one_point_b = gs.array([[2., 10.]])

        result = self.metric.inner_product(one_point_a, one_point_b)
        expected = gs.array([10.])
        self.assertAllClose(gs.shape(result), (1,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([[0., 1.]])
        one_point_b = gs.array([2., 10.])

        result = self.metric.inner_product(one_point_a, one_point_b)
        expected = gs.array([10.])
        self.assertAllClose(gs.shape(result), (1,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([[2., 10.]])

        result = self.metric.inner_product(one_point_a, one_point_b)
        expected = gs.array([10.])
        self.assertAllClose(gs.shape(result), (1,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])

        result = self.metric.inner_product(one_point_a, one_point_b)
        expected = 10.
        self.assertAllClose(gs.shape(result), ())
        self.assertAllClose(result, expected)

    def test_inner_product_vectorization_n_samples(self):
        n_samples = 3
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])

        result = self.metric.inner_product(n_points_a, one_point_b)
        expected = gs.array([14., -44., 0.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(one_point_a, n_points_b)
        expected = gs.array([10., -1., 6.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, n_points_b)
        expected = gs.array([14., -12., 21.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([[0., 1.]])
        one_point_b = gs.array([[2., 10]])

        result = self.metric.inner_product(n_points_a, one_point_b)
        expected = gs.array([14., -44., 0.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(one_point_a, n_points_b)
        expected = gs.array([10., -1., 6.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, n_points_b)
        expected = gs.array([14., -12., 21.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([[0., 1.]])
        one_point_b = gs.array([2., 10.])

        result = self.metric.inner_product(n_points_a, one_point_b)
        expected = gs.array([14., -44., 0.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(one_point_a, n_points_b)
        expected = gs.array([10., -1., 6.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, n_points_b)
        expected = gs.array([14., -12., 21.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([[2., 10.]])

        result = self.metric.inner_product(n_points_a, one_point_b)
        expected = gs.array([14., -44., 0.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(one_point_a, n_points_b)
        expected = gs.array([10., -1., 6.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, n_points_b)
        expected = gs.array([14., -12., 21.])
        self.assertAllClose(gs.shape(result), (n_samples,))
        self.assertAllClose(result, expected)

    def test_squared_norm(self):
        point = gs.array([-2., 4.])

        result = self.metric.squared_norm(point)
        expected = 20.

        self.assertAllClose(result, expected)

    def test_norm(self):
        point = gs.array([-2., 4.])
        result = self.metric.norm(point)
        expected = 4.472135955

        self.assertAllClose(result, expected)

    def test_exp(self):
        base_point = gs.array([0., 1.])
        vector = gs.array([2., 10.])

        result = self.metric.exp(tangent_vec=vector,
                                 base_point=base_point)
        expected = base_point + vector

        self.assertAllClose(result, expected)

    def test_log(self):
        base_point = gs.array([0., 1.])
        point = gs.array([2., 10.])

        result = self.metric.log(point=point, base_point=base_point)
        expected = point - base_point

        self.assertAllClose(result, expected)

    def test_squared_dist(self):
        point_a = gs.array([-1., 4.])
        point_b = gs.array([1., 1.])

        result = self.metric.squared_dist(point_a, point_b)
        vec = point_b - point_a
        expected = gs.dot(vec, vec)

        self.assertAllClose(result, expected)

    def test_dist(self):
        point_a = gs.array([0., 1.])
        point_b = gs.array([2., 10.])

        result = self.metric.dist(point_a, point_b)
        expected = gs.linalg.norm(point_b - point_a)

        self.assertAllClose(result, expected)

    def test_geodesic_and_belongs(self):
        n_geodesic_points = 100
        initial_point = gs.array([[2., -1.]])
        initial_tangent_vec = gs.array([2., 0.])
        geodesic = self.metric.geodesic(
            initial_point=initial_point,
            initial_tangent_vec=initial_tangent_vec)

        t = gs.linspace(start=0., stop=1., num=n_geodesic_points)
        points = geodesic(t)

        result = self.space.belongs(points)
        expected = gs.array(n_geodesic_points * [True])

        self.assertAllClose(expected, result)
Пример #6
0
class _SpecialEuclidean3Vectors(LieGroup):
    """Class for the special euclidean group in 3d, SE(3).

    i.e. the Lie group of rigid transformations. Elements of SE(3) can either
    be represented as vectors (in 3d) or as matrices in general. The matrix
    representation corresponds to homogeneous coordinates.This class is
    specific to the vector representation of rotations. For the matrix
    representation use the SpecialEuclidean class and set `n=3`.

    Parameter
    ---------
    epsilon : float, optional (defaults to 0)
        Precision to use for calculations involving potential
        division by 0 in rotations.
    """
    def __init__(self, epsilon=0.):
        super(_SpecialEuclidean3Vectors,
              self).__init__(dim=6, default_point_type='vector')

        self.n = 3
        self.epsilon = epsilon
        self.rotations = SpecialOrthogonal(n=3,
                                           point_type='vector',
                                           epsilon=epsilon)
        self.translations = Euclidean(dim=3)

    def get_identity(self, point_type=None):
        """Get the identity of the group.

        Parameters
        ----------
        point_type : str, {'vector', 'matrix'}, optional
            The point_type of the returned value.
            default: self.default_point_type

        Returns
        -------
        identity : array-like, shape={[dim], [n + 1, n + 1]}
        """
        if point_type is None:
            point_type = self.default_point_type
        identity = gs.zeros(self.dim)
        if point_type == 'matrix':
            identity = gs.eye(self.n + 1)
        return identity

    identity = property(get_identity)

    def get_point_type_shape(self, point_type=None):
        """Get the shape of the instance given the default_point_style."""
        return self.get_identity(point_type).shape

    def belongs(self, point):
        """Evaluate if a point belongs to SE(3).

        Parameters
        ----------
        point : array-like, shape=[..., 3]
            The point of which to check whether it belongs to SE(3).

        Returns
        -------
        belongs : array-like, shape=[..., 1]
            Boolean indicating whether point belongs to SE(3).
        """
        point_dim = point.shape[-1]
        point_ndim = point.ndim
        belongs = gs.logical_and(point_dim == self.dim, point_ndim < 3)

        belongs = gs.logical_and(belongs,
                                 self.rotations.belongs(point[..., :self.n]))
        return belongs

    def regularize(self, point):
        """Regularize a point to the default representation for SE(n).

        Parameters
        ----------
        point : array-like, shape=[..., 3]
            The point to regularize.

        Returns
        -------
        point : array-like, shape=[..., 3]
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        rot_vec = point[..., :dim_rotations]
        regularized_rot_vec = rotations.regularize(rot_vec)

        translation = point[..., dim_rotations:]

        return gs.concatenate([regularized_rot_vec, translation], axis=-1)

    @geomstats.vectorization.decorator(['else', 'vector', 'else'])
    def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None):
        """Regularize a tangent vector at the identity.

        Parameters
        ----------
        tangent_vec: array-like, shape=[..., 3]
        metric : RiemannianMetric, optional

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        return self.regularize_tangent_vec(tangent_vec, self.identity, metric)

    def regularize_tangent_vec(self, tangent_vec, base_point, metric=None):
        """Regularize a tangent vector at a base point.

        Parameters
        ----------
        tangent_vec: array-like, shape=[..., 3]
        base_point : array-like, shape=[..., 3]
        metric : RiemannianMetric, optional
            default: self.left_canonical_metric

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        if metric is None:
            metric = self.left_canonical_metric

        rotations = self.rotations
        dim_rotations = rotations.dim

        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)

        rotations_vec = rotations.regularize_tangent_vec(
            tangent_vec=rot_tangent_vec,
            base_point=rot_base_point,
            metric=rot_metric)

        return gs.concatenate(
            [rotations_vec, tangent_vec[..., dim_rotations:]], axis=-1)

    @geomstats.vectorization.decorator(['else', 'vector'])
    def matrix_from_vector(self, vec):
        """Convert point in vector point-type to matrix.

        Parameters
        ----------
        vec: array-like, shape=[..., 3]

        Returns
        -------
        mat: array-like, shape=[..., n+1, n+1]
        """
        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

    @geomstats.vectorization.decorator(['else', 'vector', 'vector'])
    def compose(self, point_a, point_b):
        r"""Compose two elements of SE(3).

        Parameters
        ----------
        point_a : array-like, shape=[..., 3]
            Point of the group.
        point_b : array-like, shape=[..., 3]
            Point of the group.

        Equation
        ---------
        (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`)

        Returns
        -------
        composition :
            The composition of point_a and point_b.

        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point_a = self.regularize(point_a)
        point_b = self.regularize(point_b)

        rot_vec_a = point_a[..., :dim_rotations]
        rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a)

        rot_vec_b = point_b[..., :dim_rotations]
        rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b)

        translation_a = point_a[..., dim_rotations:]
        translation_b = point_b[..., dim_rotations:]

        composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b)
        composition_rot_vec = rotations.rotation_vector_from_matrix(
            composition_rot_mat)

        composition_translation = gs.einsum('...j,...kj->...k', translation_b,
                                            rot_mat_a) + translation_a

        composition = gs.concatenate(
            (composition_rot_vec, composition_translation), axis=-1)
        return self.regularize(composition)

    @geomstats.vectorization.decorator(['else', 'vector'])
    def inverse(self, point):
        r"""Compute the group inverse in SE(n).

        Parameters
        ----------
        point: array-like, shape=[..., 3]

        Returns
        -------
        inverse_point : array-like, shape=[..., 3]
            The inverted point.

        Notes
        -----
        :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))`
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point = self.regularize(point)

        rot_vec = point[:, :dim_rotations]
        translation = point[:, dim_rotations:]

        inverse_rotation = -rot_vec

        inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation)

        inverse_translation = gs.einsum(
            'ni,nij->nj', -translation,
            gs.transpose(inv_rot_mat, axes=(0, 2, 1)))

        inverse_point = gs.concatenate([inverse_rotation, inverse_translation],
                                       axis=-1)
        return self.regularize(inverse_point)

    @geomstats.vectorization.decorator(['else', 'vector', 'else'])
    def jacobian_translation(self, point, left_or_right='left'):
        """Compute the Jacobian matrix resulting from translation.

        Compute the matrix of the differential of the left/right translations
        from the identity to point in SE(3).

        Parameters
        ----------
        point: array-like, shape=[..., 3]
        left_or_right: str, {'left', 'right'}, optional
            Whether to compute the jacobian of the left or right translation.

        Returns
        -------
        jacobian : array-like, shape=[..., 3]
            The jacobian of the left / right translation.
        """
        if left_or_right not in ('left', 'right'):
            raise ValueError('`left_or_right` must be `left` or `right`.')

        rotations = self.rotations
        translations = self.translations
        dim_rotations = rotations.dim
        dim_translations = translations.dim

        point = self.regularize(point)

        n_points, _ = point.shape

        rot_vec = point[:, :dim_rotations]

        jacobian_rot = self.rotations.jacobian_translation(
            point=rot_vec, left_or_right=left_or_right)
        block_zeros_1 = gs.zeros((n_points, dim_rotations, dim_translations))
        jacobian_block_line_1 = gs.concatenate([jacobian_rot, block_zeros_1],
                                               axis=2)

        if left_or_right == 'left':
            rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec)
            jacobian_trans = rot_mat
            block_zeros_2 = gs.zeros(
                (n_points, dim_translations, dim_rotations))
            jacobian_block_line_2 = gs.concatenate(
                [block_zeros_2, jacobian_trans], axis=2)

        else:
            inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec)
            eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3)
            eye = gs.tile(eye, [n_points, 1, 1])
            jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye], axis=2)

        return gs.concatenate([jacobian_block_line_1, jacobian_block_line_2],
                              axis=1)

    @geomstats.vectorization.decorator(['else', 'vector'])
    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

    @geomstats.vectorization.decorator(['else', 'vector'])
    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 random_uniform(self, n_samples=1):
        """Sample in SE(3) with the uniform distribution.

        Parameters
        ----------
        n_samples : int, optional
            default : 1

        Returns
        -------
        random_point : array-like, shape=[..., 3]
            An array of random elements in SE(3) having the given.
        """
        random_translation = self.translations.random_uniform(n_samples)
        random_rot_vec = self.rotations.random_uniform(n_samples)
        return gs.concatenate([random_rot_vec, random_translation], axis=-1)

    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 : The matrix exponential of rot_vec
        """
        # TODO(nguigs): find usecase for this method
        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 = 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
Пример #7
0
class TestEuclideanMethods(geomstats.tests.TestCase):
    def setUp(self):
        gs.random.seed(1234)

        self.dimension = 2
        self.space = Euclidean(self.dimension)
        self.metric = self.space.metric

        self.n_samples = 3

        self.one_point_a = gs.array([0., 1.])
        self.one_point_b = gs.array([2., 10.])
        self.n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        self.n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

    def test_random_uniform_and_belongs(self):
        point = self.space.random_uniform()
        result = self.space.belongs(point)
        expected = gs.array([[True]])

        self.assertAllClose(result, expected)

    def test_squared_norm_vectorization(self):
        n_samples = self.n_samples
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        result = self.metric.squared_norm(n_points)

        expected = gs.array([[5.], [20.], [26.]])

        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

    def test_norm_vectorization(self):
        n_samples = self.n_samples
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])

        result = self.metric.norm(n_points)

        expected = gs.array([[2.2360679775], [4.472135955], [5.09901951359]])

        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

    def test_exp_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension

        one_tangent_vec = gs.array([0., 1.])
        one_base_point = gs.array([2., 10.])
        n_tangent_vecs = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_base_points = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.exp(one_tangent_vec, one_base_point)
        expected = one_tangent_vec + one_base_point
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

        result = self.metric.exp(n_tangent_vecs, one_base_point)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.exp(one_tangent_vec, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.exp(n_tangent_vecs, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

    def test_log_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension

        one_point = gs.array([0., 1.])
        one_base_point = gs.array([2., 10.])
        n_points = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_base_points = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.log(one_point, one_base_point)
        expected = one_point - one_base_point
        expected = helper.to_vector(expected)
        self.assertAllClose(result, expected)

        result = self.metric.log(n_points, one_base_point)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.log(one_point, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

        result = self.metric.log(n_points, n_base_points)
        self.assertAllClose(gs.shape(result), (n_samples, dim))

    def test_squared_dist_vectorization(self):
        n_samples = self.n_samples

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.squared_dist(one_point_a, one_point_b)
        vec = one_point_a - one_point_b
        expected = gs.dot(vec, gs.transpose(vec))
        expected = helper.to_scalar(expected)
        self.assertAllClose(result, expected)

        result = self.metric.squared_dist(n_points_a, one_point_b)
        self.assertAllClose(gs.shape(result), (n_samples, 1))

        result = self.metric.squared_dist(one_point_a, n_points_b)
        self.assertAllClose(gs.shape(result), (n_samples, 1))

        result = self.metric.squared_dist(n_points_a, n_points_b)
        expected = gs.array([[81.], [109.], [29.]])
        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

    def test_dist_vectorization(self):
        n_samples = self.n_samples

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.dist(one_point_a, one_point_b)
        vec = one_point_a - one_point_b
        expected = gs.sqrt(gs.dot(vec, gs.transpose(vec)))
        expected = helper.to_scalar(expected)
        self.assertAllClose(result, expected)

        result = self.metric.dist(n_points_a, one_point_b)
        self.assertAllClose(gs.shape(result), (n_samples, 1))

        result = self.metric.dist(one_point_a, n_points_b)
        self.assertAllClose(gs.shape(result), (n_samples, 1))

        result = self.metric.dist(n_points_a, n_points_b)
        expected = gs.array([[9.], [gs.sqrt(109.)], [gs.sqrt(29.)]])

        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

    def test_belongs(self):
        point = gs.array([0., 1.])

        result = self.space.belongs(point)
        expected = gs.array([[True]])

        self.assertAllClose(result, expected)

    def test_random_uniform(self):
        result = self.space.random_uniform()

        self.assertAllClose(gs.shape(result), (1, self.dimension))

    def test_inner_product_matrix(self):
        result = self.metric.inner_product_matrix()

        expected = gs.eye(self.dimension)
        expected = helper.to_matrix(expected)

        self.assertAllClose(result, expected)

    def test_inner_product(self):
        point_a = gs.array([0., 1.])
        point_b = gs.array([2., 10.])

        result = self.metric.inner_product(point_a, point_b)
        expected = gs.array([[10.]])

        self.assertAllClose(result, expected)

    def test_inner_product_vectorization(self):
        n_samples = 3

        one_point_a = gs.array([0., 1.])
        one_point_b = gs.array([2., 10.])
        n_points_a = gs.array([
            [2., 1.],
            [-2., -4.],
            [-5., 1.]])
        n_points_b = gs.array([
            [2., 10.],
            [8., -1.],
            [-3., 6.]])

        result = self.metric.inner_product(one_point_a, one_point_b)
        expected = gs.array([[10.]])
        self.assertAllClose(gs.shape(result), (1, 1))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, one_point_b)
        expected = gs.array([[14.], [-44.], [0.]])
        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(one_point_a, n_points_b)
        expected = gs.array([[10.], [-1.], [6.]])
        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

        result = self.metric.inner_product(n_points_a, n_points_b)
        expected = gs.array([[14.], [-12.], [21.]])
        self.assertAllClose(gs.shape(result), (n_samples, 1))
        self.assertAllClose(result, expected)

    def test_squared_norm(self):
        point = gs.array([-2., 4.])

        result = self.metric.squared_norm(point)
        expected = gs.array([[20.]])

        self.assertAllClose(result, expected)

    def test_norm(self):
        point = gs.array([-2., 4.])
        result = self.metric.norm(point)
        expected = gs.array([[4.472135955]])

        self.assertAllClose(result, expected)

    def test_exp(self):
        base_point = gs.array([0., 1.])
        vector = gs.array([2., 10.])

        result = self.metric.exp(tangent_vec=vector,
                                 base_point=base_point)
        expected = base_point + vector
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    def test_log(self):
        base_point = gs.array([0., 1.])
        point = gs.array([2., 10.])

        result = self.metric.log(point=point, base_point=base_point)
        expected = point - base_point
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    def test_squared_dist(self):
        point_a = gs.array([-1., 4.])
        point_b = gs.array([1., 1.])

        result = self.metric.squared_dist(point_a, point_b)
        vec = point_b - point_a
        expected = gs.dot(vec, vec)
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    def test_dist(self):
        point_a = gs.array([0., 1.])
        point_b = gs.array([2., 10.])

        result = self.metric.dist(point_a, point_b)
        expected = gs.linalg.norm(point_b - point_a)
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    def test_geodesic_and_belongs(self):
        n_geodesic_points = 100
        initial_point = gs.array([[2., -1.]])
        initial_tangent_vec = gs.array([2., 0.])
        geodesic = self.metric.geodesic(
            initial_point=initial_point,
            initial_tangent_vec=initial_tangent_vec)

        t = gs.linspace(start=0., stop=1., num=n_geodesic_points)
        points = geodesic(t)

        result = self.space.belongs(points)
        expected = gs.array(n_geodesic_points * [[True]])

        self.assertAllClose(expected, result)

    def test_mean(self):
        point = gs.array([[1., 4.]])
        result = self.metric.mean(points=[point, point, point])
        expected = point
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

        points = gs.array([
            [1., 2.],
            [2., 3.],
            [3., 4.],
            [4., 5.]])
        weights = gs.array([1., 2., 1., 2.])

        result = self.metric.mean(points, weights)
        expected = gs.array([16. / 6., 22. / 6.])
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    def test_variance(self):
        points = gs.array([
            [1., 2.],
            [2., 3.],
            [3., 4.],
            [4., 5.]])
        weights = gs.array([1., 2., 1., 2.])
        base_point = gs.zeros(2)
        result = self.metric.variance(points, weights, base_point)
        # we expect the average of the points' sq norms.
        expected = (1 * 5. + 2 * 13. + 1 * 25. + 2 * 41.) / 6.
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)
Пример #8
0
class SpecialEuclidean(LieGroup):
    """Class for the special euclidean group SE(n).

    i.e. the Lie group of rigid transformations.
    """
    def __init__(self, n, point_type=None, epsilon=0.):
        """Initiate an object of class SpecialEuclidean.

        Parameter
        ---------
        n : int
            the dimension of the euclidean space that SE(n) acts upon
        point_type : str, {'vector', 'matrix'}, optional
            whether to represent elmenents of SE(n) by vectors or matrices
            if None is given, point_type is set to 'vector' for dimension 3
            and 'matrix' otherwise
        epsilon : float, optional
            precision to use for calculations involving potential division by
            rotations
            default: 0
        """
        assert isinstance(n, int) and n > 1

        self.n = n
        self.dimension = int((n * (n - 1)) / 2 + n)

        self.epsilon = epsilon

        self.default_point_type = point_type
        if point_type is None:
            self.default_point_type = 'vector' if n == 3 else 'matrix'

        super(SpecialEuclidean, self).__init__(dimension=self.dimension)

        self.rotations = SpecialOrthogonal(n=n, epsilon=epsilon)
        self.translations = Euclidean(dimension=n)

    def get_identity(self, point_type=None):
        """Get the identity of the group.

        Parameters
        ----------
        point_type : str, {'vector', 'matrix'}, optional
            the point_type of the returned value
            default: self.default_point_type

        Returns
        -------
        identity : array-like, shape={[dimension], [n + 1, n + 1]}
        """
        if point_type is None:
            point_type = self.default_point_type

        identity = gs.zeros(self.dimension)
        if self.default_point_type == 'matrix':
            identity = gs.eye(self.n)
        return identity

    identity = property(get_identity)

    def belongs(self, point, point_type=None):
        """Evaluate if a point belongs to SE(n).

        Parameters
        ----------
        point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
            the point of which to check whether it belongs to SE(n)
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        belongs : array-like, shape=[n_samples, 1]
            array of booleans indicating whether point belongs to 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)
            n_points, point_dim = point.shape
            belongs = point_dim == self.dimension
            belongs = gs.to_ndarray(belongs, to_ndim=1)
            belongs = gs.to_ndarray(belongs, to_ndim=2, axis=1)
            belongs = gs.tile(belongs, (n_points, 1))
        elif point_type == 'matrix':
            point = gs.to_ndarray(point, to_ndim=3)
            raise NotImplementedError()

        return belongs

    def regularize(self, point, point_type=None):
        """Regularize a point to the default representation for SE(n).

        Parameters
        ----------
        point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
            the point which should be regularized
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        """
        if point_type is None:
            point_type = self.default_point_type

        if point_type == 'vector':
            point = gs.to_ndarray(point, to_ndim=2)

            rotations = self.rotations
            dim_rotations = rotations.dimension

            rot_vec = point[:, :dim_rotations]
            regularized_rot_vec = rotations.regularize(rot_vec,
                                                       point_type=point_type)

            translation = point[:, dim_rotations:]

            regularized_point = gs.concatenate(
                [regularized_rot_vec, translation], axis=1)

        elif point_type == 'matrix':
            point = gs.to_ndarray(point, to_ndim=3)
            regularized_point = gs.copy(point)

        return regularized_point

    def regularize_tangent_vec_at_identity(self,
                                           tangent_vec,
                                           metric=None,
                                           point_type=None):
        """Regularize a tangent vector at the identity.

        Parameters
        ----------
        tangent_vec: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        metric : RiemannianMetric, optional
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        if point_type is None:
            point_type = self.default_point_type

        return self.regularize_tangent_vec(tangent_vec,
                                           self.identity,
                                           metric,
                                           point_type=point_type)

    def regularize_tangent_vec(self,
                               tangent_vec,
                               base_point,
                               metric=None,
                               point_type=None):
        """Regularize a tangent vector at a base point.

        Parameters
        ----------
        tangent_vec: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        base_point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        metric : RiemannianMetric, optional
            default: self.left_canonical_metric
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        if point_type is None:
            point_type = self.default_point_type

        if metric is None:
            metric = self.left_canonical_metric

        if point_type == 'vector':
            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)
            rotations_vec = rotations.regularize_tangent_vec(
                tangent_vec=rot_tangent_vec,
                base_point=rot_base_point,
                metric=rot_metric,
                point_type=point_type)

            regularized_vec = gs.concatenate(
                [rotations_vec, tangent_vec[:, dim_rotations:]], axis=1)

        elif point_type == 'matrix':
            regularized_vec = tangent_vec

        return regularized_vec

    def compose(self, point_1, point_2, point_type=None):
        r"""Compose two elements of SE(n).

        Parameters
        ----------
        point_1 : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        point_2 : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Equation
        ---------
        (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`)

        Returns
        -------
        composition : the composition of point_1 and point_2

        """
        if point_type is None:
            point_type = self.default_point_type

        rotations = self.rotations
        dim_rotations = rotations.dimension

        point_1 = self.regularize(point_1, point_type=point_type)
        point_2 = self.regularize(point_2, point_type=point_type)

        if point_type == 'vector':
            n_points_1, _ = point_1.shape
            n_points_2, _ = point_2.shape

            assert (point_1.shape == point_2.shape or n_points_1 == 1
                    or n_points_2 == 1)

            if n_points_1 == 1:
                point_1 = gs.stack([point_1[0]] * n_points_2)

            if n_points_2 == 1:
                point_2 = gs.stack([point_2[0]] * n_points_1)

            rot_vec_1 = point_1[:, :dim_rotations]
            rot_mat_1 = rotations.matrix_from_rotation_vector(rot_vec_1)

            rot_vec_2 = point_2[:, :dim_rotations]
            rot_mat_2 = rotations.matrix_from_rotation_vector(rot_vec_2)

            translation_1 = point_1[:, dim_rotations:]
            translation_2 = point_2[:, dim_rotations:]

            composition_rot_mat = gs.matmul(rot_mat_1, rot_mat_2)
            composition_rot_vec = rotations.rotation_vector_from_matrix(
                composition_rot_mat)

            composition_translation = gs.einsum('ij,ikj->ik', translation_2,
                                                rot_mat_1) + translation_1

            composition = gs.concatenate(
                (composition_rot_vec, composition_translation), axis=1)

        elif point_type == 'matrix':
            raise NotImplementedError()

        composition = self.regularize(composition, point_type=point_type)
        return composition

    def inverse(self, point, point_type=None):
        r"""Compute the group inverse in SE(n).

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]

        Formula
        -------
        :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))`

        Returns
        -------
        inverse_point : array-like,
            shape=[n_samples, {dimension, [n + 1, n + 1]}]
            the inverted point
        """
        if point_type is None:
            point_type = self.default_point_type

        rotations = self.rotations
        dim_rotations = rotations.dimension

        point = self.regularize(point)

        if point_type == 'vector':
            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.einsum(
                'ni,nij->nj', -translation,
                gs.transpose(inv_rot_mat, axes=(0, 2, 1)))

            inverse_point = gs.concatenate(
                [inverse_rotation, inverse_translation], axis=1)

        elif point_type == 'matrix':
            raise NotImplementedError()

        inverse_point = self.regularize(inverse_point, point_type=point_type)
        return inverse_point

    def jacobian_translation(self,
                             point,
                             left_or_right='left',
                             point_type=None):
        """Compute the Jacobian matrix resulting from translation.

        Compute the jacobian matrix of the differential
        of the left/right translations from the identity to point in SE(n).
        Currently only implemented for point_type == 'vector'.

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]

        left_or_right: str, {'left', 'right'}, optional
            default: 'left'
            whether to compute the jacobian of the left or right translation
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        jacobian : array-like, shape=[n_samples, dimension]
            The jacobian of the left / right translation
        """
        if point_type is None:
            point_type = self.default_point_type

        assert left_or_right in ('left', 'right')

        dim = self.dimension
        rotations = self.rotations
        translations = self.translations
        dim_rotations = rotations.dimension
        dim_translations = translations.dimension

        point = self.regularize(point, point_type=point_type)

        if point_type == 'vector':
            n_points, _ = point.shape

            rot_vec = point[:, :dim_rotations]

            jacobian = gs.zeros((n_points, ) + (dim, ) * 2)
            jacobian_rot = self.rotations.jacobian_translation(
                point=rot_vec,
                left_or_right=left_or_right,
                point_type=point_type)
            block_zeros_1 = gs.zeros(
                (n_points, dim_rotations, dim_translations))
            jacobian_block_line_1 = gs.concatenate(
                [jacobian_rot, block_zeros_1], axis=2)

            if left_or_right == 'left':
                rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec)
                jacobian_trans = rot_mat
                block_zeros_2 = gs.zeros(
                    (n_points, dim_translations, dim_rotations))
                jacobian_block_line_2 = gs.concatenate(
                    [block_zeros_2, jacobian_trans], axis=2)

            else:
                inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec)
                eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3)
                eye = gs.tile(eye, [n_points, 1, 1])
                jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye],
                                                       axis=2)

            jacobian = gs.concatenate(
                [jacobian_block_line_1, jacobian_block_line_2], axis=1)

            assert gs.ndim(jacobian) == 3

        elif point_type == 'matrix':
            raise NotImplementedError()

        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, {dimension, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        group_exp: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
            the group exponential of the tangent vectors calculated
            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
            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 += 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
        elif point_type == 'matrix':
            raise NotImplementedError()

    def log_from_identity(self, point, point_type=None):
        """Compute the group logarithm of the point at the identity.

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        group_log: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
            the group logarithm in the Lie algbra
        """
        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
            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 += mask_i_float * (translation_i + term_1_i +
                                                   term_2_i)

            group_log = gs.concatenate([rot_vec, log_translation], axis=1)

            assert gs.ndim(group_log) == 2

        elif point_type == 'matrix':
            raise NotImplementedError()

        return group_log

    def random_uniform(self, n_samples=1, point_type=None):
        """Sample in SE(n) with the uniform distribution.

        Parameters
        ----------
        n_samples: int, optional
            default: 1
        point_typ: str, {'vector', 'matrix'}, optional
            default: self.default_point_type


        Returns
        -------
        random_transfo: array-like,
            shape=[n_samples, {dimension, [n + 1, n + 1]}]
            an array of random elements in SE(n) having the given point_type
        """
        if point_type is None:
            point_type = self.default_point_type

        random_rot_vec = self.rotations.random_uniform(n_samples,
                                                       point_type=point_type)
        random_translation = self.translations.random_uniform(n_samples)

        if point_type == 'vector':
            random_transfo = gs.concatenate(
                [random_rot_vec, random_translation], axis=1)

        elif point_type == 'matrix':
            raise NotImplementedError()

        random_transfo = self.regularize(random_transfo, point_type=point_type)
        return random_transfo

    def exponential_matrix(self, rot_vec):
        """Compute exponential of rotation matrix represented by rot_vec.

        Parameters
        ----------
        rot_vec : array-like, shape=[n_samples, dimension]

        Returns
        -------
        exponential_mat: The matrix exponential of rot_vec
        """
        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 = 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
        assert exponential_mat.ndim == 3

        return exponential_mat

    def exponential_barycenter(self, points, weights=None, point_type=None):
        """Compute the group exponential barycenter in SE(n).

        Parameters
        ----------
        points: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}]
        weights: array-like, shape=[n_samples], optional
            default: weight 1 / n_samples for each point
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type


        Returns
        -------
        exp_bar: array-like, shape=[{dimension, [n + 1, n + 1]}]
            the exponential barycenter
        """
        if point_type is None:
            point_type = self.default_point_type

        n_points = points.shape[0]
        assert n_points > 0

        if weights is None:
            weights = gs.ones((n_points, 1))

        weights = gs.to_ndarray(weights, to_ndim=2, axis=1)
        n_weights, _ = weights.shape
        assert n_points == n_weights

        dim = self.dimension
        rotations = self.rotations
        dim_rotations = rotations.dimension

        if point_type == 'vector':
            rotation_vectors = points[:, :dim_rotations]
            translations = points[:, dim_rotations:dim]
            assert rotation_vectors.shape == (n_points, dim_rotations)
            assert translations.shape == (n_points, self.n)

            mean_rotation = rotations.exponential_barycenter(
                points=rotation_vectors, weights=weights)
            mean_rotation_mat = rotations.matrix_from_rotation_vector(
                mean_rotation)

            matrix = gs.zeros((1, ) + (self.n, ) * 2)
            translation_aux = gs.zeros((1, self.n))

            inv_rot_mats = rotations.matrix_from_rotation_vector(
                -rotation_vectors)
            matrix_aux = gs.matmul(mean_rotation_mat, inv_rot_mats)
            assert matrix_aux.shape == (n_points, ) + (dim_rotations, ) * 2

            vec_aux = rotations.rotation_vector_from_matrix(matrix_aux)
            matrix_aux = self.exponential_matrix(vec_aux)
            matrix_aux = gs.linalg.inv(matrix_aux)

            for i in range(n_points):
                matrix += weights[i] * matrix_aux[i]
                translation_aux += weights[i] * gs.dot(
                    gs.matmul(matrix_aux[i], inv_rot_mats[i]), translations[i])

            mean_translation = gs.dot(
                translation_aux,
                gs.transpose(gs.linalg.inv(matrix), axes=(0, 2, 1)))

            exp_bar = gs.zeros((1, dim))
            exp_bar[0, :dim_rotations] = mean_rotation
            exp_bar[0, dim_rotations:dim] = mean_translation

        elif point_type == 'matrix':
            vector_points = self.rotation_vector_from_matrix(points)
            vector_exp_bar = self.exponential_barycenter(vector_points,
                                                         weights,
                                                         point_type='vector')
            exp_bar = self.matrix_from_rotation_vector(vector_exp_bar)
        return exp_bar
Пример #9
0
class SpecialEuclidean(LieGroup):
    """Class for the special euclidean group SE(n).

    i.e. the Lie group of rigid transformations. Elements of SE(n) can either
    be represented as vectors (in 3d) or as matrices in general. The matrix
    representation corresponds to homogeneous coordinates.
    """
    def __init__(self, n, default_point_type=None, epsilon=0.):
        """Initiate an object of class SpecialEuclidean.

        Parameter
        ---------
        n : int
            the dimension of the euclidean space that SE(n) acts upon
        point_type : str, {'vector', 'matrix'}, optional
            whether to represent elmenents of SE(n) by vectors or matrices
            if None is given, point_type is set to 'vector' for dimension 3
            and 'matrix' otherwise
        epsilon : float, optional
            precision to use for calculations involving potential division by
            rotations
            default: 0
        """
        if not (isinstance(n, int) and n > 1):
            raise ValueError('n must be an integer > 1.')

        self.n = n
        self.dimension = int((n * (n - 1)) / 2 + n)

        self.epsilon = epsilon

        super(SpecialEuclidean,
              self).__init__(dim=self.dimension,
                             default_point_type=default_point_type)
        if default_point_type is None:
            self.default_point_type = 'vector' if n == 3 else 'matrix'

        self.rotations = SpecialOrthogonal(
            n=n, epsilon=epsilon, default_point_type=default_point_type)
        self.translations = Euclidean(dim=n)

    def get_identity(self, point_type=None):
        """Get the identity of the group.

        Parameters
        ----------
        point_type : str, {'vector', 'matrix'}, optional
            the point_type of the returned value
            default: self.default_point_type

        Returns
        -------
        identity : array-like, shape={[dim], [n + 1, n + 1]}
        """
        if point_type is None:
            point_type = self.default_point_type

        identity = gs.zeros(self.dimension)
        if point_type == 'matrix':
            identity = gs.eye(self.n + 1)
        return identity

    identity = property(get_identity)

    def get_point_type_shape(self, point_type=None):
        """Get the shape of the instance given the default_point_style."""
        return self.get_identity(point_type).shape

    @geomstats.vectorization.decorator(['else', 'point', 'point_type'])
    def belongs(self, point, point_type=None):
        """Evaluate if a point belongs to SE(n).

        Parameters
        ----------
        point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
            the point of which to check whether it belongs to SE(n)
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        belongs : array-like, shape=[n_samples, 1]
            array of booleans indicating whether point belongs to SE(n)
        """
        if point_type == 'vector':
            n_points, vec_dim = gs.shape(point)
            belongs = vec_dim == self.dimension

            belongs = gs.tile([belongs], (point.shape[0], ))

            belongs = gs.logical_and(belongs,
                                     self.rotations.belongs(point[:, :self.n]))
            return gs.flatten(belongs)
        if point_type == 'matrix':
            n_points, point_dim1, point_dim2 = point.shape
            belongs = (point_dim1 == point_dim2 == self.n + 1)
            belongs = [belongs] * n_points

            rotation = point[:, :self.n, :self.n]
            rot_belongs = self.rotations.belongs(rotation,
                                                 point_type=point_type)

            belongs = gs.logical_and(belongs, rot_belongs)

            last_line_except_last_term = point[:, self.n:, :-1]
            all_but_last_zeros = ~gs.any(last_line_except_last_term,
                                         axis=(1, 2))

            belongs = gs.logical_and(belongs, all_but_last_zeros)

            last_term = point[:, self.n:, self.n:]
            belongs = gs.logical_and(belongs,
                                     gs.all(last_term == 1, axis=(1, 2)))
            return gs.flatten(belongs)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'point_type'])
    def regularize(self, point, point_type=None):
        """Regularize a point to the default representation for SE(n).

        Parameters
        ----------
        point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
            the point which should be regularized
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        """
        if point_type == 'vector':
            rotations = self.rotations
            dim_rotations = rotations.dim

            rot_vec = point[:, :dim_rotations]
            regularized_rot_vec = rotations.regularize(rot_vec,
                                                       point_type=point_type)

            translation = point[:, dim_rotations:]

            return gs.concatenate([regularized_rot_vec, translation], axis=1)

        if point_type == 'matrix':
            return gs.to_ndarray(point, to_ndim=3)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'else', 'point_type'])
    def regularize_tangent_vec_at_identity(self,
                                           tangent_vec,
                                           metric=None,
                                           point_type=None):
        """Regularize a tangent vector at the identity.

        Parameters
        ----------
        tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        metric : RiemannianMetric, optional
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        if point_type == 'vector':
            return self.regularize_tangent_vec(tangent_vec,
                                               self.identity,
                                               metric,
                                               point_type=point_type)

        if point_type == 'matrix':
            translation_mask = gs.hstack(
                [gs.ones((self.n, ) * 2), 2 * gs.ones((self.n, 1))])
            translation_mask = gs.concatenate(
                [translation_mask, gs.zeros((1, self.n + 1))], axis=0)
            tangent_vec = tangent_vec * gs.where(translation_mask != 0.,
                                                 gs.array(1.), gs.array(0.))
            tangent_vec = (tangent_vec -
                           GeneralLinear.transpose(tangent_vec)) / 2.
            return tangent_vec * translation_mask

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(
        ['else', 'point', 'point', 'else', 'point_type'])
    def regularize_tangent_vec(self,
                               tangent_vec,
                               base_point,
                               metric=None,
                               point_type=None):
        """Regularize a tangent vector at a base point.

        Parameters
        ----------
        tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        base_point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        metric : RiemannianMetric, optional
            default: self.left_canonical_metric
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        regularized_vec : the regularized tangent vector
        """
        if metric is None:
            metric = self.left_canonical_metric

        if point_type == 'vector':
            rotations = self.rotations
            dim_rotations = rotations.dim

            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)

            rotations_vec = rotations.regularize_tangent_vec(
                tangent_vec=rot_tangent_vec,
                base_point=rot_base_point,
                metric=rot_metric,
                point_type=point_type)

            return gs.concatenate(
                [rotations_vec, tangent_vec[:, dim_rotations:]], axis=1)

        if point_type == 'matrix':
            tangent_vec_at_id = self.compose(self.inverse(base_point),
                                             tangent_vec)
            regularized = self.regularize_tangent_vec_at_identity(
                tangent_vec_at_id, point_type=point_type)
            return self.compose(base_point, regularized)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'vector'])
    def matrix_from_vector(self, vec):
        """Convert point in vector point-type to matrix.

        Parameters
        ----------
        vec: array-like, shape=[n_samples, dim]

        Returns
        -------
        mat: array-like, shape=[n_samples, {dim, [n+1, n+1]}]
        """
        vec = self.regularize(vec, point_type='vector')
        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

    @geomstats.vectorization.decorator(
        ['else', 'point', 'point', 'point_type'])
    def compose(self, point_a, point_b, point_type=None):
        r"""Compose two elements of SE(n).

        Parameters
        ----------
        point_1 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        point_2 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Equation
        ---------
        (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`)

        Returns
        -------
        composition : the composition of point_1 and point_2

        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point_a = self.regularize(point_a, point_type=point_type)
        point_b = self.regularize(point_b, point_type=point_type)

        if point_type == 'vector':
            n_points_a, _ = point_a.shape
            n_points_b, _ = point_b.shape

            if not (point_a.shape == point_b.shape or n_points_a == 1
                    or n_points_b == 1):
                raise ValueError()

            rot_vec_a = point_a[:, :dim_rotations]
            rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a)

            rot_vec_b = point_b[:, :dim_rotations]
            rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b)

            translation_a = point_a[:, dim_rotations:]
            translation_b = point_b[:, dim_rotations:]

            composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b)
            composition_rot_vec = rotations.rotation_vector_from_matrix(
                composition_rot_mat)

            composition_translation = gs.einsum(
                '...j,...kj->...k', translation_b, rot_mat_a) + translation_a

            composition = gs.concatenate(
                (composition_rot_vec, composition_translation), axis=-1)
            return self.regularize(composition, point_type=point_type)

        if point_type == 'matrix':
            return GeneralLinear.compose(point_a, point_b)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'point_type'])
    def inverse(self, point, point_type=None):
        r"""Compute the group inverse in SE(n).

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        inverse_point : array-like,
            shape=[n_samples, {dim, [n + 1, n + 1]}]
            the inverted point

        Notes
        -----
        :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))`
        """
        rotations = self.rotations
        dim_rotations = rotations.dim

        point = self.regularize(point)

        if point_type == 'vector':
            rot_vec = point[:, :dim_rotations]
            translation = point[:, dim_rotations:]

            inverse_rotation = -rot_vec

            inv_rot_mat = rotations.matrix_from_rotation_vector(
                inverse_rotation)

            inverse_translation = gs.einsum(
                'ni,nij->nj', -translation,
                gs.transpose(inv_rot_mat, axes=(0, 2, 1)))

            inverse_point = gs.concatenate(
                [inverse_rotation, inverse_translation], axis=-1)
            return self.regularize(inverse_point, point_type=point_type)

        if point_type == 'matrix':
            inv_rot = gs.transpose(point[:, :self.n, :self.n], axes=(0, 2, 1))
            inv_trans = gs.matmul(inv_rot, -point[:, :self.n, self.n:])
            last_line = point[:, self.n:, :]
            inverse_point = gs.concatenate((inv_rot, inv_trans), axis=2)
            return gs.concatenate((inverse_point, last_line), axis=1)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'else', 'point_type'])
    def jacobian_translation(self,
                             point,
                             left_or_right='left',
                             point_type=None):
        """Compute the Jacobian matrix resulting from translation.

        Compute the matrix of the differential
        of the left/right translations from the identity to point in SE(n).

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]

        left_or_right: str, {'left', 'right'}, optional
            default: 'left'
            whether to compute the jacobian of the left or right translation
        point_type : str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        jacobian : array-like, shape=[n_samples, dim]
            The jacobian of the left / right translation
        """
        if point_type is None:
            point_type = self.default_point_type

        if left_or_right not in ('left', 'right'):
            raise ValueError('`left_or_right` must be `left` or `right`.')

        rotations = self.rotations
        translations = self.translations
        dim_rotations = rotations.dim
        dim_translations = translations.dim

        point = self.regularize(point, point_type=point_type)

        if point_type == 'vector':
            n_points, _ = point.shape

            rot_vec = point[:, :dim_rotations]

            jacobian_rot = self.rotations.jacobian_translation(
                point=rot_vec,
                left_or_right=left_or_right,
                point_type=point_type)
            block_zeros_1 = gs.zeros(
                (n_points, dim_rotations, dim_translations))
            jacobian_block_line_1 = gs.concatenate(
                [jacobian_rot, block_zeros_1], axis=2)

            if left_or_right == 'left':
                rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec)
                jacobian_trans = rot_mat
                block_zeros_2 = gs.zeros(
                    (n_points, dim_translations, dim_rotations))
                jacobian_block_line_2 = gs.concatenate(
                    [block_zeros_2, jacobian_trans], axis=2)

            else:
                inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec)
                eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3)
                eye = gs.tile(eye, [n_points, 1, 1])
                jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye],
                                                       axis=2)

            return gs.concatenate(
                [jacobian_block_line_1, jacobian_block_line_2], axis=1)

        if point_type == 'matrix':
            return point

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'point_type'])
    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\'.')

    @geomstats.vectorization.decorator(['else', 'point', 'point_type'])
    def log_from_identity(self, point, point_type=None):
        """Compute the group logarithm of the point at the identity.

        Parameters
        ----------
        point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        group_log: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
            the group logarithm in the Lie algbra
        """
        point = self.regularize(point, point_type=point_type)

        rotations = self.rotations
        dim_rotations = rotations.dim

        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
            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)

        if point_type == 'matrix':
            return GeneralLinear.log(point)

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    def random_uniform(self, n_samples=1, point_type=None):
        """Sample in SE(n) with the uniform distribution.

        Parameters
        ----------
        n_samples: int, optional
            default: 1
        point_type: str, {'vector', 'matrix'}, optional
            default: self.default_point_type

        Returns
        -------
        random_point: array-like,
            shape=[n_samples, {dim, [n + 1, n + 1]}]
            An array of random elements in SE(n) having the given point_type.
        """
        if point_type is None:
            point_type = self.default_point_type

        random_translation = self.translations.random_uniform(n_samples)

        if point_type == 'vector':
            random_rot_vec = self.rotations.random_uniform(
                n_samples, point_type=point_type)
            return gs.concatenate([random_rot_vec, random_translation],
                                  axis=-1)

        if point_type == 'matrix':
            random_rotation = self.rotations.random_uniform(
                n_samples, point_type=point_type)
            random_rotation = gs.to_ndarray(random_rotation, to_ndim=3)

            random_translation = gs.to_ndarray(random_translation, to_ndim=2)
            random_translation = gs.transpose(
                gs.to_ndarray(random_translation, to_ndim=3, axis=1),
                (0, 2, 1))

            random_point = gs.concatenate(
                (random_rotation, random_translation), axis=2)
            last_line = gs.zeros((n_samples, 1, self.n + 1))
            random_point = gs.concatenate((random_point, last_line), axis=1)
            random_point = gs.assignment(random_point, 1, (-1, -1), axis=0)
            if gs.shape(random_point)[0] == 1:
                random_point = gs.squeeze(random_point, axis=0)
            return random_point

        raise ValueError('Invalid point_type, expected \'vector\' or '
                         '\'matrix\'.')

    def _exponential_matrix(self, rot_vec):
        """Compute exponential of rotation matrix represented by rot_vec.

        Parameters
        ----------
        rot_vec : array-like, shape=[n_samples, dim]

        Returns
        -------
        exponential_mat : The matrix exponential of rot_vec
        """
        # TODO(nguigs): find usecase for this method
        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 = 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
Пример #10
0
class _SpecialEuclideanMatrices(GeneralLinear, LieGroup):
    """Class for special Euclidean group.

    Parameters
    ----------
    n : int
        Integer dimension of the underlying Euclidean space. Matrices will
        be of size: (n+1) x (n+1).
    """

    def __init__(self, n):
        super().__init__(
            n=n + 1, dim=int((n * (n + 1)) / 2), default_point_type='matrix',
            lie_algebra=SpecialEuclideanMatrixLieAlgebra(n=n))
        self.rotations = SpecialOrthogonal(n=n)
        self.translations = Euclidean(dim=n)
        self.n = n

        self.left_canonical_metric = \
            SpecialEuclideanMatrixCannonicalLeftMetric(group=self)

    def get_identity(self):
        """Return the identity matrix."""
        return gs.eye(self.n + 1, self.n + 1)
    identity = property(get_identity)

    def belongs(self, point):
        """Check whether point is of the form rotation, translation.

        Parameters
        ----------
        point : array-like, shape=[..., n, n].
            Point to be checked.

        Returns
        -------
        belongs : array-like, shape=[...,]
            Boolean denoting if point belongs to the group.
        """
        point_dim1, point_dim2 = point.shape[-2:]
        belongs = (point_dim1 == point_dim2 == self.n + 1)

        rotation = point[..., :self.n, :self.n]
        rot_belongs = self.rotations.belongs(rotation)

        belongs = gs.logical_and(belongs, rot_belongs)

        last_line_except_last_term = point[..., self.n:, :-1]
        all_but_last_zeros = ~ gs.any(
            last_line_except_last_term, axis=(-2, -1))

        belongs = gs.logical_and(belongs, all_but_last_zeros)

        last_term = point[..., self.n, self.n]
        belongs = gs.logical_and(belongs, gs.isclose(last_term, 1.))

        if point.ndim == 2:
            return gs.squeeze(belongs)
        return gs.flatten(belongs)

    def random_uniform(self, n_samples=1, tol=1e-6):
        """Sample in SE(n) from the uniform distribution.

        Parameters
        ----------
        n_samples : int
            Number of samples.
            Optional, default: 1.
        tol : unused

        Returns
        -------
        samples : array-like, shape=[..., n + 1, n + 1]
            Sample in SE(n).
        """
        random_translation = self.translations.random_uniform(n_samples)
        random_rotation = self.rotations.random_uniform(n_samples)
        output_shape = (
            (n_samples, self.n + 1, self.n + 1) if n_samples != 1
            else (self.n + 1, ) * 2)
        random_point = homogeneous_representation(
            random_rotation, random_translation, output_shape)
        return random_point

    @classmethod
    def inverse(cls, point):
        """Return the inverse of a point.

        Parameters
        ----------
        point : array-like, shape=[..., n, n]
            Point to be inverted.
        """
        n = point.shape[-1] - 1
        transposed_rot = cls.transpose(point[..., :n, :n])
        translation = point[..., :n, -1]
        translation = gs.einsum(
            '...ij,...j->...i', transposed_rot, translation)
        return homogeneous_representation(
            transposed_rot, -translation, point.shape)