Ejemplo n.º 1
0
 def test_random_von_mises_one_sample_belongs(self):
     for dim in [2, 9]:
         sphere = Hypersphere(dim)
         point = sphere.random_von_mises_fisher()
         self.assertAllClose(point.shape, (dim + 1, ))
         result = sphere.belongs(point)
         self.assertTrue(result)
Ejemplo n.º 2
0
    def test_load_cities(self):
        """Test that the cities coordinates belong to the sphere."""
        sphere = Hypersphere(dim=2)
        data, _ = data_utils.load_cities()
        self.assertAllClose(gs.shape(data), (50, 3))

        tokyo = data[0]
        self.assertAllClose(tokyo,
                            gs.array([0.61993792, -0.52479018, 0.58332859]))

        result = sphere.belongs(data)
        self.assertTrue(gs.all(result))
    def test_fit_init_random_sphere(self):
        """Test fitting data into a GMM."""
        space = Hypersphere(2)
        gmm_learning = RiemannianEM(
            metric=space.metric,
            n_gaussians=2,
            initialisation_method=self.initialisation_method,
        )

        means = space.random_uniform(2)
        cluster_1 = space.random_von_mises_fisher(mu=means[0], kappa=20, n_samples=140)
        cluster_2 = space.random_von_mises_fisher(mu=means[1], kappa=20, n_samples=140)

        data = gs.concatenate((cluster_1, cluster_2), axis=0)
        means, variances, coefficients = gmm_learning.fit(data)

        self.assertTrue((coefficients < 1).all() and (coefficients > 0).all())
        self.assertTrue((variances < 1).all() and (variances > 0).all())
        self.assertTrue(space.belongs(means).all())
