From 68a78fe64ffd2bcc4c1da2d6941ee418b5e4e96d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20K=C3=B6hler?=
 <27728103+Ceyron@users.noreply.github.com>
Date: Tue, 9 Apr 2024 07:02:54 +0200
Subject: [PATCH] Fix faceted state 1d animation

---
 exponax/viz/_animate.py | 34 +++++++++++++++++++++++++++++-----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/exponax/viz/_animate.py b/exponax/viz/_animate.py
index 62ae9b4..948d798 100644
--- a/exponax/viz/_animate.py
+++ b/exponax/viz/_animate.py
@@ -72,19 +72,43 @@ def animate_state_1d_facet(
 
     fig, ax_s = plt.subplots(*grid, figsize=figsize)
 
-    for i, ax in enumerate(ax_s.flatten()):
+    num_subplots = trj.shape[0]
+
+    for j, ax in enumerate(ax_s.flatten()):
         plot_state_1d(
-            trj[i],
+            trj[j, 0],
             vlim=vlim,
             domain_extent=domain_extent,
             labels=labels,
             ax=ax,
             **kwargs,
         )
-        if titles is not None:
-            ax.set_title(titles[i])
+        if j >= num_subplots:
+            ax.remove()
+        else:
+            if titles is not None:
+                ax.set_title(titles[j])
+
+    def animate(i):
+        for j, ax in enumerate(ax_s.flatten()):
+            ax.clear()
+            plot_state_1d(
+                trj[j, i],
+                vlim=vlim,
+                domain_extent=domain_extent,
+                labels=labels,
+                ax=ax,
+                **kwargs,
+            )
+            if j >= num_subplots:
+                ax.remove()
+            else:
+                if titles is not None:
+                    ax.set_title(titles[j])
+
+    ani = FuncAnimation(fig, animate, frames=trj.shape[1], interval=100, blit=False)
 
-    return fig
+    return ani
 
 
 def animate_spatio_temporal(