geomstats · Benleeee · Jan 26, 2023 · Jan 29, 2023 · Jan 29, 2023 · Jan 30, 2023
diff --git a/geomstats/visualization/CategoricalDistributionsManifold.py b/geomstats/visualization/CategoricalDistributionsManifold.py
diff --git a/notebooks/CategoricalDistributionsManifold.ipynb b/notebooks/CategoricalDistributionsManifold.ipynb
diff --git a/notebooks/images/coins.png b/notebooks/images/coins.png
diff --git a/notebooks/images/color_naming.png b/notebooks/images/color_naming.png
diff --git a/notebooks/images/dice.jpeg b/notebooks/images/dice.jpeg
diff --git a/notebooks/images/dome.jpeg b/notebooks/images/dome.jpeg
diff --git a/notebooks/images/simplex.png b/notebooks/images/simplex.png
diff --git a/tests/tests_geomstats/test_visualization_manifold_of_categorical_distributions.py b/tests/tests_geomstats/test_visualization_manifold_of_categorical_distributions.py
@@ -0,0 +1,48 @@
+"""Unit tests for visualization."""
+
+import geomstats.visualization as visualization
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import tests.conftest
+
+from CategoricalDistributionsManifold import CategoricalDistributionsManifold
+
+matplotlib.use("Agg")  # NOQA
+
+
+class TestVisualizationManifoldOfCategoricalDistributions(
+        tests.conftest.TestCase):
+    def setup_method(self):
+        self.n_samples = 10
+        self.CD2 = CategoricalDistributionsManifold(dim=2)
+        self.CD3 = CategoricalDistributionsManifold(dim=3)
+        plt.figure()
+
+    @staticmethod
+    def test_tutorial_matplotlib():
+        visualization.tutorial_matplotlib()
+
+    def test_plot(self):
+        self.CD2.plot()
+        self.CD3.plot()
+
+    def test_scatter(self):
+        self.CD2.scatter(self.n_samples)
+        self.CD3.scatter(self.n_samples)
+
+    def test_plot_geodesic(self):
+        b2 = np.array([0.2, 0.3, 0.5])
+        v2 = np.array([1, 0, 0])
+        self.CD2.plot_geodesic(initial_point=b2, tangent_vector=v2)
+
+        b3 = np.array([0.3, 0.3, 0.2, 0.2])
+        v3 = np.array([1, 0, 0, 0])
+        self.CD3.plot_geodesic(initial_point=b3, tangent_vector=v3)
+
+    def test_plot_grid(self):
+        self.CD2.plot_grid()
+
+    @staticmethod
+    def teardown_method():
+        plt.close()