From 0dfe4b42b59b640c330322ea956fef9bda08dffd Mon Sep 17 00:00:00 2001
From: Felix Koehler <f.koehler@tum.de>
Date: Mon, 8 Apr 2024 16:22:04 +0200
Subject: [PATCH] Correctly wrap boundary conditions

---
 exponax/viz/_plot.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/exponax/viz/_plot.py b/exponax/viz/_plot.py
index c86798a..f97c3af 100644
--- a/exponax/viz/_plot.py
+++ b/exponax/viz/_plot.py
@@ -249,11 +249,13 @@ def plot_state_2d(
         # One more because we wrapped the BC
         space_range = (0, state.shape[-1])
 
+    state_wrapped = wrap_bc(state)
+
     if ax is None:
         fig, ax = plt.subplots()
 
     im = ax.imshow(
-        state.T,
+        state_wrapped.T,
         vmin=vlim[0],
         vmax=vlim[1],
         cmap="RdBu_r",
@@ -264,6 +266,7 @@ def plot_state_2d(
     )
     ax.set_xlabel("x_0")
     ax.set_ylabel("x_1")
+    ax.set_aspect("equal")
 
     return im