fix feature 3d animation bugs

Author: yifanfeng97

dhg/visualization/feature/utils.py

@@ -6,11 +6,18 @@
 from sklearn import preprocessing
 from sklearn.manifold import TSNE
 from sklearn.decomposition import PCA
+import matplotlib.animation as animation
 
 eps = 1e-5
 min_norm = 1e-15
 
 
+def on_key_press(event):
+    print(event.key)
+    if event.key == "escape" or event.key == "q":
+        plt.close(event.canvas.figure)
+
+
 def make_animation(embeddings: np.ndarray, colors: Union[np.ndarray, str], cmap="viridis"):
     r"""Make an animation of embeddings.
 
@@ -22,18 +29,20 @@ def make_animation(embeddings: np.ndarray, colors: Union[np.ndarray, str], cmap=
     embeddings = normalize(embeddings)
     x, y, z = embeddings[:, 0], embeddings[:, 1], embeddings[:, 2]
     fig = plt.figure(figsize=(8, 8))
-    ax = fig.gca(projection="3d")
-    plt.ion()
-    for i in range(30000):
-        plt.clf()
-        fig = plt.gcf()
-        ax = fig.gca(projection="3d")
+    ax = fig.add_subplot(111, projection="3d")
+
+    def init():
         if colors is not None:
             ax.scatter(x, y, z, c=colors, cmap=cmap)
         else:
             ax.scatter(x, y, z, cmap=cmap)
+        return fig
+
+    def animate(i):
         ax.view_init(elev=20, azim=i % 360)
-        plt.pause(0.001)
+
+    ani = animation.FuncAnimation(fig, animate, init_func=init, frames=360, interval=20, blit=False)
+    return ani
 
 
 def plot_2d_embedding(embeddings: np.ndarray, label: Optional[np.ndarray] = None, cmap="viridis"):

dhg/visualization/feature/visualization.py

@@ -62,43 +62,53 @@ def draw_in_poincare_ball(
 def animation_of_3d_poincare_ball(
     embeddings: np.ndarray, label: Optional[np.ndarray] = None, reduce_method: str = "pca", cmap="viridis"
 ):
-    r"""Play animation of embeddings visualization on Poincare Ball.
+    r"""Make 3D animation of embeddings visualization on Poincare Ball. 
+    This function will return the animation object `ani <https://matplotlib.org/stable/api/_as_gen/matplotlib.animation.FuncAnimation.html>`_.
+    You can save the animation by ``ani.save("animation.gif")``.
 
     Args:
         ``feature`` (``np.ndarray``): The feature matrix. Size :math:`(N, C)`.
         ``label`` (``np.ndarray``, optional): The label matrix. Size :math:`(N, )`. Defaults to ``None``.
         ``reduce_method`` (``str``): The method to project the embedding into low-dimensional space. It can be ``pca`` or ``tsne``. Defaults to ``pca``.
         ``cmap`` (``str``, optional): The `color map <https://matplotlib.org/stable/tutorials/colors/colormaps.html>`_. Defaults to ``"viridis"``.
+    
+    Example:
+        >>> import numpy as np
+        >>> import matplotlib.pyplot as plt
+        >>> from dhg.visualization import animation_of_3d_poincare_ball
+        >>> x = np.random.rand(100, 32)
+        >>> ani = animation_of_3d_poincare_ball(x)
+        >>> plt.show()
+        >>> ani.save('a.gif')
     """
     emb_low = project_to_poincare_ball(embeddings, 3, reduce_method)
-    colors = label if label is not None else "b"
-    make_animation(emb_low, colors, cmap=cmap)
+    colors = label if label is not None else "r"
+    return make_animation(emb_low, colors, cmap=cmap)
 
 
 def animation_of_3d_euclidean_space(
     embeddings: np.ndarray, label: Optional[np.ndarray] = None, cmap="viridis",
 ):
-    r"""Play animation of embeddings visualization of tSNE algorithm.
+    r"""Make 3D animation of embeddings visualization of tSNE algorithm.
+    This function will return the animation object `ani <https://matplotlib.org/stable/api/_as_gen/matplotlib.animation.FuncAnimation.html>`_.
+    You can save the animation by ``ani.save("animation.gif")``.
 
     Args:
         ``feature`` (``np.ndarray``): The feature matrix. Size :math:`(N, C)`.
         ``label`` (``np.ndarray``, optional): The label matrix. Size :math:`(N, )`. Defaults to ``None``.
         ``cmap`` (``str``, optional): The `color map <https://matplotlib.org/stable/tutorials/colors/colormaps.html>`_. Defaults to ``"viridis"``.
+    
+    Example:
+        >>> import numpy as np
+        >>> import matplotlib.pyplot as plt
+        >>> from dhg.visualization import animation_of_3d_euclidean_space
+        >>> x = np.random.rand(100, 32)
+        >>> ani = animation_of_3d_euclidean_space(x)
+        >>> plt.show()
+        >>> ani.save('a.gif')
     """
     tsne = TSNE(n_components=3, init="pca")
     emb_low = tsne.fit_transform(embeddings)
-    colors = label if label is not None else "b"
-    make_animation(emb_low, colors, cmap=cmap)
+    colors = label if label is not None else "r"
+    return make_animation(emb_low, colors, cmap=cmap)
 
-
-if __name__ == "__main__":
-    file_dir = "data/modelnet40/train_img_feat_4.npy"
-    # save_dir = "./tmp/figure"
-    save_dir = None  # None for show now or file name to save
-    low_demen_method = "TSNE"  # vis for poincare_ball, PCA or TSNE
-    show_method = "Rotation"  # None for 2d or Rotation and Drag for 3d
-    label = np.load("data/modelnet40/train_label.npy")
-    ft = np.load(file_dir)
-    d = 3
-    # vis_tsne(ft, save_dir,d)
-    draw_in_poincare_ball(ft, save_dir, d, label, reduce_method=low_demen_method, auto_play=show_method)

docs/source/api/visualization.rst

@@ -45,8 +45,6 @@ Feature Visualization in Poincare Ball
 Make Animations
 ~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: dhg.visualization.make_animation
-
 .. autofunction:: dhg.visualization.animation_of_3d_euclidean_space
 
 .. autofunction:: dhg.visualization.animation_of_3d_poincare_ball