defer plot close until end

Author: JoelLucaAdams

src/sdf_xarray/plotting.py

@@ -366,8 +366,6 @@ def animate(
     # Create plot if no ax is provided
     if ax is None:
         fig, ax = plt.subplots()
-        # Prevents figure from prematurely displaying in Jupyter notebook
-        plt.close(fig)
 
     animation = _generate_animation(
         data,
@@ -382,14 +380,19 @@ def animate(
         kwargs=kwargs,
     )
 
-    return FuncAnimation(
-        ax.get_figure(),
+    anim = FuncAnimation(
+        fig,
         animation.update,
         frames=range(animation.n_frames),
         interval=1000 / fps,
         repeat=True,
     )
 
+    # Prevents figure from prematurely displaying in Jupyter notebook
+    plt.close(fig)
+
+    return anim
+
 
 def animate_multiple(
     *datasets: xr.DataArray,
@@ -458,8 +461,6 @@ def animate_multiple(
     # Create plot if no ax is provided
     if ax is None:
         fig, ax = plt.subplots()
-        # Prevents figure from prematurely displaying in Jupyter notebook
-        plt.close(fig)
 
     n_datasets = len(datasets)
     if datasets_kwargs is None:
@@ -508,14 +509,19 @@ def update(frame):
         if show_legend:
             ax.legend(loc="upper right")
 
-    return FuncAnimation(
-        ax.get_figure(),
+    anim = FuncAnimation(
+        fig,
         update,
         frames=range(n_frames),
         interval=1000 / fps,
         repeat=True,
     )
 
+    # Prevents figure from prematurely displaying in Jupyter notebook
+    plt.close(fig)
+
+    return anim
+
 
 def show(anim):
     """Shows the FuncAnimation in a Jupyter notebook.