Start writing tests for viz routines

tests/test_viz.py

@@ -0,0 +1,38 @@
+import jax
+import matplotlib.pyplot as plt
+
+import exponax as ex
+
+
+def test_plot_state_1d():
+    state = jax.random.normal(
+        jax.random.PRNGKey(0),
+        (
+            10,
+            100,
+        ),
+    )
+
+    fig = ex.viz.plot_state_1d(state)
+    plt.close(fig)
+
+
+def test_plot_spatio_temporal():
+    trj = jax.random.normal(jax.random.PRNGKey(0), (100, 1, 64))
+
+    fig = ex.viz.plot_spatio_temporal(trj)
+    plt.close(fig)
+
+
+def test_plot_state_2d():
+    state = jax.random.normal(jax.random.PRNGKey(0), (1, 100, 100))
+
+    fig = ex.viz.plot_state_2d(state)
+    plt.close(fig)
+
+
+def test_plot_state_3d():
+    state = jax.random.normal(jax.random.PRNGKey(0), (1, 32, 32, 32))
+
+    fig = ex.viz.plot_state_3d(state)
+    plt.close(fig)