def test_geodesic_invalid_initial_conditions(self): space = SpecialOrthogonal(n=4) initial_point = space.random_uniform(2) vector = gs.random.rand(2, 4, 4) initial_tangent_vec = space.to_tangent(vector=vector, base_point=initial_point) end_point = space.random_uniform(2) self.assertRaises( RuntimeError, lambda: space.bi_invariant_metric.geodesic( initial_point=initial_point, initial_tangent_vec=initial_tangent_vec, end_point=end_point))
class TestTangentPCA(geomstats.tests.TestCase): _multiprocess_can_split_ = True def setUp(self): self.so3 = SpecialOrthogonal(n=3) self.n_samples = 10 self.X = self.so3.random_uniform(n_samples=self.n_samples) self.metric = self.so3.bi_invariant_metric self.n_components = 2 @geomstats.tests.np_only def test_tangent_pca_error(self): X = self.X trans = TangentPCA(self.metric, n_components=self.n_components) trans.fit(X) X_diff_size = gs.ones((self.n_samples, gs.shape(X)[1] + 1)) self.assertRaises(ValueError, trans.transform, X_diff_size) @geomstats.tests.np_only def test_tangent_pca(self): X = self.X trans = TangentPCA(self.metric, n_components=gs.shape(X)[1]) trans.fit(X) self.assertEquals(trans.n_features_, gs.shape(X)[1])
def test_geodesic_and_coincides_exp_son(self): n_geodesic_points = 10 space = SpecialOrthogonal(n=4) initial_point = space.random_uniform(2) vector = gs.random.rand(2, 4, 4) initial_tangent_vec = space.to_tangent(vector=vector, base_point=initial_point) geodesic = space.bi_invariant_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 = points[:, -1] expected = space.bi_invariant_metric.exp(initial_tangent_vec, initial_point) self.assertAllClose(result, expected) initial_point = initial_point[0] initial_tangent_vec = initial_tangent_vec[0] geodesic = space.bi_invariant_metric.geodesic( initial_point=initial_point, initial_tangent_vec=initial_tangent_vec) points = geodesic(t) result = points[-1] expected = space.bi_invariant_metric.exp(initial_tangent_vec, initial_point) self.assertAllClose(expected, result)
class TestVisualization(geomstats.tests.TestCase): def setUp(self): self.n_samples = 10 self.SO3_GROUP = SpecialOrthogonal(n=3, point_type='vector') self.SE3_GROUP = SpecialEuclidean(n=3, point_type='vector') self.S1 = Hypersphere(dim=1) self.S2 = Hypersphere(dim=2) self.H2 = Hyperbolic(dim=2) self.H2_half_plane = PoincareHalfSpace(dim=2) plt.figure() @staticmethod def test_tutorial_matplotlib(): visualization.tutorial_matplotlib() def test_plot_points_so3(self): points = self.SO3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space='SO3_GROUP') def test_plot_points_se3(self): points = self.SE3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space='SE3_GROUP') @geomstats.tests.np_and_pytorch_only def test_plot_points_s1(self): points = self.S1.random_uniform(self.n_samples) visualization.plot(points, space='S1') def test_plot_points_s2(self): points = self.S2.random_uniform(self.n_samples) visualization.plot(points, space='S2') def test_plot_points_h2_poincare_disk(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_disk') def test_plot_points_h2_poincare_half_plane_ext(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane', point_type='extrinsic') def test_plot_points_h2_poincare_half_plane_none(self): points = self.H2_half_plane.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane') def test_plot_points_h2_poincare_half_plane_hs(self): points = self.H2_half_plane.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane', point_type='half_space') def test_plot_points_h2_klein_disk(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_klein_disk')
class TestVisualizationMethods(geomstats.tests.TestCase): def setUp(self): self.n_samples = 10 self.SO3_GROUP = SpecialOrthogonal(n=3) self.SE3_GROUP = SpecialEuclidean(n=3) self.S1 = Hypersphere(dim=1) self.S2 = Hypersphere(dim=2) self.H2 = Hyperbolic(dim=2) plt.figure() @geomstats.tests.np_only def test_plot_points_so3(self): points = self.SO3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space='SO3_GROUP') @geomstats.tests.np_only def test_plot_points_se3(self): points = self.SE3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space='SE3_GROUP') @geomstats.tests.np_only def test_plot_points_s1(self): points = self.S1.random_uniform(self.n_samples) visualization.plot(points, space='S1') @geomstats.tests.np_only def test_plot_points_s2(self): points = self.S2.random_uniform(self.n_samples) visualization.plot(points, space='S2') @geomstats.tests.np_only def test_plot_points_h2_poincare_disk(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_disk') @geomstats.tests.np_only def test_plot_points_h2_poincare_half_plane(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane') @geomstats.tests.np_only def test_plot_points_h2_klein_disk(self): points = self.H2.random_uniform(self.n_samples) visualization.plot(points, space='H2_klein_disk')
def test_double_cluster_riemannian_mean_shift(self): gs.random.seed(10) number_of_samples = 20 sphere = Hypersphere(dim=2) metric = HypersphereMetric(2) cluster = sphere.random_von_mises_fisher(kappa=20, n_samples=number_of_samples) special_orthogonal = SpecialOrthogonal(3) rotation1 = special_orthogonal.random_uniform() rotation2 = special_orthogonal.random_uniform() cluster_1 = cluster @ rotation1 cluster_2 = cluster @ rotation2 combined_cluster = gs.concatenate((cluster_1, cluster_2)) rms = riemannian_mean_shift(manifold=sphere, metric=metric, bandwidth=0.3, tol=1e-4, n_centers=2) rms.fit(combined_cluster) closest_centers = rms.predict(combined_cluster) count_in_first_cluster = 0 for point in closest_centers: if gs.allclose(point, rms.centers[0]): count_in_first_cluster += 1 count_in_second_cluster = 0 for point in closest_centers: if gs.allclose(point, rms.centers[1]): count_in_second_cluster += 1 self.assertEqual(combined_cluster.shape[0], count_in_first_cluster + count_in_second_cluster)
def _init_double_cluster( seed=10, num_of_samples=20, size_of_dim=2, kappa_value=20, orthogonality_of_sphere=3, bandwidth=0.3, tol=1e-4, num_of_centers=2, ): gs.random.seed(seed) number_of_samples = num_of_samples sphere = Hypersphere(size_of_dim) metric = sphere.metric cluster = sphere.random_von_mises_fisher( kappa=kappa_value, n_samples=number_of_samples ) special_orthogonal = SpecialOrthogonal(orthogonality_of_sphere) rotation1 = special_orthogonal.random_uniform() rotation2 = special_orthogonal.random_uniform() cluster_1 = cluster @ rotation1 cluster_2 = cluster @ rotation2 combined_cluster = gs.concatenate((cluster_1, cluster_2)) rms = riemannian_mean_shift( manifold=sphere, metric=metric, bandwidth=bandwidth, tol=tol, n_centers=num_of_centers, ) rms.fit(combined_cluster) return combined_cluster, rms
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 TestTangentPCA(geomstats.tests.TestCase): _multiprocess_can_split_ = True def setUp(self): self.so3 = SpecialOrthogonal(n=3) self.spd = SPDMatrices(3) self.spd_metric = SPDMetricAffine(3) self.n_samples = 10 self.X = self.so3.random_uniform(n_samples=self.n_samples) self.metric = self.so3.bi_invariant_metric self.n_components = 2 @geomstats.tests.np_only def test_tangent_pca_error(self): X = self.X tpca = TangentPCA(self.metric, n_components=self.n_components) tpca.fit(X) X_diff_size = gs.ones((self.n_samples, gs.shape(X)[1] + 1)) self.assertRaises(ValueError, tpca.transform, X_diff_size) @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]) @geomstats.tests.np_only def test_fit_mle(self): X = self.X tpca = TangentPCA(self.metric, n_components='mle') tpca.fit(X) self.assertEqual(tpca.n_features_, gs.shape(X)[1]) @geomstats.tests.np_only def test_fit_to_target_explained_variance(self): X = self.spd.random_uniform(n_samples=5) target = 0.90 tpca = TangentPCA(self.spd_metric, point_type='matrix', n_components=target) tpca.fit(X) result = gs.cumsum(tpca.explained_variance_ratio_)[-1] > target expected = True self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_matrix(self): expected = 2 X = self.spd.random_uniform(n_samples=5) tpca = TangentPCA(metric=self.spd_metric, point_type='matrix', n_components=expected) tpca.fit(X) result = tpca.n_components_ self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_transform_matrix(self): expected = 2 X = self.spd.random_uniform(n_samples=5) tpca = TangentPCA(metric=self.spd_metric, point_type='matrix', n_components=expected) tangent_projected_data = tpca.fit_transform(X) result = tangent_projected_data.shape[-1] self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_inverse_transform_matrix(self): X = self.spd.random_uniform(n_samples=5) tpca = TangentPCA(metric=self.spd_metric, point_type='matrix') tangent_projected_data = tpca.fit_transform(X) result = tpca.inverse_transform(tangent_projected_data) expected = X self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_transform_vector(self): expected = 2 tpca = TangentPCA(metric=self.metric, point_type='vector', n_components=expected) tangent_projected_data = tpca.fit_transform(self.X) result = tangent_projected_data.shape[-1] self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_inverse_transform_vector(self): tpca = TangentPCA(metric=self.metric, point_type='vector') tangent_projected_data = tpca.fit_transform(self.X) result = tpca.inverse_transform(tangent_projected_data) expected = self.X self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_fit_transform_matrix(self): X = self.spd.random_uniform(n_samples=5) tpca = TangentPCA(metric=self.spd_metric, point_type='matrix') expected = tpca.fit_transform(X) result = tpca.fit(X).transform(X) self.assertAllClose(result, expected) @geomstats.tests.np_only def test_fit_matrix_se(self): se_mat = SpecialEuclidean(n=3, default_point_type='matrix') X = se_mat.random_uniform(self.n_samples) estimator = ExponentialBarycenter(se_mat) estimator.fit(X) mean = estimator.estimate_ tpca = TangentPCA(metric=se_mat, point_type='matrix') tangent_projected_data = tpca.fit_transform(X, base_point=mean) result = tpca.inverse_transform(tangent_projected_data) expected = X self.assertAllClose(result, expected)
class TestBackends(tests.conftest.TestCase): def setup_method(self): warnings.simplefilter("ignore", category=ImportWarning) self.so3_group = SpecialOrthogonal(n=3) self.n_samples = 2 def test_array(self): gs_mat = gs.array(1.5) np_mat = _np.array(1.5) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)]) np_mat = _np.array([_np.ones(3), _np.ones(3)]) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)], dtype=gs.float64) np_mat = _np.array([_np.ones(3), _np.ones(3)], dtype=_np.float64) self.assertTrue(gs_mat.dtype == gs.float64) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([[gs.ones(3)], [gs.ones(3)]], dtype=gs.uint8) np_mat = _np.array([[_np.ones(3)], [_np.ones(3)]], dtype=_np.uint8) self.assertTrue(gs_mat.dtype == gs.uint8) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), [0, 0, 0]], dtype=gs.int32) np_mat = _np.array([_np.ones(3), [0, 0, 0]], dtype=_np.int32) self.assertTrue(gs_mat.dtype == gs.int32) self.assertAllCloseToNp(gs_mat, np_mat) def test_matmul(self): mat_a = [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [7.0, 0.0, 4.0]] mat_b = [[1.0, 0.0, 2.0], [0.0, 3.0, 0.0], [0.0, 0.0, 1.0]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) gs_result = gs.matmul(gs_mat_a, gs_mat_b) np_result = _np.matmul(np_mat_a, np_mat_b) self.assertAllCloseToNp(gs_result, np_result) @tests.conftest.np_autograd_and_tf_only def test_matmul_vectorization(self): mat_a = [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [7.0, 0.0, 4.0]] mat_b = [[1.0, 0.0, 2.0], [0.0, 3.0, 0.0], [0.0, 0.0, 1.0]] mat_c = [[1.0, 4.0, 2.0], [4.0, 3.0, 4.0], [0.0, 0.0, 4.0]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) gs_mat_c = gs.array(mat_c) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) np_mat_c = _np.array(mat_c) gs_result = gs.matmul(gs_mat_a, [gs_mat_b, gs_mat_c]) np_result = _np.matmul(np_mat_a, [np_mat_b, np_mat_c]) self.assertAllCloseToNp(gs_result, np_result) def test_logm(self): point = gs.array([[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]]) result = gs.linalg.logm(point) expected = gs.array([[0.693147180, 0.0, 0.0], [0.0, 1.098612288, 0.0], [0.0, 0.0, 1.38629436]]) self.assertAllClose(result, expected) np_point = _np.array([[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]]) scipy_result = scipy.linalg.logm(np_point) self.assertAllCloseToNp(result, scipy_result) def test_expm_and_logm(self): point = gs.array([[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) np_point = _np.array([[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]]) scipy_result = scipy.linalg.expm(scipy.linalg.logm(np_point)) self.assertAllCloseToNp(result, scipy_result) @tests.conftest.np_and_autograd_only def test_expm_vectorization(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([ [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]], [[1.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 6.0]], ]) expected = gs.array([ [ [7.38905609, 0.0, 0.0], [0.0, 20.0855369, 0.0], [0.0, 0.0, 54.5981500], ], [ [2.718281828, 0.0, 0.0], [0.0, 148.413159, 0.0], [0.0, 0.0, 403.42879349], ], ]) result = gs.linalg.expm(point) self.assertAllClose(result, expected) def test_logm_vectorization_diagonal(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([ [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]], [[1.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 6.0]], ]) expected = gs.array([ [ [0.693147180, 0.0, 0.0], [0.0, 1.09861228866, 0.0], [0.0, 0.0, 1.38629436], ], [[0.0, 0.0, 0.0], [0.0, 1.609437912, 0.0], [0.0, 0.0, 1.79175946]], ]) result = gs.linalg.logm(point) self.assertAllClose(result, expected) def test_expm_and_logm_vectorization_random_rotation(self): point = self.so3_group.random_uniform(self.n_samples) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected, atol=gs.atol * 100) def test_expm_and_logm_vectorization(self): point = gs.array([ [[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]], [[1.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 6.0]], ]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @tests.conftest.tf_only def test_vstack(self): import tensorflow as tf tensor_1 = tf.convert_to_tensor([1.0, 2.0, 3.0]) tensor_2 = tf.convert_to_tensor([7.0, 8.0, 9.0]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([[1.0, 2.0, 3.0], [7.0, 8.0, 9.0]]) self.assertAllClose(result, expected) tensor_1 = tf.convert_to_tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) tensor_2 = tf.convert_to_tensor([7.0, 8.0, 9.0]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) self.assertAllClose(result, expected) def test_cumprod(self): result = gs.cumprod(gs.arange(1, 10)) expected = gs.array(([1, 2, 6, 24, 120, 720, 5040, 40320, 362880])) self.assertAllClose(result, expected) result = gs.reshape(gs.arange(1, 11), (2, 5)) result = gs.cumprod(result, axis=1) expected = gs.array(([[1, 2, 6, 24, 120], [6, 42, 336, 3024, 30240]])) self.assertAllClose(result, expected) @tests.conftest.torch_only def test_cumsum(self): result = gs.cumsum(gs.arange(10)) expected = gs.array(([0, 1, 3, 6, 10, 15, 21, 28, 36, 45])) self.assertAllClose(result, expected) result = gs.cumsum(gs.arange(10).reshape(2, 5), axis=1) expected = gs.array(([[0, 1, 3, 6, 10], [5, 11, 18, 26, 35]])) self.assertAllClose(result, expected) def test_array_from_sparse(self): expected = gs.array([[0, 1, 0], [0, 0, 2]]) result = gs.array_from_sparse([(0, 1), (1, 2)], [1, 2], (2, 3)) self.assertAllClose(result, expected) def test_einsum(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum("...i,...i->...", np_array_1, np_array_2) gs_result = gs.einsum("...i,...i->...", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4], [-1, 5]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4], [-1, 5]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum("...i,...i->...", np_array_1, np_array_2) gs_result = gs.einsum("...i,...i->...", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3], [5, 6]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3], [5, 6]]) np_result = _np.einsum("...i,...i->...", np_array_1, np_array_2) gs_result = gs.einsum("...i,...i->...", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array([5]) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum("...,...i->...i", np_array_1, np_array_2) gs_result = gs.einsum("...,...i->...i", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array(5) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum("...,...i->...i", np_array_1, np_array_2) gs_result = gs.einsum("...,...i->...i", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array([5]) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum("...,...i->...i", np_array_1, np_array_2) gs_result = gs.einsum("...,...i->...i", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array(5) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum("...,...i->...i", np_array_1, np_array_2) gs_result = gs.einsum("...,...i->...i", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_einsum_dtypes(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2.0, 3.0]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2.0, 3.0]]) np_result = _np.einsum("...i,...i->...", np_array_1, np_array_2) gs_result = gs.einsum("...i,...i->...", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1.0, 4.0], [-1.0, 5.0]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1.0, 4.0], [-1.0, 5.0]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum("...i,...i->...", np_array_1, np_array_2) gs_result = gs.einsum("...i,...i->...", array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_matrices(self): np_array = _np.zeros((2, 3, 3)) gs_array = gs.zeros((2, 3, 3)) np_array[:, 0, 1] = 44.0 gs_array = gs.assignment(gs_array, 44.0, (0, 1), axis=0) self.assertAllCloseToNp(gs_array, np_array) n_samples = 3 theta = _np.random.rand(5) phi = _np.random.rand(5) np_array = _np.zeros((n_samples, 5, 4)) gs_array = gs.array(np_array) np_array[0, :, 0] = gs.cos(theta) * gs.cos(phi) np_array[0, :, 1] = -gs.sin(theta) * gs.sin(phi) gs_array = gs.assignment(gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment(gs_array, -gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) self.assertAllCloseToNp(gs_array, np_array) def test_assignment_with_booleans_single_index(self): np_array = _np.array([[2.0, 5.0]]) gs_array = gs.array([[2.0, 5.0]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2.0, 5.0]]) gs_array = gs.array([[2.0, 5.0]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2.0, 5.0]]) gs_array = gs.array([[2.0, 5.0]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2.0, 5.0]]) gs_array = gs.array([[2.0, 5.0]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_booleans_many_indices(self): np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([True, False, True]) gs_mask = gs.array([True, False, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([False, True, True]) gs_mask = gs.array([False, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([False, False, False]) gs_mask = gs.array([False, False, False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment(self): gs_array_1 = gs.ones(3) with pytest.raises(ValueError): gs.assignment(gs_array_1, [0.1, 2.0, 1.0], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] = 1.5 gs_result = gs.assignment(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] = 1.5 gs_result = gs.assignment(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] = 1 gs_result = gs.assignment(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] = 1 gs_result = gs.assignment(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] = 1 gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) gs_array_4_arr = gs.zeros_like(gs.array(np_array_4)) gs_result = gs.assignment(gs_array_4_arr, 1, gs.array((0, 1)), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] = 1 gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) def test_assignment_by_sum(self): gs_array_1 = gs.ones(3) with pytest.raises(ValueError): gs.assignment_by_sum(gs_array_1, [0.1, 2.0, 1.0], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) gs_result_list = gs.assignment_by_sum(gs_array_1, [2.0, 1.5], [0, 2]) np_array_1[0] += 2.0 self.assertAllCloseToNp(gs_result_list, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] += 1 gs_result = gs.assignment_by_sum(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] += 1 gs_result = gs.assignment_by_sum(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) n_samples = 3 theta = _np.array([0.1, 0.2, 0.3, 0.4, 5.5]) phi = _np.array([0.11, 0.22, 0.33, 0.44, -0.55]) np_array = _np.ones((n_samples, 5, 4)) gs_array = gs.array(np_array) gs_array = gs.assignment_by_sum(gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment_by_sum(gs_array, -gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) np_array[0, :, 0] += _np.cos(theta) * _np.cos(phi) np_array[0, :, 1] -= _np.sin(theta) * _np.sin(phi) # TODO (ninamiolane): This test fails 15% of the time, # when gs and _np computations are in the reverse order. # We should investigate this. self.assertAllCloseToNp(gs_array, np_array) np_array = _np.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) gs_array = gs.array([[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] += _np.zeros_like(np_array[np_mask]) np_array[~np_mask] += 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment_by_sum(gs_array, values_mask, gs_mask) gs_result = gs.assignment_by_sum(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_any(self): base_list = [ [[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]], [[34.0, 12.0], [2.0, -3.0], [67.0, 35.0]], ] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.any(np_array > 30.0) gs_result = gs.any(gs_array > 30.0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30.0, axis=0) gs_result = gs.any(gs_array > 30.0, axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30.0, axis=-2) gs_result = gs.any(gs_array > 30.0, axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30.0, axis=(-2, -1)) gs_result = gs.any(gs_array > 30.0, axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_all(self): base_list = [ [[22.0, 55.0], [33.0, 88.0], [77.0, 99.0]], [[34.0, 12.0], [2.0, -3.0], [67.0, 35.0]], ] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.all(np_array > 30.0) gs_result = gs.all(gs_array > 30.0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30.0, axis=0) gs_result = gs.all(gs_array > 30.0, axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30.0, axis=-2) gs_result = gs.all(gs_array > 30.0, axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30.0, axis=(-2, -1)) gs_result = gs.all(gs_array > 30.0, axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_trace(self): base_list = [[[22.0, 55.0], [33.0, 88.0]], [[34.0, 12.0], [67.0, 35.0]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.trace(np_array) gs_result = gs.trace(gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=1, axis2=2) gs_result = gs.trace(gs_array, axis1=1, axis2=2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=-1, axis2=-2) gs_result = gs.trace(gs_array, axis1=-1, axis2=-2) self.assertAllCloseToNp(gs_result, np_result) def test_isclose(self): base_list = [[[22.0 + 1e-5, 22.0 + 1e-7], [22.0 + 1e-6, 88.0 + 1e-4]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.isclose(np_array, 22.0) gs_result = gs.isclose(gs_array, 22.0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22.0) gs_result = gs.isclose(gs_array, 22.0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22.0, rtol=1e-8, atol=1e-7) gs_result = gs.isclose(gs_array, 22.0, rtol=1e-8, atol=1e-7) self.assertAllCloseToNp(gs_result, np_result) @tests.conftest.np_autograd_and_torch_only def test_where(self): # TODO (ninamiolane): Make tf behavior consistent with np # Currently, tf returns array, while np returns tuple base_list = [[[22.0, 55.0], [33.0, 88.0]], [[34.0, 12.0], [67.0, 35.0]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.where(np_array > 20.0, 0.0, np_array) gs_result = gs.where(gs_array > 20.0, 0.0, gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.where(np_array > 20, np_array**2, 4.0) gs_result = gs.where(gs_array > 20, gs_array**2, 4.0) self.assertAllCloseToNp(gs_result, np_result) base_list = [[0, 1, 0, 1, 0, 1, 0, 1, 0, 1]] np_array = _np.array(base_list) gs_array = gs.array(base_list) result = gs.where(gs_array == 0) expected = _np.where(np_array == 0) self.assertAllCloseToNp(*result, *expected) result = gs.where(gs_array == 0, -1, gs_array) expected = _np.where(np_array == 0, -1, np_array) self.assertAllCloseToNp(result, expected) expected = _np.where(np_array == 1, _np.ones(10), np_array) result = gs.where(gs_array == 1, gs.ones(10), gs_array) self.assertAllCloseToNp(result, expected) def test_convert_to_wider_dtype(self): gs_list = [gs.array([1, 2]), gs.array([2.2, 3.3], dtype=gs.float32)] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float32 for a in gs_result] self.assertTrue(gs.all(result)) gs_list = [gs.array([1, 2]), gs.array([2.2, 3.3], dtype=gs.float64)] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float64 for a in gs_result] self.assertTrue(gs.all(result)) gs_list = [ gs.array([11.11, 222.2], dtype=gs.float64), gs.array([2.2, 3.3], dtype=gs.float32), ] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float64 for a in gs_result] self.assertTrue(gs.all(result)) def test_broadcast_arrays(self): array_1 = gs.array([[1, 2, 3]]) array_2 = gs.array([[4], [5]]) result = gs.broadcast_arrays(array_1, array_2) result_verdict = [ gs.array([[1, 2, 3], [1, 2, 3]]), gs.array([[4, 4, 4], [5, 5, 5]]), ] self.assertAllClose(result[0], result_verdict[0]) self.assertAllClose(result[1], result_verdict[1]) with pytest.raises((ValueError, RuntimeError)): gs.broadcast_arrays(gs.array([1, 2]), gs.array([3, 4, 5])) def test_choice(self): x = gs.array([0.1, 0.2, 0.3, 0.4, 0.5]) a = 4 result = gs.random.choice(x, a) result_bool = True for i in result: if i in x: continue result_bool = False self.assertTrue(result_bool) self.assertEqual(len(result), a) def test_split(self): x = gs.array([0.1, 0.2, 0.3, 0.4]) result = gs.split(x, 2) expected = _np.split(x, 2) for res, exp in zip(result, expected): self.assertAllClose(res, exp) @tests.conftest.autograd_and_torch_only def test_expm_backward(self): mat = gs.array([[0, 1, 0.5], [-1, 0, 0.2], [-0.5, -0.2, 0]]) mat = gs.cast(mat, gs.float64) def loss(p): return gs.sum((gs.linalg.expm(p) - gs.eye(3))**2) value_and_grad = gs.autodiff.value_and_grad(loss) value, grad = value_and_grad(mat) expected_value = 2.31430522 expected_grad = gs.array([ [1.12127191, 1.68659998, 0.61904561], [-1.50719647, 0.93289823, 0.76788841], [-0.97785262, 0.12912912, 0.26013508], ]) self.assertAllClose(value, expected_value) self.assertAllClose(grad, expected_grad) def test_svd(self): gs_point = gs.reshape(gs.arange(12), (4, 3)) gs_point = gs.cast(gs_point, gs.float64) np_point = _np.arange(12).reshape(4, 3) reconstruction = gs.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]]) u, s, v = _np.linalg.svd(np_point) u_r, s_r, v_r = gs.linalg.svd(gs_point) s_r_reconstructed = gs.einsum("kl,l->kl", reconstruction, s_r) gs_a_approx = gs.matmul(gs.matmul(u_r, s_r_reconstructed), v_r) s_reconstructed = _np.einsum("kl,l->kl", reconstruction, s) np_a_approx = _np.dot(u, _np.dot(s_reconstructed, v)) self.assertAllClose(gs_a_approx, np_a_approx) full_matrices = False u, s, v = _np.linalg.svd(np_point, full_matrices=full_matrices) u_r, s_r, v_r = gs.linalg.svd(gs_point, full_matrices) reconstruction = gs.eye(3) s_r_reconstructed = gs.einsum("kl,l->kl", reconstruction, s_r) gs_a_approx = gs.matmul(gs.matmul(u_r, s_r_reconstructed), v_r) s_reconstructed = _np.einsum("kl,l->kl", reconstruction, s) np_a_approx = _np.dot(u, _np.dot(s_reconstructed, v)) self.assertAllClose(gs_a_approx, np_a_approx) compute_uv = False s = _np.linalg.svd(np_point, compute_uv=compute_uv) s_r = gs.linalg.svd(gs_point, compute_uv=compute_uv) self.assertAllClose(s, s_r) @tests.conftest.np_and_autograd_only def test_sylvester_solve(self): mat = gs.random.rand(4, 3) spd = gs.matmul(gs.transpose(mat), mat) mat = gs.random.rand(3, 3) skew = mat - gs.transpose(mat) solution = gs.linalg.solve_sylvester(spd, spd, skew) result = gs.matmul(spd, solution) result += gs.matmul(solution, spd) self.assertAllClose(result, skew) def test_sylvester_solve_psd(self): psd = gs.array([[1.0, 0.0, 0.0], [0.0, 0.5, 0.0], [0.0, 0.0, 0.0]]) mat = gs.random.rand(3, 3) skew = mat - gs.transpose(mat) solution = gs.linalg.solve_sylvester(psd, psd, skew) result = gs.matmul(psd, solution) result += gs.matmul(solution, psd) self.assertAllClose(result, skew) @tests.conftest.np_autograd_and_torch_only def test_general_sylvester_solve(self): a = gs.array([[-3.0, -2.0, 0.0], [-1.0, -1.0, 3.0], [3.0, -5.0, -1.0]]) b = gs.array([[1.0]]) q = gs.array([[1.0], [2.0], [3.0]]) sol = gs.linalg.solve_sylvester(a, b, q) result = gs.matmul(a, sol) + gs.matmul(sol, b) self.assertAllClose(result, q) def test_sylvester_solve_vectorization(self): gs.random.seed(0) mat = gs.random.rand(2, 4, 3) spd = gs.matmul(gs.transpose(mat, (0, 2, 1)), mat) mat = gs.random.rand(2, 3, 3) skew = mat - gs.transpose(mat, (0, 2, 1)) solution = gs.linalg.solve_sylvester(spd, spd, skew) result = gs.matmul(spd, solution) result += gs.matmul(solution, spd) self.assertAllClose(result, skew) def test_eigvalsh(self): mat = gs.array([[2.0, 1.0], [1.0, -1.0]]) result = gs.linalg.eigvalsh(mat, UPLO="U") expected = _np.linalg.eigvalsh(mat) self.assertAllCloseToNp(result, expected) def test_cholesky(self): mat = SPDMatrices(3).random_point(2) result = gs.linalg.cholesky(mat) expected = _np.linalg.cholesky(mat) self.assertAllClose(result, expected) def test_triu(self): mat = gs.array([[2.0, 1.0, 1.0], [1.0, -1.5, 2.0], [-1.0, 10.0, 2.0]]) result = gs.triu(mat) expected = gs.array([[2.0, 1.0, 1.0], [0.0, -1.5, 2.0], [0.0, 0.0, 2.0]]) self.assertAllClose(result, expected) def test_mat_from_diag_triu_tril(self): diag = gs.array([9.0, 9.0, 9.0]) triu = gs.array([1.0, 2.0, 3.0]) tril = -1 * triu mat = gs.mat_from_diag_triu_tril(diag, triu, tril) expected = gs.array([ [9.0, 1.0, 2.0], [ -1.0, 9.0, 3.0, ], [-2.0, -3.0, 9.0], ]) result = mat batch_diag = gs.eye(3) batch_triu = gs.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) batch_tril = -1 * batch_triu batch_mat = gs.mat_from_diag_triu_tril(batch_diag, batch_triu, batch_tril) batch_expected = gs.array([ [[1.0, 1.0, 2.0], [-1.0, 0.0, 3.0], [-2.0, -3.0, 0.0]], [[0.0, 4.0, 5.0], [-4.0, 1.0, 6.0], [-5.0, -6.0, 0.0]], [[0.0, 7.0, 8.0], [-7.0, 0.0, 9.0], [-8.0, -9.0, 1.0]], ]) batch_result = batch_mat self.assertAllClose(expected, result) self.assertAllClose(batch_expected, batch_result) def test_prod(self): vec = gs.random.rand(10) result = gs.prod(vec) expected = gs.cumprod(vec)[-1] self.assertAllClose(result, expected) def test_is_single_matrix_pd(self): pd = gs.eye(3) not_pd_1 = -1 * gs.eye(3) not_pd_2 = gs.ones((3, 3)) pd_result = gs.linalg.is_single_matrix_pd(pd) not_pd_1_result = gs.linalg.is_single_matrix_pd(not_pd_1) not_pd_2_result = gs.linalg.is_single_matrix_pd(not_pd_2) pd_expected = gs.array(True) not_pd_1_expected = gs.array(False) not_pd_2_expected = gs.array(False) self.assertAllClose(pd_expected, pd_result) self.assertAllClose(not_pd_1_expected, not_pd_1_result) self.assertAllClose(not_pd_2_expected, not_pd_2_result) def test_unique(self): vec = gs.array([-1, 0, 1, 1, 0, -1]) result = gs.unique(vec) expected = gs.array([-1, 0, 1]) self.assertAllClose(result, expected)
class TestBackends(geomstats.tests.TestCase): def setUp(self): warnings.simplefilter('ignore', category=ImportWarning) self.so3_group = SpecialOrthogonal(n=3) self.n_samples = 2 def test_array(self): gs_mat = gs.array(1.5) np_mat = _np.array(1.5) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)]) np_mat = _np.array([_np.ones(3), _np.ones(3)]) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)], dtype=gs.float64) np_mat = _np.array([_np.ones(3), _np.ones(3)], dtype=_np.float64) self.assertTrue(gs_mat.dtype == gs.float64) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([[gs.ones(3)], [gs.ones(3)]], dtype=gs.uint8) np_mat = _np.array([[_np.ones(3)], [_np.ones(3)]], dtype=_np.uint8) self.assertTrue(gs_mat.dtype == gs.uint8) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), [0, 0, 0]], dtype=gs.int32) np_mat = _np.array([_np.ones(3), [0, 0, 0]], dtype=_np.int32) self.assertTrue(gs_mat.dtype == gs.int32) self.assertAllCloseToNp(gs_mat, np_mat) def test_matmul(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) gs_result = gs.matmul(gs_mat_a, gs_mat_b) np_result = _np.matmul(np_mat_a, np_mat_b) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_matmul_vectorization(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] mat_c = [[1., 4., 2.], [4., 3., 4.], [0., 0., 4.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) gs_mat_c = gs.array(mat_c) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) np_mat_c = _np.array(mat_c) gs_result = gs.matmul(gs_mat_a, [gs_mat_b, gs_mat_c]) np_result = _np.matmul(np_mat_a, [np_mat_b, np_mat_c]) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.logm(point) expected = gs.array([[0.693147180, 0., 0.], [0., 1.098612288, 0.], [0., 0., 1.38629436]]) self.assertAllClose(result, expected) np_point = _np.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.logm(np_point) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_and_tf_only def test_expm_and_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) np_point = _np.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.expm(scipy.linalg.logm(np_point)) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_only def test_expm_vectorization(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[7.38905609, 0., 0.], [0., 20.0855369, 0.], [0., 0., 54.5981500]], [[2.718281828, 0., 0.], [0., 148.413159, 0.], [0., 0., 403.42879349]]]) result = gs.linalg.expm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_logm_vectorization_diagonal(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[0.693147180, 0., 0.], [0., 1.09861228866, 0.], [0., 0., 1.38629436]], [[0., 0., 0.], [0., 1.609437912, 0.], [0., 0., 1.79175946]]]) result = gs.linalg.logm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization_random_rotation(self): point = self.so3_group.random_uniform(self.n_samples) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization(self): point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.tf_only def test_vstack(self): import tensorflow as tf tensor_1 = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.]]) tensor_2 = tf.convert_to_tensor([[7., 8., 9.]]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) self.assertAllClose(result, expected) def test_cumprod(self): result = gs.cumprod(gs.arange(1, 10)) expected = gs.array(([1, 2, 6, 24, 120, 720, 5040, 40320, 362880])) self.assertAllClose(result, expected) result = gs.reshape(gs.arange(1, 11), (2, 5)) result = gs.cumprod(result, axis=1) expected = gs.array(([[1, 2, 6, 24, 120], [6, 42, 336, 3024, 30240]])) self.assertAllClose(result, expected) @geomstats.tests.pytorch_only def test_cumsum(self): result = gs.cumsum(gs.arange(10)) expected = gs.array(([0, 1, 3, 6, 10, 15, 21, 28, 36, 45])) self.assertAllClose(result, expected) result = gs.cumsum(gs.arange(10).reshape(2, 5), axis=1) expected = gs.array(([[0, 1, 3, 6, 10], [5, 11, 18, 26, 35]])) self.assertAllClose(result, expected) def test_array_from_sparse(self): expected = gs.array([[0, 1, 0], [0, 0, 2]]) result = gs.array_from_sparse([(0, 1), (1, 2)], [1, 2], (2, 3)) self.assertAllClose(result, expected) def test_einsum(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4], [-1, 5]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4], [-1, 5]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3], [5, 6]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3], [5, 6]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array([5]) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array(5) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array([5]) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array(5) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_einsum_dtypes(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2., 3.]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2., 3.]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1., 4.], [-1., 5.]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1., 4.], [-1., 5.]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_matrices(self): np_array = _np.zeros((2, 3, 3)) gs_array = gs.zeros((2, 3, 3)) np_array[:, 0, 1] = 44. gs_array = gs.assignment(gs_array, 44., (0, 1), axis=0) self.assertAllCloseToNp(gs_array, np_array) n_samples = 3 theta = _np.random.rand(5) phi = _np.random.rand(5) np_array = _np.zeros((n_samples, 5, 4)) gs_array = gs.array(np_array) np_array[0, :, 0] = gs.cos(theta) * gs.cos(phi) np_array[0, :, 1] = -gs.sin(theta) * gs.sin(phi) gs_array = gs.assignment(gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment(gs_array, -gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) self.assertAllCloseToNp(gs_array, np_array) def test_assignment_with_booleans_single_index(self): np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_booleans_many_indices(self): np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, False, True]) gs_mask = gs.array([True, False, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([False, True, True]) gs_mask = gs.array([False, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([False, False, False]) gs_mask = gs.array([False, False, False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment(gs_array, values_mask, gs_mask) gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment(self): gs_array_1 = gs.ones(3) self.assertRaises(ValueError, gs.assignment, gs_array_1, [.1, 2., 1.], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] = 1.5 gs_result = gs.assignment(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] = 1.5 gs_result = gs.assignment(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] = 1 gs_result = gs.assignment(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] = 1 gs_result = gs.assignment(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] = 1 gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) gs_array_4_arr = gs.zeros_like(gs.array(np_array_4)) gs_result = gs.assignment(gs_array_4_arr, 1, gs.array((0, 1)), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] = 1 gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) def test_assignment_by_sum(self): gs_array_1 = gs.ones(3) self.assertRaises(ValueError, gs.assignment_by_sum, gs_array_1, [.1, 2., 1.], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) gs_result_list = gs.assignment_by_sum(gs_array_1, [2., 1.5], [0, 2]) np_array_1[0] += 2. self.assertAllCloseToNp(gs_result_list, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] += 1 gs_result = gs.assignment_by_sum(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] += 1 gs_result = gs.assignment_by_sum(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) n_samples = 3 theta = _np.array([0.1, 0.2, 0.3, 0.4, 5.5]) phi = _np.array([0.11, 0.22, 0.33, 0.44, -.55]) np_array = _np.ones((n_samples, 5, 4)) gs_array = gs.array(np_array) gs_array = gs.assignment_by_sum(gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment_by_sum(gs_array, -gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) np_array[0, :, 0] += _np.cos(theta) * _np.cos(phi) np_array[0, :, 1] -= _np.sin(theta) * _np.sin(phi) # TODO (ninamiolane): This test fails 15% of the time, # when gs and _np computations are in the reverse order. # We should investigate this. self.assertAllCloseToNp(gs_array, np_array) np_array = _np.array([[22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([[22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] += _np.zeros_like(np_array[np_mask]) np_array[~np_mask] += 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment_by_sum(gs_array, values_mask, gs_mask) gs_result = gs.assignment_by_sum(gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_any(self): base_list = [[[22., 55.], [33., 88.], [77., 99.]], [[34., 12.], [2., -3.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.any(np_array > 30.) gs_result = gs.any(gs_array > 30.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=0) gs_result = gs.any(gs_array > 30., axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=-2) gs_result = gs.any(gs_array > 30., axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=(-2, -1)) gs_result = gs.any(gs_array > 30., axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_all(self): base_list = [[[22., 55.], [33., 88.], [77., 99.]], [[34., 12.], [2., -3.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.all(np_array > 30.) gs_result = gs.all(gs_array > 30.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=0) gs_result = gs.all(gs_array > 30., axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=-2) gs_result = gs.all(gs_array > 30., axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=(-2, -1)) gs_result = gs.all(gs_array > 30., axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_trace(self): base_list = [[[22., 55.], [33., 88.]], [[34., 12.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.trace(np_array) gs_result = gs.trace(gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=1, axis2=2) gs_result = gs.trace(gs_array, axis1=1, axis2=2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=-1, axis2=-2) gs_result = gs.trace(gs_array, axis1=-1, axis2=-2) self.assertAllCloseToNp(gs_result, np_result) def test_isclose(self): base_list = [[[22. + 1e-5, 22. + 1e-7], [22. + 1e-6, 88. + 1e-4]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.isclose(np_array, 22.) gs_result = gs.isclose(gs_array, 22.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22., atol=1e-8) gs_result = gs.isclose(gs_array, 22., atol=1e-8) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22., rtol=1e-8, atol=1e-7) gs_result = gs.isclose(gs_array, 22., rtol=1e-8, atol=1e-7) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_pytorch_only def test_where(self): # TODO (ninamiolane): Make tf behavior consistent with np # Currently, tf returns array, while np returns tuple base_list = [[[22., 55.], [33., 88.]], [[34., 12.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.where(np_array > 20., 0., np_array) gs_result = gs.where(gs_array > 20., 0., gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.where(np_array > 20, np_array**2, 4.) gs_result = gs.where(gs_array > 20, gs_array**2, 4.) self.assertAllCloseToNp(gs_result, np_result) base_list = [[0, 1, 0, 1, 0, 1, 0, 1, 0, 1]] np_array = _np.array(base_list) gs_array = gs.array(base_list) result = gs.where(gs_array == 0) expected = _np.where(np_array == 0) self.assertAllCloseToNp(*result, *expected) result = gs.where(gs_array == 0, -1, gs_array) expected = _np.where(np_array == 0, -1, np_array) self.assertAllCloseToNp(result, expected) expected = _np.where(np_array == 1, _np.ones(10), np_array) result = gs.where(gs_array == 1, gs.ones(10), gs_array) self.assertAllCloseToNp(result, expected) def test_convert_to_wider_dtype(self): gs_list = [gs.array([1, 2]), gs.array([2.2, 3.3], dtype=gs.float32)] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float32 for a in gs_result] self.assertTrue(gs.all(result)) gs_list = [gs.array([1, 2]), gs.array([2.2, 3.3], dtype=gs.float64)] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float64 for a in gs_result] self.assertTrue(gs.all(result)) gs_list = [ gs.array([11.11, 222.2], dtype=gs.float64), gs.array([2.2, 3.3], dtype=gs.float32) ] gs_result = gs.convert_to_wider_dtype(gs_list) result = [a.dtype == gs.float64 for a in gs_result] self.assertTrue(gs.all(result)) def test_broadcast_arrays(self): array_1 = gs.array([[1, 2, 3]]) array_2 = gs.array([[4], [5]]) result = gs.broadcast_arrays(array_1, array_2) result_verdict = [ gs.array([[1, 2, 3], [1, 2, 3]]), gs.array([[4, 4, 4], [5, 5, 5]]) ] self.assertAllClose(result[0], result_verdict[0]) self.assertAllClose(result[1], result_verdict[1]) with self.assertRaises((ValueError, RuntimeError)): gs.broadcast_arrays(gs.array([1, 2]), gs.array([3, 4, 5])) def test_value_and_grad(self): n = 10 vector = gs.ones(n) result_loss, result_grad = gs.autograd.value_and_grad( lambda v: gs.sum(v**2))(vector) expected_loss = n expected_grad = 2 * vector self.assertAllClose(result_loss, expected_loss) self.assertAllClose(result_grad, expected_grad) def test_value_and_grad_numpy_input(self): n = 10 vector = _np.ones(n) result_loss, result_grad = gs.autograd.value_and_grad( lambda v: gs.sum(v**2))(vector) expected_loss = n expected_grad = 2 * vector self.assertAllClose(result_loss, expected_loss) self.assertAllClose(result_grad, expected_grad) def test_choice(self): x = gs.array([0.1, 0.2, 0.3, 0.4, 0.5]) a = 4 result = gs.random.choice(x, a) result_bool = True for i in result: if i in x: continue result_bool = False self.assertTrue(result_bool) self.assertEqual(len(result), a) def test_split(self): x = gs.array([0.1, 0.2, 0.3, 0.4]) result = gs.split(x, 2) expected = _np.split(x, 2) for res, exp in zip(result, expected): self.assertAllClose(res, exp) def test_svd(self): gs_point = gs.reshape(gs.arange(12), (4, 3)) gs_point = gs.cast(gs_point, gs.float64) np_point = _np.arange(12).reshape(4, 3) reconstruction = gs.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.], [0., 0., 0.]]) u, s, v = _np.linalg.svd(np_point) u_r, s_r, v_r = gs.linalg.svd(gs_point) s_r_reconstructed = gs.einsum('kl,l->kl', reconstruction, s_r) gs_a_approx = gs.matmul(gs.matmul(u_r, s_r_reconstructed), v_r) s_reconstructed = _np.einsum('kl,l->kl', reconstruction, s) np_a_approx = _np.dot(u, _np.dot(s_reconstructed, v)) self.assertAllClose(gs_a_approx, np_a_approx) full_matrices = False u, s, v = _np.linalg.svd(np_point, full_matrices=full_matrices) u_r, s_r, v_r = gs.linalg.svd(gs_point, full_matrices) reconstruction = gs.eye(3) s_r_reconstructed = gs.einsum('kl,l->kl', reconstruction, s_r) gs_a_approx = gs.matmul(gs.matmul(u_r, s_r_reconstructed), v_r) s_reconstructed = _np.einsum('kl,l->kl', reconstruction, s) np_a_approx = _np.dot(u, _np.dot(s_reconstructed, v)) self.assertAllClose(gs_a_approx, np_a_approx) compute_uv = False s = _np.linalg.svd(np_point, compute_uv=compute_uv) s_r = gs.linalg.svd(gs_point, compute_uv=compute_uv) self.assertAllClose(s, s_r) def test_sylvester_solve(self): mat = gs.random.rand(4, 3) spd = gs.matmul(gs.transpose(mat), mat) mat = gs.random.rand(3, 3) skew = mat - gs.transpose(mat) solution = gs.linalg.solve_sylvester(spd, spd, skew) result = gs.matmul(spd, solution) result += gs.matmul(solution, spd) self.assertAllClose(result, skew) @geomstats.tests.np_and_pytorch_only def test_general_sylvester_solve(self): a = gs.array([[-3., -2., 0.], [-1., -1., 3.], [3., -5., -1.]]) b = gs.array([[1.]]) q = gs.array([[1.], [2.], [3.]]) sol = gs.linalg.solve_sylvester(a, b, q) result = gs.matmul(a, sol) + gs.matmul(sol, b) self.assertAllClose(result, q) def test_sylvester_solve_vectorization(self): gs.random.seed(0) mat = gs.random.rand(2, 4, 3) spd = gs.matmul(gs.transpose(mat, (0, 2, 1)), mat) mat = gs.random.rand(2, 3, 3) skew = mat - gs.transpose(mat, (0, 2, 1)) solution = gs.linalg.solve_sylvester(spd, spd, skew) result = gs.matmul(spd, solution) result += gs.matmul(solution, spd) self.assertAllClose(result, skew) def test_eigvalsh(self): mat = gs.array([[2., 1.], [1., -1.]]) result = gs.linalg.eigvalsh(mat, UPLO='U') expected = _np.linalg.eigvalsh(mat) self.assertAllCloseToNp(result, expected) def test_cholesky(self): mat = SPDMatrices(3).random_point(2) result = gs.linalg.cholesky(mat) expected = _np.linalg.cholesky(mat) self.assertAllClose(result, expected) @geomstats.tests.np_and_pytorch_only def test_expm_backward(self): mat = gs.array([[0, 1, .5], [-1, 0, 0.2], [-.5, -.2, 0]]) def loss(p): return gs.sum((gs.linalg.expm(p) - gs.eye(3))**2) value_and_grad = gs.autograd.value_and_grad(loss) result = value_and_grad(mat) def loss_torch(p): return torch.sum((torch.matrix_exp(p) - torch.eye(3))**2) torch_mat = torch.tensor([[0, 1, .5], [-1, 0, 0.2], [-.5, -.2, 0]], requires_grad=True) value = loss_torch(torch_mat) value.backward() grad = torch_mat.grad self.assertAllClose(result[0], value.detach()) self.assertAllClose(result[1], grad)
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)
class _SpecialEuclideanMatrices(GeneralLinear, LieGroup): """Class for special orthogonal groups. Parameters ---------- n : int Integer representing the shape of the matrices: n x n. """ def __init__(self, n): super(_SpecialEuclideanMatrices, self).__init__(default_point_type='matrix', n=n + 1) self.rotations = SpecialOrthogonal(n=n) self.translations = Euclidean(dim=n) self.n = n self.dim = int((n * (n + 1)) / 2) def get_identity(self): """Return the identity matrix.""" return gs.eye(self.n + 1, self.n + 1) identity = property(get_identity) def belongs(self, point): """Check whether point is of the form rotation, translation. Parameters ---------- point : array-like, shape=[..., n, n]. Point to be checked. Returns ------- belongs : array-like, shape=[...,] Boolean denoting if point belongs to the group. """ point_dim1, point_dim2 = point.shape[-2:] belongs = (point_dim1 == point_dim2 == self.n + 1) rotation = point[..., :self.n, :self.n] rot_belongs = self.rotations.belongs(rotation) belongs = gs.logical_and(belongs, rot_belongs) last_line_except_last_term = point[..., self.n:, :-1] all_but_last_zeros = ~gs.any(last_line_except_last_term, axis=(-2, -1)) belongs = gs.logical_and(belongs, all_but_last_zeros) last_term = point[..., self.n:, self.n:] belongs = gs.logical_and(belongs, gs.all(last_term == 1, axis=(-2, -1))) if point.ndim == 2: return gs.squeeze(belongs) return gs.flatten(belongs) def _is_in_lie_algebra(self, tangent_vec, atol=TOLERANCE): """Project vector rotation part onto skew-symmetric matrices.""" point_dim1, point_dim2 = tangent_vec.shape[-2:] belongs = (point_dim1 == point_dim2 == self.n + 1) rotation = tangent_vec[..., :self.n, :self.n] rot_belongs = self.is_skew_symmetric(rotation, atol=atol) belongs = gs.logical_and(belongs, rot_belongs) last_line = tangent_vec[..., -1, :] all_zeros = ~gs.any(last_line, axis=-1) belongs = gs.logical_and(belongs, all_zeros) return belongs def _to_lie_algebra(self, tangent_vec): """Project vector rotation part onto skew-symmetric matrices.""" translation_mask = gs.hstack( [gs.ones((self.n, ) * 2), 2 * gs.ones((self.n, 1))]) translation_mask = gs.concatenate( [translation_mask, gs.zeros((1, self.n + 1))], axis=0) tangent_vec = tangent_vec * gs.where(translation_mask != 0., gs.array(1.), gs.array(0.)) tangent_vec = (tangent_vec - GeneralLinear.transpose(tangent_vec)) / 2. return tangent_vec * translation_mask def random_uniform(self, n_samples=1, tol=1e-6): """Sample in SE(n) from the uniform distribution. Parameters ---------- n_samples : int Number of samples. Optional, default: 1. tol : unused Returns ------- samples : array-like, shape=[..., n + 1, n + 1] Sample in SE(n). """ random_translation = self.translations.random_uniform(n_samples) random_rotation = self.rotations.random_uniform(n_samples) random_rotation = gs.to_ndarray(random_rotation, to_ndim=3) random_translation = gs.to_ndarray(random_translation, to_ndim=2) random_translation = gs.transpose( gs.to_ndarray(random_translation, to_ndim=3, axis=1), (0, 2, 1)) random_point = gs.concatenate((random_rotation, random_translation), axis=2) last_line = gs.zeros((n_samples, 1, self.n + 1)) random_point = gs.concatenate((random_point, last_line), axis=1) random_point = gs.assignment(random_point, 1, (-1, -1), axis=0) if gs.shape(random_point)[0] == 1: random_point = gs.squeeze(random_point, axis=0) return random_point
class TestSpecialOrthogonal(geomstats.tests.TestCase): def setUp(self): self.n = 2 self.group = SpecialOrthogonal(n=self.n) self.n_samples = 4 def test_dim(self): for n in [2, 3, 4, 5, 6]: group = SpecialOrthogonal(n=n) result = group.dim expected = n * (n - 1) / 2 self.assertAllClose(result, expected) def test_belongs(self): theta = gs.pi / 3 point_1 = gs.array([[gs.cos(theta), - gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_1) self.assertTrue(result) point_2 = gs.array([[gs.cos(theta), gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_2) self.assertFalse(result) point = gs.array([point_1, point_2]) expected = gs.array([True, False]) result = self.group.belongs(point) self.assertAllClose(result, expected) point = point_1[0] result = self.group.belongs(point) self.assertFalse(result) point = gs.zeros((2, 3)) result = self.group.belongs(point) self.assertFalse(result) point = gs.zeros((2, 2, 3)) result = self.group.belongs(point) self.assertFalse(gs.all(result)) def test_random_uniform_and_belongs(self): point = self.group.random_uniform() result = self.group.belongs(point) expected = True self.assertAllClose(result, expected) point = self.group.random_uniform(self.n_samples) result = self.group.belongs(point) expected = gs.array([True] * self.n_samples) self.assertAllClose(result, expected) def test_identity(self): result = self.group.identity expected = gs.eye(self.n) self.assertAllClose(result, expected) def test_is_in_lie_algebra(self): theta = gs.pi / 3 vec_1 = gs.array([[0., - theta], [theta, 0.]]) result = self.group.is_tangent(vec_1) self.assertTrue(result) vec_2 = gs.array([[0., - theta], [theta, 1.]]) result = self.group.is_tangent(vec_2) self.assertFalse(result) vec = gs.array([vec_1, vec_2]) result = self.group.is_tangent(vec) expected = gs.array([True, False]) self.assertAllClose(result, expected) def test_is_tangent(self): point = self.group.random_uniform() theta = 1. vec_1 = gs.array([[0., - theta], [theta, 0.]]) vec_1 = self.group.compose(point, vec_1) result = self.group.is_tangent(vec_1, point) self.assertTrue(result) vec_2 = gs.array([[0., - theta], [theta, 1.]]) vec_2 = self.group.compose(point, vec_2) result = self.group.is_tangent(vec_2, point) self.assertFalse(result) vec = gs.array([vec_1, vec_2]) point = gs.array([point, point]) expected = gs.array([True, False]) result = self.group.is_tangent(vec, point) self.assertAllClose(result, expected) def test_to_tangent(self): theta = 1. vec_1 = gs.array([[0., - theta], [theta, 0.]]) result = self.group.to_tangent(vec_1) expected = vec_1 self.assertAllClose(result, expected) n_samples = self.n_samples base_points = self.group.random_uniform(n_samples=n_samples) tangent_vecs = self.group.compose(base_points, vec_1) result = self.group.to_tangent(tangent_vecs, base_points) expected = tangent_vecs self.assertAllClose(result, expected) def test_projection_and_belongs(self): gs.random.seed(4) shape = (self.n_samples, self.n, self.n) result = helper.test_projection_and_belongs( self.group, shape, gs.atol * 100) for res in result: self.assertTrue(res) def test_skew_to_vec_and_back(self): group = SpecialOrthogonal(n=4) vec = gs.random.rand(group.dim) mat = group.skew_matrix_from_vector(vec) result = group.vector_from_skew_matrix(mat) self.assertAllClose(result, vec) def test_parallel_transport(self): metric = self.group.bi_invariant_metric shape = (self.n_samples, self.group.n, self.group.n) results = helper.test_parallel_transport(self.group, metric, shape) for res in results: self.assertTrue(res) def test_metric_left_invariant(self): group = self.group point = group.random_point() tangent_vec = self.group.lie_algebra.basis[0] expected = group.bi_invariant_metric.norm( tangent_vec) translated = group.tangent_translation_map(point)(tangent_vec) result = group.bi_invariant_metric.norm(translated) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_distance_broadcast(self): group = self.group point = group.random_point(5) result = group.bi_invariant_metric.dist_broadcast(point[:3], point) expected = [] for a in point[:3]: expected.append(group.bi_invariant_metric.dist(a, point)) expected = gs.stack(expected) self.assertAllClose(result, expected)
class TestSpecialOrthogonal2(geomstats.tests.TestCase): def setup_method(self): warnings.simplefilter("ignore", category=ImportWarning) warnings.simplefilter("ignore", category=UserWarning) gs.random.seed(1234) self.group = SpecialOrthogonal(n=2, point_type="vector") # -- Set attributes self.n_samples = 4 def test_projection(self): # Test 2D and nD cases rot_mat = gs.eye(2) delta = 1e-12 * gs.ones((2, 2)) rot_mat_plus_delta = rot_mat + delta result = self.group.projection(rot_mat_plus_delta) expected = rot_mat self.assertAllClose(result, expected) def test_projection_vectorization(self): n_samples = self.n_samples mats = gs.ones((n_samples, 2, 2)) result = self.group.projection(mats) self.assertAllClose(gs.shape(result), (n_samples, 2, 2)) def test_skew_matrix_from_vector(self): rot_vec = gs.array([0.9]) skew_matrix = self.group.skew_matrix_from_vector(rot_vec) result = gs.matmul(skew_matrix, skew_matrix) diag = gs.array([-0.81, -0.81]) expected = algebra_utils.from_vector_to_diagonal_matrix(diag) self.assertAllClose(result, expected) def test_skew_matrix_and_vector(self): rot_vec = gs.array([0.8]) skew_mat = self.group.skew_matrix_from_vector(rot_vec) result = self.group.vector_from_skew_matrix(skew_mat) expected = rot_vec self.assertAllClose(result, expected) def test_skew_matrix_from_vector_vectorization(self): n_samples = self.n_samples rot_vecs = self.group.random_uniform(n_samples=n_samples) result = self.group.skew_matrix_from_vector(rot_vecs) self.assertAllClose(gs.shape(result), (n_samples, 2, 2)) def test_random_uniform_shape(self): result = self.group.random_uniform() self.assertAllClose(gs.shape(result), (self.group.dim, )) def test_random_and_belongs(self): point = self.group.random_uniform() result = self.group.belongs(point) expected = True self.assertAllClose(result, expected) 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_regularize(self): angle = 2 * gs.pi + 1 result = self.group.regularize(gs.array([angle])) expected = gs.array([1.0]) self.assertAllClose(result, expected) def test_regularize_vectorization(self): n_samples = self.n_samples rot_vecs = self.group.random_uniform(n_samples=n_samples) result = self.group.regularize(rot_vecs) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) def test_matrix_from_rotation_vector(self): angle = gs.pi / 3 expected = gs.array([[1.0 / 2, -gs.sqrt(3.0) / 2], [gs.sqrt(3.0) / 2, 1.0 / 2]]) result = self.group.matrix_from_rotation_vector(gs.array([angle])) self.assertAllClose(result, expected) def test_matrix_from_rotation_vector_vectorization(self): n_samples = self.n_samples rot_vecs = self.group.random_uniform(n_samples=n_samples) rot_mats = self.group.matrix_from_rotation_vector(rot_vecs) self.assertAllClose(gs.shape(rot_mats), (n_samples, self.group.n, self.group.n)) def test_rotation_vector_from_matrix(self): angle = 0.12 rot_mat = gs.array([[gs.cos(angle), -gs.sin(angle)], [gs.sin(angle), gs.cos(angle)]]) result = self.group.rotation_vector_from_matrix(rot_mat) expected = gs.array([0.12]) self.assertAllClose(result, expected) def test_rotation_vector_and_rotation_matrix(self): """ This tests that the composition of rotation_vector_from_matrix and matrix_from_rotation_vector is the identity. """ # TODO(nguigs): bring back a 1d representation of SO2 point = gs.array([0.78]) rot_mat = self.group.matrix_from_rotation_vector(point) result = self.group.rotation_vector_from_matrix(rot_mat) expected = point self.assertAllClose(result, expected) def test_rotation_vector_and_rotation_matrix_vectorization(self): rot_vecs = gs.array([[2.0], [1.3], [0.8], [0.03]]) rot_mats = self.group.matrix_from_rotation_vector(rot_vecs) result = self.group.rotation_vector_from_matrix(rot_mats) expected = self.group.regularize(rot_vecs) self.assertAllClose(result, expected) def test_compose(self): point_a = gs.array([0.12]) point_b = gs.array([-0.15]) result = self.group.compose(point_a, point_b) expected = self.group.regularize(gs.array([-0.03])) self.assertAllClose(result, expected) def test_compose_and_inverse(self): angle = 0.986 point = gs.array([angle]) inv_point = self.group.inverse(point) result = self.group.compose(point, inv_point) expected = self.group.identity self.assertAllClose(result, expected) result = self.group.compose(inv_point, point) expected = self.group.identity self.assertAllClose(result, expected) def test_compose_vectorization(self): point_type = "vector" self.group.default_point_type = point_type n_samples = self.n_samples n_points_a = self.group.random_uniform(n_samples=n_samples) n_points_b = self.group.random_uniform(n_samples=n_samples) one_point = self.group.random_uniform(n_samples=1) result = self.group.compose(one_point, n_points_a) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) result = self.group.compose(n_points_a, one_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) result = self.group.compose(n_points_a, n_points_b) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) def test_inverse_vectorization(self): n_samples = self.n_samples points = self.group.random_uniform(n_samples=n_samples) result = self.group.inverse(points) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) def test_group_exp(self): """ The Riemannian exp and log are inverse functions of each other. This test is the inverse of test_log's. """ rot_vec_base_point = gs.array([gs.pi / 5]) rot_vec = gs.array([2 * gs.pi / 5]) expected = gs.array([3 * gs.pi / 5]) result = self.group.exp(base_point=rot_vec_base_point, tangent_vec=rot_vec) self.assertAllClose(result, expected) def test_group_exp_vectorization(self): n_samples = self.n_samples one_tangent_vec = self.group.random_uniform(n_samples=1) one_base_point = self.group.random_uniform(n_samples=1) n_tangent_vec = self.group.random_uniform(n_samples=n_samples) n_base_point = self.group.random_uniform(n_samples=n_samples) # Test with the 1 base point, and n tangent vecs result = self.group.exp(n_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) # Test with the several base point, and one tangent vec result = self.group.exp(one_tangent_vec, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) # Test with the same number n of base point and n tangent vec result = self.group.exp(n_tangent_vec, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) def test_group_log(self): """ The Riemannian exp and log are inverse functions of each other. This test is the inverse of test_exp's. """ rot_vec_base_point = gs.array([gs.pi / 5]) rot_vec = gs.array([2 * gs.pi / 5]) expected = gs.array([1 * gs.pi / 5]) result = self.group.log(point=rot_vec, base_point=rot_vec_base_point) self.assertAllClose(result, expected) def test_group_log_vectorization(self): n_samples = self.n_samples one_point = self.group.random_uniform(n_samples=1) one_base_point = self.group.random_uniform(n_samples=1) n_point = self.group.random_uniform(n_samples=n_samples) n_base_point = self.group.random_uniform(n_samples=n_samples) # Test with the 1 base point, and several different points result = self.group.log(n_point, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) # Test with the several base point, and 1 point result = self.group.log(one_point, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) # Test with the same number n of base point and point result = self.group.log(n_point, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.group.dim)) def test_group_exp_then_log_from_identity(self): """ Test that the group exponential and the group logarithm are inverse. Expect their composition to give the identity function. """ tangent_vec = gs.array([0.12]) result = helper.group_exp_then_log_from_identity( group=self.group, tangent_vec=tangent_vec) expected = self.group.regularize(tangent_vec) self.assertAllClose(result, expected) def test_group_log_then_exp_from_identity(self): """ Test that the group exponential and the group logarithm are inverse. Expect their composition to give the identity function. """ point = gs.array([0.12]) result = helper.group_log_then_exp_from_identity(group=self.group, point=point) expected = self.group.regularize(point) self.assertAllClose(result, expected) def test_group_exp_then_log(self): """ This tests that the composition of log and exp gives identity. """ base_point = gs.array([0.12]) tangent_vec = gs.array([0.35]) result = helper.group_exp_then_log(group=self.group, tangent_vec=tangent_vec, base_point=base_point) expected = self.group.regularize_tangent_vec(tangent_vec=tangent_vec, base_point=base_point) self.assertAllClose(result, expected) def test_group_log_then_exp(self): """ This tests that the composition of log and exp gives identity. """ base_point = gs.array([0.12]) point = gs.array([0.35]) result = helper.group_log_then_exp(group=self.group, point=point, base_point=base_point) expected = self.group.regularize(point) self.assertAllClose(result, expected)
class TestSpecialOrthogonal(geomstats.tests.TestCase): def setUp(self): self.n = 2 self.group = SpecialOrthogonal(n=self.n) self.n_samples = 4 def test_belongs(self): theta = gs.pi / 3 point_1 = gs.array([[gs.cos(theta), -gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_1) expected = True self.assertAllClose(result, expected) point_2 = gs.array([[gs.cos(theta), gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_2) expected = False self.assertAllClose(result, expected) point = gs.array([point_1, point_2]) expected = gs.array([True, False]) result = self.group.belongs(point) self.assertAllClose(result, expected) def test_random_uniform_and_belongs(self): point = self.group.random_uniform() result = self.group.belongs(point) expected = True self.assertAllClose(result, expected) point = self.group.random_uniform(self.n_samples) result = self.group.belongs(point) expected = gs.array([True] * self.n_samples) self.assertAllClose(result, expected) def test_identity(self): result = self.group.identity expected = gs.eye(self.n) self.assertAllClose(result, expected) def test_is_in_lie_algebra(self): theta = gs.pi / 3 vec_1 = gs.array([[0., -theta], [theta, 0.]]) result = self.group.is_tangent(vec_1) expected = True self.assertAllClose(result, expected) vec_2 = gs.array([[0., -theta], [theta, 1.]]) result = self.group.is_tangent(vec_2) expected = False self.assertAllClose(result, expected) vec = gs.array([vec_1, vec_2]) expected = gs.array([True, False]) result = self.group.is_tangent(vec) self.assertAllClose(result, expected) def test_is_tangent(self): point = self.group.random_uniform() theta = 1. vec_1 = gs.array([[0., -theta], [theta, 0.]]) vec_1 = self.group.compose(point, vec_1) result = self.group.is_tangent(vec_1, point) expected = True self.assertAllClose(result, expected) vec_2 = gs.array([[0., -theta], [theta, 1.]]) vec_2 = self.group.compose(point, vec_2) result = self.group.is_tangent(vec_2, point) expected = False self.assertAllClose(result, expected) vec = gs.array([vec_1, vec_2]) point = gs.array([point, point]) expected = gs.array([True, False]) result = self.group.is_tangent(vec, point) self.assertAllClose(result, expected) def test_to_tangent(self): theta = 1. vec_1 = gs.array([[0., -theta], [theta, 0.]]) result = self.group.to_tangent(vec_1) expected = vec_1 self.assertAllClose(result, expected) n_samples = self.n_samples base_points = self.group.random_uniform(n_samples=n_samples) tangent_vecs = self.group.compose(base_points, vec_1) result = self.group.to_tangent(tangent_vecs, base_points) expected = tangent_vecs self.assertAllClose(result, expected) def test_projection_and_belongs(self): gs.random.seed(3) group = SpecialOrthogonal(n=4) mat = gs.random.rand(4, 4) point = group.projection(mat) result = group.belongs(point) self.assertTrue(result) mat = gs.random.rand(2, 4, 4) point = group.projection(mat) result = group.belongs(point, atol=1e-4) self.assertTrue(gs.all(result)) def test_skew_to_vec_and_back(self): group = SpecialOrthogonal(n=4) vec = gs.random.rand(group.dim) mat = group.skew_matrix_from_vector(vec) result = group.vector_from_skew_matrix(mat) self.assertAllClose(result, vec)
class TestExponentialBarycenter(geomstats.tests.TestCase): def setUp(self): logger = logging.getLogger() logger.disabled = True self.se_mat = SpecialEuclidean(n=3) self.so_vec = SpecialOrthogonal(n=3, point_type='vector') self.so = SpecialOrthogonal(n=3) self.n_samples = 4 @geomstats.tests.np_only def test_estimate_and_belongs_se(self): point = self.se_mat.random_point(self.n_samples) estimator = ExponentialBarycenter(self.se_mat) estimator.fit(point) barexp = estimator.estimate_ result = self.se_mat.belongs(barexp) expected = True self.assertAllClose(result, expected) point = self.so_vec.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so_vec) estimator.fit(point) barexp = estimator.estimate_ result = self.so_vec.belongs(barexp) expected = True self.assertAllClose(result, expected) def test_estimate_one_sample_se(self): point = self.se_mat.random_point() estimator = ExponentialBarycenter(self.se_mat) estimator.fit(point) result = estimator.estimate_ expected = point self.assertAllClose(result, expected) point = self.so_vec.random_uniform(1) estimator = ExponentialBarycenter(self.so_vec) estimator.fit(point) result = estimator.estimate_ expected = point self.assertAllClose(result, expected) @geomstats.tests.np_only def test_estimate_and_reach_max_iter_se(self): point = self.se_mat.random_point(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_estimate_so_matrix(self): points = self.so.random_uniform(2) mean_vec = ExponentialBarycenter(group=self.so) mean_vec.fit(points) logs = self.so.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_belongs_so(self): point = self.so.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so) estimator.fit(point) barexp = estimator.estimate_ result = self.so.belongs(barexp) expected = True self.assertAllClose(result, expected) point = self.so_vec.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so_vec) estimator.fit(point) barexp = estimator.estimate_ result = self.so_vec.belongs(barexp) expected = True self.assertAllClose(result, expected) @geomstats.tests.np_only def test_estimate_one_sample_so(self): point = self.so.random_uniform(1) estimator = ExponentialBarycenter(self.so) estimator.fit(point) result = estimator.estimate_ expected = point self.assertAllClose(result, expected) point = self.so_vec.random_uniform(1) estimator = ExponentialBarycenter(self.so_vec) estimator.fit(point) result = estimator.estimate_ expected = point self.assertAllClose(result, expected) @geomstats.tests.np_only def test_estimate_and_reach_max_iter_so(self): point = self.so.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so, max_iter=2) estimator.fit(point) barexp = estimator.estimate_ result = self.so.belongs(barexp) expected = True self.assertAllClose(result, expected) point = self.so_vec.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so_vec, max_iter=2) estimator.fit(point) barexp = estimator.estimate_ result = self.so_vec.belongs(barexp) expected = True self.assertAllClose(result, expected) @geomstats.tests.np_only def test_coincides_with_frechet_so(self): gs.random.seed(0) point = self.so.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so, max_iter=40, epsilon=1e-10) estimator.fit(point) result = estimator.estimate_ frechet_estimator = FrechetMean(self.so.bi_invariant_metric, max_iter=40, epsilon=1e-10, lr=1., method='adaptive') frechet_estimator.fit(point) expected = frechet_estimator.estimate_ self.assertAllClose(result, expected) @geomstats.tests.np_only def test_estimate_weights(self): point = self.so.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so, verbose=True) weights = gs.arange(self.n_samples) estimator.fit(point, weights=weights) barexp = estimator.estimate_ result = self.so.belongs(barexp) expected = True self.assertAllClose(result, expected) point = self.so_vec.random_uniform(self.n_samples) estimator = ExponentialBarycenter(self.so_vec) estimator.fit(point, weights=weights) barexp = estimator.estimate_ result = self.so_vec.belongs(barexp) expected = True self.assertAllClose(result, expected) def test_linear_mean(self): euclidean = Euclidean(3) point = euclidean.random_point(self.n_samples) estimator = ExponentialBarycenter(euclidean) estimator.fit(point) result = estimator.estimate_ expected = gs.mean(point, axis=0) self.assertAllClose(result, expected)
class SpecialEuclidean(LieGroup): """Class for the special euclidean group SE(n). i.e. the Lie group of rigid transformations. """ def __init__(self, n, point_type=None, epsilon=0.): """Initiate an object of class SpecialEuclidean. Parameter --------- n : int the dimension of the euclidean space that SE(n) acts upon point_type : str, {'vector', 'matrix'}, optional whether to represent elmenents of SE(n) by vectors or matrices if None is given, point_type is set to 'vector' for dimension 3 and 'matrix' otherwise epsilon : float, optional precision to use for calculations involving potential division by rotations default: 0 """ assert isinstance(n, int) and n > 1 self.n = n self.dimension = int((n * (n - 1)) / 2 + n) self.epsilon = epsilon self.default_point_type = point_type if point_type is None: self.default_point_type = 'vector' if n == 3 else 'matrix' super(SpecialEuclidean, self).__init__(dimension=self.dimension) self.rotations = SpecialOrthogonal(n=n, epsilon=epsilon) self.translations = Euclidean(dimension=n) def get_identity(self, point_type=None): """Get the identity of the group. Parameters ---------- point_type : str, {'vector', 'matrix'}, optional the point_type of the returned value default: self.default_point_type Returns ------- identity : array-like, shape={[dimension], [n + 1, n + 1]} """ if point_type is None: point_type = self.default_point_type identity = gs.zeros(self.dimension) if self.default_point_type == 'matrix': identity = gs.eye(self.n) return identity identity = property(get_identity) def belongs(self, point, point_type=None): """Evaluate if a point belongs to SE(n). Parameters ---------- point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] the point of which to check whether it belongs to SE(n) point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- belongs : array-like, shape=[n_samples, 1] array of booleans indicating whether point belongs to SE(n) """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': point = gs.to_ndarray(point, to_ndim=2) n_points, point_dim = point.shape belongs = point_dim == self.dimension belongs = gs.to_ndarray(belongs, to_ndim=1) belongs = gs.to_ndarray(belongs, to_ndim=2, axis=1) belongs = gs.tile(belongs, (n_points, 1)) elif point_type == 'matrix': point = gs.to_ndarray(point, to_ndim=3) raise NotImplementedError() return belongs def regularize(self, point, point_type=None): """Regularize a point to the default representation for SE(n). Parameters ---------- point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] the point which should be regularized point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': point = gs.to_ndarray(point, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_vec = point[:, :dim_rotations] regularized_rot_vec = rotations.regularize(rot_vec, point_type=point_type) translation = point[:, dim_rotations:] regularized_point = gs.concatenate( [regularized_rot_vec, translation], axis=1) elif point_type == 'matrix': point = gs.to_ndarray(point, to_ndim=3) regularized_point = gs.copy(point) return regularized_point def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None, point_type=None): """Regularize a tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] metric : RiemannianMetric, optional point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- regularized_vec : the regularized tangent vector """ if point_type is None: point_type = self.default_point_type return self.regularize_tangent_vec(tangent_vec, self.identity, metric, point_type=point_type) def regularize_tangent_vec(self, tangent_vec, base_point, metric=None, point_type=None): """Regularize a tangent vector at a base point. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] base_point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] metric : RiemannianMetric, optional default: self.left_canonical_metric point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- regularized_vec : the regularized tangent vector """ if point_type is None: point_type = self.default_point_type if metric is None: metric = self.left_canonical_metric if point_type == 'vector': tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) base_point = gs.to_ndarray(base_point, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_tangent_vec = tangent_vec[:, :dim_rotations] rot_base_point = base_point[:, :dim_rotations] metric_mat = metric.inner_product_mat_at_identity rot_metric_mat = metric_mat[:, :dim_rotations, :dim_rotations] rot_metric = InvariantMetric( group=rotations, inner_product_mat_at_identity=rot_metric_mat, left_or_right=metric.left_or_right) regularized_vec = gs.zeros_like(tangent_vec) rotations_vec = rotations.regularize_tangent_vec( tangent_vec=rot_tangent_vec, base_point=rot_base_point, metric=rot_metric, point_type=point_type) regularized_vec = gs.concatenate( [rotations_vec, tangent_vec[:, dim_rotations:]], axis=1) elif point_type == 'matrix': regularized_vec = tangent_vec return regularized_vec def compose(self, point_1, point_2, point_type=None): r"""Compose two elements of SE(n). Parameters ---------- point_1 : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] point_2 : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Equation --------- (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`) Returns ------- composition : the composition of point_1 and point_2 """ if point_type is None: point_type = self.default_point_type rotations = self.rotations dim_rotations = rotations.dimension point_1 = self.regularize(point_1, point_type=point_type) point_2 = self.regularize(point_2, point_type=point_type) if point_type == 'vector': n_points_1, _ = point_1.shape n_points_2, _ = point_2.shape assert (point_1.shape == point_2.shape or n_points_1 == 1 or n_points_2 == 1) if n_points_1 == 1: point_1 = gs.stack([point_1[0]] * n_points_2) if n_points_2 == 1: point_2 = gs.stack([point_2[0]] * n_points_1) rot_vec_1 = point_1[:, :dim_rotations] rot_mat_1 = rotations.matrix_from_rotation_vector(rot_vec_1) rot_vec_2 = point_2[:, :dim_rotations] rot_mat_2 = rotations.matrix_from_rotation_vector(rot_vec_2) translation_1 = point_1[:, dim_rotations:] translation_2 = point_2[:, dim_rotations:] composition_rot_mat = gs.matmul(rot_mat_1, rot_mat_2) composition_rot_vec = rotations.rotation_vector_from_matrix( composition_rot_mat) composition_translation = gs.einsum('ij,ikj->ik', translation_2, rot_mat_1) + translation_1 composition = gs.concatenate( (composition_rot_vec, composition_translation), axis=1) elif point_type == 'matrix': raise NotImplementedError() composition = self.regularize(composition, point_type=point_type) return composition def inverse(self, point, point_type=None): r"""Compute the group inverse in SE(n). Parameters ---------- point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] Formula ------- :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))` Returns ------- inverse_point : array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] the inverted point """ if point_type is None: point_type = self.default_point_type rotations = self.rotations dim_rotations = rotations.dimension point = self.regularize(point) if point_type == 'vector': n_points, _ = point.shape rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] inverse_point = gs.zeros_like(point) inverse_rotation = -rot_vec inv_rot_mat = rotations.matrix_from_rotation_vector( inverse_rotation) inverse_translation = gs.einsum( 'ni,nij->nj', -translation, gs.transpose(inv_rot_mat, axes=(0, 2, 1))) inverse_point = gs.concatenate( [inverse_rotation, inverse_translation], axis=1) elif point_type == 'matrix': raise NotImplementedError() inverse_point = self.regularize(inverse_point, point_type=point_type) return inverse_point def jacobian_translation(self, point, left_or_right='left', point_type=None): """Compute the Jacobian matrix resulting from translation. Compute the jacobian matrix of the differential of the left/right translations from the identity to point in SE(n). Currently only implemented for point_type == 'vector'. Parameters ---------- point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] left_or_right: str, {'left', 'right'}, optional default: 'left' whether to compute the jacobian of the left or right translation point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- jacobian : array-like, shape=[n_samples, dimension] The jacobian of the left / right translation """ if point_type is None: point_type = self.default_point_type assert left_or_right in ('left', 'right') dim = self.dimension rotations = self.rotations translations = self.translations dim_rotations = rotations.dimension dim_translations = translations.dimension point = self.regularize(point, point_type=point_type) if point_type == 'vector': n_points, _ = point.shape rot_vec = point[:, :dim_rotations] jacobian = gs.zeros((n_points, ) + (dim, ) * 2) jacobian_rot = self.rotations.jacobian_translation( point=rot_vec, left_or_right=left_or_right, point_type=point_type) block_zeros_1 = gs.zeros( (n_points, dim_rotations, dim_translations)) jacobian_block_line_1 = gs.concatenate( [jacobian_rot, block_zeros_1], axis=2) if left_or_right == 'left': rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) jacobian_trans = rot_mat block_zeros_2 = gs.zeros( (n_points, dim_translations, dim_rotations)) jacobian_block_line_2 = gs.concatenate( [block_zeros_2, jacobian_trans], axis=2) else: inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec) eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3) eye = gs.tile(eye, [n_points, 1, 1]) jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye], axis=2) jacobian = gs.concatenate( [jacobian_block_line_1, jacobian_block_line_2], axis=1) assert gs.ndim(jacobian) == 3 elif point_type == 'matrix': raise NotImplementedError() return jacobian def exp_from_identity(self, tangent_vec, point_type=None): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- group_exp: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] the group exponential of the tangent vectors calculated at the identity """ if point_type is None: point_type = self.default_point_type if point_type == 'vector': tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2) rotations = self.rotations dim_rotations = rotations.dimension rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec, point_type=point_type) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) exp_translation += mask_i_float * (translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp, point_type=point_type) return group_exp elif point_type == 'matrix': raise NotImplementedError() def log_from_identity(self, point, point_type=None): """Compute the group logarithm of the point at the identity. Parameters ---------- point: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- group_log: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] the group logarithm in the Lie algbra """ if point_type is None: point_type = self.default_point_type point = self.regularize(point, point_type=point_type) rotations = self.rotations dim_rotations = rotations.dimension if point_type == 'vector': rot_vec = point[:, :dim_rotations] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) translation = point[:, dim_rotations:] skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec) sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec) mask_close_0 = gs.isclose(angle, 0.) mask_close_pi = gs.isclose(angle, gs.pi) mask_else = ~mask_close_0 & ~mask_close_pi mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_close_pi_float = gs.cast(mask_close_pi, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) mask_0 = gs.isclose(angle, 0., atol=1e-6) mask_0_float = gs.cast(mask_0, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = -0.5 * gs.ones_like(angle) coef_2 = gs.zeros_like(angle) coef_2 += mask_close_0_float * (1. / 12. + angle**2 / 720. + angle**4 / 30240. + angle**6 / 1209600.) delta_angle = angle - gs.pi coef_2 += mask_close_pi_float * ( 1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) - ((PI2 - 12.) * delta_angle**2 / (4. * PI4)) + ((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) - ((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) + ((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 / (480. * PI7)) - ((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 / (480. * PI8))) psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle)) coef_2 += mask_else_float * (1 - psi) / (angle**2) n_points, _ = point.shape log_translation = gs.zeros((n_points, self.n)) for i in range(n_points): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_rot_vec[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_rot_vec[i])) mask_i_float = gs.get_mask_i_float(i, n_points) log_translation += mask_i_float * (translation_i + term_1_i + term_2_i) group_log = gs.concatenate([rot_vec, log_translation], axis=1) assert gs.ndim(group_log) == 2 elif point_type == 'matrix': raise NotImplementedError() return group_log def random_uniform(self, n_samples=1, point_type=None): """Sample in SE(n) with the uniform distribution. Parameters ---------- n_samples: int, optional default: 1 point_typ: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- random_transfo: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] an array of random elements in SE(n) having the given point_type """ if point_type is None: point_type = self.default_point_type random_rot_vec = self.rotations.random_uniform(n_samples, point_type=point_type) random_translation = self.translations.random_uniform(n_samples) if point_type == 'vector': random_transfo = gs.concatenate( [random_rot_vec, random_translation], axis=1) elif point_type == 'matrix': raise NotImplementedError() random_transfo = self.regularize(random_transfo, point_type=point_type) return random_transfo def exponential_matrix(self, rot_vec): """Compute exponential of rotation matrix represented by rot_vec. Parameters ---------- rot_vec : array-like, shape=[n_samples, dimension] Returns ------- exponential_mat: The matrix exponential of rot_vec """ rot_vec = self.rotations.regularize(rot_vec) n_rot_vecs, _ = rot_vec.shape angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_rot_vec = self.rotations.skew_matrix_from_vector(rot_vec) coef_1 = gs.empty_like(angle) coef_2 = gs.empty_like(coef_1) mask_0 = gs.equal(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_close_to_0 = gs.isclose(angle, 0) mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1) mask_else = ~mask_0 & ~mask_close_to_0 coef_1[mask_close_to_0] = (1. / 2. - angle[mask_close_to_0]**2 / 24.) coef_2[mask_close_to_0] = (1. / 6. - angle[mask_close_to_0]**3 / 120.) # TODO(nina): Check if the discontinuity at 0 is expected. coef_1[mask_0] = 0 coef_2[mask_0] = 0 coef_1[mask_else] = (angle[mask_else]**(-2) * (1. - gs.cos(angle[mask_else]))) coef_2[mask_else] = (angle[mask_else]**(-2) * (1. - (gs.sin(angle[mask_else]) / angle[mask_else]))) term_1 = gs.zeros((n_rot_vecs, self.n, self.n)) term_2 = gs.zeros_like(term_1) for i in range(n_rot_vecs): term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i] term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i] exponential_mat = term_1 + term_2 assert exponential_mat.ndim == 3 return exponential_mat def exponential_barycenter(self, points, weights=None, point_type=None): """Compute the group exponential barycenter in SE(n). Parameters ---------- points: array-like, shape=[n_samples, {dimension, [n + 1, n + 1]}] weights: array-like, shape=[n_samples], optional default: weight 1 / n_samples for each point point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- exp_bar: array-like, shape=[{dimension, [n + 1, n + 1]}] the exponential barycenter """ if point_type is None: point_type = self.default_point_type n_points = points.shape[0] assert n_points > 0 if weights is None: weights = gs.ones((n_points, 1)) weights = gs.to_ndarray(weights, to_ndim=2, axis=1) n_weights, _ = weights.shape assert n_points == n_weights dim = self.dimension rotations = self.rotations dim_rotations = rotations.dimension if point_type == 'vector': rotation_vectors = points[:, :dim_rotations] translations = points[:, dim_rotations:dim] assert rotation_vectors.shape == (n_points, dim_rotations) assert translations.shape == (n_points, self.n) mean_rotation = rotations.exponential_barycenter( points=rotation_vectors, weights=weights) mean_rotation_mat = rotations.matrix_from_rotation_vector( mean_rotation) matrix = gs.zeros((1, ) + (self.n, ) * 2) translation_aux = gs.zeros((1, self.n)) inv_rot_mats = rotations.matrix_from_rotation_vector( -rotation_vectors) matrix_aux = gs.matmul(mean_rotation_mat, inv_rot_mats) assert matrix_aux.shape == (n_points, ) + (dim_rotations, ) * 2 vec_aux = rotations.rotation_vector_from_matrix(matrix_aux) matrix_aux = self.exponential_matrix(vec_aux) matrix_aux = gs.linalg.inv(matrix_aux) for i in range(n_points): matrix += weights[i] * matrix_aux[i] translation_aux += weights[i] * gs.dot( gs.matmul(matrix_aux[i], inv_rot_mats[i]), translations[i]) mean_translation = gs.dot( translation_aux, gs.transpose(gs.linalg.inv(matrix), axes=(0, 2, 1))) exp_bar = gs.zeros((1, dim)) exp_bar[0, :dim_rotations] = mean_rotation exp_bar[0, dim_rotations:dim] = mean_translation elif point_type == 'matrix': vector_points = self.rotation_vector_from_matrix(points) vector_exp_bar = self.exponential_barycenter(vector_points, weights, point_type='vector') exp_bar = self.matrix_from_rotation_vector(vector_exp_bar) return exp_bar
class TestVisualization(geomstats.tests.TestCase): def setup_method(self): self.n_samples = 10 self.SO3_GROUP = SpecialOrthogonal(n=3, point_type="vector") self.SE3_GROUP = SpecialEuclidean(n=3, point_type="vector") self.S1 = Hypersphere(dim=1) self.S2 = Hypersphere(dim=2) self.H2 = Hyperbolic(dim=2) self.H2_half_plane = PoincareHalfSpace(dim=2) self.M32 = Matrices(m=3, n=2) self.S32 = PreShapeSpace(k_landmarks=3, m_ambient=2) self.KS = visualization.KendallSphere() self.M33 = Matrices(m=3, n=3) self.S33 = PreShapeSpace(k_landmarks=3, m_ambient=3) self.KD = visualization.KendallDisk() self.spd = SPDMatrices(n=2) plt.figure() @staticmethod def test_tutorial_matplotlib(): visualization.tutorial_matplotlib() def test_plot_points_so3(self): points = self.SO3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space="SO3_GROUP") def test_plot_points_se3(self): points = self.SE3_GROUP.random_point(self.n_samples) visualization.plot(points, space="SE3_GROUP") def test_draw_pre_shape_2d(self): self.KS.draw() def test_draw_points_pre_shape_2d(self): points = self.S32.random_point(self.n_samples) visualization.plot(points, space="S32") points = self.M32.random_point(self.n_samples) visualization.plot(points, space="M32") self.KS.clear_points() def test_draw_curve_pre_shape_2d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) times = gs.linspace(0.0, 1.0, 1000) speeds = gs.array([-t * tangent_vec for t in times]) points = self.S32.ambient_metric.exp(speeds, base_point) self.KS.add_points(points) self.KS.draw_curve() self.KS.clear_points() def test_draw_vector_pre_shape_2d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) self.KS.draw_vector(tangent_vec, base_point) def test_convert_to_spherical_coordinates_pre_shape_2d(self): points = self.S32.random_point(self.n_samples) coords = self.KS.convert_to_spherical_coordinates(points) x = coords[:, 0] y = coords[:, 1] z = coords[:, 2] result = x**2 + y**2 + z**2 expected = 0.25 * gs.ones(self.n_samples) self.assertAllClose(result, expected) def test_rotation_pre_shape_2d(self): theta = gs.random.rand(1)[0] phi = gs.random.rand(1)[0] rot = self.KS.rotation(theta, phi) result = _SpecialOrthogonalMatrices(3).belongs(rot) expected = True self.assertAllClose(result, expected) def test_draw_pre_shape_3d(self): self.KD.draw() def test_draw_points_pre_shape_3d(self): points = self.S33.random_point(self.n_samples) visualization.plot(points, space="S33") points = self.M33.random_point(self.n_samples) visualization.plot(points, space="M33") self.KD.clear_points() def test_draw_curve_pre_shape_3d(self): self.KD.draw() base_point = self.S33.random_point() vec = self.S33.random_point() tangent_vec = self.S33.to_tangent(vec, base_point) tangent_vec = 0.5 * tangent_vec / self.S33.ambient_metric.norm( tangent_vec) times = gs.linspace(0.0, 1.0, 1000) speeds = gs.array([-t * tangent_vec for t in times]) points = self.S33.ambient_metric.exp(speeds, base_point) self.KD.add_points(points) self.KD.draw_curve() self.KD.clear_points() def test_draw_vector_pre_shape_3d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) self.KS.draw_vector(tangent_vec, base_point) def test_convert_to_planar_coordinates_pre_shape_3d(self): points = self.S33.random_point(self.n_samples) coords = self.KD.convert_to_planar_coordinates(points) x = coords[:, 0] y = coords[:, 1] radius = x**2 + y**2 result = [r <= 1.0 for r in radius] self.assertTrue(gs.all(result)) @geomstats.tests.np_autograd_and_torch_only def test_plot_points_s1(self): points = self.S1.random_uniform(self.n_samples) visualization.plot(points, space="S1") def test_plot_points_s2(self): points = self.S2.random_uniform(self.n_samples) visualization.plot(points, space="S2") def test_plot_points_h2_poincare_disk(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space="H2_poincare_disk") def test_plot_points_h2_poincare_half_plane_ext(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space="H2_poincare_half_plane", point_type="extrinsic") def test_plot_points_h2_poincare_half_plane_none(self): points = self.H2_half_plane.random_point(self.n_samples) visualization.plot(points, space="H2_poincare_half_plane") def test_plot_points_h2_poincare_half_plane_hs(self): points = self.H2_half_plane.random_point(self.n_samples) visualization.plot(points, space="H2_poincare_half_plane", point_type="half_space") def test_plot_points_h2_klein_disk(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space="H2_klein_disk") @staticmethod def test_plot_points_se2(): points = SpecialEuclidean(n=2, point_type="vector").random_point(4) visu = visualization.SpecialEuclidean2(points, point_type="vector") ax = visu.set_ax() visu.draw_points(ax) def test_plot_points_spd2(self): one_point = self.spd.random_point() visualization.plot(one_point, space="SPD2") points = self.spd.random_point(4) visualization.plot(points, space="SPD2") def test_compute_coordinates_spd2(self): point = gs.eye(2) ellipsis = visualization.Ellipses(n_sampling_points=4) x, y = ellipsis.compute_coordinates(point) self.assertAllClose(x, gs.array([1, 0, -1, 0, 1])) self.assertAllClose(y, gs.array([0, 1, 0, -1, 0])) @staticmethod def teardown_method(): plt.close()
class _SpecialEuclidean3Vectors(LieGroup): """Class for the special euclidean group in 3d, SE(3). i.e. the Lie group of rigid transformations. Elements of SE(3) can either be represented as vectors (in 3d) or as matrices in general. The matrix representation corresponds to homogeneous coordinates.This class is specific to the vector representation of rotations. For the matrix representation use the SpecialEuclidean class and set `n=3`. Parameter --------- epsilon : float, optional (defaults to 0) Precision to use for calculations involving potential division by 0 in rotations. """ def __init__(self, epsilon=0.): super(_SpecialEuclidean3Vectors, self).__init__(dim=6, default_point_type='vector') self.n = 3 self.epsilon = epsilon self.rotations = SpecialOrthogonal(n=3, point_type='vector', epsilon=epsilon) self.translations = Euclidean(dim=3) def get_identity(self, point_type=None): """Get the identity of the group. Parameters ---------- point_type : str, {'vector', 'matrix'}, optional The point_type of the returned value. default: self.default_point_type Returns ------- identity : array-like, shape={[dim], [n + 1, n + 1]} """ if point_type is None: point_type = self.default_point_type identity = gs.zeros(self.dim) if point_type == 'matrix': identity = gs.eye(self.n + 1) return identity identity = property(get_identity) def get_point_type_shape(self, point_type=None): """Get the shape of the instance given the default_point_style.""" return self.get_identity(point_type).shape def belongs(self, point): """Evaluate if a point belongs to SE(3). Parameters ---------- point : array-like, shape=[..., 3] The point of which to check whether it belongs to SE(3). Returns ------- belongs : array-like, shape=[..., 1] Boolean indicating whether point belongs to SE(3). """ point_dim = point.shape[-1] point_ndim = point.ndim belongs = gs.logical_and(point_dim == self.dim, point_ndim < 3) belongs = gs.logical_and(belongs, self.rotations.belongs(point[..., :self.n])) return belongs def regularize(self, point): """Regularize a point to the default representation for SE(n). Parameters ---------- point : array-like, shape=[..., 3] The point to regularize. Returns ------- point : array-like, shape=[..., 3] """ rotations = self.rotations dim_rotations = rotations.dim rot_vec = point[..., :dim_rotations] regularized_rot_vec = rotations.regularize(rot_vec) translation = point[..., dim_rotations:] return gs.concatenate([regularized_rot_vec, translation], axis=-1) @geomstats.vectorization.decorator(['else', 'vector', 'else']) def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None): """Regularize a tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[..., 3] metric : RiemannianMetric, optional Returns ------- regularized_vec : the regularized tangent vector """ return self.regularize_tangent_vec(tangent_vec, self.identity, metric) def regularize_tangent_vec(self, tangent_vec, base_point, metric=None): """Regularize a tangent vector at a base point. Parameters ---------- tangent_vec: array-like, shape=[..., 3] base_point : array-like, shape=[..., 3] metric : RiemannianMetric, optional default: self.left_canonical_metric Returns ------- regularized_vec : the regularized tangent vector """ if metric is None: metric = self.left_canonical_metric rotations = self.rotations dim_rotations = rotations.dim rot_tangent_vec = tangent_vec[..., :dim_rotations] rot_base_point = base_point[..., :dim_rotations] metric_mat = metric.inner_product_mat_at_identity rot_metric_mat = metric_mat[:dim_rotations, :dim_rotations] rot_metric = InvariantMetric( group=rotations, inner_product_mat_at_identity=rot_metric_mat, left_or_right=metric.left_or_right) rotations_vec = rotations.regularize_tangent_vec( tangent_vec=rot_tangent_vec, base_point=rot_base_point, metric=rot_metric) return gs.concatenate( [rotations_vec, tangent_vec[..., dim_rotations:]], axis=-1) @geomstats.vectorization.decorator(['else', 'vector']) def matrix_from_vector(self, vec): """Convert point in vector point-type to matrix. Parameters ---------- vec: array-like, shape=[..., 3] Returns ------- mat: array-like, shape=[..., n+1, n+1] """ vec = self.regularize(vec) n_vecs, _ = vec.shape rot_vec = vec[:, :self.rotations.dim] trans_vec = vec[:, self.rotations.dim:] rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) trans_vec = gs.reshape(trans_vec, (n_vecs, self.n, 1)) mat = gs.concatenate((rot_mat, trans_vec), axis=2) last_lines = gs.array(gs.get_mask_i_float(self.n, self.n + 1)) last_lines = gs.to_ndarray(last_lines, to_ndim=2) last_lines = gs.to_ndarray(last_lines, to_ndim=3) mat = gs.concatenate((mat, last_lines), axis=1) return mat @geomstats.vectorization.decorator(['else', 'vector', 'vector']) def compose(self, point_a, point_b): r"""Compose two elements of SE(3). Parameters ---------- point_a : array-like, shape=[..., 3] Point of the group. point_b : array-like, shape=[..., 3] Point of the group. Equation --------- (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`) Returns ------- composition : The composition of point_a and point_b. """ rotations = self.rotations dim_rotations = rotations.dim point_a = self.regularize(point_a) point_b = self.regularize(point_b) rot_vec_a = point_a[..., :dim_rotations] rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a) rot_vec_b = point_b[..., :dim_rotations] rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b) translation_a = point_a[..., dim_rotations:] translation_b = point_b[..., dim_rotations:] composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b) composition_rot_vec = rotations.rotation_vector_from_matrix( composition_rot_mat) composition_translation = gs.einsum('...j,...kj->...k', translation_b, rot_mat_a) + translation_a composition = gs.concatenate( (composition_rot_vec, composition_translation), axis=-1) return self.regularize(composition) @geomstats.vectorization.decorator(['else', 'vector']) def inverse(self, point): r"""Compute the group inverse in SE(n). Parameters ---------- point: array-like, shape=[..., 3] Returns ------- inverse_point : array-like, shape=[..., 3] The inverted point. Notes ----- :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))` """ rotations = self.rotations dim_rotations = rotations.dim point = self.regularize(point) rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] inverse_rotation = -rot_vec inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation) inverse_translation = gs.einsum( 'ni,nij->nj', -translation, gs.transpose(inv_rot_mat, axes=(0, 2, 1))) inverse_point = gs.concatenate([inverse_rotation, inverse_translation], axis=-1) return self.regularize(inverse_point) @geomstats.vectorization.decorator(['else', 'vector', 'else']) def jacobian_translation(self, point, left_or_right='left'): """Compute the Jacobian matrix resulting from translation. Compute the matrix of the differential of the left/right translations from the identity to point in SE(3). Parameters ---------- point: array-like, shape=[..., 3] left_or_right: str, {'left', 'right'}, optional Whether to compute the jacobian of the left or right translation. Returns ------- jacobian : array-like, shape=[..., 3] The jacobian of the left / right translation. """ if left_or_right not in ('left', 'right'): raise ValueError('`left_or_right` must be `left` or `right`.') rotations = self.rotations translations = self.translations dim_rotations = rotations.dim dim_translations = translations.dim point = self.regularize(point) n_points, _ = point.shape rot_vec = point[:, :dim_rotations] jacobian_rot = self.rotations.jacobian_translation( point=rot_vec, left_or_right=left_or_right) block_zeros_1 = gs.zeros((n_points, dim_rotations, dim_translations)) jacobian_block_line_1 = gs.concatenate([jacobian_rot, block_zeros_1], axis=2) if left_or_right == 'left': rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) jacobian_trans = rot_mat block_zeros_2 = gs.zeros( (n_points, dim_translations, dim_rotations)) jacobian_block_line_2 = gs.concatenate( [block_zeros_2, jacobian_trans], axis=2) else: inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec) eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3) eye = gs.tile(eye, [n_points, 1, 1]) jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye], axis=2) return gs.concatenate([jacobian_block_line_1, jacobian_block_line_2], axis=1) @geomstats.vectorization.decorator(['else', 'vector']) def exp_from_identity(self, tangent_vec): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[..., 3] Returns ------- group_exp: array-like, shape=[..., 3] The group exponential of the tangent vectors calculated at the identity. """ rotations = self.rotations dim_rotations = rotations.dim rot_vec = tangent_vec[..., :dim_rotations] rot_vec = self.rotations.regularize(rot_vec) translation = tangent_vec[..., dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=-1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * (TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * (TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) exp_translation += gs.outer(mask_i_float, translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp) return group_exp @geomstats.vectorization.decorator(['else', 'vector']) def log_from_identity(self, point): """Compute the group logarithm of the point at the identity. Parameters ---------- point: array-like, shape=[..., 3] Returns ------- group_log: array-like, shape=[..., 3] the group logarithm in the Lie algbra """ point = self.regularize(point) rotations = self.rotations dim_rotations = rotations.dim rot_vec = point[:, :dim_rotations] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) translation = point[:, dim_rotations:] skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec) sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec) mask_close_0 = gs.isclose(angle, 0.) mask_close_pi = gs.isclose(angle, gs.pi) mask_else = ~mask_close_0 & ~mask_close_pi mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_close_pi_float = gs.cast(mask_close_pi, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) mask_0 = gs.isclose(angle, 0., atol=1e-6) mask_0_float = gs.cast(mask_0, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = -0.5 * gs.ones_like(angle) coef_2 = gs.zeros_like(angle) coef_2 += mask_close_0_float * (1. / 12. + angle**2 / 720. + angle**4 / 30240. + angle**6 / 1209600.) delta_angle = angle - gs.pi coef_2 += mask_close_pi_float * ( 1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) - ((PI2 - 12.) * delta_angle**2 / (4. * PI4)) + ((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) - ((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) + ((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 / (480. * PI7)) - ((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 / (480. * PI8))) psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle)) coef_2 += mask_else_float * (1 - psi) / (angle**2) n_points, _ = point.shape log_translation = gs.zeros((n_points, self.n)) for i in range(n_points): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_rot_vec[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_rot_vec[i])) mask_i_float = gs.get_mask_i_float(i, n_points) log_translation += gs.outer(mask_i_float, translation_i + term_1_i + term_2_i) return gs.concatenate([rot_vec, log_translation], axis=1) def random_uniform(self, n_samples=1): """Sample in SE(3) with the uniform distribution. Parameters ---------- n_samples : int, optional default : 1 Returns ------- random_point : array-like, shape=[..., 3] An array of random elements in SE(3) having the given. """ random_translation = self.translations.random_uniform(n_samples) random_rot_vec = self.rotations.random_uniform(n_samples) return gs.concatenate([random_rot_vec, random_translation], axis=-1) def _exponential_matrix(self, rot_vec): """Compute exponential of rotation matrix represented by rot_vec. Parameters ---------- rot_vec : array-like, shape=[..., 3] Returns ------- exponential_mat : The matrix exponential of rot_vec """ # TODO(nguigs): find usecase for this method rot_vec = self.rotations.regularize(rot_vec) n_rot_vecs, _ = rot_vec.shape angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_rot_vec = self.rotations.skew_matrix_from_vector(rot_vec) coef_1 = gs.empty_like(angle) coef_2 = gs.empty_like(coef_1) mask_0 = gs.equal(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_close_to_0 = gs.isclose(angle, 0) mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1) mask_else = ~mask_0 & ~mask_close_to_0 coef_1[mask_close_to_0] = (1. / 2. - angle[mask_close_to_0]**2 / 24.) coef_2[mask_close_to_0] = (1. / 6. - angle[mask_close_to_0]**3 / 120.) # TODO(nina): Check if the discontinuity at 0 is expected. coef_1[mask_0] = 0 coef_2[mask_0] = 0 coef_1[mask_else] = (angle[mask_else]**(-2) * (1. - gs.cos(angle[mask_else]))) coef_2[mask_else] = (angle[mask_else]**(-2) * (1. - (gs.sin(angle[mask_else]) / angle[mask_else]))) term_1 = gs.zeros((n_rot_vecs, self.n, self.n)) term_2 = gs.zeros_like(term_1) for i in range(n_rot_vecs): term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i] term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i] exponential_mat = term_1 + term_2 return exponential_mat
class _SpecialEuclideanMatrices(GeneralLinear, LieGroup): """Class for special Euclidean group. Parameters ---------- n : int Integer dimension of the underlying Euclidean space. Matrices will be of size: (n+1) x (n+1). Attributes ---------- rotations : SpecialOrthogonal Subgroup of rotations of size n. translations : Euclidean Subgroup of translations of size n. left_canonical_metric : InvariantMetric The left invariant metric that corresponds to the Frobenius inner product at the identity. right_canonical_metric : InvariantMetric The right invariant metric that corresponds to the Frobenius inner product at the identity. metric : MatricesMetric The Euclidean (Frobenius) inner product. """ def __init__(self, n): super().__init__(n=n + 1, dim=int((n * (n + 1)) / 2), default_point_type='matrix', lie_algebra=SpecialEuclideanMatrixLieAlgebra(n=n)) self.rotations = SpecialOrthogonal(n=n) self.translations = Euclidean(dim=n) self.n = n self.left_canonical_metric = \ SpecialEuclideanMatrixCannonicalLeftMetric(group=self) def get_identity(self): """Return the identity matrix.""" return gs.eye(self.n + 1, self.n + 1) identity = property(get_identity) def belongs(self, point, atol=gs.atol): """Check whether point is of the form rotation, translation. Parameters ---------- point : array-like, shape=[..., n, n]. Point to be checked. atol : float Tolerance threshold. Returns ------- belongs : array-like, shape=[...,] Boolean denoting if point belongs to the group. """ n = self.n belongs = Matrices(n + 1, n + 1).belongs(point) if gs.all(belongs): rotation = point[..., :n, :n] belongs = self.rotations.belongs(rotation, atol=atol) last_line_except_last_term = point[..., n:, :-1] all_but_last_zeros = ~gs.any(last_line_except_last_term, axis=(-2, -1)) belongs = gs.logical_and(belongs, all_but_last_zeros) last_term = point[..., n, n] belongs = gs.logical_and(belongs, gs.isclose(last_term, 1., atol=atol)) return belongs def random_point(self, n_samples=1, bound=1.): """Sample in SE(n) from the uniform distribution. Parameters ---------- n_samples : int Number of samples. Optional, default: 1. bound: float Bound of the interval in which to sample each entry of the translation part. Optional, default: 1. Returns ------- samples : array-like, shape=[..., n + 1, n + 1] Sample in SE(n). """ random_translation = self.translations.random_point(n_samples) random_rotation = self.rotations.random_uniform(n_samples) output_shape = ((n_samples, self.n + 1, self.n + 1) if n_samples != 1 else (self.n + 1, ) * 2) random_point = homogeneous_representation(random_rotation, random_translation, output_shape) return random_point @classmethod def inverse(cls, point): """Return the inverse of a point. Parameters ---------- point : array-like, shape=[..., n, n] Point to be inverted. """ n = point.shape[-1] - 1 transposed_rot = cls.transpose(point[..., :n, :n]) translation = point[..., :n, -1] translation = gs.einsum('...ij,...j->...i', transposed_rot, translation) return homogeneous_representation(transposed_rot, -translation, point.shape)
class TestToTangentSpace(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, point_type='matrix') def test_estimate_transform_sphere(self): point = gs.array([0., 0., 0., 0., 1.]) points = gs.array([point, point]) transformer = ToTangentSpace(geometry=self.sphere) transformer.fit(X=points) result = transformer.transform(points) expected = gs.zeros_like(points) self.assertAllClose(expected, result) def test_inverse_transform_no_fit_sphere(self): point = self.sphere.random_uniform(3) base_point = point[0] point = point[1:] transformer = ToTangentSpace(geometry=self.sphere) X = transformer.transform(point, base_point=base_point) result = transformer.inverse_transform(X, base_point=base_point) expected = point self.assertAllClose(expected, result) @geomstats.tests.np_and_tf_only def test_estimate_transform_so_group(self): point = self.so_matrix.random_uniform() points = gs.array([point, point]) transformer = ToTangentSpace(geometry=self.so_matrix) transformer.fit(X=points) result = transformer.transform(points) expected = gs.zeros((2, 6)) self.assertAllClose(expected, result) def test_estimate_transform_spd(self): point = spd.SPDMatrices(3).random_uniform() points = gs.stack([point, point]) transformer = ToTangentSpace(geometry=spd.SPDMetricAffine(3)) transformer.fit(X=points) result = transformer.transform(points) expected = gs.zeros((2, 6)) self.assertAllClose(expected, result, atol=1e-5) def test_fit_transform_hyperbolic(self): point = gs.array([2., 1., 1., 1.]) points = gs.array([point, point]) transformer = ToTangentSpace(geometry=self.hyperbolic.metric) result = transformer.fit_transform(X=points) expected = gs.zeros_like(points) self.assertAllClose(expected, result) def test_inverse_transform_hyperbolic(self): points = self.hyperbolic.random_uniform(10) transformer = ToTangentSpace(geometry=self.hyperbolic.metric) X = transformer.fit_transform(X=points) result = transformer.inverse_transform(X) expected = points self.assertAllClose(expected, result) def test_inverse_transform_spd(self): point = spd.SPDMatrices(3).random_uniform(10) transformer = ToTangentSpace(geometry=spd.SPDMetricLogEuclidean(3)) X = transformer.fit_transform(X=point) result = transformer.inverse_transform(X) expected = point self.assertAllClose(expected, result, atol=1e-4) transformer = ToTangentSpace(geometry=spd.SPDMetricAffine(3)) X = transformer.fit_transform(X=point) result = transformer.inverse_transform(X) expected = point self.assertAllClose(expected, result, atol=1e-4) @geomstats.tests.np_only def test_inverse_transform_so(self): # FIXME: einsum vectorization error for invariant_metric log in tf point = self.so_matrix.random_uniform(10) transformer = ToTangentSpace( geometry=self.so_matrix.bi_invariant_metric) X = transformer.transform(X=point, base_point=self.so_matrix.identity) result = transformer.inverse_transform( X, base_point=self.so_matrix.identity) expected = point self.assertAllClose(expected, result)
class TestBackends(geomstats.tests.TestCase): def setUp(self): warnings.simplefilter('ignore', category=ImportWarning) self.so3_group = SpecialOrthogonal(n=3) self.n_samples = 2 def test_array(self): gs_mat = gs.array(1.5) np_mat = _np.array(1.5) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)]) np_mat = _np.array([_np.ones(3), _np.ones(3)]) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)], dtype=gs.float64) np_mat = _np.array([_np.ones(3), _np.ones(3)], dtype=_np.float64) assert gs_mat.dtype == gs.float64 self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([[gs.ones(3)], [gs.ones(3)]], dtype=gs.uint8) np_mat = _np.array([[_np.ones(3)], [_np.ones(3)]], dtype=_np.uint8) assert gs_mat.dtype == gs.uint8 self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), [0, 0, 0]], dtype=gs.int32) np_mat = _np.array([_np.ones(3), [0, 0, 0]], dtype=_np.int32) assert gs_mat.dtype == gs.int32 self.assertAllCloseToNp(gs_mat, np_mat) def test_matmul(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) gs_result = gs.matmul(gs_mat_a, gs_mat_b) np_result = _np.matmul(np_mat_a, np_mat_b) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_matmul_vectorization(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] mat_c = [[1., 4., 2.], [4., 3., 4.], [0., 0., 4.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) gs_mat_c = gs.array(mat_c) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) np_mat_c = _np.array(mat_c) gs_result = gs.matmul(gs_mat_a, [gs_mat_b, gs_mat_c]) np_result = _np.matmul(np_mat_a, [np_mat_b, np_mat_c]) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.logm(point) expected = gs.array([[0.693147180, 0., 0.], [0., 1.098612288, 0.], [0., 0., 1.38629436]]) self.assertAllClose(result, expected) np_point = _np.array( [[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.logm(np_point) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_and_tf_only def test_expm_and_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) np_point = _np.array( [[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.expm(scipy.linalg.logm(np_point)) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_only def test_expm_vectorization(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[7.38905609, 0., 0.], [0., 20.0855369, 0.], [0., 0., 54.5981500]], [[2.718281828, 0., 0.], [0., 148.413159, 0.], [0., 0., 403.42879349]]]) result = gs.linalg.expm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_logm_vectorization_diagonal(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[0.693147180, 0., 0.], [0., 1.09861228866, 0.], [0., 0., 1.38629436]], [[0., 0., 0.], [0., 1.609437912, 0.], [0., 0., 1.79175946]]]) result = gs.linalg.logm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization_random_rotation(self): point = self.so3_group.random_uniform(self.n_samples) point = self.so3_group.matrix_from_rotation_vector(point) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization(self): point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm_diagonal(self): power = .5 point = gs.array([[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]]) result = gs.linalg.powerm(point, power) expected = gs.array([[1., 0., 0.], [0., 2., 0.], [0., 0., 3.]]) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm(self): power = 2.4 point = gs.array([[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]) result = gs.linalg.powerm(point, power) result = gs.linalg.powerm(result, 1 / power) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_only def test_powerm_vectorization(self): power = 2.4 points = gs.array([[[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]], [[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]]) result = gs.linalg.powerm(points, power) result = gs.linalg.powerm(result, 1. / power) expected = points self.assertAllClose(result, expected) @geomstats.tests.tf_only def test_vstack(self): import tensorflow as tf tensor_1 = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.]]) tensor_2 = tf.convert_to_tensor([[7., 8., 9.]]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([ [1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) self.assertAllClose(result, expected) @geomstats.tests.tf_only def test_tensor_addition(self): tensor_1 = gs.ones((1, 1)) tensor_2 = gs.ones((0, 1)) tensor_1 + tensor_2 @geomstats.tests.pytorch_only def test_cumsum(self): result = gs.cumsum(gs.arange(10)) expected = gs.array(([0, 1, 3, 6, 10, 15, 21, 28, 36, 45])) self.assertAllClose(result, expected) result = gs.cumsum(gs.arange(10).reshape(2, 5), axis=1) expected = gs.array(([[0, 1, 3, 6, 10], [5, 11, 18, 26, 35]])) self.assertAllClose(result, expected) def test_array_from_sparse(self): expected = gs.array([[0, 1, 0], [0, 0, 2]]) result = gs.array_from_sparse([(0, 1), (1, 2)], [1, 2], (2, 3)) self.assertAllClose(result, expected) def test_einsum(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4], [-1, 5]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4], [-1, 5]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3], [5, 6]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3], [5, 6]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array([5]) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array(5) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array([5]) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array(5) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_assignment(self): np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] = 1.5 gs_result = gs.assignment(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] = 1.5 gs_result = gs.assignment(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] = 1 gs_result = gs.assignment(gs_array_2, 1, 0, axis=0) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] = 1 gs_result = gs.assignment(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] = 1 gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] = 1 gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4) def test_assignment_by_sum(self): np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] += 1 gs_result = gs.assignment_by_sum(gs_array_2, 1, 0, axis=0) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] += 1 gs_result = gs.assignment_by_sum( gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4)
class _SpecialEuclideanVectors(LieGroup): """Base Class for the special Euclidean groups in 2d and 3d in vector form. i.e. the Lie group of rigid transformations. Elements of SE(2), SE(3) can either be represented as vectors (in 2d or 3d) or as matrices in general. The matrix representation corresponds to homogeneous coordinates. This class is specific to the vector representation of rotations. For the matrix representation use the SpecialEuclidean class and set `n=2` or `n=3`. Parameter --------- epsilon : float Precision to use for calculations involving potential division by 0 in rotations. Optional, default: 0. """ def __init__(self, n, epsilon=0.0): dim = n * (n + 1) // 2 LieGroup.__init__( self, dim=dim, shape=(dim, ), default_point_type="vector", lie_algebra=Euclidean(dim), ) self.n = n self.epsilon = epsilon self.rotations = SpecialOrthogonal(n=n, point_type="vector", epsilon=epsilon) self.translations = Euclidean(dim=n) def get_identity(self, point_type=None): """Get the identity of the group. Parameters ---------- point_type : str, {'vector', 'matrix'} The point_type of the returned value. Optional, default: self.default_point_type Returns ------- identity : array-like, shape={[dim], [n + 1, n + 1]} """ if point_type is None: point_type = self.default_point_type identity = gs.zeros(self.dim) return identity identity = property(get_identity) def get_point_type_shape(self, point_type=None): """Get the shape of the instance given the default_point_style.""" return self.get_identity(point_type).shape def belongs(self, point, atol=gs.atol): """Evaluate if a point belongs to SE(2) or SE(3). Parameters ---------- point : array-like, shape=[..., dim] Point to check. Returns ------- belongs : array-like, shape=[...,] Boolean indicating whether point belongs to SE(2) or SE(3). """ point_dim = point.shape[-1] point_ndim = point.ndim belongs = gs.logical_and(point_dim == self.dim, point_ndim < 3) belongs = gs.logical_and( belongs, self.rotations.belongs(point[..., :self.rotations.dim], atol=atol)) return belongs def projection(self, point): """Project a point to the group. The point is regularized, so that the norm of the rotation part lie in [0, pi). Parameters ---------- point: array-like, shape[..., dim] Point. Returns ------- projected: array-like, shape[..., dim] Regularized point. """ return self.regularize(point) def regularize(self, point): """Regularize a point to the default representation for SE(n). Parameters ---------- point : array-like, shape=[..., dim] Point to regularize. Returns ------- point : array-like, shape=[..., dim] Regularized point. """ rotations = self.rotations dim_rotations = rotations.dim regularized_point = gs.copy(point) rot_vec = regularized_point[..., :dim_rotations] regularized_rot_vec = rotations.regularize(rot_vec) translation = regularized_point[..., dim_rotations:] return gs.concatenate([regularized_rot_vec, translation], axis=-1) @geomstats.vectorization.decorator(["else", "vector", "else"]) def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None): """Regularize a tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[..., dim] Tangent vector at base point. metric : RiemannianMetric Metric. Optional, default: None. Returns ------- regularized_vec : array-like, shape=[..., dim] Regularized vector. """ return self.regularize_tangent_vec(tangent_vec, self.identity, metric) @geomstats.vectorization.decorator(["else", "vector"]) def matrix_from_vector(self, vec): """Convert point in vector point-type to matrix. Parameters ---------- vec : array-like, shape=[..., dim] Vector. Returns ------- mat : array-like, shape=[..., n+1, n+1] Matrix. """ vec = self.regularize(vec) output_shape = ((vec.shape[0], self.n + 1, self.n + 1) if vec.ndim == 2 else (self.n + 1, ) * 2) rot_vec = vec[..., :self.rotations.dim] trans_vec = vec[..., self.rotations.dim:] rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) return homogeneous_representation(rot_mat, trans_vec, output_shape) @geomstats.vectorization.decorator(["else", "vector", "vector"]) def compose(self, point_a, point_b): r"""Compose two elements of SE(2) or SE(3). Parameters ---------- point_a : array-like, shape=[..., dim] Point of the group. point_b : array-like, shape=[..., dim] Point of the group. Equation -------- (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`) Returns ------- composition : array-like, shape=[..., dim] Composition of point_a and point_b. """ rotations = self.rotations dim_rotations = rotations.dim point_a = self.regularize(point_a) point_b = self.regularize(point_b) rot_vec_a = point_a[..., :dim_rotations] rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a) rot_vec_b = point_b[..., :dim_rotations] rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b) translation_a = point_a[..., dim_rotations:] translation_b = point_b[..., dim_rotations:] composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b) composition_rot_vec = rotations.rotation_vector_from_matrix( composition_rot_mat) composition_translation = ( gs.einsum("...j,...kj->...k", translation_b, rot_mat_a) + translation_a) composition = gs.concatenate( (composition_rot_vec, composition_translation), axis=-1) return self.regularize(composition) @geomstats.vectorization.decorator(["else", "vector"]) def inverse(self, point): r"""Compute the group inverse in SE(n). Parameters ---------- point: array-like, shape=[..., dim] Point. Returns ------- inverse_point : array-like, shape=[..., dim] Inverted point. Notes ----- :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))` """ rotations = self.rotations dim_rotations = rotations.dim point = self.regularize(point) rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] inverse_rotation = -rot_vec inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation) inverse_translation = gs.einsum( "ni,nij->nj", -translation, gs.transpose(inv_rot_mat, axes=(0, 2, 1))) inverse_point = gs.concatenate([inverse_rotation, inverse_translation], axis=-1) return self.regularize(inverse_point) @geomstats.vectorization.decorator(["else", "vector"]) def exp_from_identity(self, tangent_vec): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[..., 3] Tangent vector at base point. Returns ------- group_exp: array-like, shape=[..., 3] Group exponential of the tangent vectors computed at the identity. """ rotations = self.rotations dim_rotations = rotations.dim rot_vec = tangent_vec[..., :dim_rotations] rot_vec_regul = self.rotations.regularize(rot_vec) rot_vec_regul = gs.to_ndarray(rot_vec_regul, to_ndim=2, axis=1) transform = self._exp_translation_transform(rot_vec_regul) translation = tangent_vec[..., dim_rotations:] exp_translation = gs.einsum("ijk, ik -> ij", transform, translation) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp) return group_exp @geomstats.vectorization.decorator(["else", "vector"]) def log_from_identity(self, point): """Compute the group logarithm of the point at the identity. Parameters ---------- point: array-like, shape=[..., 3] Point. Returns ------- group_log: array-like, shape=[..., 3] Group logarithm in the Lie algebra. """ point = self.regularize(point) rotations = self.rotations dim_rotations = rotations.dim rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] transform = self._log_translation_transform(rot_vec) log_translation = gs.einsum("ijk, ik -> ij", transform, translation) return gs.concatenate([rot_vec, log_translation], axis=1) def random_point(self, n_samples=1, bound=1.0, **kwargs): r"""Sample in SE(n) with the uniform distribution. Parameters ---------- n_samples : int Number of samples. Optional, default: 1. bound : float Upper bound for the translation part of the sample. Optional, default: 1. Returns ------- random_point : array-like, shape=[..., dim] Sample. """ random_translation = self.translations.random_point(n_samples, bound) random_rot_vec = self.rotations.random_uniform(n_samples) return gs.concatenate([random_rot_vec, random_translation], axis=-1)
class TestBackends(geomstats.tests.TestCase): def setUp(self): warnings.simplefilter('ignore', category=ImportWarning) self.so3_group = SpecialOrthogonal(n=3) self.n_samples = 2 def test_matmul(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) np_mat_a = np.array(mat_a) np_mat_b = np.array(mat_b) gs_result = gs.matmul(gs_mat_a, gs_mat_b) np_result = np.matmul(np_mat_a, np_mat_b) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_matmul_vectorization(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] mat_c = [[1., 4., 2.], [4., 3., 4.], [0., 0., 4.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) gs_mat_c = gs.array(mat_c) np_mat_a = np.array(mat_a) np_mat_b = np.array(mat_b) np_mat_c = np.array(mat_c) gs_result = gs.matmul(gs_mat_a, [gs_mat_b, gs_mat_c]) np_result = np.matmul(np_mat_a, [np_mat_b, np_mat_c]) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.logm(point) expected = gs.array([[0.693147180, 0., 0.], [0., 1.098612288, 0.], [0., 0., 1.38629436]]) self.assertAllClose(result, expected) np_point = np.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.logm(np_point) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_and_tf_only def test_expm_and_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) np_point = np.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.expm(scipy.linalg.logm(np_point)) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_only def test_expm_vectorization(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[7.38905609, 0., 0.], [0., 20.0855369, 0.], [0., 0., 54.5981500]], [[2.718281828, 0., 0.], [0., 148.413159, 0.], [0., 0., 403.42879349]]]) result = gs.linalg.expm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_logm_vectorization_diagonal(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[0.693147180, 0., 0.], [0., 1.09861228866, 0.], [0., 0., 1.38629436]], [[0., 0., 0.], [0., 1.609437912, 0.], [0., 0., 1.79175946]]]) result = gs.linalg.logm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization_random_rotation(self): point = self.so3_group.random_uniform(self.n_samples) point = self.so3_group.matrix_from_rotation_vector(point) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization(self): point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm_diagonal(self): power = .5 point = gs.array([[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]]) result = gs.linalg.powerm(point, power) expected = gs.array([[1., 0., 0.], [0., 2., 0.], [0., 0., 3.]]) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm(self): power = 2.4 point = gs.array([[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]) result = gs.linalg.powerm(point, power) result = gs.linalg.powerm(result, 1 / power) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_only def test_powerm_vectorization(self): power = 2.4 points = gs.array([[[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]], [[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]]) result = gs.linalg.powerm(points, power) result = gs.linalg.powerm(result, 1. / power) expected = points self.assertAllClose(result, expected) @geomstats.tests.pytorch_only def test_sampling_choice(self): res = gs.random.choice(10, (5, 1, 3)) self.assertAllClose(res.shape, [5, 1, 3]) @geomstats.tests.tf_only def test_vstack(self): import tensorflow as tf with self.test_session(): tensor_1 = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.]]) tensor_2 = tf.convert_to_tensor([[7., 8., 9.]]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) self.assertAllClose(result, expected) @geomstats.tests.tf_only def test_tensor_addition(self): with self.test_session(): tensor_1 = gs.ones((1, 1)) tensor_2 = gs.ones((0, 1)) tensor_1 + tensor_2
class TestSpecialOrthogonal(geomstats.tests.TestCase): def setUp(self): self.n = 2 self.group = SpecialOrthogonal(n=self.n) self.n_samples = 4 def test_belongs(self): theta = gs.pi / 3 point_1 = gs.array([[gs.cos(theta), -gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_1) expected = True self.assertAllClose(result, expected) point_2 = gs.array([[gs.cos(theta), gs.sin(theta)], [gs.sin(theta), gs.cos(theta)]]) result = self.group.belongs(point_2) expected = False self.assertAllClose(result, expected) point = gs.array([point_1, point_2]) expected = gs.array([True, False]) result = self.group.belongs(point) self.assertAllClose(result, expected) def test_random_uniform_and_belongs(self): point = self.group.random_uniform() result = self.group.belongs(point) expected = True self.assertAllClose(result, expected) point = self.group.random_uniform(self.n_samples) result = self.group.belongs(point) expected = gs.array([True] * self.n_samples) self.assertAllClose(result, expected) def test_identity(self): result = self.group.identity expected = gs.eye(self.n) self.assertAllClose(result, expected) def test_is_in_lie_algebra(self): theta = gs.pi / 3 vec_1 = gs.array([[0., -theta], [theta, 0.]]) result = self.group.is_tangent(vec_1) expected = True self.assertAllClose(result, expected) vec_2 = gs.array([[0., -theta], [theta, 1.]]) result = self.group.is_tangent(vec_2) expected = False self.assertAllClose(result, expected) vec = gs.array([vec_1, vec_2]) expected = gs.array([True, False]) result = self.group.is_tangent(vec) self.assertAllClose(result, expected) def test_is_tangent(self): point = self.group.random_uniform() theta = 1. vec_1 = gs.array([[0., -theta], [theta, 0.]]) vec_1 = self.group.compose(point, vec_1) result = self.group.is_tangent(vec_1, point, atol=1e-6) expected = True self.assertAllClose(result, expected) vec_2 = gs.array([[0., -theta], [theta, 1.]]) vec_2 = self.group.compose(point, vec_2) result = self.group.is_tangent(vec_2, point, atol=1e-6) expected = False self.assertAllClose(result, expected) vec = gs.array([vec_1, vec_2]) point = gs.array([point, point]) expected = gs.array([True, False]) result = self.group.is_tangent(vec, point, atol=1e-6) self.assertAllClose(result, expected) def test_to_tangent(self): theta = 1. vec_1 = gs.array([[0., -theta], [theta, 0.]]) result = self.group.to_tangent(vec_1) expected = vec_1 self.assertAllClose(result, expected) n_samples = self.n_samples base_points = self.group.random_uniform(n_samples=n_samples) tangent_vecs = self.group.compose(base_points, vec_1) result = self.group.to_tangent(tangent_vecs, base_points) expected = tangent_vecs self.assertAllClose(result, expected)
class _SpecialEuclideanMatrices(MatrixLieGroup, LevelSet): """Class for special Euclidean group. Parameters ---------- n : int Integer dimension of the underlying Euclidean space. Matrices will be of size: (n+1) x (n+1). Attributes ---------- rotations : SpecialOrthogonal Subgroup of rotations of size n. translations : Euclidean Subgroup of translations of size n. left_canonical_metric : InvariantMetric The left invariant metric that corresponds to the Frobenius inner product at the identity. right_canonical_metric : InvariantMetric The right invariant metric that corresponds to the Frobenius inner product at the identity. metric : MatricesMetric The Euclidean (Frobenius) inner product. """ def __init__(self, n, **kwargs): super().__init__(n=n + 1, dim=int((n * (n + 1)) / 2), embedding_space=GeneralLinear(n + 1, positive_det=True), submersion=submersion, value=gs.eye(n + 1), tangent_submersion=tangent_submersion, lie_algebra=SpecialEuclideanMatrixLieAlgebra(n=n), **kwargs) self.rotations = SpecialOrthogonal(n=n) self.translations = Euclidean(dim=n) self.n = n self.left_canonical_metric = SpecialEuclideanMatrixCannonicalLeftMetric( group=self) if self._metric is None: self._metric = self.left_canonical_metric @property def identity(self): """Return the identity matrix.""" return gs.eye(self.n + 1, self.n + 1) def random_point(self, n_samples=1, bound=1.0): """Sample in SE(n) from the uniform distribution. Parameters ---------- n_samples : int Number of samples. Optional, default: 1. bound: float Bound of the interval in which to sample each entry of the translation part. Optional, default: 1. Returns ------- samples : array-like, shape=[..., n + 1, n + 1] Sample in SE(n). """ random_translation = self.translations.random_point(n_samples) random_rotation = self.rotations.random_uniform(n_samples) output_shape = ((n_samples, self.n + 1, self.n + 1) if n_samples != 1 else (self.n + 1, ) * 2) random_point = homogeneous_representation(random_rotation, random_translation, output_shape) return random_point @classmethod def inverse(cls, point): """Return the inverse of a point. Parameters ---------- point : array-like, shape=[..., n + 1, n + 1] Point to be inverted. Returns ------- inverse : array-like, shape=[..., n + 1, n + 1] Inverse of point. """ n = point.shape[-1] - 1 transposed_rot = Matrices.transpose(point[..., :n, :n]) translation = point[..., :n, -1] translation = gs.einsum("...ij,...j->...i", transposed_rot, translation) return homogeneous_representation(transposed_rot, -translation, point.shape) def projection(self, mat): """Project a matrix on SE(n). The upper-left n x n block is projected to SO(n) by minimizing the Frobenius norm. The last columns is kept unchanged and used as the translation part. The last row is discarded. Parameters ---------- mat : array-like, shape=[..., n + 1, n + 1] Matrix. Returns ------- projected : array-like, shape=[..., n + 1, n + 1] Rotation-translation matrix in homogeneous representation. """ n = mat.shape[-1] - 1 projected_rot = self.rotations.projection(mat[..., :n, :n]) translation = mat[..., :n, -1] return homogeneous_representation(projected_rot, translation, mat.shape)
class TestBackends(geomstats.tests.TestCase): def setUp(self): warnings.simplefilter('ignore', category=ImportWarning) self.so3_group = SpecialOrthogonal(n=3) self.n_samples = 2 def test_array(self): gs_mat = gs.array(1.5) np_mat = _np.array(1.5) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)]) np_mat = _np.array([_np.ones(3), _np.ones(3)]) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), gs.ones(3)], dtype=gs.float64) np_mat = _np.array([_np.ones(3), _np.ones(3)], dtype=_np.float64) self.assertTrue(gs_mat.dtype == gs.float64) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([[gs.ones(3)], [gs.ones(3)]], dtype=gs.uint8) np_mat = _np.array([[_np.ones(3)], [_np.ones(3)]], dtype=_np.uint8) self.assertTrue(gs_mat.dtype == gs.uint8) self.assertAllCloseToNp(gs_mat, np_mat) gs_mat = gs.array([gs.ones(3), [0, 0, 0]], dtype=gs.int32) np_mat = _np.array([_np.ones(3), [0, 0, 0]], dtype=_np.int32) self.assertTrue(gs_mat.dtype == gs.int32) self.assertAllCloseToNp(gs_mat, np_mat) def test_matmul(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) gs_result = gs.matmul(gs_mat_a, gs_mat_b) np_result = _np.matmul(np_mat_a, np_mat_b) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_matmul_vectorization(self): mat_a = [[2., 0., 0.], [0., 3., 0.], [7., 0., 4.]] mat_b = [[1., 0., 2.], [0., 3., 0.], [0., 0., 1.]] mat_c = [[1., 4., 2.], [4., 3., 4.], [0., 0., 4.]] gs_mat_a = gs.array(mat_a) gs_mat_b = gs.array(mat_b) gs_mat_c = gs.array(mat_c) np_mat_a = _np.array(mat_a) np_mat_b = _np.array(mat_b) np_mat_c = _np.array(mat_c) gs_result = gs.matmul(gs_mat_a, [gs_mat_b, gs_mat_c]) np_result = _np.matmul(np_mat_a, [np_mat_b, np_mat_c]) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_tf_only def test_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.logm(point) expected = gs.array([[0.693147180, 0., 0.], [0., 1.098612288, 0.], [0., 0., 1.38629436]]) self.assertAllClose(result, expected) np_point = _np.array( [[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.logm(np_point) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_and_tf_only def test_expm_and_logm(self): point = gs.array([[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) np_point = _np.array( [[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]]) scipy_result = scipy.linalg.expm(scipy.linalg.logm(np_point)) self.assertAllCloseToNp(result, scipy_result) @geomstats.tests.np_only def test_expm_vectorization(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[7.38905609, 0., 0.], [0., 20.0855369, 0.], [0., 0., 54.5981500]], [[2.718281828, 0., 0.], [0., 148.413159, 0.], [0., 0., 403.42879349]]]) result = gs.linalg.expm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_logm_vectorization_diagonal(self): # Note: scipy.linalg.expm is not vectorized point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) expected = gs.array([[[0.693147180, 0., 0.], [0., 1.09861228866, 0.], [0., 0., 1.38629436]], [[0., 0., 0.], [0., 1.609437912, 0.], [0., 0., 1.79175946]]]) result = gs.linalg.logm(point) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization_random_rotation(self): point = self.so3_group.random_uniform(self.n_samples) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_expm_and_logm_vectorization(self): point = gs.array([[[2., 0., 0.], [0., 3., 0.], [0., 0., 4.]], [[1., 0., 0.], [0., 5., 0.], [0., 0., 6.]]]) result = gs.linalg.expm(gs.linalg.logm(point)) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm_diagonal(self): power = .5 point = gs.array([[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]]) result = gs.linalg.powerm(point, power) expected = gs.array([[1., 0., 0.], [0., 2., 0.], [0., 0., 3.]]) self.assertAllClose(result, expected) @geomstats.tests.np_and_tf_only def test_powerm(self): power = 2.4 point = gs.array([[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]) result = gs.linalg.powerm(point, power) result = gs.linalg.powerm(result, 1 / power) expected = point self.assertAllClose(result, expected) @geomstats.tests.np_only def test_powerm_vectorization(self): power = 2.4 points = gs.array([[[1., 0., 0.], [0., 4., 0.], [0., 0., 9.]], [[1., 0., 0.], [0., 2.5, 1.5], [0., 1.5, 2.5]]]) result = gs.linalg.powerm(points, power) result = gs.linalg.powerm(result, 1. / power) expected = points self.assertAllClose(result, expected) @geomstats.tests.tf_only def test_vstack(self): import tensorflow as tf tensor_1 = tf.convert_to_tensor([[1., 2., 3.], [4., 5., 6.]]) tensor_2 = tf.convert_to_tensor([[7., 8., 9.]]) result = gs.vstack([tensor_1, tensor_2]) expected = tf.convert_to_tensor([ [1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) self.assertAllClose(result, expected) @geomstats.tests.pytorch_only def test_cumsum(self): result = gs.cumsum(gs.arange(10)) expected = gs.array(([0, 1, 3, 6, 10, 15, 21, 28, 36, 45])) self.assertAllClose(result, expected) result = gs.cumsum(gs.arange(10).reshape(2, 5), axis=1) expected = gs.array(([[0, 1, 3, 6, 10], [5, 11, 18, 26, 35]])) self.assertAllClose(result, expected) def test_array_from_sparse(self): expected = gs.array([[0, 1, 0], [0, 0, 2]]) result = gs.array_from_sparse([(0, 1), (1, 2)], [1, 2], (2, 3)) self.assertAllClose(result, expected) def test_einsum(self): np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4], [-1, 5]]) np_array_2 = _np.array([[2, 3]]) array_1 = gs.array([[1, 4], [-1, 5]]) array_2 = gs.array([[2, 3]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([[1, 4]]) np_array_2 = _np.array([[2, 3], [5, 6]]) array_1 = gs.array([[1, 4]]) array_2 = gs.array([[2, 3], [5, 6]]) np_result = _np.einsum('...i,...i->...', np_array_1, np_array_2) gs_result = gs.einsum('...i,...i->...', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array([5]) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([[1, 2, 3]]) array_1 = gs.array(5) array_2 = gs.array([[1, 2, 3]]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array([5]) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array([5]) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) np_array_1 = _np.array(5) np_array_2 = _np.array([1, 2, 3]) array_1 = gs.array(5) array_2 = gs.array([1, 2, 3]) np_result = _np.einsum('...,...i->...i', np_array_1, np_array_2) gs_result = gs.einsum('...,...i->...i', array_1, array_2) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_matrices(self): np_array = _np.zeros((2, 3, 3)) gs_array = gs.zeros((2, 3, 3)) np_array[:, 0, 1] = 44. gs_array = gs.assignment( gs_array, 44., (0, 1), axis=0) self.assertAllCloseToNp(gs_array, np_array) n_samples = 3 theta = _np.random.rand(5) phi = _np.random.rand(5) np_array = _np.zeros((n_samples, 5, 4)) gs_array = gs.array(np_array) np_array[0, :, 0] = gs.cos(theta) * gs.cos(phi) np_array[0, :, 1] = - gs.sin(theta) * gs.sin(phi) gs_array = gs.assignment( gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment( gs_array, - gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) self.assertAllCloseToNp(gs_array, np_array) def test_assignment_with_booleans_single_index(self): np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([True]) gs_mask = gs.array([True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([[2., 5.]]) gs_array = gs.array([[2., 5.]]) np_mask = _np.array([False]) gs_mask = gs.array([False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment_with_booleans_many_indices(self): np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, False, True]) gs_mask = gs.array([True, False, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([False, True, True]) gs_mask = gs.array([False, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask]) np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs.ones_like(gs_array[~gs_mask]), ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([True, True, True]) gs_mask = gs.array([True, True, True]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([False, False, False]) gs_mask = gs.array([False, False, False]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] = _np.zeros_like(np_array[np_mask]) np_array[~np_mask] = 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment( gs_array, values_mask, gs_mask) gs_result = gs.assignment( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_assignment(self): gs_array_1 = gs.ones(3) self.assertRaises( ValueError, gs.assignment, gs_array_1, [.1, 2., 1.], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] = 1.5 gs_result = gs.assignment(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] = 1.5 gs_result = gs.assignment(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] = 1 gs_result = gs.assignment(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] = 1 gs_result = gs.assignment(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] = 1 gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) gs_array_4_arr = gs.zeros_like(gs.array(np_array_4)) gs_result = gs.assignment(gs_array_4_arr, 1, gs.array((0, 1)), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] = 1 gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) def test_assignment_by_sum(self): gs_array_1 = gs.ones(3) self.assertRaises( ValueError, gs.assignment_by_sum, gs_array_1, [.1, 2., 1.], [0, 1]) np_array_1 = _np.ones(3) gs_array_1 = gs.ones_like(gs.array(np_array_1)) np_array_1[2] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1, 1.5, 2) self.assertAllCloseToNp(gs_result, np_array_1) gs_result_list = gs.assignment_by_sum(gs_array_1, [2., 1.5], [0, 2]) np_array_1[0] += 2. self.assertAllCloseToNp(gs_result_list, np_array_1) np_array_1_list = _np.ones(3) gs_array_1_list = gs.ones_like(gs.array(np_array_1_list)) indices = [1, 2] np_array_1_list[indices] += 1.5 gs_result = gs.assignment_by_sum(gs_array_1_list, 1.5, indices) self.assertAllCloseToNp(gs_result, np_array_1_list) np_array_2 = _np.zeros((3, 2)) gs_array_2 = gs.zeros_like(gs.array(np_array_2)) np_array_2[0, :] += 1 gs_result = gs.assignment_by_sum(gs_array_2, 1, 0, axis=1) self.assertAllCloseToNp(gs_result, np_array_2) np_array_3 = _np.zeros((3, 3)) gs_array_3 = gs.zeros_like(gs.array(np_array_3)) np_array_3[0, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_3, 1, (0, 1)) self.assertAllCloseToNp(gs_result, np_array_3) np_array_4 = _np.zeros((3, 3, 2)) gs_array_4 = gs.zeros_like(gs.array(np_array_4)) np_array_4[0, :, 1] += 1 gs_result = gs.assignment_by_sum(gs_array_4, 1, (0, 1), axis=1) self.assertAllCloseToNp(gs_result, np_array_4) np_array_4_list = _np.zeros((3, 3, 2)) gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list)) np_array_4_list[(0, 1), :, (1, 1)] += 1 gs_result = gs.assignment_by_sum( gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1) self.assertAllCloseToNp(gs_result, np_array_4_list) n_samples = 3 theta = _np.array([0.1, 0.2, 0.3, 0.4, 5.5]) phi = _np.array([0.11, 0.22, 0.33, 0.44, -.55]) np_array = _np.ones((n_samples, 5, 4)) gs_array = gs.array(np_array) gs_array = gs.assignment_by_sum( gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1) gs_array = gs.assignment_by_sum( gs_array, - gs.sin(theta) * gs.sin(phi), (0, 1), axis=1) np_array[0, :, 0] += _np.cos(theta) * _np.cos(phi) np_array[0, :, 1] -= _np.sin(theta) * _np.sin(phi) # TODO (ninamiolane): This test fails 15% of the time, # when gs and _np computations are in the reverse order. # We should investigate this. self.assertAllCloseToNp(gs_array, np_array) np_array = _np.array([ [22., 55.], [33., 88.], [77., 99.]]) gs_array = gs.array([ [22., 55.], [33., 88.], [77., 99.]]) np_mask = _np.array([[False, False], [False, True], [True, True]]) gs_mask = gs.array([[False, False], [False, True], [True, True]]) np_array[np_mask] += _np.zeros_like(np_array[np_mask]) np_array[~np_mask] += 4 * np_array[~np_mask] np_result = np_array values_mask = gs.zeros_like(gs_array[gs_mask]) gs_result = gs.assignment_by_sum( gs_array, values_mask, gs_mask) gs_result = gs.assignment_by_sum( gs_result, 4 * gs_array[~gs_mask], ~gs_mask) self.assertAllCloseToNp(gs_result, np_result) def test_any(self): base_list = [ [[22., 55.], [33., 88.], [77., 99.]], [[34., 12.], [2., -3.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.any(np_array > 30.) gs_result = gs.any(gs_array > 30.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=0) gs_result = gs.any(gs_array > 30., axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=-2) gs_result = gs.any(gs_array > 30., axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.any(np_array > 30., axis=(-2, -1)) gs_result = gs.any(gs_array > 30., axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_all(self): base_list = [ [[22., 55.], [33., 88.], [77., 99.]], [[34., 12.], [2., -3.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.all(np_array > 30.) gs_result = gs.all(gs_array > 30.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=0) gs_result = gs.all(gs_array > 30., axis=0) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=-2) gs_result = gs.all(gs_array > 30., axis=-2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.all(np_array > 30., axis=(-2, -1)) gs_result = gs.all(gs_array > 30., axis=(-2, -1)) self.assertAllCloseToNp(gs_result, np_result) def test_trace(self): base_list = [ [[22., 55.], [33., 88.]], [[34., 12.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.trace(np_array) gs_result = gs.trace(gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=1, axis2=2) gs_result = gs.trace(gs_array, axis1=1, axis2=2) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.trace(np_array, axis1=-1, axis2=-2) gs_result = gs.trace(gs_array, axis1=-1, axis2=-2) self.assertAllCloseToNp(gs_result, np_result) def test_isclose(self): base_list = [ [[22. + 1e-5, 22. + 1e-7], [22. + 1e-6, 88. + 1e-4]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.isclose(np_array, 22.) gs_result = gs.isclose(gs_array, 22.) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22., atol=1e-8) gs_result = gs.isclose(gs_array, 22., atol=1e-8) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.isclose(np_array, 22., rtol=1e-8, atol=1e-7) gs_result = gs.isclose(gs_array, 22., rtol=1e-8, atol=1e-7) self.assertAllCloseToNp(gs_result, np_result) @geomstats.tests.np_and_pytorch_only def test_where(self): # TODO (ninamiolane): Make tf behavior consistent with np # Currently, tf returns array, while np returns tuple base_list = [ [[22., 55.], [33., 88.]], [[34., 12.], [67., 35.]]] np_array = _np.array(base_list) gs_array = gs.array(base_list) np_result = _np.where(np_array > 20., 0., np_array) gs_result = gs.where(gs_array > 20., 0., gs_array) self.assertAllCloseToNp(gs_result, np_result) np_result = _np.where(np_array > 20, np_array**2, 4.) gs_result = gs.where(gs_array > 20, gs_array**2, 4.) self.assertAllCloseToNp(gs_result, np_result) base_list = [[0, 1, 0, 1, 0, 1, 0, 1, 0, 1]] np_array = _np.array(base_list) gs_array = gs.array(base_list) result = gs.where(gs_array == 0) expected = _np.where(np_array == 0) self.assertAllCloseToNp(*result, *expected) result = gs.where(gs_array == 0, - 1, gs_array) expected = _np.where(np_array == 0, - 1, np_array) self.assertAllCloseToNp(result, expected) expected = _np.where(np_array == 1, _np.ones(10), np_array) result = gs.where(gs_array == 1, gs.ones(10), gs_array) self.assertAllCloseToNp(result, expected)
class TestVisualization(geomstats.tests.TestCase): def setUp(self): self.n_samples = 10 self.SO3_GROUP = SpecialOrthogonal(n=3, point_type='vector') self.SE3_GROUP = SpecialEuclidean(n=3, point_type='vector') self.S1 = Hypersphere(dim=1) self.S2 = Hypersphere(dim=2) self.H2 = Hyperbolic(dim=2) self.H2_half_plane = PoincareHalfSpace(dim=2) self.M32 = Matrices(m=3, n=2) self.S32 = PreShapeSpace(k_landmarks=3, m_ambient=2) self.KS = visualization.KendallSphere() self.M33 = Matrices(m=3, n=3) self.S33 = PreShapeSpace(k_landmarks=3, m_ambient=3) self.KD = visualization.KendallDisk() plt.figure() @staticmethod def test_tutorial_matplotlib(): visualization.tutorial_matplotlib() def test_plot_points_so3(self): points = self.SO3_GROUP.random_uniform(self.n_samples) visualization.plot(points, space='SO3_GROUP') def test_plot_points_se3(self): points = self.SE3_GROUP.random_point(self.n_samples) visualization.plot(points, space='SE3_GROUP') def test_draw_pre_shape_2d(self): self.KS.draw() def test_draw_points_pre_shape_2d(self): points = self.S32.random_point(self.n_samples) visualization.plot(points, space='S32') points = self.M32.random_point(self.n_samples) visualization.plot(points, space='M32') self.KS.clear_points() def test_draw_curve_pre_shape_2d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) times = gs.linspace(0., 1., 1000) speeds = gs.array([-t * tangent_vec for t in times]) points = self.S32.ambient_metric.exp(speeds, base_point) self.KS.add_points(points) self.KS.draw_curve() self.KS.clear_points() def test_draw_vector_pre_shape_2d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) self.KS.draw_vector(tangent_vec, base_point) def test_convert_to_spherical_coordinates_pre_shape_2d(self): points = self.S32.random_point(self.n_samples) coords = self.KS.convert_to_spherical_coordinates(points) x = coords[:, 0] y = coords[:, 1] z = coords[:, 2] result = x**2 + y**2 + z**2 expected = .25 * gs.ones(self.n_samples) self.assertAllClose(result, expected) def test_rotation_pre_shape_2d(self): theta = gs.random.rand(1)[0] phi = gs.random.rand(1)[0] rot = self.KS.rotation(theta, phi) result = _SpecialOrthogonalMatrices(3).belongs(rot) expected = True self.assertAllClose(result, expected) def test_draw_pre_shape_3d(self): self.KD.draw() def test_draw_points_pre_shape_3d(self): points = self.S33.random_point(self.n_samples) visualization.plot(points, space='S33') points = self.M33.random_point(self.n_samples) visualization.plot(points, space='M33') self.KD.clear_points() def test_draw_curve_pre_shape_3d(self): self.KD.draw() base_point = self.S33.random_point() vec = self.S33.random_point() tangent_vec = self.S33.to_tangent(vec, base_point) tangent_vec = .5 * tangent_vec / self.S33.ambient_metric.norm( tangent_vec) times = gs.linspace(0., 1., 1000) speeds = gs.array([-t * tangent_vec for t in times]) points = self.S33.ambient_metric.exp(speeds, base_point) self.KD.add_points(points) self.KD.draw_curve() self.KD.clear_points() def test_draw_vector_pre_shape_3d(self): self.KS.draw() base_point = self.S32.random_point() vec = self.S32.random_point() tangent_vec = self.S32.to_tangent(vec, base_point) self.KS.draw_vector(tangent_vec, base_point) def test_convert_to_planar_coordinates_pre_shape_3d(self): points = self.S33.random_point(self.n_samples) coords = self.KD.convert_to_planar_coordinates(points) x = coords[:, 0] y = coords[:, 1] radius = x**2 + y**2 result = [r <= 1. for r in radius] self.assertTrue(gs.all(result)) @geomstats.tests.np_and_pytorch_only def test_plot_points_s1(self): points = self.S1.random_uniform(self.n_samples) visualization.plot(points, space='S1') def test_plot_points_s2(self): points = self.S2.random_uniform(self.n_samples) visualization.plot(points, space='S2') def test_plot_points_h2_poincare_disk(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space='H2_poincare_disk') def test_plot_points_h2_poincare_half_plane_ext(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane', point_type='extrinsic') def test_plot_points_h2_poincare_half_plane_none(self): points = self.H2_half_plane.random_point(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane') def test_plot_points_h2_poincare_half_plane_hs(self): points = self.H2_half_plane.random_point(self.n_samples) visualization.plot(points, space='H2_poincare_half_plane', point_type='half_space') def test_plot_points_h2_klein_disk(self): points = self.H2.random_point(self.n_samples) visualization.plot(points, space='H2_klein_disk') @staticmethod def test_plot_points_se2(): points = SpecialEuclidean(n=2, point_type='vector').random_point(4) visu = visualization.SpecialEuclidean2(points, point_type='vector') ax = visu.set_ax() visu.draw(ax)
class SpecialEuclidean(LieGroup): """Class for the special euclidean group SE(n). i.e. the Lie group of rigid transformations. Elements of SE(n) can either be represented as vectors (in 3d) or as matrices in general. The matrix representation corresponds to homogeneous coordinates. """ def __init__(self, n, default_point_type=None, epsilon=0.): """Initiate an object of class SpecialEuclidean. Parameter --------- n : int the dimension of the euclidean space that SE(n) acts upon point_type : str, {'vector', 'matrix'}, optional whether to represent elmenents of SE(n) by vectors or matrices if None is given, point_type is set to 'vector' for dimension 3 and 'matrix' otherwise epsilon : float, optional precision to use for calculations involving potential division by rotations default: 0 """ if not (isinstance(n, int) and n > 1): raise ValueError('n must be an integer > 1.') self.n = n self.dimension = int((n * (n - 1)) / 2 + n) self.epsilon = epsilon super(SpecialEuclidean, self).__init__(dim=self.dimension, default_point_type=default_point_type) if default_point_type is None: self.default_point_type = 'vector' if n == 3 else 'matrix' self.rotations = SpecialOrthogonal( n=n, epsilon=epsilon, default_point_type=default_point_type) self.translations = Euclidean(dim=n) def get_identity(self, point_type=None): """Get the identity of the group. Parameters ---------- point_type : str, {'vector', 'matrix'}, optional the point_type of the returned value default: self.default_point_type Returns ------- identity : array-like, shape={[dim], [n + 1, n + 1]} """ if point_type is None: point_type = self.default_point_type identity = gs.zeros(self.dimension) if point_type == 'matrix': identity = gs.eye(self.n + 1) return identity identity = property(get_identity) def get_point_type_shape(self, point_type=None): """Get the shape of the instance given the default_point_style.""" return self.get_identity(point_type).shape @geomstats.vectorization.decorator(['else', 'point', 'point_type']) def belongs(self, point, point_type=None): """Evaluate if a point belongs to SE(n). Parameters ---------- point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the point of which to check whether it belongs to SE(n) point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- belongs : array-like, shape=[n_samples, 1] array of booleans indicating whether point belongs to SE(n) """ if point_type == 'vector': n_points, vec_dim = gs.shape(point) belongs = vec_dim == self.dimension belongs = gs.tile([belongs], (point.shape[0], )) belongs = gs.logical_and(belongs, self.rotations.belongs(point[:, :self.n])) return gs.flatten(belongs) if point_type == 'matrix': n_points, point_dim1, point_dim2 = point.shape belongs = (point_dim1 == point_dim2 == self.n + 1) belongs = [belongs] * n_points rotation = point[:, :self.n, :self.n] rot_belongs = self.rotations.belongs(rotation, point_type=point_type) belongs = gs.logical_and(belongs, rot_belongs) last_line_except_last_term = point[:, self.n:, :-1] all_but_last_zeros = ~gs.any(last_line_except_last_term, axis=(1, 2)) belongs = gs.logical_and(belongs, all_but_last_zeros) last_term = point[:, self.n:, self.n:] belongs = gs.logical_and(belongs, gs.all(last_term == 1, axis=(1, 2))) return gs.flatten(belongs) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'point_type']) def regularize(self, point, point_type=None): """Regularize a point to the default representation for SE(n). Parameters ---------- point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the point which should be regularized point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] """ if point_type == 'vector': rotations = self.rotations dim_rotations = rotations.dim rot_vec = point[:, :dim_rotations] regularized_rot_vec = rotations.regularize(rot_vec, point_type=point_type) translation = point[:, dim_rotations:] return gs.concatenate([regularized_rot_vec, translation], axis=1) if point_type == 'matrix': return gs.to_ndarray(point, to_ndim=3) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'else', 'point_type']) def regularize_tangent_vec_at_identity(self, tangent_vec, metric=None, point_type=None): """Regularize a tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] metric : RiemannianMetric, optional point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- regularized_vec : the regularized tangent vector """ if point_type == 'vector': return self.regularize_tangent_vec(tangent_vec, self.identity, metric, point_type=point_type) if point_type == 'matrix': translation_mask = gs.hstack( [gs.ones((self.n, ) * 2), 2 * gs.ones((self.n, 1))]) translation_mask = gs.concatenate( [translation_mask, gs.zeros((1, self.n + 1))], axis=0) tangent_vec = tangent_vec * gs.where(translation_mask != 0., gs.array(1.), gs.array(0.)) tangent_vec = (tangent_vec - GeneralLinear.transpose(tangent_vec)) / 2. return tangent_vec * translation_mask raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator( ['else', 'point', 'point', 'else', 'point_type']) def regularize_tangent_vec(self, tangent_vec, base_point, metric=None, point_type=None): """Regularize a tangent vector at a base point. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] base_point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] metric : RiemannianMetric, optional default: self.left_canonical_metric point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- regularized_vec : the regularized tangent vector """ if metric is None: metric = self.left_canonical_metric if point_type == 'vector': rotations = self.rotations dim_rotations = rotations.dim rot_tangent_vec = tangent_vec[:, :dim_rotations] rot_base_point = base_point[:, :dim_rotations] metric_mat = metric.inner_product_mat_at_identity rot_metric_mat = metric_mat[:dim_rotations, :dim_rotations] rot_metric = InvariantMetric( group=rotations, inner_product_mat_at_identity=rot_metric_mat, left_or_right=metric.left_or_right) rotations_vec = rotations.regularize_tangent_vec( tangent_vec=rot_tangent_vec, base_point=rot_base_point, metric=rot_metric, point_type=point_type) return gs.concatenate( [rotations_vec, tangent_vec[:, dim_rotations:]], axis=1) if point_type == 'matrix': tangent_vec_at_id = self.compose(self.inverse(base_point), tangent_vec) regularized = self.regularize_tangent_vec_at_identity( tangent_vec_at_id, point_type=point_type) return self.compose(base_point, regularized) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'vector']) def matrix_from_vector(self, vec): """Convert point in vector point-type to matrix. Parameters ---------- vec: array-like, shape=[n_samples, dim] Returns ------- mat: array-like, shape=[n_samples, {dim, [n+1, n+1]}] """ vec = self.regularize(vec, point_type='vector') n_vecs, _ = vec.shape rot_vec = vec[:, :self.rotations.dim] trans_vec = vec[:, self.rotations.dim:] rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) trans_vec = gs.reshape(trans_vec, (n_vecs, self.n, 1)) mat = gs.concatenate((rot_mat, trans_vec), axis=2) last_lines = gs.array(gs.get_mask_i_float(self.n, self.n + 1)) last_lines = gs.to_ndarray(last_lines, to_ndim=2) last_lines = gs.to_ndarray(last_lines, to_ndim=3) mat = gs.concatenate((mat, last_lines), axis=1) return mat @geomstats.vectorization.decorator( ['else', 'point', 'point', 'point_type']) def compose(self, point_a, point_b, point_type=None): r"""Compose two elements of SE(n). Parameters ---------- point_1 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_2 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Equation --------- (:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`) Returns ------- composition : the composition of point_1 and point_2 """ rotations = self.rotations dim_rotations = rotations.dim point_a = self.regularize(point_a, point_type=point_type) point_b = self.regularize(point_b, point_type=point_type) if point_type == 'vector': n_points_a, _ = point_a.shape n_points_b, _ = point_b.shape if not (point_a.shape == point_b.shape or n_points_a == 1 or n_points_b == 1): raise ValueError() rot_vec_a = point_a[:, :dim_rotations] rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a) rot_vec_b = point_b[:, :dim_rotations] rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b) translation_a = point_a[:, dim_rotations:] translation_b = point_b[:, dim_rotations:] composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b) composition_rot_vec = rotations.rotation_vector_from_matrix( composition_rot_mat) composition_translation = gs.einsum( '...j,...kj->...k', translation_b, rot_mat_a) + translation_a composition = gs.concatenate( (composition_rot_vec, composition_translation), axis=-1) return self.regularize(composition, point_type=point_type) if point_type == 'matrix': return GeneralLinear.compose(point_a, point_b) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'point_type']) def inverse(self, point, point_type=None): r"""Compute the group inverse in SE(n). Parameters ---------- point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- inverse_point : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the inverted point Notes ----- :math:`(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))` """ rotations = self.rotations dim_rotations = rotations.dim point = self.regularize(point) if point_type == 'vector': rot_vec = point[:, :dim_rotations] translation = point[:, dim_rotations:] inverse_rotation = -rot_vec inv_rot_mat = rotations.matrix_from_rotation_vector( inverse_rotation) inverse_translation = gs.einsum( 'ni,nij->nj', -translation, gs.transpose(inv_rot_mat, axes=(0, 2, 1))) inverse_point = gs.concatenate( [inverse_rotation, inverse_translation], axis=-1) return self.regularize(inverse_point, point_type=point_type) if point_type == 'matrix': inv_rot = gs.transpose(point[:, :self.n, :self.n], axes=(0, 2, 1)) inv_trans = gs.matmul(inv_rot, -point[:, :self.n, self.n:]) last_line = point[:, self.n:, :] inverse_point = gs.concatenate((inv_rot, inv_trans), axis=2) return gs.concatenate((inverse_point, last_line), axis=1) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'else', 'point_type']) def jacobian_translation(self, point, left_or_right='left', point_type=None): """Compute the Jacobian matrix resulting from translation. Compute the matrix of the differential of the left/right translations from the identity to point in SE(n). Parameters ---------- point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] left_or_right: str, {'left', 'right'}, optional default: 'left' whether to compute the jacobian of the left or right translation point_type : str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- jacobian : array-like, shape=[n_samples, dim] The jacobian of the left / right translation """ if point_type is None: point_type = self.default_point_type if left_or_right not in ('left', 'right'): raise ValueError('`left_or_right` must be `left` or `right`.') rotations = self.rotations translations = self.translations dim_rotations = rotations.dim dim_translations = translations.dim point = self.regularize(point, point_type=point_type) if point_type == 'vector': n_points, _ = point.shape rot_vec = point[:, :dim_rotations] jacobian_rot = self.rotations.jacobian_translation( point=rot_vec, left_or_right=left_or_right, point_type=point_type) block_zeros_1 = gs.zeros( (n_points, dim_rotations, dim_translations)) jacobian_block_line_1 = gs.concatenate( [jacobian_rot, block_zeros_1], axis=2) if left_or_right == 'left': rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec) jacobian_trans = rot_mat block_zeros_2 = gs.zeros( (n_points, dim_translations, dim_rotations)) jacobian_block_line_2 = gs.concatenate( [block_zeros_2, jacobian_trans], axis=2) else: inv_skew_mat = -self.rotations.skew_matrix_from_vector(rot_vec) eye = gs.to_ndarray(gs.eye(self.n), to_ndim=3) eye = gs.tile(eye, [n_points, 1, 1]) jacobian_block_line_2 = gs.concatenate([inv_skew_mat, eye], axis=2) return gs.concatenate( [jacobian_block_line_1, jacobian_block_line_2], axis=1) if point_type == 'matrix': return point raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'point_type']) def exp_from_identity(self, tangent_vec, point_type=None): """Compute group exponential of the tangent vector at the identity. Parameters ---------- tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- group_exp: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the group exponential of the tangent vectors calculated at the identity """ if point_type == 'vector': rotations = self.rotations dim_rotations = rotations.dim rot_vec = tangent_vec[:, :dim_rotations] rot_vec = self.rotations.regularize(rot_vec, point_type=point_type) translation = tangent_vec[:, dim_rotations:] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_mat = self.rotations.skew_matrix_from_vector(rot_vec) sq_skew_mat = gs.matmul(skew_mat, skew_mat) mask_0 = gs.equal(angle, 0.) mask_close_0 = gs.isclose(angle, 0.) & ~mask_0 mask_else = ~mask_0 & ~mask_close_0 mask_0_float = gs.cast(mask_0, gs.float32) mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = gs.zeros_like(angle) coef_2 = gs.zeros_like(angle) coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle) coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle) coef_1 += mask_close_0_float * ( TAYLOR_COEFFS_1_AT_0[0] + TAYLOR_COEFFS_1_AT_0[2] * angle**2 + TAYLOR_COEFFS_1_AT_0[4] * angle**4 + TAYLOR_COEFFS_1_AT_0[6] * angle**6) coef_2 += mask_close_0_float * ( TAYLOR_COEFFS_2_AT_0[0] + TAYLOR_COEFFS_2_AT_0[2] * angle**2 + TAYLOR_COEFFS_2_AT_0[4] * angle**4 + TAYLOR_COEFFS_2_AT_0[6] * angle**6) coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2) coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3) n_tangent_vecs, _ = tangent_vec.shape exp_translation = gs.zeros((n_tangent_vecs, self.n)) for i in range(n_tangent_vecs): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_mat[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_mat[i])) mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs) exp_translation += gs.outer( mask_i_float, translation_i + term_1_i + term_2_i) group_exp = gs.concatenate([rot_vec, exp_translation], axis=1) group_exp = self.regularize(group_exp, point_type=point_type) return group_exp if point_type == 'matrix': return GeneralLinear.exp(tangent_vec) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') @geomstats.vectorization.decorator(['else', 'point', 'point_type']) def log_from_identity(self, point, point_type=None): """Compute the group logarithm of the point at the identity. Parameters ---------- point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- group_log: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] the group logarithm in the Lie algbra """ point = self.regularize(point, point_type=point_type) rotations = self.rotations dim_rotations = rotations.dim if point_type == 'vector': rot_vec = point[:, :dim_rotations] angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) translation = point[:, dim_rotations:] skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec) sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec) mask_close_0 = gs.isclose(angle, 0.) mask_close_pi = gs.isclose(angle, gs.pi) mask_else = ~mask_close_0 & ~mask_close_pi mask_close_0_float = gs.cast(mask_close_0, gs.float32) mask_close_pi_float = gs.cast(mask_close_pi, gs.float32) mask_else_float = gs.cast(mask_else, gs.float32) mask_0 = gs.isclose(angle, 0., atol=1e-6) mask_0_float = gs.cast(mask_0, gs.float32) angle += mask_0_float * gs.ones_like(angle) coef_1 = -0.5 * gs.ones_like(angle) coef_2 = gs.zeros_like(angle) coef_2 += mask_close_0_float * (1. / 12. + angle**2 / 720. + angle**4 / 30240. + angle**6 / 1209600.) delta_angle = angle - gs.pi coef_2 += mask_close_pi_float * ( 1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) - ((PI2 - 12.) * delta_angle**2 / (4. * PI4)) + ((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) - ((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) + ((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 / (480. * PI7)) - ((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 / (480. * PI8))) psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle)) coef_2 += mask_else_float * (1 - psi) / (angle**2) n_points, _ = point.shape log_translation = gs.zeros((n_points, self.n)) for i in range(n_points): translation_i = translation[i] term_1_i = coef_1[i] * gs.dot(translation_i, gs.transpose(skew_rot_vec[i])) term_2_i = coef_2[i] * gs.dot(translation_i, gs.transpose(sq_skew_rot_vec[i])) mask_i_float = gs.get_mask_i_float(i, n_points) log_translation += gs.outer( mask_i_float, translation_i + term_1_i + term_2_i) return gs.concatenate([rot_vec, log_translation], axis=1) if point_type == 'matrix': return GeneralLinear.log(point) raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') def random_uniform(self, n_samples=1, point_type=None): """Sample in SE(n) with the uniform distribution. Parameters ---------- n_samples: int, optional default: 1 point_type: str, {'vector', 'matrix'}, optional default: self.default_point_type Returns ------- random_point: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}] An array of random elements in SE(n) having the given point_type. """ if point_type is None: point_type = self.default_point_type random_translation = self.translations.random_uniform(n_samples) if point_type == 'vector': random_rot_vec = self.rotations.random_uniform( n_samples, point_type=point_type) return gs.concatenate([random_rot_vec, random_translation], axis=-1) if point_type == 'matrix': random_rotation = self.rotations.random_uniform( n_samples, point_type=point_type) random_rotation = gs.to_ndarray(random_rotation, to_ndim=3) random_translation = gs.to_ndarray(random_translation, to_ndim=2) random_translation = gs.transpose( gs.to_ndarray(random_translation, to_ndim=3, axis=1), (0, 2, 1)) random_point = gs.concatenate( (random_rotation, random_translation), axis=2) last_line = gs.zeros((n_samples, 1, self.n + 1)) random_point = gs.concatenate((random_point, last_line), axis=1) random_point = gs.assignment(random_point, 1, (-1, -1), axis=0) if gs.shape(random_point)[0] == 1: random_point = gs.squeeze(random_point, axis=0) return random_point raise ValueError('Invalid point_type, expected \'vector\' or ' '\'matrix\'.') def _exponential_matrix(self, rot_vec): """Compute exponential of rotation matrix represented by rot_vec. Parameters ---------- rot_vec : array-like, shape=[n_samples, dim] Returns ------- exponential_mat : The matrix exponential of rot_vec """ # TODO(nguigs): find usecase for this method rot_vec = self.rotations.regularize(rot_vec) n_rot_vecs, _ = rot_vec.shape angle = gs.linalg.norm(rot_vec, axis=1) angle = gs.to_ndarray(angle, to_ndim=2, axis=1) skew_rot_vec = self.rotations.skew_matrix_from_vector(rot_vec) coef_1 = gs.empty_like(angle) coef_2 = gs.empty_like(coef_1) mask_0 = gs.equal(angle, 0) mask_0 = gs.squeeze(mask_0, axis=1) mask_close_to_0 = gs.isclose(angle, 0) mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1) mask_else = ~mask_0 & ~mask_close_to_0 coef_1[mask_close_to_0] = (1. / 2. - angle[mask_close_to_0]**2 / 24.) coef_2[mask_close_to_0] = (1. / 6. - angle[mask_close_to_0]**3 / 120.) # TODO(nina): Check if the discontinuity at 0 is expected. coef_1[mask_0] = 0 coef_2[mask_0] = 0 coef_1[mask_else] = (angle[mask_else]**(-2) * (1. - gs.cos(angle[mask_else]))) coef_2[mask_else] = (angle[mask_else]**(-2) * (1. - (gs.sin(angle[mask_else]) / angle[mask_else]))) term_1 = gs.zeros((n_rot_vecs, self.n, self.n)) term_2 = gs.zeros_like(term_1) for i in range(n_rot_vecs): term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i] term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i] exponential_mat = term_1 + term_2 return exponential_mat