Ejemplo n.º 4
0
class TestHypersphereMethods(geomstats.tests.TestCase):
    def setUp(self):
        gs.random.seed(1234)

        self.dimension = 4
        self.space = Hypersphere(dimension=self.dimension)
        self.metric = self.space.metric
        self.n_samples = 10

    @geomstats.tests.np_and_pytorch_only
    def test_random_uniform_and_belongs(self):
        """
        Test that the random uniform method samples
        on the hypersphere space.
        """
        n_samples = self.n_samples
        point = self.space.random_uniform(n_samples)
        result = self.space.belongs(point)
        expected = gs.array([[True]] * n_samples)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_random_uniform(self):
        point = self.space.random_uniform()

        self.assertAllClose(gs.shape(point), (1, self.dimension + 1))

    def test_projection_and_belongs(self):
        point = gs.array([1., 2., 3., 4., 5.])
        proj = self.space.projection(point)
        result = self.space.belongs(proj)
        expected = gs.array([[True]])

        self.assertAllClose(expected, result)

    def test_intrinsic_and_extrinsic_coords(self):
        """
        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        point_int = gs.array([.1, 0., 0., .1])
        point_ext = self.space.intrinsic_to_extrinsic_coords(point_int)
        result = self.space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

        point_ext = (1. / (gs.sqrt(6.)) * gs.array([1., 0., 0., 1., 2.]))
        point_int = self.space.extrinsic_to_intrinsic_coords(point_ext)
        result = self.space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    def test_intrinsic_and_extrinsic_coords_vectorization(self):
        """
        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        point_int = gs.array([[.1, 0., 0., .1], [.1, .1, .1, .4],
                              [.1, .3, 0., .1], [-0.1, .1, -.4, .1],
                              [0., 0., .1, .1], [.1, .1, .1, .1]])
        point_ext = self.space.intrinsic_to_extrinsic_coords(point_int)
        result = self.space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

        point_int = self.space.extrinsic_to_intrinsic_coords(point_ext)
        result = self.space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_log_and_exp_general_case(self):
        """
        Test that the riemannian exponential
        and the riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # General case
        base_point = gs.array([1., 2., 3., 4., 6.])
        base_point = base_point / gs.linalg.norm(base_point)
        point = gs.array([0., 5., 6., 2., -1.])
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected, atol=1e-6)

    @geomstats.tests.np_and_pytorch_only
    def test_log_and_exp_edge_case(self):
        """
        Test that the riemannian exponential
        and the riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # Edge case: two very close points, base_point_2 and point_2,
        # form an angle < epsilon
        base_point = gs.array([1., 2., 3., 4., 6.])
        base_point = base_point / gs.linalg.norm(base_point)
        point = (base_point + 1e-12 * gs.array([-1., -2., 1., 1., .1]))
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_exp_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_vec = self.space.random_uniform()
        one_base_point = self.space.random_uniform()
        n_vecs = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

        one_tangent_vec = self.space.projection_to_tangent_space(
            one_vec, base_point=one_base_point)
        result = self.metric.exp(one_tangent_vec, one_base_point)

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

        n_tangent_vecs = self.space.projection_to_tangent_space(
            n_vecs, base_point=one_base_point)
        result = self.metric.exp(n_tangent_vecs, one_base_point)

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

        one_tangent_vec = self.space.projection_to_tangent_space(
            one_vec, base_point=n_base_points)
        result = self.metric.exp(one_tangent_vec, n_base_points)

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

        n_tangent_vecs = self.space.projection_to_tangent_space(
            n_vecs, base_point=n_base_points)
        result = self.metric.exp(n_tangent_vecs, n_base_points)

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

    @geomstats.tests.np_and_pytorch_only
    def test_log_vectorization(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_base_point = self.space.random_uniform()
        one_point = self.space.random_uniform()
        n_points = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

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

    @geomstats.tests.np_and_pytorch_only
    def test_exp_and_log_and_projection_to_tangent_space_general_case(self):
        """
        Test that the riemannian exponential
        and the riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # TODO(nina): Fix that this test fails, also in numpy
        # Riemannian Exp then Riemannian Log
        # General case
        # NB: Riemannian log gives a regularized tangent vector,
        # so we take the norm modulo 2 * pi.
        base_point = gs.array([0., -3., 0., 3., 4.])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = gs.array([9., 5., 0., 0., -1.])
        vector = self.space.projection_to_tangent_space(vector=vector,
                                                        base_point=base_point)

        # exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        # result = self.metric.log(point=exp, base_point=base_point)

        expected = vector
        norm_expected = gs.linalg.norm(expected)
        regularized_norm_expected = gs.mod(norm_expected, 2 * gs.pi)
        expected = expected / norm_expected * regularized_norm_expected
        expected = helper.to_vector(expected)

    @geomstats.tests.np_and_pytorch_only
    def test_exp_and_log_and_projection_to_tangent_space_edge_case(self):
        """
        Test that the riemannian exponential
        and the riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Exp then Riemannian Log
        # Edge case: tangent vector has norm < epsilon
        base_point = gs.array([10., -2., -.5, 34., 3.])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = 1e-10 * gs.array([.06, -51., 6., 5., 3.])
        vector = self.space.projection_to_tangent_space(vector=vector,
                                                        base_point=base_point)

        exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        result = self.metric.log(point=exp, base_point=base_point)
        expected = self.space.projection_to_tangent_space(
            vector=vector, base_point=base_point)
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected, atol=1e-8)

    def test_squared_norm_and_squared_dist(self):
        """
        Test that the squared distance between two points is
        the squared norm of their logarithm.
        """
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = (1. / gs.sqrt(435.) * gs.array([1., -20., -5., 0., 3.]))
        log = self.metric.log(point=point_a, base_point=point_b)
        result = self.metric.squared_norm(vector=log)
        expected = self.metric.squared_dist(point_a, point_b)
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_squared_dist_vectorization(self):
        n_samples = self.n_samples

        one_point_a = self.space.random_uniform()
        one_point_b = self.space.random_uniform()
        n_points_a = self.space.random_uniform(n_samples=n_samples)
        n_points_b = self.space.random_uniform(n_samples=n_samples)

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

        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)
        self.assertAllClose(gs.shape(result), (n_samples, 1))

    def test_norm_and_dist(self):
        """
        Test that the distance between two points is
        the norm of their logarithm.
        """
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = (1. / gs.sqrt(435.) * gs.array([1., -20., -5., 0., 3.]))
        log = self.metric.log(point=point_a, base_point=point_b)
        result = self.metric.norm(vector=log)
        expected = self.metric.dist(point_a, point_b)
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    def test_dist_point_and_itself(self):
        # Distance between a point and itself is 0
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = point_a
        result = self.metric.dist(point_a, point_b)
        expected = 0.
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    def test_dist_orthogonal_points(self):
        # Distance between two orthogonal points is pi / 2.
        point_a = gs.array([10., -2., -.5, 0., 0.])
        point_a = point_a / gs.linalg.norm(point_a)
        point_b = gs.array([2., 10, 0., 0., 0.])
        point_b = point_b / gs.linalg.norm(point_b)
        result = gs.dot(point_a, point_b)
        result = helper.to_scalar(result)
        expected = 0
        expected = helper.to_scalar(expected)
        self.assertAllClose(result, expected)

        result = self.metric.dist(point_a, point_b)
        expected = gs.pi / 2
        expected = helper.to_scalar(expected)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_exp_and_dist_and_projection_to_tangent_space(self):
        base_point = gs.array([16., -2., -2.5, 84., 3.])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = gs.array([9., 0., -1., -2., 1.])
        tangent_vec = self.space.projection_to_tangent_space(
            vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)
        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec) % (2 * gs.pi)
        expected = helper.to_scalar(expected)
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_exp_and_dist_and_projection_to_tangent_space_vec(self):
        base_point = gs.array([[16., -2., -2.5, 84., 3.],
                               [16., -2., -2.5, 84., 3.]])

        base_single_point = gs.array([16., -2., -2.5, 84., 3.])
        scalar_norm = gs.linalg.norm(base_single_point)

        base_point = base_point / scalar_norm
        vector = gs.array([[9., 0., -1., -2., 1.], [9., 0., -1., -2., 1]])

        tangent_vec = self.space.projection_to_tangent_space(
            vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)

        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec, axis=-1) % (2 * gs.pi)

        expected = helper.to_scalar(expected)
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_geodesic_and_belongs(self):
        n_geodesic_points = 100
        initial_point = self.space.random_uniform()
        vector = gs.array([2., 0., -1., -2., 1.])
        initial_tangent_vec = self.space.projection_to_tangent_space(
            vector=vector, base_point=initial_point)
        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_inner_product(self):
        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([0., 0., 0., 0., 1.])
        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([[0.]])

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_variance(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.zeros((2, point.shape[0]))
        points[0, :] = point
        points[1, :] = point
        result = self.metric.variance(points)
        expected = helper.to_scalar(0.)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_mean(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.zeros((2, point.shape[0]))
        points[0, :] = point
        points[1, :] = point
        result = self.metric.mean(points)
        expected = helper.to_vector(point)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_only
    def test_adaptive_gradientdescent_mean(self):
        n_tests = 100
        result = gs.zeros(n_tests)
        expected = gs.zeros(n_tests)

        for i in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.space.random_uniform(n_samples=2)
            mean = self.metric.adaptive_gradientdescent_mean(points)
            logs = self.metric.log(point=points, base_point=mean)
            result[i] = gs.linalg.norm(logs[1, :] + logs[0, :])

        self.assertAllClose(expected, result, rtol=1e-10, atol=1e-10)

    @geomstats.tests.np_and_pytorch_only
    def test_mean_and_belongs(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.zeros((2, point_a.shape[0]))
        points[0, :] = point_a
        points[1, :] = point_b
        mean = self.metric.mean(points)
        result = self.space.belongs(mean)
        expected = gs.array([[True]])
        self.assertAllClose(result, expected)

    def test_diameter(self):
        dim = 2
        sphere = Hypersphere(dim)
        point_a = gs.array([[0., 0., 1.]])
        point_b = gs.array([[1., 0., 0.]])
        point_c = gs.array([[0., 0., -1.]])
        result = sphere.metric.diameter(gs.vstack((point_a, point_b, point_c)))
        expected = gs.pi
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_closest_neighbor_index(self):
        """
        Check that the closest neighbor is one of neighbors.
        """
        n_samples = 10
        points = self.space.random_uniform(n_samples=n_samples)
        point = points[0, :]
        neighbors = points[1:, :]
        index = self.metric.closest_neighbor_index(point, neighbors)
        closest_neighbor = points[index, :]

        test = gs.sum(gs.all(points == closest_neighbor, axis=1))
        result = test > 0
        self.assertTrue(result)

    @geomstats.tests.np_and_pytorch_only
    def test_sample_von_mises_fisher(self):
        """
        Check that the maximum likelihood estimates of the mean and
        concentration parameter are close to the real values. A first
        estimation of the concentration parameter is obtained by a
        closed-form expression and improved through the Newton method.
        """
        dim = 2
        n_points = 1000000
        sphere = Hypersphere(dim)

        # check mean value for concentrated distribution
        kappa = 10000000
        points = sphere.random_von_mises_fisher(kappa, n_points)
        sum_points = gs.sum(points, axis=0)
        mean = gs.array([0., 0., 1.])
        mean_estimate = sum_points / gs.linalg.norm(sum_points)
        expected = mean
        result = mean_estimate
        self.assertTrue(gs.allclose(result, expected,
                                    atol=MEAN_ESTIMATION_TOL))
        # check concentration parameter for dispersed distribution
        kappa = 1
        points = sphere.random_von_mises_fisher(kappa, n_points)
        sum_points = gs.sum(points, axis=0)
        mean_norm = gs.linalg.norm(sum_points) / n_points
        kappa_estimate = (mean_norm * (dim + 1. - mean_norm**2) /
                          (1. - mean_norm**2))
        kappa_estimate = gs.cast(kappa_estimate, gs.float64)
        p = dim + 1
        n_steps = 100
        for i in range(n_steps):
            bessel_func_1 = scipy.special.iv(p / 2., kappa_estimate)
            bessel_func_2 = scipy.special.iv(p / 2. - 1., kappa_estimate)
            ratio = bessel_func_1 / bessel_func_2
            denominator = 1. - ratio**2 - (p - 1.) * ratio / kappa_estimate
            mean_norm = gs.cast(mean_norm, gs.float64)
            kappa_estimate = kappa_estimate - (ratio - mean_norm) / denominator
        expected = kappa
        result = kappa_estimate
        self.assertTrue(
            gs.allclose(result, expected, atol=KAPPA_ESTIMATION_TOL))

    @geomstats.tests.np_and_pytorch_only
    def test_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        result = sphere.spherical_to_extrinsic(points_spherical)
        expected = gs.array([[1., 0., 0.],
                             [gs.sqrt(2) / 4,
                              gs.sqrt(2) / 4,
                              gs.sqrt(3) / 2]])
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_tangent_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates for tangent vectors to the
        2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        base_points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 2, 0]])
        tangent_vecs_spherical = gs.array([[0.25, 0.5], [0.3, 0.2]])
        result = sphere.tangent_spherical_to_extrinsic(tangent_vecs_spherical,
                                                       base_points_spherical)
        expected = gs.array([[0, 0.5, -0.25], [0, 0.2, -0.3]])
        self.assertAllClose(result, expected)

    def test_christoffels_vectorization(self):
        """
        Check vectorization of Christoffel symbols in
        spherical coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        christoffel = sphere.metric.christoffels(points_spherical)
        result = christoffel.shape
        expected = gs.array([2, dim, dim, dim])
        self.assertAllClose(result, expected)
Ejemplo n.º 5
0
class TestFrechetMean(geomstats.tests.TestCase):
    _multiprocess_can_split_ = True

    def setUp(self):
        gs.random.seed(123)
        self.sphere = Hypersphere(dim=4)
        self.hyperbolic = Hyperboloid(dim=3)
        self.euclidean = Euclidean(dim=2)
        self.minkowski = Minkowski(dim=2)
        self.so3 = SpecialOrthogonal(n=3, point_type='vector')
        self.so_matrix = SpecialOrthogonal(n=3)

    def test_logs_at_mean_default_gradient_descent_sphere(self):
        n_tests = 10
        estimator = FrechetMean(
            metric=self.sphere.metric, method='default', lr=1.)

        result = []
        for _ in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result.append(gs.linalg.norm(logs[1, :] + logs[0, :]))
        result = gs.stack(result)
        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result)

    def test_logs_at_mean_adaptive_gradient_descent_sphere(self):
        n_tests = 10
        estimator = FrechetMean(metric=self.sphere.metric, method='adaptive')

        result = []
        for _ in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result.append(gs.linalg.norm(logs[1, :] + logs[0, :]))
        result = gs.stack(result)

        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result)

    def test_estimate_shape_default_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(
            metric=self.sphere.metric, method='default', verbose=True)
        mean.fit(points)
        result = mean.estimate_

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

    def test_estimate_shape_adaptive_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(points)
        result = mean.estimate_

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

    def test_estimate_and_belongs_default_gradient_descent_sphere(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    def test_estimate_default_gradient_descent_so3(self):
        points = self.so3.random_uniform(2)

        mean_vec = FrechetMean(
            metric=self.so3.bi_invariant_metric, method='default', lr=1.)
        mean_vec.fit(points)

        logs = self.so3.bi_invariant_metric.log(points, mean_vec.estimate_)
        result = gs.sum(logs, axis=0)
        expected = gs.zeros_like(points[0])
        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_default_gradient_descent_so3(self):
        point = self.so3.random_uniform(10)

        mean_vec = FrechetMean(
            metric=self.so3.bi_invariant_metric, method='default')
        mean_vec.fit(point)

        result = self.so3.belongs(mean_vec.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_default_gradient_descent_so_matrix(self):
        points = self.so_matrix.random_uniform(2)
        mean_vec = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric, method='default',
            lr=1.)
        mean_vec.fit(points)
        logs = self.so_matrix.bi_invariant_metric.log(
            points, mean_vec.estimate_)
        result = gs.sum(logs, axis=0)
        expected = gs.zeros_like(points[0])

        self.assertAllClose(result, expected, atol=1e-5)

    @geomstats.tests.np_and_tf_only
    def test_estimate_and_belongs_default_gradient_descent_so_matrix(self):
        point = self.so_matrix.random_uniform(10)

        mean = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric, method='default')
        mean.fit(point)

        result = self.so_matrix.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_and_belongs_adaptive_gradient_descent_so_matrix(self):
        point = self.so_matrix.random_uniform(10)

        mean = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric, method='adaptive',
            verbose=True, lr=.5)
        mean.fit(point)

        result = self.so_matrix.belongs(mean.estimate_)
        self.assertTrue(result)

    @geomstats.tests.np_and_tf_only
    def test_estimate_and_coincide_default_so_vec_and_mat(self):
        point = self.so_matrix.random_uniform(3)

        mean = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric, method='default')
        mean.fit(point)
        expected = mean.estimate_

        mean_vec = FrechetMean(
            metric=self.so3.bi_invariant_metric, method='default')
        point_vec = self.so3.rotation_vector_from_matrix(point)
        mean_vec.fit(point_vec)
        result_vec = mean_vec.estimate_
        result = self.so3.matrix_from_rotation_vector(result_vec)

        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_adaptive_gradient_descent_sphere(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

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

        result = variance(
            points, base_point=point, metric=self.sphere.metric)
        expected = gs.array(0.)

        self.assertAllClose(expected, result)

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

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

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

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    def test_estimate_spd(self):
        point = SPDMatrices(3).random_point()
        points = gs.array([point, point])
        mean = FrechetMean(metric=SPDMetricAffine(3), point_type='matrix')
        mean.fit(X=points)
        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)

    def test_variance_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])
        result = variance(
            points, base_point=point, metric=self.hyperbolic.metric)
        expected = gs.array(0.)

        self.assertAllClose(result, expected)

    def test_estimate_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)
        expected = point

        result = mean.estimate_

        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_hyperbolic(self):
        point_a = self.hyperbolic.random_point()
        point_b = self.hyperbolic.random_point()
        point_c = self.hyperbolic.random_point()
        points = gs.stack([point_a, point_b, point_c], axis=0)

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)

        result = self.hyperbolic.belongs(mean.estimate_)
        expected = True

        self.assertAllClose(result, expected)

    def test_mean_euclidean_shape(self):
        dim = 2
        point = gs.array([1., 4.])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

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

    def test_mean_euclidean(self):
        point = gs.array([1., 4.])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

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

        mean = FrechetMean(metric=self.euclidean.metric)
        mean.fit(points, weights=weights)

        result = mean.estimate_
        expected = gs.array([16. / 6., 22. / 6.])

        self.assertAllClose(result, expected)

    def test_variance_euclidean(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 = variance(
            points, weights=weights, base_point=base_point,
            metric=self.euclidean.metric)
        # we expect the average of the points' sq norms.
        expected = gs.array((1 * 5. + 2 * 13. + 1 * 25. + 2 * 41.) / 6.)

        self.assertAllClose(result, expected)

    def test_mean_matrices_shape(self):
        m, n = (2, 2)
        point = gs.array([
            [1., 4.],
            [2., 3.]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type='matrix')
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

        self.assertAllClose(gs.shape(result), (m, n))

    def test_mean_matrices(self):
        m, n = (2, 2)
        point = gs.array([
            [1., 4.],
            [2., 3.]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type='matrix')
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

    def test_mean_minkowski_shape(self):
        dim = 2
        point = gs.array([2., -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

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

    def test_mean_minkowski(self):
        point = gs.array([2., -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

        expected = point

        self.assertAllClose(result, expected)

        points = gs.array([
            [1., 0.],
            [2., math.sqrt(3)],
            [3., math.sqrt(8)],
            [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points, weights=weights)
        result = mean.estimate_
        result = self.minkowski.belongs(result)
        expected = gs.array(True)

        self.assertAllClose(result, expected)

    def test_variance_minkowski(self):
        points = gs.array([
            [1., 0.],
            [2., math.sqrt(3)],
            [3., math.sqrt(8)],
            [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])
        base_point = gs.array([-1., 0.])
        var = variance(
            points, weights=weights, base_point=base_point,
            metric=self.minkowski.metric)
        result = var != 0
        # we expect the average of the points' Minkowski sq norms.
        expected = True
        self.assertAllClose(result, expected)

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

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(X=point)

        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)

        mean = FrechetMean(
            metric=self.sphere.metric, method='frechet-poincare-ball')
        mean.fit(X=point)

        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)
Ejemplo n.º 6
0
class TestFrechetMean(geomstats.tests.TestCase):
    _multiprocess_can_split_ = True

    def setUp(self):
        self.sphere = Hypersphere(dimension=4)
        self.hyperbolic = Hyperbolic(dimension=3)
        self.euclidean = Euclidean(dimension=2)
        self.minkowski = Minkowski(dimension=2)

    @geomstats.tests.np_only
    def test_adaptive_gradient_descent_sphere(self):
        n_tests = 100
        result = gs.zeros(n_tests)
        expected = gs.zeros(n_tests)

        for i in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            mean = _adaptive_gradient_descent(points=points,
                                              metric=self.sphere.metric)

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result[i] = gs.linalg.norm(logs[1, :] + logs[0, :])

        self.assertAllClose(expected, result, rtol=1e-10, atol=1e-10)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_and_belongs_sphere(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.zeros((2, point_a.shape[0]))
        points[0, :] = point_a
        points[1, :] = point_b

        mean = FrechetMean(metric=self.sphere.metric)
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = gs.array([[True]])
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_variance_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.zeros((2, point.shape[0]))
        points[0, :] = point
        points[1, :] = point

        result = variance(points, base_point=point, metric=self.sphere.metric)
        expected = helper.to_scalar(0.)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.zeros((2, point.shape[0]))
        points[0, :] = point
        points[1, :] = point

        mean = FrechetMean(metric=self.sphere.metric)
        mean.fit(X=points)

        result = mean.estimate_
        expected = helper.to_vector(point)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_tf_only
    def test_variance_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])
        result = variance(points,
                          base_point=point,
                          metric=self.hyperbolic.metric)
        expected = helper.to_scalar(0.)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)

        result = mean.estimate_
        expected = helper.to_vector(point)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_and_belongs_hyperbolic(self):
        point_a = self.hyperbolic.random_uniform()
        point_b = self.hyperbolic.random_uniform()
        point_c = self.hyperbolic.random_uniform()
        points = gs.concatenate([point_a, point_b, point_c], axis=0)

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)

        result = self.hyperbolic.belongs(mean.estimate_)
        expected = gs.array([[True]])

        self.assertAllClose(result, expected)

    def test_mean_euclidean(self):
        point = gs.array([[1., 4.]])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        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.])

        mean = FrechetMean(metric=self.euclidean.metric)
        mean.fit(points, weights=weights)

        result = mean.estimate_
        expected = gs.array([16. / 6., 22. / 6.])
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

    def test_variance_euclidean(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 = variance(points,
                          weights=weights,
                          base_point=base_point,
                          metric=self.euclidean.metric)
        # 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)

    def test_mean_minkowski(self):
        point = gs.array([[2., -math.sqrt(3)]])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

        expected = point
        expected = helper.to_vector(expected)

        self.assertAllClose(result, expected)

        points = gs.array([[1., 0.], [2., math.sqrt(3)], [3., math.sqrt(8)],
                           [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points, weights=weights)
        result = mean.estimate_
        result = self.minkowski.belongs(result)
        expected = gs.array([[True]])

        self.assertAllClose(result, expected)

    def test_variance_minkowski(self):
        points = gs.array([[1., 0.], [2., math.sqrt(3)], [3., math.sqrt(8)],
                           [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])
        base_point = gs.array([-1., 0.])
        var = variance(points,
                       weights=weights,
                       base_point=base_point,
                       metric=self.minkowski.metric)
        result = helper.to_scalar(var != 0)
        # we expect the average of the points' Minkowski sq norms.
        expected = helper.to_scalar(gs.array([True]))
        self.assertAllClose(result, expected)
Ejemplo n.º 7
0
class TestHypersphere(geomstats.tests.TestCase):
    def setUp(self):
        gs.random.seed(1234)

        self.dimension = 4
        self.space = Hypersphere(dim=self.dimension)
        self.metric = self.space.metric
        self.n_samples = 10

    def test_random_uniform_and_belongs(self):
        """Test random uniform and belongs.

        Test that the random uniform method samples
        on the hypersphere space.
        """
        n_samples = self.n_samples
        point = self.space.random_uniform(n_samples)
        result = self.space.belongs(point)
        expected = gs.array([True] * n_samples)

        self.assertAllClose(expected, result)

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

        self.assertAllClose(gs.shape(point), (self.dimension + 1, ))

    def test_replace_values(self):
        points = gs.ones((3, 5))
        new_points = gs.zeros((2, 5))
        indcs = [True, False, True]
        update = self.space._replace_values(points, new_points, indcs)
        self.assertAllClose(update,
                            gs.stack([gs.zeros(5),
                                      gs.ones(5),
                                      gs.zeros(5)]))

    def test_projection_and_belongs(self):
        shape = (self.n_samples, self.dimension + 1)
        result = helper.test_projection_and_belongs(self.space, shape)
        for res in result:
            self.assertTrue(res)

    def test_intrinsic_and_extrinsic_coords(self):
        """
        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        space = Hypersphere(dim=2)
        point_int = gs.array([0.1, 0.0])
        point_ext = space.intrinsic_to_extrinsic_coords(point_int)
        result = space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int

        self.assertAllClose(result, expected)

        point_ext = 1. / (gs.sqrt(2.)) * gs.array([1.0, 0.0, 1.0])
        point_int = space.extrinsic_to_intrinsic_coords(point_ext)
        result = space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext

        self.assertAllClose(result, expected)

    def test_intrinsic_and_extrinsic_coords_vectorization(self):
        """Test change of coordinates.

        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        space = Hypersphere(dim=2)
        point_int = gs.array([
            [0.1, 0.1],
            [0.1, 0.4],
            [0.1, 0.3],
            [0.0, 0.0],
            [0.1, 0.5],
        ])
        point_ext = space.intrinsic_to_extrinsic_coords(point_int)
        result = space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int

        self.assertAllClose(result, expected)

        point_int = space.extrinsic_to_intrinsic_coords(point_ext)
        result = space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext

        self.assertAllClose(result, expected)

    def test_log_and_exp_general_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # General case
        base_point = gs.array([1.0, 2.0, 3.0, 4.0, 6.0])
        base_point = base_point / gs.linalg.norm(base_point)
        point = gs.array([0.0, 5.0, 6.0, 2.0, -1.0])
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected)

    def test_log_and_exp_edge_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # Edge case: two very close points, base_point_2 and point_2,
        # form an angle < epsilon
        base_point = gs.array([1.0, 2.0, 3.0, 4.0, 6.0])
        base_point = base_point / gs.linalg.norm(base_point)
        point = base_point + 1e-4 * gs.array([-1.0, -2.0, 1.0, 1.0, 0.1])
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected)

    def test_exp_vectorization_single_samples(self):
        dim = self.dimension + 1

        one_vec = self.space.random_uniform()
        one_base_point = self.space.random_uniform()
        one_tangent_vec = self.space.to_tangent(one_vec,
                                                base_point=one_base_point)

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

        one_base_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_tangent_vec = gs.to_ndarray(one_tangent_vec, to_ndim=2)
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_base_point = self.space.random_uniform()
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

    def test_exp_vectorization_n_samples(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_vec = self.space.random_uniform()
        one_base_point = self.space.random_uniform()
        n_vecs = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

        n_tangent_vecs = self.space.to_tangent(n_vecs,
                                               base_point=one_base_point)
        result = self.metric.exp(n_tangent_vecs, one_base_point)

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

        one_tangent_vec = self.space.to_tangent(one_vec,
                                                base_point=n_base_points)
        result = self.metric.exp(one_tangent_vec, n_base_points)

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

        n_tangent_vecs = self.space.to_tangent(n_vecs,
                                               base_point=n_base_points)
        result = self.metric.exp(n_tangent_vecs, n_base_points)

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

    def test_log_vectorization_single_samples(self):
        dim = self.dimension + 1

        one_base_point = self.space.random_uniform()
        one_point = self.space.random_uniform()

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

        one_base_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_base_point = self.space.random_uniform()
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

    def test_log_vectorization_n_samples(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_base_point = self.space.random_uniform()
        one_point = self.space.random_uniform()
        n_points = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

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

        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_exp_log_are_inverse(self):
        initial_point = self.space.random_uniform(2)
        end_point = self.space.random_uniform(2)
        vec = self.space.metric.log(point=end_point, base_point=initial_point)
        result = self.space.metric.exp(vec, initial_point)
        self.assertAllClose(end_point, result)

    def test_log_extreme_case(self):
        initial_point = self.space.random_uniform(2)
        vec = 1e-4 * gs.random.rand(*initial_point.shape)
        vec = self.space.to_tangent(vec, initial_point)
        point = self.space.metric.exp(vec, base_point=initial_point)
        result = self.space.metric.log(point, initial_point)
        self.assertAllClose(vec, result)

    def test_exp_and_log_and_projection_to_tangent_space_general_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Exp then Riemannian Log
        # General case
        # NB: Riemannian log gives a regularized tangent vector,
        # so we take the norm modulo 2 * pi.
        base_point = gs.array([0.0, -3.0, 0.0, 3.0, 4.0])
        base_point = base_point / gs.linalg.norm(base_point)

        vector = gs.array([3.0, 2.0, 0.0, 0.0, -1.0])
        vector = self.space.to_tangent(vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        result = self.metric.log(point=exp, base_point=base_point)

        expected = vector
        norm_expected = gs.linalg.norm(expected)
        regularized_norm_expected = gs.mod(norm_expected, 2 * gs.pi)
        expected = expected / norm_expected * regularized_norm_expected

        # The Log can be the opposite vector on the tangent space,
        # whose Exp gives the base_point
        are_close = gs.allclose(result, expected)
        norm_2pi = gs.isclose(gs.linalg.norm(result - expected), 2 * gs.pi)
        self.assertTrue(are_close or norm_2pi)

    def test_exp_and_log_and_projection_to_tangent_space_edge_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Exp then Riemannian Log
        # Edge case: tangent vector has norm < epsilon
        base_point = gs.array([10.0, -2.0, -0.5, 34.0, 3.0])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = 1e-4 * gs.array([0.06, -51.0, 6.0, 5.0, 3.0])
        vector = self.space.to_tangent(vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        result = self.metric.log(point=exp, base_point=base_point)
        self.assertAllClose(result, vector)

    def test_squared_norm_and_squared_dist(self):
        """
        Test that the squared distance between two points is
        the squared norm of their logarithm.
        """
        point_a = 1.0 / gs.sqrt(129.0) * gs.array([10.0, -2.0, -5.0, 0.0, 0.0])
        point_b = 1.0 / gs.sqrt(435.0) * gs.array([1.0, -20.0, -5.0, 0.0, 3.0])
        log = self.metric.log(point=point_a, base_point=point_b)
        result = self.metric.squared_norm(vector=log)
        expected = self.metric.squared_dist(point_a, point_b)

        self.assertAllClose(result, expected)

    def test_squared_dist_vectorization_single_sample(self):
        one_point_a = self.space.random_uniform()
        one_point_b = self.space.random_uniform()

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

        one_point_a = gs.to_ndarray(one_point_a, to_ndim=2)
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

        one_point_b = gs.to_ndarray(one_point_b, to_ndim=2)
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

        one_point_a = self.space.random_uniform()
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

    def test_squared_dist_vectorization_n_samples(self):
        n_samples = self.n_samples

        one_point_a = self.space.random_uniform()
        one_point_b = self.space.random_uniform()
        n_points_a = self.space.random_uniform(n_samples=n_samples)
        n_points_b = self.space.random_uniform(n_samples=n_samples)

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

        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)
        self.assertAllClose(gs.shape(result), (n_samples, ))

        one_point_a = gs.to_ndarray(one_point_a, to_ndim=2)
        one_point_b = gs.to_ndarray(one_point_b, to_ndim=2)

        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)
        self.assertAllClose(gs.shape(result), (n_samples, ))

    def test_norm_and_dist(self):
        """
        Test that the distance between two points is
        the norm of their logarithm.
        """
        point_a = 1.0 / gs.sqrt(129.0) * gs.array([10.0, -2.0, -5.0, 0.0, 0.0])
        point_b = 1.0 / gs.sqrt(435.0) * gs.array([1.0, -20.0, -5.0, 0.0, 3.0])
        log = self.metric.log(point=point_a, base_point=point_b)

        self.assertAllClose(gs.shape(log), (5, ))

        result = self.metric.norm(vector=log)
        self.assertAllClose(gs.shape(result), ())

        expected = self.metric.dist(point_a, point_b)
        self.assertAllClose(gs.shape(expected), ())

        self.assertAllClose(result, expected)

    def test_dist_point_and_itself(self):
        # Distance between a point and itself is 0
        point_a = 1.0 / gs.sqrt(129.0) * gs.array([10.0, -2.0, -5.0, 0.0, 0.0])
        point_b = point_a
        result = self.metric.dist(point_a, point_b)
        expected = 0.0

        self.assertAllClose(result, expected)

    def test_dist_pairwise(self):

        point_a = 1.0 / gs.sqrt(129.0) * gs.array([10.0, -2.0, -5.0, 0.0, 0.0])
        point_b = 1.0 / gs.sqrt(435.0) * gs.array([1.0, -20.0, -5.0, 0.0, 3.0])

        point = gs.array([point_a, point_b])
        result = self.metric.dist_pairwise(point)

        expected = gs.array([[0.0, 1.24864502], [1.24864502, 0.0]])

        self.assertAllClose(result, expected, rtol=1e-3)

    def test_dist_pairwise_parallel(self):
        n_samples = 15
        points = self.space.random_uniform(n_samples)
        result = self.metric.dist_pairwise(points, n_jobs=2, prefer="threads")
        is_sym = Matrices.is_symmetric(result)
        belongs = Matrices(n_samples, n_samples).belongs(result)
        self.assertTrue(is_sym)
        self.assertTrue(belongs)

    def test_dist_orthogonal_points(self):
        # Distance between two orthogonal points is pi / 2.
        point_a = gs.array([10.0, -2.0, -0.5, 0.0, 0.0])
        point_a = point_a / gs.linalg.norm(point_a)
        point_b = gs.array([2.0, 10, 0.0, 0.0, 0.0])
        point_b = point_b / gs.linalg.norm(point_b)
        result = gs.dot(point_a, point_b)
        expected = 0
        self.assertAllClose(result, expected)

        result = self.metric.dist(point_a, point_b)
        expected = gs.pi / 2

        self.assertAllClose(result, expected)

    def test_exp_and_dist_and_projection_to_tangent_space(self):
        base_point = gs.array([16.0, -2.0, -2.5, 84.0, 3.0])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = gs.array([9.0, 0.0, -1.0, -2.0, 1.0])
        tangent_vec = self.space.to_tangent(vector=vector,
                                            base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)
        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec) % (2 * gs.pi)
        self.assertAllClose(result, expected)

    def test_exp_and_dist_and_projection_to_tangent_space_vec(self):
        base_point = gs.array([[16.0, -2.0, -2.5, 84.0, 3.0],
                               [16.0, -2.0, -2.5, 84.0, 3.0]])

        base_single_point = gs.array([16.0, -2.0, -2.5, 84.0, 3.0])
        scalar_norm = gs.linalg.norm(base_single_point)

        base_point = base_point / scalar_norm
        vector = gs.array([[9.0, 0.0, -1.0, -2.0, 1.0],
                           [9.0, 0.0, -1.0, -2.0, 1]])

        tangent_vec = self.space.to_tangent(vector=vector,
                                            base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)

        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec, axis=-1) % (2 * gs.pi)

        self.assertAllClose(result, expected)

    def test_geodesic_and_belongs(self):
        n_geodesic_points = 10
        initial_point = self.space.random_uniform(2)
        vector = gs.array([[2.0, 0.0, -1.0, -2.0, 1.0]] * 2)
        initial_tangent_vec = self.space.to_tangent(vector=vector,
                                                    base_point=initial_point)
        geodesic = self.metric.geodesic(
            initial_point=initial_point,
            initial_tangent_vec=initial_tangent_vec)
        t = gs.linspace(start=0.0, stop=1.0, num=n_geodesic_points)
        points = geodesic(t)
        result = gs.stack([self.space.belongs(pt) for pt in points])
        self.assertTrue(gs.all(result))

        initial_point = initial_point[0]
        initial_tangent_vec = initial_tangent_vec[0]
        geodesic = self.metric.geodesic(
            initial_point=initial_point,
            initial_tangent_vec=initial_tangent_vec)
        points = geodesic(t)
        result = self.space.belongs(points)
        expected = gs.array(n_geodesic_points * [True])
        self.assertAllClose(expected, result)

    def test_geodesic_end_point(self):
        n_geodesic_points = 10
        initial_point = self.space.random_uniform(4)
        geodesic = self.metric.geodesic(initial_point=initial_point[:2],
                                        end_point=initial_point[2:])
        t = gs.linspace(start=0.0, stop=1.0, num=n_geodesic_points)
        points = geodesic(t)
        result = points[:, -1]
        expected = initial_point[2:]
        self.assertAllClose(expected, result)

    def test_inner_product(self):
        tangent_vec_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        tangent_vec_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        base_point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])
        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = 0.0

        self.assertAllClose(expected, result)

    def test_inner_product_vectorization_single_samples(self):
        tangent_vec_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        tangent_vec_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        base_point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = 0.0
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1.0, 0.0, 0.0, 0.0, 0.0]])
        tangent_vec_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        base_point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.0])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        tangent_vec_b = gs.array([[0.0, 1.0, 0.0, 0.0, 0.0]])
        base_point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.0])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1.0, 0.0, 0.0, 0.0, 0.0]])
        tangent_vec_b = gs.array([[0.0, 1.0, 0.0, 0.0, 0.0]])
        base_point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.0])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1.0, 0.0, 0.0, 0.0, 0.0]])
        tangent_vec_b = gs.array([[0.0, 1.0, 0.0, 0.0, 0.0]])
        base_point = gs.array([[0.0, 0.0, 0.0, 0.0, 1.0]])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.0])
        self.assertAllClose(expected, result)

    def test_diameter(self):
        dim = 2
        sphere = Hypersphere(dim)
        point_a = gs.array([[0.0, 0.0, 1.0]])
        point_b = gs.array([[1.0, 0.0, 0.0]])
        point_c = gs.array([[0.0, 0.0, -1.0]])
        result = sphere.metric.diameter(gs.vstack((point_a, point_b, point_c)))
        expected = gs.pi
        self.assertAllClose(expected, result)

    def test_closest_neighbor_index(self):
        """Check that the closest neighbor is one of neighbors."""
        n_samples = 10
        points = self.space.random_uniform(n_samples=n_samples)
        point = points[0, :]
        neighbors = points[1:, :]
        index = self.metric.closest_neighbor_index(point, neighbors)
        closest_neighbor = points[index, :]

        test = gs.sum(gs.all(points == closest_neighbor, axis=1))
        result = test > 0
        self.assertTrue(result)

    def test_sample_von_mises_fisher_arbitrary_mean(self):
        """
        Check that the maximum likelihood estimates of the mean and
        concentration parameter are close to the real values. A first
        estimation of the concentration parameter is obtained by a
        closed-form expression and improved through the Newton method.
        """
        for dim in [2, 9]:
            n_points = 10000
            sphere = Hypersphere(dim)

            # check mean value for concentrated distribution for different mean
            kappa = 1000.0
            mean = sphere.random_uniform()
            points = sphere.random_von_mises_fisher(mu=mean,
                                                    kappa=kappa,
                                                    n_samples=n_points)
            sum_points = gs.sum(points, axis=0)
            result = sum_points / gs.linalg.norm(sum_points)
            expected = mean
            self.assertAllClose(result, expected, atol=MEAN_ESTIMATION_TOL)

    def test_random_von_mises_kappa(self):
        # check concentration parameter for dispersed distribution
        kappa = 1.0
        n_points = 100000
        for dim in [2, 9]:
            sphere = Hypersphere(dim)
            points = sphere.random_von_mises_fisher(kappa=kappa,
                                                    n_samples=n_points)
            sum_points = gs.sum(points, axis=0)
            mean_norm = gs.linalg.norm(sum_points) / n_points
            kappa_estimate = (mean_norm * (dim + 1.0 - mean_norm**2) /
                              (1.0 - mean_norm**2))
            kappa_estimate = gs.cast(kappa_estimate, gs.float64)
            p = dim + 1
            n_steps = 100
            for _ in range(n_steps):
                bessel_func_1 = scipy.special.iv(p / 2.0, kappa_estimate)
                bessel_func_2 = scipy.special.iv(p / 2.0 - 1.0, kappa_estimate)
                ratio = bessel_func_1 / bessel_func_2
                denominator = 1.0 - ratio**2 - (p -
                                                1.0) * ratio / kappa_estimate
                mean_norm = gs.cast(mean_norm, gs.float64)
                kappa_estimate = kappa_estimate - (ratio -
                                                   mean_norm) / denominator
            result = kappa_estimate
            expected = kappa
            self.assertAllClose(result, expected, atol=KAPPA_ESTIMATION_TOL)

    def test_random_von_mises_general_dim_mean(self):
        for dim in [2, 9]:
            sphere = Hypersphere(dim)
            n_points = 100000

            # check mean value for concentrated distribution
            kappa = 10
            points = sphere.random_von_mises_fisher(kappa=kappa,
                                                    n_samples=n_points)
            sum_points = gs.sum(points, axis=0)
            expected = gs.array([1.0] + [0.0] * dim)
            result = sum_points / gs.linalg.norm(sum_points)
            self.assertAllClose(result, expected, atol=KAPPA_ESTIMATION_TOL)

    def test_random_von_mises_one_sample_belongs(self):
        for dim in [2, 9]:
            sphere = Hypersphere(dim)
            point = sphere.random_von_mises_fisher()
            self.assertAllClose(point.shape, (dim + 1, ))
            result = sphere.belongs(point)
            self.assertTrue(result)

    def test_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)

        points_spherical = gs.array([gs.pi / 2, 0])
        result = sphere.spherical_to_extrinsic(points_spherical)
        expected = gs.array([1.0, 0.0, 0.0])
        self.assertAllClose(result, expected)

    def test_extrinsic_to_spherical(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)

        points_extrinsic = gs.array([1.0, 0.0, 0.0])
        result = sphere.extrinsic_to_spherical(points_extrinsic)
        expected = gs.array([gs.pi / 2, 0])
        self.assertAllClose(result, expected)

    def test_spherical_to_extrinsic_vectorization(self):
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        result = sphere.spherical_to_extrinsic(points_spherical)
        expected = gs.array([
            [1.0, 0.0, 0.0],
            [gs.sqrt(2.0) / 4.0,
             gs.sqrt(2.0) / 4.0,
             gs.sqrt(3.0) / 2.0],
        ])
        self.assertAllClose(result, expected)

    def test_extrinsic_to_spherical_vectorization(self):
        dim = 2
        sphere = Hypersphere(dim)
        expected = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        point_extrinsic = gs.array([
            [1.0, 0.0, 0.0],
            [gs.sqrt(2.0) / 4.0,
             gs.sqrt(2.0) / 4.0,
             gs.sqrt(3.0) / 2.0],
        ])
        result = sphere.extrinsic_to_spherical(point_extrinsic)
        self.assertAllClose(result, expected)

    def test_spherical_to_extrinsic_and_inverse(self):
        dim = 2
        n_samples = 5
        sphere = Hypersphere(dim)
        points = gs.random.rand(n_samples, 2) * gs.pi * gs.array([1., 2.
                                                                  ])[None, :]
        extrinsic = sphere.spherical_to_extrinsic(points)
        result = sphere.extrinsic_to_spherical(extrinsic)
        self.assertAllClose(result, points)

        points_extrinsic = sphere.random_uniform(n_samples)
        spherical = sphere.extrinsic_to_spherical(points_extrinsic)
        result = sphere.spherical_to_extrinsic(spherical)
        self.assertAllClose(result, points_extrinsic)

    def test_tangent_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates for tangent vectors to the
        2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        base_points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 2, 0]])
        tangent_vecs_spherical = gs.array([[0.25, 0.5], [0.3, 0.2]])
        result = sphere.tangent_spherical_to_extrinsic(tangent_vecs_spherical,
                                                       base_points_spherical)
        expected = gs.array([[0, 0.5, -0.25], [0, 0.2, -0.3]])
        self.assertAllClose(result, expected)

        result = sphere.tangent_spherical_to_extrinsic(
            tangent_vecs_spherical[0], base_points_spherical[0])
        self.assertAllClose(result, expected[0])

    def test_tangent_extrinsic_to_spherical(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates for tangent vectors to the
        2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        base_points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 2, 0]])
        expected = gs.array([[0.25, 0.5], [0.3, 0.2]])
        tangent_vecs = gs.array([[0, 0.5, -0.25], [0, 0.2, -0.3]])
        result = sphere.tangent_extrinsic_to_spherical(
            tangent_vecs, base_point_spherical=base_points_spherical)
        self.assertAllClose(result, expected)

        result = sphere.tangent_extrinsic_to_spherical(tangent_vecs[0],
                                                       base_point=gs.array(
                                                           [1., 0., 0.]))
        self.assertAllClose(result, expected[0])

    def test_tangent_spherical_and_extrinsic_inverse(self):
        dim = 2
        n_samples = 5
        sphere = Hypersphere(dim)
        points = gs.random.rand(n_samples, 2) * gs.pi * gs.array([1., 2.
                                                                  ])[None, :]
        tangent_spherical = gs.random.rand(n_samples, 2)
        tangent_extrinsic = sphere.tangent_spherical_to_extrinsic(
            tangent_spherical, points)
        result = sphere.tangent_extrinsic_to_spherical(
            tangent_extrinsic, base_point_spherical=points)
        self.assertAllClose(result, tangent_spherical)

        points_extrinsic = sphere.random_uniform(n_samples)
        vector = gs.random.rand(n_samples, dim + 1)
        tangent_extrinsic = sphere.to_tangent(vector, points_extrinsic)
        tangent_spherical = sphere.tangent_extrinsic_to_spherical(
            tangent_extrinsic, base_point=points_extrinsic)
        spherical = sphere.extrinsic_to_spherical(points_extrinsic)
        result = sphere.tangent_spherical_to_extrinsic(tangent_spherical,
                                                       spherical)
        self.assertAllClose(result, tangent_extrinsic)

    def test_christoffels_vectorization(self):
        """
        Check vectorization of Christoffel symbols in
        spherical coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        christoffel = sphere.metric.christoffels(points_spherical)
        result = christoffel.shape
        expected = gs.array([2, dim, dim, dim])
        self.assertAllClose(result, expected)

    def test_parallel_transport_vectorization(self):
        sphere = Hypersphere(2)
        metric = sphere.metric
        shape = (4, 3)

        results = helper.test_parallel_transport(sphere, metric, shape)
        for res in results:
            self.assertTrue(res)

    def test_is_tangent(self):
        space = self.space
        vec = space.random_uniform()
        result = space.is_tangent(vec, vec)
        self.assertFalse(result)

        base_point = space.random_uniform()
        tangent_vec = space.to_tangent(vec, base_point)
        result = space.is_tangent(tangent_vec, base_point)
        self.assertTrue(result)

        base_point = space.random_uniform(2)
        vec = space.random_uniform(2)
        tangent_vec = space.to_tangent(vec, base_point)
        result = space.is_tangent(tangent_vec, base_point)
        self.assertAllClose(gs.shape(result), (2, ))
        self.assertTrue(gs.all(result))

    def test_sectional_curvature(self):
        n_samples = 4
        sphere = self.space
        base_point = sphere.random_uniform(n_samples)
        tan_vec_a = sphere.to_tangent(
            gs.random.rand(n_samples, sphere.dim + 1), base_point)
        tan_vec_b = sphere.to_tangent(
            gs.random.rand(n_samples, sphere.dim + 1), base_point)
        result = sphere.metric.sectional_curvature(tan_vec_a, tan_vec_b,
                                                   base_point)
        expected = gs.ones(result.shape)
        self.assertAllClose(result, expected)

    @geomstats.tests.np_autograd_and_torch_only
    def test_riemannian_normal_and_belongs(self):
        mean = self.space.random_uniform()
        cov = gs.eye(self.space.dim)
        sample = self.space.random_riemannian_normal(mean, cov, 10)
        result = self.space.belongs(sample)
        self.assertTrue(gs.all(result))

    @geomstats.tests.np_autograd_and_torch_only
    def test_riemannian_normal_mean(self):
        space = self.space
        mean = space.random_uniform()
        precision = gs.eye(space.dim) * 10
        sample = space.random_riemannian_normal(mean, precision, 10000)
        estimator = FrechetMean(space.metric, method="adaptive")
        estimator.fit(sample)
        estimate = estimator.estimate_
        self.assertAllClose(estimate, mean, atol=1e-2)

    def test_raises(self):
        space = self.space
        point = space.random_uniform()
        self.assertRaises(NotImplementedError,
                          lambda: space.extrinsic_to_spherical(point))

        self.assertRaises(
            NotImplementedError,
            lambda: space.tangent_extrinsic_to_spherical(point, point))

        sphere = Hypersphere(2)
        self.assertRaises(ValueError,
                          lambda: sphere.tangent_extrinsic_to_spherical(point))

    def test_angle_to_extrinsic(self):
        space = Hypersphere(1)
        point = gs.pi / 4
        result = space.angle_to_extrinsic(point)
        expected = gs.array([1., 1.]) / gs.sqrt(2.)
        self.assertAllClose(result, expected)

        point = gs.array([1. / 3, 0.]) * gs.pi
        result = space.angle_to_extrinsic(point)
        expected = gs.array([[1. / 2, gs.sqrt(3.) / 2], [1., 0.]])
        self.assertAllClose(result, expected)

    def test_extrinsic_to_angle(self):
        space = Hypersphere(1)
        point = gs.array([1., 1.]) / gs.sqrt(2.)
        result = space.extrinsic_to_angle(point)
        expected = gs.pi / 4
        self.assertAllClose(result, expected)

        point = gs.array([[1. / 2, gs.sqrt(3.) / 2], [1., 0.]])
        result = space.extrinsic_to_angle(point)
        expected = gs.array([1. / 3, 0.]) * gs.pi
        self.assertAllClose(result, expected)

    def test_extrinsic_to_angle_inverse(self):
        space = Hypersphere(1)
        point = space.random_uniform()
        angle = space.extrinsic_to_angle(point)
        result = space.angle_to_extrinsic(angle)
        self.assertAllClose(result, point)

        space = Hypersphere(1, default_coords_type='intrinsic')
        angle = space.random_uniform()
        extrinsic = space.angle_to_extrinsic(angle)
        result = space.extrinsic_to_angle(extrinsic)
        self.assertAllClose(result, angle)
