From 8d3d26236040567253c80bd5fd0ae2cee1b5ad56 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20K=C3=B6hler?=
 <27728103+Ceyron@users.noreply.github.com>
Date: Thu, 31 Oct 2024 12:39:53 +0100
Subject: [PATCH] Fix Streamlit (#28)

* Fix 1D spatio-temporal

* Fix distance_scale default

* Remove no longer needed color map hack

* Vmap periodic extension before using volume renderer
---
 explore_sample_data_streamlit.py | 40 ++++++++------------------------
 1 file changed, 10 insertions(+), 30 deletions(-)

diff --git a/explore_sample_data_streamlit.py b/explore_sample_data_streamlit.py
index 795871f..ab4a5db 100644
--- a/explore_sample_data_streamlit.py
+++ b/explore_sample_data_streamlit.py
@@ -13,12 +13,12 @@
 from dataclasses import dataclass
 from typing import Optional
 
+import jax
 import matplotlib.pyplot as plt
 import numpy as np
 import streamlit as st
 import streamlit.components.v1 as components
 from IPython.display import DisplayObject
-from matplotlib.colors import Colormap, LinearSegmentedColormap, ListedColormap
 
 import apebench
 
@@ -88,7 +88,7 @@
 
 TEMPLATE_IFRAME = """
     <div>
-        <iframe id="{canvas_id}" src="https://keksboter.github.io/v4dv/?inline" width="{canvas_width}" height="{canvas_height}" frameBorder="0" sandbox="allow-same-origin allow-scripts"></iframe>
+        <iframe id="{canvas_id}" src="https://vape.niedermayr.dev/?inline" width="{canvas_width}" height="{canvas_height}" frameBorder="0" sandbox="allow-same-origin allow-scripts"></iframe>
     </div>
     <script>
 
@@ -124,6 +124,7 @@ class ViewerSettings:
     show_volume_info: bool
     vmin: Optional[float]
     vmax: Optional[float]
+    distance_scale: float
 
 
 def show(
@@ -136,6 +137,7 @@ def show(
     show_volume_info=False,
     vmin=None,
     vmax=None,
+    distance_scale=1.0,
 ):
     return VolumeRenderer(
         data,
@@ -148,6 +150,7 @@ def show(
             show_volume_info,
             vmin,
             vmax,
+            distance_scale,
         ),
     )
 
@@ -191,33 +194,8 @@ def __html__(self):
         return self._repr_html_()
 
 
-def felix_cmap_hack(cmap: Colormap) -> Colormap:
-    """changes the alpha channel of a colormap to be diverging (0->1, 0.5 > 0, 1->1)
-
-    Args:
-        cmap (Colormap): colormap
-
-    Returns:
-        Colormap: new colormap
-    """
-    cmap = cmap.copy()
-    if isinstance(cmap, ListedColormap):
-        for i, a in enumerate(cmap.colors):
-            a.append(2 * abs(i / cmap.N - 0.5))
-    elif isinstance(cmap, LinearSegmentedColormap):
-        cmap._segmentdata["alpha"] = np.array(
-            [[0.0, 1.0, 1.0], [0.5, 0.0, 0.0], [1.0, 1.0, 1.0]]
-        )
-    else:
-        raise TypeError(
-            "cmap must be either a ListedColormap or a LinearSegmentedColormap"
-        )
-    return cmap
-
-
 if dimension_type == "1d ST":
-    apebench.exponax.viz.plot_spatio_temporal(trj, vlim=scenario.vlim)
-    fig = plt.gcf()
+    fig = apebench.exponax.viz.plot_spatio_temporal(trj, vlim=scenario.vlim)
     st.pyplot(fig)
 elif dimension_type == "1d":
     ani = apebench.exponax.viz.animate_state_1d(trj, vlim=scenario.vlim)
@@ -226,7 +204,8 @@ def felix_cmap_hack(cmap: Colormap) -> Colormap:
     ani = apebench.exponax.viz.animate_state_2d(trj, vlim=scenario.vlim)
     components.html(ani.to_jshtml(), height=800)
 elif dimension_type == "2d ST":
-    trj_rearranged = trj.transpose(1, 0, 2, 3)[None]
+    trj_wrapped = jax.vmap(apebench.exponax.wrap_bc)(trj)
+    trj_rearranged = trj_wrapped.transpose(1, 0, 2, 3)[None]
     components.html(
         show(
             trj_rearranged,
@@ -239,9 +218,10 @@ def felix_cmap_hack(cmap: Colormap) -> Colormap:
         height=800,
     )
 elif dimension_type == "3d":
+    trj_wrapped = jax.vmap(apebench.exponax.wrap_bc)(trj)
     components.html(
         show(
-            trj,
+            trj_wrapped,
             plt.get_cmap("RdBu_r"),
             width=1500,
             height=800,