Small ui cleanups

tabs/compare.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass
 from itertools import combinations
 from pathlib import Path
 
+import plotly.graph_objects as go
 import streamlit as st
 from persona_data.environment import get_artifacts_dir
 from persona_data.synth_persona import BASELINE_PERSONA_ID
@@ -44,6 +45,7 @@ from utils.helpers import (
     slugify,
     widget_key,
 )
+from utils.theme import style_plotly_layer_controls
 
 
 def _filename(*parts: str) -> str:
@@ -295,12 +297,25 @@ def _render_save_buttons(
     """Render the Save HTML button for one or more figures."""
     if st.button("Save HTML", key=widget_key("load", "save_html", key_suffix)):
         try:
+            _style_plotly_figures(figs)
             paths = [save_plot_html(fig, fn) for fig, fn in zip(figs, filenames)]
             st.success(f"Saved {len(paths)} HTML file(s) to `artifacts/plots`.")
         except Exception as exc:
             st.error(f"Could not save HTML: {exc}")
 
 
+def _style_plotly_figures(figs: list[object]) -> None:
+    base = st.get_option("theme.base")
+    for fig in figs:
+        if isinstance(fig, go.Figure):
+            style_plotly_layer_controls(fig, base)
+
+
+def _plotly_chart(fig: object) -> None:
+    _style_plotly_figures([fig])
+    st.plotly_chart(fig, width="stretch")
+
+
 def _render_mask_strategy_select(scope: str) -> MaskStrategy:
     return render_mask_strategy_select(
         key=widget_key("load", "mask_strategy", scope),
@@ -494,12 +509,12 @@ def _render_cosine_similarity(
 
     if cosine_fig_key in st.session_state:
         fig, pair_fig, n_traces, n_pair_traces = st.session_state[cosine_fig_key]
-        st.plotly_chart(fig, width="stretch")
+        _plotly_chart(fig)
         figs = [fig]
         filenames = [filename]
         if pair_fig is not None:
             st.subheader("Variant pairs")
-            st.plotly_chart(pair_fig, width="stretch")
+            _plotly_chart(pair_fig)
             figs.append(pair_fig)
             filenames.append(pairs_filename)
         _render_save_buttons(figs, filenames, "cosine")
@@ -662,12 +677,12 @@ def _render_layered_figure_analysis(
 
     if fig_key in st.session_state:
         main_fig, extra_fig, n_samples = st.session_state[fig_key]
-        st.plotly_chart(main_fig, width="stretch")
+        _plotly_chart(main_fig)
         figs = [main_fig]
         filenames = [filename]
         if extra_fig is not None:
             st.subheader("Pair trajectories")
-            st.plotly_chart(extra_fig, width="stretch")
+            _plotly_chart(extra_fig)
             figs.append(extra_fig)
             filenames.append(f"{filename}__pair_trajectories")
         _render_save_buttons(figs, filenames, scope)
@@ -816,12 +831,12 @@ def _render_dendrogram_analysis(
             col_a, col_b = st.columns(2)
             with col_a:
                 st.subheader(prompt_variant_label(va))
-                st.plotly_chart(fig_a, width="stretch")
+                _plotly_chart(fig_a)
             with col_b:
                 st.subheader(prompt_variant_label(vb))
-                st.plotly_chart(fig_b, width="stretch")
+                _plotly_chart(fig_b)
         else:
-            st.plotly_chart(fig_a, width="stretch")
+            _plotly_chart(fig_a)
 
         figs = [fig_a] + ([fig_b] if fig_b else [])
         filenames = [

utils/theme.py

@@ -5,6 +5,10 @@ import plotly.io as pio
 from catppuccin import PALETTE
 
 
+def _flavor(base: str | None):
+    return PALETTE.latte if base == "light" else PALETTE.mocha
+
+
 def install_catppuccin_theme(base: str | None = None) -> None:
     """Register a Catppuccin template and alias it as ``plotly_white``.
 
@@ -12,8 +16,7 @@ def install_catppuccin_theme(base: str | None = None) -> None:
     every figure, so replacing that entry themes all plots without any
     per-figure code.
     """
-    flavor = PALETTE.latte if base == "light" else PALETTE.mocha
-    c = flavor.colors
+    c = _flavor(base).colors
     bg, surface, line = c.base.hex, c.surface0.hex, c.surface1.hex
     text, subtext = c.text.hex, c.subtext1.hex
 
@@ -65,3 +68,33 @@ def install_catppuccin_theme(base: str | None = None) -> None:
     pio.templates["catppuccin"] = template
     pio.templates["plotly_white"] = template
     pio.templates.default = "catppuccin"
+
+
+def style_plotly_layer_controls(fig: go.Figure, base: str | None = None) -> go.Figure:
+    """Theme persona-vectors layer sliders/buttons for this Streamlit app."""
+    c = _flavor(base).colors
+    surface = c.surface0.hex
+    overlay = c.surface1.hex
+    text = c.text.hex
+    primary = c.blue.hex
+
+    for slider in fig.layout.sliders:
+        slider.bgcolor = surface
+        slider.activebgcolor = primary
+        slider.bordercolor = overlay
+        slider.borderwidth = 1
+        slider.font = dict(color=text, size=11)
+        slider.tickcolor = primary
+        slider.currentvalue = dict(
+            slider.currentvalue.to_plotly_json(),
+            font=dict(color=text, size=13),
+        )
+
+    for menu in fig.layout.updatemenus:
+        if menu.type != "buttons":
+            continue
+        menu.bgcolor = surface
+        menu.bordercolor = overlay
+        menu.font = dict(color=text, size=13)
+
+    return fig