Ejemplo n.º 8
0
class TestFrechetMean(geomstats.tests.TestCase):
    _multiprocess_can_split_ = True

    def setup_method(self):
        gs.random.seed(123)
        self.sphere = Hypersphere(dim=4)
        self.hyperbolic = Hyperboloid(dim=3)
        self.euclidean = Euclidean(dim=2)
        self.minkowski = Minkowski(dim=2)
        self.so3 = SpecialOrthogonal(n=3, point_type="vector")
        self.so_matrix = SpecialOrthogonal(n=3)
        self.curves_2d = DiscreteCurves(R2)
        self.elastic_metric = ElasticMetric(a=1, b=1, ambient_manifold=R2)

    def test_logs_at_mean_curves_2d(self):
        n_tests = 10
        metric = self.curves_2d.srv_metric
        estimator = FrechetMean(metric=metric, init_step_size=1.0)

        result = []
        for _ in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.curves_2d.random_point(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            # Expand and tile are added because of GitHub issue 1575
            mean = gs.expand_dims(mean, axis=0)
            mean = gs.tile(mean, (2, 1, 1))

            logs = metric.log(point=points, base_point=mean)
            logs_srv = metric.aux_differential_srv_transform(logs, curve=mean)
            # Note that the logs are NOT inverse, only the logs_srv are.
            result.append(gs.linalg.norm(logs_srv[1, :] + logs_srv[0, :]))
        result = gs.stack(result)
        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result, atol=1e-5)

    def test_logs_at_mean_default_gradient_descent_sphere(self):
        n_tests = 10
        estimator = FrechetMean(metric=self.sphere.metric,
                                method="default",
                                init_step_size=1.0)

        result = []
        for _ in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result.append(gs.linalg.norm(logs[1, :] + logs[0, :]))
        result = gs.stack(result)
        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result)

    def test_logs_at_mean_adaptive_gradient_descent_sphere(self):
        n_tests = 10
        estimator = FrechetMean(metric=self.sphere.metric, method="adaptive")

        result = []
        for _ in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result.append(gs.linalg.norm(logs[1, :] + logs[0, :]))
        result = gs.stack(result)

        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result)

    def test_estimate_shape_default_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        point_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric,
                           method="default",
                           verbose=True)
        mean.fit(points)
        result = mean.estimate_

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

    def test_estimate_shape_adaptive_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        point_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method="adaptive")
        mean.fit(points)
        result = mean.estimate_

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

    def test_estimate_shape_elastic_metric(self):
        points = self.curves_2d.random_point(n_samples=2)

        mean = FrechetMean(metric=self.elastic_metric)
        mean.fit(points)
        result = mean.estimate_

        self.assertAllClose(gs.shape(result), (points.shape[1:]))

    def test_estimate_shape_metric(self):
        points = self.curves_2d.random_point(n_samples=2)

        mean = FrechetMean(metric=self.curves_2d.srv_metric)
        mean.fit(points)
        result = mean.estimate_

        self.assertAllClose(gs.shape(result), (points.shape[1:]))

    def test_estimate_and_belongs_default_gradient_descent_sphere(self):
        point_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        point_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method="default")
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_curves_2d(self):
        points = self.curves_2d.random_point(n_samples=2)

        mean = FrechetMean(metric=self.curves_2d.srv_metric)
        mean.fit(points)

        result = self.curves_2d.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    def test_estimate_default_gradient_descent_so3(self):
        points = self.so3.random_uniform(2)

        mean_vec = FrechetMean(metric=self.so3.bi_invariant_metric,
                               method="default",
                               init_step_size=1.0)
        mean_vec.fit(points)

        logs = self.so3.bi_invariant_metric.log(points, mean_vec.estimate_)
        result = gs.sum(logs, axis=0)
        expected = gs.zeros_like(points[0])
        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_default_gradient_descent_so3(self):
        point = self.so3.random_uniform(10)

        mean_vec = FrechetMean(metric=self.so3.bi_invariant_metric,
                               method="default")
        mean_vec.fit(point)

        result = self.so3.belongs(mean_vec.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_autograd_and_tf_only
    def test_estimate_default_gradient_descent_so_matrix(self):
        points = self.so_matrix.random_uniform(2)
        mean_vec = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric,
            method="default",
            init_step_size=1.0,
        )
        mean_vec.fit(points)
        logs = self.so_matrix.bi_invariant_metric.log(points,
                                                      mean_vec.estimate_)
        result = gs.sum(logs, axis=0)
        expected = gs.zeros_like(points[0])

        self.assertAllClose(result, expected, atol=1e-5)

    @geomstats.tests.np_autograd_and_tf_only
    def test_estimate_and_belongs_default_gradient_descent_so_matrix(self):
        point = self.so_matrix.random_uniform(10)

        mean = FrechetMean(metric=self.so_matrix.bi_invariant_metric,
                           method="default")
        mean.fit(point)

        result = self.so_matrix.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_autograd_and_tf_only
    def test_estimate_and_belongs_adaptive_gradient_descent_so_matrix(self):
        point = self.so_matrix.random_uniform(10)

        mean = FrechetMean(
            metric=self.so_matrix.bi_invariant_metric,
            method="adaptive",
            init_step_size=0.5,
            verbose=True,
        )
        mean.fit(point)

        result = self.so_matrix.belongs(mean.estimate_)
        self.assertTrue(result)

    @geomstats.tests.np_autograd_and_tf_only
    def test_estimate_and_coincide_default_so_vec_and_mat(self):
        point = self.so_matrix.random_uniform(3)

        mean = FrechetMean(metric=self.so_matrix.bi_invariant_metric,
                           method="default")
        mean.fit(point)
        expected = mean.estimate_

        mean_vec = FrechetMean(metric=self.so3.bi_invariant_metric,
                               method="default")
        point_vec = self.so3.rotation_vector_from_matrix(point)
        mean_vec.fit(point_vec)
        result_vec = mean_vec.estimate_
        result = self.so3.matrix_from_rotation_vector(result_vec)

        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_adaptive_gradient_descent_sphere(self):
        point_a = gs.array([1.0, 0.0, 0.0, 0.0, 0.0])
        point_b = gs.array([0.0, 1.0, 0.0, 0.0, 0.0])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method="adaptive")
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    def test_variance_sphere(self):
        point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])
        points = gs.array([point, point])

        result = variance(points, base_point=point, metric=self.sphere.metric)
        expected = gs.array(0.0)

        self.assertAllClose(expected, result)

    def test_estimate_default_gradient_descent_sphere(self):
        point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.sphere.metric, method="default")
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    def test_estimate_elastic_metric(self):
        point = self.curves_2d.random_point(n_samples=1)
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.elastic_metric)
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    def test_estimate_curves_2d(self):
        point = self.curves_2d.random_point(n_samples=1)
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.curves_2d.srv_metric)
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    def test_estimate_adaptive_gradient_descent_sphere(self):
        point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.sphere.metric, method="adaptive")
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    def test_estimate_spd(self):
        point = SPDMatrices(3).random_point()
        points = gs.array([point, point])
        mean = FrechetMean(metric=SPDMetricAffine(3), point_type="matrix")
        mean.fit(X=points)
        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)

    def test_estimate_spd_two_samples(self):
        space = SPDMatrices(3)
        metric = SPDMetricAffine(3)
        point = space.random_point(2)
        mean = FrechetMean(metric)
        mean.fit(point)
        result = mean.estimate_
        expected = metric.exp(metric.log(point[0], point[1]) / 2, point[1])
        self.assertAllClose(expected, result)

    def test_variance_hyperbolic(self):
        point = gs.array([2.0, 1.0, 1.0, 1.0])
        points = gs.array([point, point])
        result = variance(points,
                          base_point=point,
                          metric=self.hyperbolic.metric)
        expected = gs.array(0.0)

        self.assertAllClose(result, expected)

    def test_estimate_hyperbolic(self):
        point = gs.array([2.0, 1.0, 1.0, 1.0])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)
        expected = point

        result = mean.estimate_

        self.assertAllClose(result, expected)

    def test_estimate_and_belongs_hyperbolic(self):
        point_a = self.hyperbolic.random_point()
        point_b = self.hyperbolic.random_point()
        point_c = self.hyperbolic.random_point()
        points = gs.stack([point_a, point_b, point_c], axis=0)

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)

        result = self.hyperbolic.belongs(mean.estimate_)
        expected = True

        self.assertAllClose(result, expected)

    def test_mean_euclidean_shape(self):
        dim = 2
        point = gs.array([1.0, 4.0])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

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

    def test_mean_euclidean(self):
        point = gs.array([1.0, 4.0])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

        points = gs.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0]])
        weights = [1.0, 2.0, 1.0, 2.0]

        mean = FrechetMean(metric=self.euclidean.metric)
        mean.fit(points, weights=weights)

        result = mean.estimate_
        expected = gs.array([16.0 / 6.0, 22.0 / 6.0])

        self.assertAllClose(result, expected)

    def test_variance_euclidean(self):
        points = gs.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0]])
        weights = gs.array([1.0, 2.0, 1.0, 2.0])
        base_point = gs.zeros(2)
        result = variance(points,
                          weights=weights,
                          base_point=base_point,
                          metric=self.euclidean.metric)
        # we expect the average of the points' sq norms.
        expected = gs.array((1 * 5.0 + 2 * 13.0 + 1 * 25.0 + 2 * 41.0) / 6.0)

        self.assertAllClose(result, expected)

    def test_mean_matrices_shape(self):
        m, n = (2, 2)
        point = gs.array([[1.0, 4.0], [2.0, 3.0]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type="matrix")
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

        self.assertAllClose(gs.shape(result), (m, n))

    def test_mean_matrices(self):
        m, n = (2, 2)
        point = gs.array([[1.0, 4.0], [2.0, 3.0]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type="matrix")
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

    def test_mean_minkowski_shape(self):
        dim = 2
        point = gs.array([2.0, -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

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

    def test_mean_minkowski(self):
        point = gs.array([2.0, -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

        expected = point

        self.assertAllClose(result, expected)

        points = gs.array([[1.0, 0.0], [2.0, math.sqrt(3)],
                           [3.0, math.sqrt(8)], [4.0, math.sqrt(24)]])
        weights = gs.array([1.0, 2.0, 1.0, 2.0])

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points, weights=weights)
        result = self.minkowski.belongs(mean.estimate_)

        self.assertTrue(result)

    def test_variance_minkowski(self):
        points = gs.array([[1.0, 0.0], [2.0, math.sqrt(3)],
                           [3.0, math.sqrt(8)], [4.0, math.sqrt(24)]])
        weights = gs.array([1.0, 2.0, 1.0, 2.0])
        base_point = gs.array([-1.0, 0.0])
        var = variance(points,
                       weights=weights,
                       base_point=base_point,
                       metric=self.minkowski.metric)
        result = var != 0
        # we expect the average of the points' Minkowski sq norms.
        expected = True
        self.assertAllClose(result, expected)

    def test_one_point(self):
        point = gs.array([0.0, 0.0, 0.0, 0.0, 1.0])

        mean = FrechetMean(metric=self.sphere.metric, method="default")
        mean.fit(X=point)

        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)

        mean = FrechetMean(metric=self.sphere.metric, method="default")
        mean.fit(X=point)

        result = mean.estimate_
        expected = point
        self.assertAllClose(expected, result)

    def test_batched(self):
        space = SPDMatrices(3)
        metric = SPDMetricAffine(3)
        point = space.random_point(4)
        mean_batch = FrechetMean(metric, method="batch", verbose=True)
        data = gs.stack([point[:2], point[2:]], axis=1)
        mean_batch.fit(data)
        result = mean_batch.estimate_

        mean = FrechetMean(metric)
        mean.fit(data[:, 0])
        expected_1 = mean.estimate_
        mean.fit(data[:, 1])
        expected_2 = mean.estimate_
        expected = gs.stack([expected_1, expected_2])
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_autograd_only
    def test_stiefel_two_samples(self):
        space = Stiefel(3, 2)
        metric = space.metric
        point = space.random_point(2)
        mean = FrechetMean(metric)
        mean.fit(point)
        result = mean.estimate_
        expected = metric.exp(metric.log(point[0], point[1]) / 2, point[1])
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_autograd_only
    def test_stiefel_n_samples(self):
        space = Stiefel(3, 2)
        metric = space.metric
        point = space.random_point(2)
        mean = FrechetMean(metric,
                           method="default",
                           init_step_size=0.5,
                           verbose=True)
        mean.fit(point)
        result = space.belongs(mean.estimate_)
        self.assertTrue(result)

    def test_circle_mean(self):
        space = Hypersphere(1)
        points = space.random_uniform(10)
        mean_circle = FrechetMean(space.metric)
        mean_circle.fit(points)
        estimate_circle = mean_circle.estimate_

        # set a wrong dimension so that the extrinsic coordinates are used
        metric = space.metric
        metric.dim = 2
        mean_extrinsic = FrechetMean(metric)
        mean_extrinsic.fit(points)
        estimate_extrinsic = mean_extrinsic.estimate_
        self.assertAllClose(estimate_circle, estimate_extrinsic)
class TestWrappedGaussianProcess(geomstats.tests.TestCase):
    def setup_method(self):
        gs.random.seed(1234)
        self.n_samples = 20

        # Set up for hypersphere
        self.dim_sphere = 2
        self.shape_sphere = (self.dim_sphere + 1, )
        self.sphere = Hypersphere(dim=self.dim_sphere)

        self.intercept_sphere_true = gs.array([0.0, -1.0, 0.0])
        self.coef_sphere_true = gs.array([1.0, 0.0, 0.5])

        # set up the prior
        self.prior = lambda x: self.sphere.metric.exp(
            x * self.coef_sphere_true,
            base_point=self.intercept_sphere_true,
        )

        self.kernel = ConstantKernel(1.0, (1e-3, 1e3)) * RBF(10.0, (1e-2, 1e2))

        # generate data
        X = gs.linspace(0.0, 1.5 * gs.pi, self.n_samples)
        self.X_sphere = gs.reshape((X - gs.mean(X)), (-1, 1))
        # generate the geodesic
        y = self.prior(self.X_sphere)
        # Then add orthogonal sinusoidal oscillations

        o = (1.0 / 20.0) * gs.array([-0.5, 0.0, 1.0])
        o = self.sphere.to_tangent(o, base_point=y)
        s = self.X_sphere * gs.sin(5.0 * gs.pi * self.X_sphere)
        self.y_sphere = self.sphere.metric.exp(s * o, base_point=y)

    def test_fit_hypersphere(self):
        """Test the fit method"""
        wgpr = WrappedGaussianProcess(self.sphere,
                                      metric=self.sphere.metric,
                                      prior=self.prior,
                                      kernel=self.kernel)
        wgpr.fit(self.X_sphere, self.y_sphere)
        self.assertAllClose(wgpr.score(self.X_sphere, self.y_sphere), 1)

    def test_predict_hypersphere(self):
        """Test the predict method"""
        wgpr = WrappedGaussianProcess(self.sphere,
                                      metric=self.sphere.metric,
                                      prior=self.prior,
                                      kernel=self.kernel)
        wgpr.fit(self.X_sphere, self.y_sphere)
        y, std = wgpr.predict(self.X_sphere, return_tangent_std=True)
        self.assertAllClose(std, gs.zeros(std.shape), atol=1e-4)
        self.assertAllClose(y, self.y_sphere, atol=1e-4)

    def test_samples_y_hypersphere(self):
        """Test the samples_y method"""
        wgpr = WrappedGaussianProcess(self.sphere,
                                      metric=self.sphere.metric,
                                      prior=self.prior,
                                      kernel=self.kernel)
        wgpr.fit(self.X_sphere, self.y_sphere)
        y = wgpr.sample_y(self.X_sphere, n_samples=100)
        y_ = gs.reshape(gs.transpose(y, [0, 2, 1]), (-1, y.shape[1]))
        self.assertTrue(gs.all(self.sphere.belongs(y_)))
Ejemplo n.º 10
0
class TestFrechetMean(geomstats.tests.TestCase):
    _multiprocess_can_split_ = True

    def setUp(self):
        gs.random.seed(123)
        self.sphere = Hypersphere(dim=4)
        self.hyperbolic = Hyperboloid(dim=3)
        self.euclidean = Euclidean(dim=2)
        self.minkowski = Minkowski(dim=2)

    @geomstats.tests.np_only
    def test_logs_at_mean_default_gradient_descent_sphere(self):
        n_tests = 100
        estimator = FrechetMean(metric=self.sphere.metric, method='default')

        result = gs.zeros(n_tests)
        for i in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result[i] = gs.linalg.norm(logs[1, :] + logs[0, :])

        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result, rtol=1e-10, atol=1e-10)

    @geomstats.tests.np_only
    def test_logs_at_mean_adaptive_gradient_descent_sphere(self):
        n_tests = 100
        estimator = FrechetMean(metric=self.sphere.metric, method='adaptive')

        result = gs.zeros(n_tests)
        for i in range(n_tests):
            # take 2 random points, compute their mean, and verify that
            # log of each at the mean is opposite
            points = self.sphere.random_uniform(n_samples=2)
            estimator.fit(points)
            mean = estimator.estimate_

            logs = self.sphere.metric.log(point=points, base_point=mean)
            result[i] = gs.linalg.norm(logs[1, :] + logs[0, :])

        expected = gs.zeros(n_tests)
        self.assertAllClose(expected, result, rtol=1e-10, atol=1e-10)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_shape_default_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(points)
        result = mean.estimate_

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

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_shape_adaptive_gradient_descent_sphere(self):
        dim = 5
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(points)
        result = mean.estimate_

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

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_and_belongs_default_gradient_descent_sphere(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_and_belongs_adaptive_gradient_descent_sphere(self):
        point_a = gs.array([1., 0., 0., 0., 0.])
        point_b = gs.array([0., 1., 0., 0., 0.])
        points = gs.array([point_a, point_b])

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(points)

        result = self.sphere.belongs(mean.estimate_)
        expected = True
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_variance_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.array([point, point])

        result = variance(points, base_point=point, metric=self.sphere.metric)
        expected = gs.array(0.)

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_default_gradient_descent_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.sphere.metric, method='default')
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_adaptive_gradient_descent_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.sphere.metric, method='adaptive')
        mean.fit(X=points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_transform_sphere(self):
        point = gs.array([0., 0., 0., 0., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.sphere.metric)
        mean.fit(X=points)
        result = mean.transform(points)
        expected = gs.zeros_like(points)
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_estimate_transform_spd(self):
        point = SPDMatrices(3).random_uniform()
        points = gs.array([point, point])
        mean = FrechetMean(metric=SPDMetricAffine(3), point_type='matrix')
        mean.fit(X=points)
        result = mean.transform(points)
        expected = gs.zeros((2, 6))
        self.assertAllClose(expected, result)

    def test_fit_transform_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])
        mean = FrechetMean(metric=self.hyperbolic.metric)
        result = mean.fit_transform(X=points)
        expected = gs.zeros_like(points)
        self.assertAllClose(expected, result)

    def test_inverse_transform_hyperbolic(self):
        points = self.hyperbolic.random_uniform(10)
        mean = FrechetMean(metric=self.hyperbolic.metric)
        X = mean.fit_transform(X=points)
        result = mean.inverse_transform(X)
        expected = points
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_pytorch_only
    def test_inverse_transform_spd(self):
        point = SPDMatrices(3).random_uniform(10)
        mean = FrechetMean(metric=SPDMetricAffine(3), point_type='matrix')
        X = mean.fit_transform(X=point)
        result = mean.inverse_transform(X)
        expected = point
        self.assertAllClose(expected, result)

    @geomstats.tests.np_and_tf_only
    def test_variance_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])
        result = variance(points,
                          base_point=point,
                          metric=self.hyperbolic.metric)
        expected = gs.array(0.)

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_hyperbolic(self):
        point = gs.array([2., 1., 1., 1.])
        points = gs.array([point, point])

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)
        expected = point

        result = mean.estimate_

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_estimate_and_belongs_hyperbolic(self):
        point_a = self.hyperbolic.random_uniform()
        point_b = self.hyperbolic.random_uniform()
        point_c = self.hyperbolic.random_uniform()
        points = gs.stack([point_a, point_b, point_c], axis=0)

        mean = FrechetMean(metric=self.hyperbolic.metric)
        mean.fit(X=points)

        result = self.hyperbolic.belongs(mean.estimate_)
        expected = True

        self.assertAllClose(result, expected)

    def test_mean_euclidean_shape(self):
        dim = 2
        point = gs.array([1., 4.])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

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

    def test_mean_euclidean(self):
        point = gs.array([1., 4.])

        mean = FrechetMean(metric=self.euclidean.metric)
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

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

        mean = FrechetMean(metric=self.euclidean.metric)
        mean.fit(points, weights=weights)

        result = mean.estimate_
        expected = gs.array([16. / 6., 22. / 6.])

        self.assertAllClose(result, expected)

    def test_variance_euclidean(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 = variance(points,
                          weights=weights,
                          base_point=base_point,
                          metric=self.euclidean.metric)
        # we expect the average of the points' sq norms.
        expected = gs.array((1 * 5. + 2 * 13. + 1 * 25. + 2 * 41.) / 6.)

        self.assertAllClose(result, expected)

    def test_mean_matrices_shape(self):
        m, n = (2, 2)
        point = gs.array([[1., 4.], [2., 3.]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type='matrix')
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_

        self.assertAllClose(gs.shape(result), (m, n))

    def test_mean_matrices(self):
        m, n = (2, 2)
        point = gs.array([[1., 4.], [2., 3.]])

        metric = MatricesMetric(m, n)
        mean = FrechetMean(metric=metric, point_type='matrix')
        points = [point, point, point]
        mean.fit(points)

        result = mean.estimate_
        expected = point

        self.assertAllClose(result, expected)

    def test_mean_minkowski_shape(self):
        dim = 2
        point = gs.array([2., -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

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

    def test_mean_minkowski(self):
        point = gs.array([2., -math.sqrt(3)])
        points = [point, point, point]

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points)
        result = mean.estimate_

        expected = point

        self.assertAllClose(result, expected)

        points = gs.array([[1., 0.], [2., math.sqrt(3)], [3., math.sqrt(8)],
                           [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])

        mean = FrechetMean(metric=self.minkowski.metric)
        mean.fit(points, weights=weights)
        result = mean.estimate_
        result = self.minkowski.belongs(result)
        expected = gs.array(True)

        self.assertAllClose(result, expected)

    def test_variance_minkowski(self):
        points = gs.array([[1., 0.], [2., math.sqrt(3)], [3., math.sqrt(8)],
                           [4., math.sqrt(24)]])
        weights = gs.array([1., 2., 1., 2.])
        base_point = gs.array([-1., 0.])
        var = variance(points,
                       weights=weights,
                       base_point=base_point,
                       metric=self.minkowski.metric)
        result = var != 0
        # we expect the average of the points' Minkowski sq norms.
        expected = True
        self.assertAllClose(result, expected)
Ejemplo n.º 11
0
class TestHypersphere(geomstats.tests.TestCase):
    def setUp(self):
        gs.random.seed(1234)

        self.dimension = 4
        self.space = Hypersphere(dim=self.dimension)
        self.metric = self.space.metric
        self.n_samples = 10

    def test_random_uniform_and_belongs(self):
        """Test random uniform and belongs.

        Test that the random uniform method samples
        on the hypersphere space.
        """
        n_samples = self.n_samples
        point = self.space.random_uniform(n_samples)
        result = self.space.belongs(point)
        expected = gs.array([True] * n_samples)

        self.assertAllClose(expected, result)

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

        self.assertAllClose(gs.shape(point), (self.dimension + 1, ))

    def test_replace_values(self):
        points = gs.ones((3, 5))
        new_points = gs.zeros((2, 5))
        indcs = [True, False, True]
        update = self.space._replace_values(points, new_points, indcs)
        self.assertAllClose(update,
                            gs.stack([gs.zeros(5),
                                      gs.ones(5),
                                      gs.zeros(5)]))

    def test_projection_and_belongs(self):
        point = gs.array([1., 2., 3., 4., 5.])
        proj = self.space.projection(point)
        result = self.space.belongs(proj)
        expected = True

        self.assertAllClose(expected, result)

    def test_intrinsic_and_extrinsic_coords(self):
        """
        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        point_int = gs.array([.1, 0., 0., .1])
        point_ext = self.space.intrinsic_to_extrinsic_coords(point_int)
        result = self.space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int

        self.assertAllClose(result, expected)

        point_ext = (1. / (gs.sqrt(6.)) * gs.array([1., 0., 0., 1., 2.]))
        point_int = self.space.extrinsic_to_intrinsic_coords(point_ext)
        result = self.space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext

        self.assertAllClose(result, expected)

    def test_intrinsic_and_extrinsic_coords_vectorization(self):
        """Test change of coordinates.

        Test that the composition of
        intrinsic_to_extrinsic_coords and
        extrinsic_to_intrinsic_coords
        gives the identity.
        """
        point_int = gs.array([[.1, 0., 0., .1], [.1, .1, .1, .4],
                              [.1, .3, 0., .1], [-0.1, .1, -.4, .1],
                              [0., 0., .1, .1], [.1, .1, .1, .1]])
        point_ext = self.space.intrinsic_to_extrinsic_coords(point_int)
        result = self.space.extrinsic_to_intrinsic_coords(point_ext)
        expected = point_int

        self.assertAllClose(result, expected)

        point_int = self.space.extrinsic_to_intrinsic_coords(point_ext)
        result = self.space.intrinsic_to_extrinsic_coords(point_int)
        expected = point_ext

        self.assertAllClose(result, expected)

    def test_log_and_exp_general_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # General case
        base_point = gs.array([1., 2., 3., 4., 6.])
        base_point = base_point / gs.linalg.norm(base_point)
        point = gs.array([0., 5., 6., 2., -1.])
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected, atol=1e-6)

    def test_log_and_exp_edge_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Log then Riemannian Exp
        # Edge case: two very close points, base_point_2 and point_2,
        # form an angle < epsilon
        base_point = gs.array([1., 2., 3., 4., 6.])
        base_point = base_point / gs.linalg.norm(base_point)
        point = (base_point + 1e-4 * gs.array([-1., -2., 1., 1., .1]))
        point = point / gs.linalg.norm(point)

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

        self.assertAllClose(result, expected)

    def test_exp_vectorization_single_samples(self):
        dim = self.dimension + 1

        one_vec = self.space.random_uniform()
        one_base_point = self.space.random_uniform()
        one_tangent_vec = self.space.to_tangent(one_vec,
                                                base_point=one_base_point)

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

        one_base_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_tangent_vec = gs.to_ndarray(one_tangent_vec, to_ndim=2)
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_base_point = self.space.random_uniform()
        result = self.metric.exp(one_tangent_vec, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

    def test_exp_vectorization_n_samples(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_vec = self.space.random_uniform()
        one_base_point = self.space.random_uniform()
        n_vecs = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

        n_tangent_vecs = self.space.to_tangent(n_vecs,
                                               base_point=one_base_point)
        result = self.metric.exp(n_tangent_vecs, one_base_point)

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

        one_tangent_vec = self.space.to_tangent(one_vec,
                                                base_point=n_base_points)
        result = self.metric.exp(one_tangent_vec, n_base_points)

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

        n_tangent_vecs = self.space.to_tangent(n_vecs,
                                               base_point=n_base_points)
        result = self.metric.exp(n_tangent_vecs, n_base_points)

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

    def test_log_vectorization_single_samples(self):
        dim = self.dimension + 1

        one_base_point = self.space.random_uniform()
        one_point = self.space.random_uniform()

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

        one_base_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_point = gs.to_ndarray(one_base_point, to_ndim=2)
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

        one_base_point = self.space.random_uniform()
        result = self.metric.log(one_point, one_base_point)
        self.assertAllClose(gs.shape(result), (1, dim))

    def test_log_vectorization_n_samples(self):
        n_samples = self.n_samples
        dim = self.dimension + 1

        one_base_point = self.space.random_uniform()
        one_point = self.space.random_uniform()
        n_points = self.space.random_uniform(n_samples=n_samples)
        n_base_points = self.space.random_uniform(n_samples=n_samples)

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

        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_exp_log_are_inverse(self):
        initial_point = self.space.random_uniform(2)
        end_point = self.space.random_uniform(2)
        vec = self.space.metric.log(point=end_point, base_point=initial_point)
        result = self.space.metric.exp(vec, initial_point)
        self.assertAllClose(end_point, result)

    def test_log_extreme_case(self):
        initial_point = self.space.random_uniform(2)
        vec = 1e-4 * gs.random.rand(*initial_point.shape)
        vec = self.space.to_tangent(vec, initial_point)
        point = self.space.metric.exp(vec, base_point=initial_point)
        result = self.space.metric.log(point, initial_point)
        self.assertAllClose(vec, result)

    def test_exp_and_log_and_projection_to_tangent_space_general_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Exp then Riemannian Log
        # General case
        # NB: Riemannian log gives a regularized tangent vector,
        # so we take the norm modulo 2 * pi.
        base_point = gs.array([0., -3., 0., 3., 4.])
        base_point = base_point / gs.linalg.norm(base_point)

        vector = gs.array([3., 2., 0., 0., -1.])
        vector = self.space.to_tangent(vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        result = self.metric.log(point=exp, base_point=base_point)

        expected = vector
        norm_expected = gs.linalg.norm(expected)
        regularized_norm_expected = gs.mod(norm_expected, 2 * gs.pi)
        expected = expected / norm_expected * regularized_norm_expected

        # The Log can be the opposite vector on the tangent space,
        # whose Exp gives the base_point
        are_close = gs.allclose(result, expected)
        norm_2pi = gs.isclose(gs.linalg.norm(result - expected), 2 * gs.pi)
        self.assertTrue(are_close or norm_2pi)

    def test_exp_and_log_and_projection_to_tangent_space_edge_case(self):
        """Test Log and Exp.

        Test that the Riemannian exponential
        and the Riemannian logarithm are inverse.

        Expect their composition to give the identity function.

        NB: points on the n-dimensional sphere are
        (n+1)-D vectors of norm 1.
        """
        # Riemannian Exp then Riemannian Log
        # Edge case: tangent vector has norm < epsilon
        base_point = gs.array([10., -2., -.5, 34., 3.])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = 1e-4 * gs.array([.06, -51., 6., 5., 3.])
        vector = self.space.to_tangent(vector=vector, base_point=base_point)

        exp = self.metric.exp(tangent_vec=vector, base_point=base_point)
        result = self.metric.log(point=exp, base_point=base_point)
        self.assertAllClose(result, vector, atol=1e-7)

    def test_squared_norm_and_squared_dist(self):
        """
        Test that the squared distance between two points is
        the squared norm of their logarithm.
        """
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = (1. / gs.sqrt(435.) * gs.array([1., -20., -5., 0., 3.]))
        log = self.metric.log(point=point_a, base_point=point_b)
        result = self.metric.squared_norm(vector=log)
        expected = self.metric.squared_dist(point_a, point_b)

        self.assertAllClose(result, expected)

    def test_squared_dist_vectorization_single_sample(self):
        one_point_a = self.space.random_uniform()
        one_point_b = self.space.random_uniform()

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

        one_point_a = gs.to_ndarray(one_point_a, to_ndim=2)
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

        one_point_b = gs.to_ndarray(one_point_b, to_ndim=2)
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

        one_point_a = self.space.random_uniform()
        result = self.metric.squared_dist(one_point_a, one_point_b)
        self.assertAllClose(gs.shape(result), (1, ))

    def test_squared_dist_vectorization_n_samples(self):
        n_samples = self.n_samples

        one_point_a = self.space.random_uniform()
        one_point_b = self.space.random_uniform()
        n_points_a = self.space.random_uniform(n_samples=n_samples)
        n_points_b = self.space.random_uniform(n_samples=n_samples)

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

        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)
        self.assertAllClose(gs.shape(result), (n_samples, ))

        one_point_a = gs.to_ndarray(one_point_a, to_ndim=2)
        one_point_b = gs.to_ndarray(one_point_b, to_ndim=2)

        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)
        self.assertAllClose(gs.shape(result), (n_samples, ))

    def test_norm_and_dist(self):
        """
        Test that the distance between two points is
        the norm of their logarithm.
        """
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = (1. / gs.sqrt(435.) * gs.array([1., -20., -5., 0., 3.]))
        log = self.metric.log(point=point_a, base_point=point_b)

        self.assertAllClose(gs.shape(log), (5, ))

        result = self.metric.norm(vector=log)
        self.assertAllClose(gs.shape(result), ())

        expected = self.metric.dist(point_a, point_b)
        self.assertAllClose(gs.shape(expected), ())

        self.assertAllClose(result, expected)

    def test_dist_point_and_itself(self):
        # Distance between a point and itself is 0
        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = point_a
        result = self.metric.dist(point_a, point_b)
        expected = 0.

        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_pytorch_only
    def test_dist_pairwise(self):

        point_a = (1. / gs.sqrt(129.) * gs.array([10., -2., -5., 0., 0.]))
        point_b = (1. / gs.sqrt(435.) * gs.array([1., -20., -5., 0., 3.]))

        point = gs.array([point_a, point_b])
        result = self.metric.dist_pairwise(point)

        expected = gs.array([[0., 1.24864502], [1.24864502, 0.]])

        self.assertAllClose(result, expected, rtol=1e-3)

    def test_dist_orthogonal_points(self):
        # Distance between two orthogonal points is pi / 2.
        point_a = gs.array([10., -2., -.5, 0., 0.])
        point_a = point_a / gs.linalg.norm(point_a)
        point_b = gs.array([2., 10, 0., 0., 0.])
        point_b = point_b / gs.linalg.norm(point_b)
        result = gs.dot(point_a, point_b)
        expected = 0
        self.assertAllClose(result, expected)

        result = self.metric.dist(point_a, point_b)
        expected = gs.pi / 2

        self.assertAllClose(result, expected)

    def test_exp_and_dist_and_projection_to_tangent_space(self):
        base_point = gs.array([16., -2., -2.5, 84., 3.])
        base_point = base_point / gs.linalg.norm(base_point)
        vector = gs.array([9., 0., -1., -2., 1.])
        tangent_vec = self.space.to_tangent(vector=vector,
                                            base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)
        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec) % (2 * gs.pi)
        self.assertAllClose(result, expected)

    def test_exp_and_dist_and_projection_to_tangent_space_vec(self):
        base_point = gs.array([[16., -2., -2.5, 84., 3.],
                               [16., -2., -2.5, 84., 3.]])

        base_single_point = gs.array([16., -2., -2.5, 84., 3.])
        scalar_norm = gs.linalg.norm(base_single_point)

        base_point = base_point / scalar_norm
        vector = gs.array([[9., 0., -1., -2., 1.], [9., 0., -1., -2., 1]])

        tangent_vec = self.space.to_tangent(vector=vector,
                                            base_point=base_point)

        exp = self.metric.exp(tangent_vec=tangent_vec, base_point=base_point)

        result = self.metric.dist(base_point, exp)
        expected = gs.linalg.norm(tangent_vec, axis=-1) % (2 * gs.pi)

        self.assertAllClose(result, expected)

    def test_geodesic_and_belongs(self):
        n_geodesic_points = 10
        initial_point = self.space.random_uniform(2)
        vector = gs.array([[2., 0., -1., -2., 1.]] * 2)
        initial_tangent_vec = self.space.to_tangent(vector=vector,
                                                    base_point=initial_point)
        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 = gs.stack([self.space.belongs(pt) for pt in points])
        self.assertTrue(gs.all(result))

        initial_point = initial_point[0]
        initial_tangent_vec = initial_tangent_vec[0]
        geodesic = self.metric.geodesic(
            initial_point=initial_point,
            initial_tangent_vec=initial_tangent_vec)
        points = geodesic(t)
        result = self.space.belongs(points)
        expected = gs.array(n_geodesic_points * [True])
        self.assertAllClose(expected, result)

    def test_geodesic_end_point(self):
        n_geodesic_points = 10
        initial_point = self.space.random_uniform(4)
        geodesic = self.metric.geodesic(initial_point=initial_point[:2],
                                        end_point=initial_point[2:])
        t = gs.linspace(start=0., stop=1., num=n_geodesic_points)
        points = geodesic(t)
        result = points[-1]
        expected = initial_point[2:]
        self.assertAllClose(expected, result)

    def test_inner_product(self):
        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([0., 0., 0., 0., 1.])
        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = 0.

        self.assertAllClose(expected, result)

    def test_inner_product_vectorization_single_samples(self):
        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([0., 0., 0., 0., 1.])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = 0.
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1., 0., 0., 0., 0.]])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([0., 0., 0., 0., 1.])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([[0., 1., 0., 0., 0.]])
        base_point = gs.array([0., 0., 0., 0., 1.])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([[0., 0., 0., 0., 1.]])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1., 0., 0., 0., 0.]])
        tangent_vec_b = gs.array([[0., 1., 0., 0., 0.]])
        base_point = gs.array([0., 0., 0., 0., 1.])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([1., 0., 0., 0., 0.])
        tangent_vec_b = gs.array([[0., 1., 0., 0., 0.]])
        base_point = gs.array([[0., 0., 0., 0., 1.]])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1., 0., 0., 0., 0.]])
        tangent_vec_b = gs.array([0., 1., 0., 0., 0.])
        base_point = gs.array([[0., 0., 0., 0., 1.]])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

        tangent_vec_a = gs.array([[1., 0., 0., 0., 0.]])
        tangent_vec_b = gs.array([[0., 1., 0., 0., 0.]])
        base_point = gs.array([[0., 0., 0., 0., 1.]])

        result = self.metric.inner_product(tangent_vec_a, tangent_vec_b,
                                           base_point)
        expected = gs.array([0.])
        self.assertAllClose(expected, result)

    def test_diameter(self):
        dim = 2
        sphere = Hypersphere(dim)
        point_a = gs.array([[0., 0., 1.]])
        point_b = gs.array([[1., 0., 0.]])
        point_c = gs.array([[0., 0., -1.]])
        result = sphere.metric.diameter(gs.vstack((point_a, point_b, point_c)))
        expected = gs.pi
        self.assertAllClose(expected, result)

    def test_closest_neighbor_index(self):
        """Check that the closest neighbor is one of neighbors."""
        n_samples = 10
        points = self.space.random_uniform(n_samples=n_samples)
        point = points[0, :]
        neighbors = points[1:, :]
        index = self.metric.closest_neighbor_index(point, neighbors)
        closest_neighbor = points[index, :]

        test = gs.sum(gs.all(points == closest_neighbor, axis=1))
        result = test > 0
        self.assertTrue(result)

    def test_sample_von_mises_fisher(self):
        """
        Check that the maximum likelihood estimates of the mean and
        concentration parameter are close to the real values. A first
        estimation of the concentration parameter is obtained by a
        closed-form expression and improved through the Newton method.
        """
        dim = 2
        n_points = 1000000
        sphere = Hypersphere(dim)

        # check mean value for concentrated distribution
        kappa = 10000000
        points = sphere.random_von_mises_fisher(kappa, n_points)
        sum_points = gs.sum(points, axis=0)
        mean = gs.array([0., 0., 1.])
        mean_estimate = sum_points / gs.linalg.norm(sum_points)
        expected = mean
        result = mean_estimate
        self.assertTrue(gs.allclose(result, expected,
                                    atol=MEAN_ESTIMATION_TOL))
        # check concentration parameter for dispersed distribution
        kappa = 1.
        points = sphere.random_von_mises_fisher(kappa, n_points)
        sum_points = gs.sum(points, axis=0)
        mean_norm = gs.linalg.norm(sum_points) / n_points
        kappa_estimate = (mean_norm * (dim + 1. - mean_norm**2) /
                          (1. - mean_norm**2))
        kappa_estimate = gs.cast(kappa_estimate, gs.float64)
        p = dim + 1
        n_steps = 100
        for _ in range(n_steps):
            bessel_func_1 = scipy.special.iv(p / 2., kappa_estimate)
            bessel_func_2 = scipy.special.iv(p / 2. - 1., kappa_estimate)
            ratio = bessel_func_1 / bessel_func_2
            denominator = 1. - ratio**2 - (p - 1.) * ratio / kappa_estimate
            mean_norm = gs.cast(mean_norm, gs.float64)
            kappa_estimate = kappa_estimate - (ratio - mean_norm) / denominator
        result = kappa_estimate
        expected = kappa
        self.assertAllClose(result, expected, atol=KAPPA_ESTIMATION_TOL)

    def test_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)

        points_spherical = gs.array([gs.pi / 2, 0])
        result = sphere.spherical_to_extrinsic(points_spherical)
        expected = gs.array([1., 0., 0.])
        self.assertAllClose(result, expected)

    def test_spherical_to_extrinsic_vectorization(self):
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        result = sphere.spherical_to_extrinsic(points_spherical)
        expected = gs.array(
            [[1., 0., 0.],
             [gs.sqrt(2.) / 4.,
              gs.sqrt(2.) / 4.,
              gs.sqrt(3.) / 2.]])
        self.assertAllClose(result, expected)

    def test_tangent_spherical_to_extrinsic(self):
        """
        Check vectorization of conversion from spherical
        to extrinsic coordinates for tangent vectors to the
        2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        base_points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 2, 0]])
        tangent_vecs_spherical = gs.array([[0.25, 0.5], [0.3, 0.2]])
        result = sphere.tangent_spherical_to_extrinsic(tangent_vecs_spherical,
                                                       base_points_spherical)
        expected = gs.array([[0, 0.5, -0.25], [0, 0.2, -0.3]])
        self.assertAllClose(result, expected)

    def test_christoffels_vectorization(self):
        """
        Check vectorization of Christoffel symbols in
        spherical coordinates on the 2-sphere.
        """
        dim = 2
        sphere = Hypersphere(dim)
        points_spherical = gs.array([[gs.pi / 2, 0], [gs.pi / 6, gs.pi / 4]])
        christoffel = sphere.metric.christoffels(points_spherical)
        result = christoffel.shape
        expected = gs.array([2, dim, dim, dim])
        self.assertAllClose(result, expected)

    @geomstats.tests.np_and_tf_only
    def test_parallel_transport_vectorization(self):
        sphere = Hypersphere(2)
        n_samples = 4

        def is_isometry(tan_a, trans_a, endpoint):
            is_tangent = gs.isclose(sphere.metric.inner_product(
                endpoint, trans_a),
                                    0.,
                                    atol=1e-6)
            is_equinormal = gs.isclose(gs.linalg.norm(trans_a, axis=-1),
                                       gs.linalg.norm(tan_a, axis=-1))
            return gs.logical_and(is_tangent, is_equinormal)

        base_point = sphere.random_uniform(n_samples)
        tan_vec_a = sphere.to_tangent(gs.random.rand(n_samples, 3), base_point)
        tan_vec_b = sphere.to_tangent(gs.random.rand(n_samples, 3), base_point)
        end_point = sphere.metric.exp(tan_vec_b, base_point)

        transported = sphere.metric.parallel_transport(tan_vec_a, tan_vec_b,
                                                       base_point)
        result = is_isometry(tan_vec_a, transported, end_point)
        self.assertTrue(gs.all(result))

        base_point = base_point[0]
        tan_vec_a = sphere.to_tangent(tan_vec_a, base_point)
        tan_vec_b = sphere.to_tangent(tan_vec_b, base_point)
        end_point = sphere.metric.exp(tan_vec_b, base_point)
        transported = sphere.metric.parallel_transport(tan_vec_a, tan_vec_b,
                                                       base_point)
        result = is_isometry(tan_vec_a, transported, end_point)
        self.assertTrue(gs.all(result))

        one_tan_vec_a = tan_vec_a[0]
        transported = sphere.metric.parallel_transport(one_tan_vec_a,
                                                       tan_vec_b, base_point)
        result = is_isometry(one_tan_vec_a, transported, end_point)
        self.assertTrue(gs.all(result))

        one_tan_vec_b = tan_vec_b[0]
        end_point = end_point[0]
        transported = sphere.metric.parallel_transport(tan_vec_a,
                                                       one_tan_vec_b,
                                                       base_point)
        result = is_isometry(tan_vec_a, transported, end_point)
        self.assertTrue(gs.all(result))

        transported = sphere.metric.parallel_transport(one_tan_vec_a,
                                                       one_tan_vec_b,
                                                       base_point)
        result = is_isometry(one_tan_vec_a, transported, end_point)
        self.assertTrue(result)