How to use geomstats - 10 common examples

def test_predict_hypersphere_distance(self):
        """Test the 'predict' class method using the hypersphere distance."""
        dim = 2
        space = Hypersphere(dim=dim)
        distance = space.metric.dist
        training_dataset = gs.array(
            [[1, 0, 0],
             [3 ** (1 / 2) / 2, 1 / 2, 0],
             [3 ** (1 / 2) / 2, - 1 / 2, 0],
             [0, 0, 1],
             [0, 1 / 2, 3 ** (1 / 2) / 2],
             [0, - 1 / 2, 3 ** (1 / 2) / 2]])
        labels = [0, 0, 0, 1, 1, 1]
        kde = KernelDensityEstimationClassifier(
            distance=distance)
        kde.fit(training_dataset, labels)
        target_dataset = gs.array(
            [[2 ** (1 / 2) / 2, 2 ** (1 / 2) / 2, 0],
             [0, 1 / 2, - 3 ** (1 / 2) / 2],
             [0, - 1 / 2, - 3 ** (1 / 2) / 2],
             [- 3 ** (1 / 2) / 2, 1 / 2, 0],
             [- 3 ** (1 / 2) / 2, - 1 / 2, 0],

def test_predict_proba(self):
        """Test the 'predict_proba' class method."""
        training_dataset = gs.array([[0], [1], [2], [3]])
        labels = [0, 0, 1, 1]
        neigh = KNearestNeighborsClassifier(n_neighbors=self.n_neighbors,
                                            distance=self.distance)
        neigh.fit(training_dataset, labels)
        result = neigh.predict_proba([[0.9]])
        expected = gs.array([[2 / 3, 1 / 3]])
        self.assertAllClose(expected, result, atol=TOLERANCE)

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

    @geomstats.tests.np_only
    def test_estimate_and_reach_max_iter_se(self):
        point = self.se_mat.random_uniform(1)
        estimator = ExponentialBarycenter(self.se_mat, max_iter=2)
        points = gs.array([point, point])
        estimator.fit(points)
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

        point = self.so_vec.random_uniform(1)
        estimator = ExponentialBarycenter(self.so_vec, max_iter=2)
        points = gs.array([point, point])
        estimator.fit(points)
        result = estimator.estimate_
        expected = point
        self.assertAllClose(result, expected)

    @geomstats.tests.np_only
    def test_spd_kmeans_fit(self):
        gs.random.seed(0)
        n_points = 100
        space = spd_matrices.SPDMatrices(10)
        data = space.random_uniform(n_samples=n_points)
        metric = spd_matrices.SPDMetricAffine(10)

        kmeans = RiemannianKMeans(metric, 1, point_type='matrix')
        kmeans.fit(data)
        result = kmeans.centroids
        mean = FrechetMean(metric=metric, point_type='matrix', max_iter=100)
        mean.fit(data)
        expected = mean.estimate_
        self.assertAllClose(result, expected, atol=1e-2, rtol=1e-2)

    @geomstats.tests.np_only
    def test_tangent_pca(self):
        X = self.X
        tpca = TangentPCA(self.metric, n_components=gs.shape(X)[1])
        tpca.fit(X)
        self.assertEqual(tpca.n_features_, gs.shape(X)[1])

def test_random_and_belongs_vectorization(self):
        n_samples = self.n_samples
        points = self.group.random_uniform(n_samples=n_samples)
        result = self.group.belongs(points)
        expected = gs.array([True] * n_samples)
        self.assertAllClose(result, expected)

def test_Localization_propagate(self):
        initial_state = gs.array([0.5, 1., 2.])
        time_step = gs.array([0.5])
        linear_vel = gs.array([1., 0.5])
        angular_vel = gs.array([0.])
        increment = gs.concatenate((
            time_step, linear_vel, angular_vel), axis=0)

        angle = initial_state[0]
        rotation = gs.array([[gs.cos(angle), -gs.sin(angle)],
                             [gs.sin(angle), gs.cos(angle)]])
        next_position = initial_state[1:] + time_step * gs.matmul(
            rotation, linear_vel)
        expected = gs.concatenate((gs.array([angle]), next_position), axis=0)
        result = self.nonlinear_model.propagate(initial_state, increment)
        self.assertAllClose(expected, result)

def test_triangular_radial_kernel(self):
        """Test the triangular radial kernel."""
        distance = gs.array([[1], [2]], dtype=float)
        bandwidth = 2
        weight = triangular_radial_kernel(
            distance=distance,
            bandwidth=bandwidth)
        result = weight
        expected = gs.array([[1 / 2], [0]], dtype=float)
        self.assertAllClose(expected, result, atol=TOLERANCE)

def test_metric_matrix(self):
        base_point = gs.array([0., 1., 0., 0.])

        result = self.euc_metric.inner_product_matrix(base_point)
        expected = gs.eye(self.dim)

        self.assertAllClose(result, expected)