Refactoring cleanups

pyproject.toml

@@ -15,6 +15,14 @@ dependencies = [
     "safetensors>=0.7.0",
 ]
 
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
 # Local development:
 # [tool.uv.sources]
 # persona-vectors = { path = "../persona-vectors", editable = true }

tabs/analysis/_shared.py

@@ -0,0 +1,449 @@
+import gc
+
+import plotly.graph_objects as go
+import streamlit as st
+from persona_data.synth_persona import BASELINE_PERSONA_ID
+from persona_vectors.extraction import MaskStrategy
+from persona_vectors.plots import save_plot_html
+
+from utils.analysis_sources import (
+    Store,
+    available_variants,
+    load_persona_vectors_cached,
+    load_variant_vectors_cached,
+    persona_names_cached,
+    personas_cached,
+    release_hf_store_cache,
+    store_cache_parts,
+    store_id,
+    store_layers_cached,
+)
+from utils.controls import render_mask_strategy_select
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
+from utils.theme import active_base, style_plotly_layer_controls
+
+from tabs.analysis._state import (
+    _DEFAULT_LAYER_FRAMES,
+    _HIGHLIGHT_OTHER_COLOR,
+    _HIGHLIGHT_OTHER_LABEL,
+    _LAST_LAYER_FRAMES_KEY,
+    _LAST_MASK_STRATEGY_KEY,
+    PersonaOptions,
+    _is_assistant_persona,
+    _persona_names_state_key,
+    _personas_empty_message,
+    _remembered_selectbox,
+    _sequence_to_list,
+)
+
+
+def _gray_out_unselected_personas(fig: go.Figure) -> None:
+    def _gray_trace(trace: object) -> None:
+        marker = getattr(trace, "marker", None)
+        if marker is None:
+            return
+
+        colors = _sequence_to_list(getattr(marker, "color", None))
+        labels = _sequence_to_list(getattr(trace, "customdata", None))
+        if colors is not None and labels is not None and len(colors) == len(labels):
+            trace.marker.color = [
+                (
+                    _HIGHLIGHT_OTHER_COLOR
+                    if str(label) == _HIGHLIGHT_OTHER_LABEL
+                    else color
+                )
+                for label, color in zip(labels, colors, strict=True)
+            ]
+            return
+
+        if getattr(trace, "name", None) == _HIGHLIGHT_OTHER_LABEL:
+            trace.marker.color = _HIGHLIGHT_OTHER_COLOR
+            trace.opacity = 0.28
+
+    for trace in fig.data:
+        _gray_trace(trace)
+    for frame in fig.frames:
+        for trace in frame.data:
+            _gray_trace(trace)
+
+
+def _layers_for_variant(
+    store: Store,
+    variant: str,
+    persona_ids: list[str],
+    mask_strategy: MaskStrategy,
+) -> list[int]:
+    source, location, model_name = store_cache_parts(store)
+    return store_layers_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        (variant,),
+        tuple(persona_ids),
+    )
+
+
+def _load_persona_vectors(
+    store: Store,
+    variant: str,
+    mask_strategy: MaskStrategy,
+    persona_ids: list[str],
+):
+    source, location, model_name = store_cache_parts(store)
+    return load_persona_vectors_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        variant,
+        tuple(persona_ids),
+    )
+
+
+def _load_variant_vectors(
+    store: Store,
+    variants: list[str] | tuple[str, ...],
+    mask_strategy: MaskStrategy,
+    persona_ids: list[str],
+):
+    source, location, model_name = store_cache_parts(store)
+    return load_variant_vectors_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        tuple(variants),
+        tuple(persona_ids),
+    )
+
+
+def _release_vector_memory(store: Store, variants: list[str] | tuple[str, ...]) -> None:
+    release_hf_store_cache(store, variants)
+    gc.collect()
+
+
+def _evenly_spaced_layers(layers: list[int], max_count: int) -> list[int]:
+    if max_count >= len(layers):
+        return layers
+    if max_count <= 1:
+        return [layers[0]]
+
+    last = len(layers) - 1
+    indices = [round(i * last / (max_count - 1)) for i in range(max_count)]
+    return [layers[index] for index in dict.fromkeys(indices)]
+
+
+def _render_layer_frame_controls(
+    store: Store,
+    scope: str,
+    layers: list[int],
+) -> list[int]:
+    if len(layers) <= _DEFAULT_LAYER_FRAMES:
+        st.caption(f"Using all {len(layers)} available layer(s).")
+        return layers
+
+    frame_count = st.slider(
+        "Layer frames",
+        min_value=2,
+        max_value=len(layers),
+        value=min(
+            max(
+                int(
+                    st.session_state.get(
+                        _LAST_LAYER_FRAMES_KEY,
+                        _DEFAULT_LAYER_FRAMES,
+                    )
+                ),
+                2,
+            ),
+            len(layers),
+        ),
+        key=widget_key("load", "layer_frames", scope, store_id(store)),
+        help="Limit animated Plotly frames to keep browser and RAM usage bounded.",
+    )
+    st.session_state[_LAST_LAYER_FRAMES_KEY] = frame_count
+    selected = _evenly_spaced_layers(layers, frame_count)
+    st.caption(f"Using {len(selected)} of {len(layers)} layers.")
+    return selected
+
+
+def _load_persona_options(
+    store: Store,
+    variants: list[str],
+    mask_strategy: MaskStrategy,
+    *,
+    empty_message: str,
+) -> PersonaOptions | None:
+    source, location, model_name = store_cache_parts(store)
+    variant_key = tuple(variants)
+    persona_ids = personas_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        variant_key,
+        include_baseline=True,
+    )
+    if not persona_ids:
+        st.info(empty_message)
+        return None
+
+    persona_names = persona_names_cached(
+        source,
+        location,
+        model_name,
+        mask_strategy.value,
+        variant_key,
+        tuple(persona_ids),
+    )
+    assistant_ids = [
+        persona_id
+        for persona_id in persona_ids
+        if _is_assistant_persona(persona_id, persona_names.get(persona_id))
+    ]
+    assistant_id = next(
+        (
+            persona_id
+            for persona_id in assistant_ids
+            if persona_id == BASELINE_PERSONA_ID
+        ),
+        assistant_ids[0] if assistant_ids else None,
+    )
+    regular_ids = [
+        persona_id for persona_id in persona_ids if persona_id not in assistant_ids
+    ]
+    if not regular_ids and assistant_id is None:
+        st.info("No personas found for this model and variant.")
+        return None
+    return PersonaOptions(
+        regular_ids=regular_ids,
+        assistant_id=assistant_id,
+        persona_names=persona_names,
+    )
+
+
+def _seed_persona_memory(
+    remember_key: str,
+    options: PersonaOptions,
+    *,
+    default_all: bool,
+    default_count_limit: int | None = None,
+) -> tuple[int, bool]:
+    remembered_count_key = f"{remember_key}:count"
+    remembered_assistant_key = f"{remember_key}:include_assistant"
+    legacy_ids = st.session_state.get(remember_key, [])
+    if isinstance(legacy_ids, list) and legacy_ids:
+        st.session_state.setdefault(
+            remembered_count_key,
+            sum(persona_id in options.regular_ids for persona_id in legacy_ids),
+        )
+        st.session_state.setdefault(
+            remembered_assistant_key,
+            options.assistant_id in legacy_ids,
+        )
+
+    if default_count_limit is not None:
+        default_count = min(default_count_limit, len(options.regular_ids))
+    elif default_all:
+        default_count = len(options.regular_ids)
+    else:
+        default_count = min(1, len(options.regular_ids))
+    remembered_count = int(st.session_state.get(remembered_count_key, default_count))
+    persona_count = min(max(remembered_count, 0), len(options.regular_ids))
+    include_assistant = bool(st.session_state.get(remembered_assistant_key, False))
+    return persona_count, include_assistant
+
+
+def _render_persona_count_controls(
+    store: Store,
+    variants: list[str],
+    mask_strategy: MaskStrategy,
+    widget_scope: str,
+    options: PersonaOptions,
+    *,
+    default_count: int,
+    include_assistant_default: bool,
+) -> tuple[int, bool]:
+    count_key = widget_key(
+        "load",
+        "persona_count",
+        widget_scope,
+        store.model_name,
+        mask_strategy.value,
+        *variants,
+    )
+    assistant_key = widget_key(
+        "load",
+        "include_assistant",
+        widget_scope,
+        store.model_name,
+        mask_strategy.value,
+        *variants,
+    )
+
+    if options.regular_ids:
+        persona_count = st.slider(
+            "Personas",
+            min_value=0 if options.assistant_id is not None else 1,
+            max_value=len(options.regular_ids),
+            value=default_count,
+            key=count_key,
+            help="Use the first N available non-assistant personas.",
+        )
+    else:
+        persona_count = 0
+        st.caption("No non-assistant personas are available for this selection.")
+    include_assistant = False
+    if options.assistant_id is not None:
+        include_assistant = st.checkbox(
+            "Include Assistant persona",
+            value=include_assistant_default,
+            key=assistant_key,
+        )
+    return persona_count, include_assistant
+
+
+def _select_artifact_personas(
+    store: Store,
+    variants: list[str],
+    mask_strategy: MaskStrategy,
+    *,
+    widget_scope: str,
+    remember_key: str,
+    default_all: bool = False,
+    default_count_limit: int | None = None,
+) -> list[str]:
+    empty_message = _personas_empty_message(variants)
+    options = _load_persona_options(
+        store,
+        variants,
+        mask_strategy,
+        empty_message=empty_message,
+    )
+    if options is None:
+        st.session_state.pop(_persona_names_state_key(widget_scope), None)
+        return []
+
+    default_count, include_assistant_default = _seed_persona_memory(
+        remember_key,
+        options,
+        default_all=default_all,
+        default_count_limit=default_count_limit,
+    )
+    persona_count, include_assistant = _render_persona_count_controls(
+        store,
+        variants,
+        mask_strategy,
+        widget_scope,
+        options,
+        default_count=default_count,
+        include_assistant_default=include_assistant_default,
+    )
+
+    persona_ids = options.regular_ids[:persona_count]
+    if include_assistant and options.assistant_id is not None:
+        persona_ids.append(options.assistant_id)
+
+    remembered_count_key = f"{remember_key}:count"
+    remembered_assistant_key = f"{remember_key}:include_assistant"
+    st.session_state[remembered_count_key] = persona_count
+    st.session_state[remembered_assistant_key] = include_assistant
+    st.session_state[remember_key] = persona_ids
+    st.session_state[_persona_names_state_key(widget_scope)] = options.persona_names
+
+    if not persona_ids:
+        st.info("Select at least one persona or include the Assistant persona.")
+        return []
+
+    regular_label = f"{persona_count} persona{'s' if persona_count != 1 else ''}"
+    assistant_label = (
+        " plus Assistant" if include_assistant and options.assistant_id else ""
+    )
+    st.caption(f"Using {regular_label}{assistant_label}.")
+    return persona_ids
+
+
+def _render_save_buttons(
+    figs: list[object],
+    filenames: list[str],
+    key_suffix: str,
+) -> None:
+    """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, strict=True)
+            ]
+            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 = active_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",
+        config={"responsive": True, "displaylogo": False},
+    )
+
+
+def _render_mask_strategy_select(scope: str) -> MaskStrategy:
+    return render_mask_strategy_select(
+        key=widget_key("load", "mask_strategy", scope),
+        last_key=_LAST_MASK_STRATEGY_KEY,
+        help_text="Which extracted activation set to load.",
+    )
+
+
+def _select_single_variant_samples(
+    store: Store,
+    mask_strategy: MaskStrategy,
+    scope: str,
+    *,
+    remember_key: str,
+    variant_remember_key: str,
+    default_count_limit: int,
+) -> tuple[str, list[str], str, list[int]] | None:
+    variants = available_variants(store, mask_strategy)
+    if not variants:
+        st.info("No variants with saved vectors for this model.")
+        return None
+    variant_key = widget_key("load", "variant", scope, store_id(store))
+    default_variant = "biography" if "biography" in variants else variants[0]
+    variant = _remembered_selectbox(
+        "Variant",
+        key=variant_key,
+        remember_key=variant_remember_key,
+        options=variants,
+        default=default_variant,
+        format_func=prompt_variant_label,
+    )
+    persona_ids = _select_artifact_personas(
+        store,
+        [variant],
+        mask_strategy,
+        widget_scope=f"{scope}:{store_id(store)}",
+        remember_key=remember_key,
+        default_count_limit=default_count_limit,
+    )
+    if not persona_ids:
+        return None
+
+    persona_key = personas_fingerprint(persona_ids)
+    layer_options = _layers_for_variant(store, variant, persona_ids, mask_strategy)
+    if not layer_options:
+        st.info("No shared layers are available for the selected personas.")
+        return None
+
+    selected_layers = _render_layer_frame_controls(store, scope, layer_options)
+    return variant, persona_ids, persona_key, selected_layers

tabs/analysis/_state.py

@@ -0,0 +1,225 @@
+from dataclasses import dataclass
+
+import streamlit as st
+from persona_data.synth_persona import BASELINE_PERSONA_ID
+from persona_vectors.attributes import DEFAULT_MAX_ATTRIBUTE_CATEGORIES
+
+from utils.helpers import slugify, widget_key
+
+
+def _filename(*parts: str) -> str:
+    return "__".join(slugify(part) for part in parts if part)
+
+
+# Keep analysis-tab selection state separate so projection defaults do not
+# overwrite cosine similarity defaults.
+_LAST_COSINE_PERSONAS_KEY = "analysis:last_personas:cosine"
+_LAST_PROJECTION_PERSONAS_KEY = "analysis:last_personas:projection"
+_LAST_SIMILARITY_PERSONAS_KEY = "analysis:last_personas:similarity"
+_LAST_MASK_STRATEGY_KEY = "analysis:last_mask_strategy"
+_LAST_SOURCE_KEY = "analysis:last_source"
+_LAST_PROJECTION_VARIANT_KEY = "analysis:last_projection_variant"
+_LAST_SIMILARITY_VARIANT_KEY = "analysis:last_similarity_variant"
+_LAST_PROJECTION_COLOR_MODE_KEY = "analysis:last_projection_color_mode"
+_LAST_PROJECTION_ATTRIBUTE_KEY = "analysis:last_projection_attribute"
+_LAST_PROJECTION_CLUSTER_K_KEY = "analysis:last_projection_cluster_k"
+_LAST_PROJECTION_CLUSTER_MODE_KEY = "analysis:last_projection_cluster_mode"
+_LAST_PROJECTION_HIGHLIGHTS_KEY = "analysis:last_projection_highlights"
+_LAST_PROJECTION_DIMS_KEY = "analysis:last_projection_dims"
+_LAST_LAYER_FRAMES_KEY = "analysis:last_layer_frames"
+
+_DEFAULT_LAYER_FRAMES = 16
+_DEFAULT_PERSONA_LIMITS = {
+    "similarity": 120,
+    "pca": 500,
+    "umap": 500,
+    "isomap": 500,
+    "dendro": 160,
+}
+_MAX_SIMILARITY_CELLS = 4_000_000
+_MAX_PAIR_TRAJECTORY_TRACES = 500
+_DEFAULT_GRAPH_NEIGHBORS = 5
+_PROJECTION_KINDS = {"pca", "umap", "isomap"}
+_CLUSTER_MODES = {
+    "Mean across layers": "mean_across_layers",
+    "First selected layer": "first_layer",
+    "Per layer": "per_layer",
+}
+_PROJECTION_COLOR_MODES = ["Persona", "K-means clusters", "Persona attribute"]
+_MAX_ATTRIBUTE_CATEGORIES = DEFAULT_MAX_ATTRIBUTE_CATEGORIES
+
+
+def _is_assistant_persona(persona_id: str, persona_name: str | None = None) -> bool:
+    persona_id_normalized = persona_id.strip().lower()
+    persona_name_normalized = (persona_name or "").strip().lower()
+    return (
+        persona_id_normalized in {"assistant", BASELINE_PERSONA_ID.lower()}
+        or persona_name_normalized == "assistant"
+    )
+
+
+@dataclass(frozen=True)
+class CosineSelection:
+    variants: list[str]
+    variant_a: str
+    variant_b: str
+    persona_ids: list[str]
+    persona_key: str
+
+
+@dataclass(frozen=True)
+class PersonaOptions:
+    regular_ids: list[str]
+    assistant_id: str | None
+    persona_names: dict[str, str]
+
+
+@dataclass(frozen=True)
+class ProjectionColorConfig:
+    color_mode: str = "Persona"
+    n_clusters: int | None = None
+    cluster_mode: str | None = None
+    attribute_name: str | None = None
+    highlight_persona_ids: tuple[str, ...] = ()
+    highlight_persona_key: str = ""
+
+
+@dataclass(frozen=True)
+class LayeredFigureStateKeys:
+    figure: str
+    projection: str | None = None
+
+
+_HIGHLIGHT_OTHER_LABEL = "Other"
+_HIGHLIGHT_OTHER_COLOR = "rgba(148, 163, 184, 0.35)"
+
+
+def _persona_names_state_key(widget_scope: str) -> str:
+    return widget_key("load", "persona_names", widget_scope)
+
+
+def _persona_display_label(persona_names: dict[str, str], persona_id: str) -> str:
+    name = persona_names.get(persona_id, persona_id)
+    return f"{name} ({persona_id})" if name != persona_id else persona_id
+
+
+def _highlight_persona_groups(
+    persona_ids: list[str],
+    persona_names: dict[str, str],
+    highlight_persona_ids: tuple[str, ...],
+) -> list[str] | None:
+    if not highlight_persona_ids:
+        return None
+
+    highlighted = set(highlight_persona_ids)
+    return [
+        (
+            _persona_display_label(persona_names, persona_id)
+            if persona_id in highlighted
+            else _HIGHLIGHT_OTHER_LABEL
+        )
+        for persona_id in persona_ids
+    ]
+
+
+def _sequence_to_list(value: object) -> list[object] | None:
+    if value is None or isinstance(value, (str, bytes)):
+        return None
+    if isinstance(value, list):
+        return value
+    if isinstance(value, tuple):
+        return list(value)
+    try:
+        return list(value)
+    except TypeError:
+        return None
+
+
+_TRACKED_STATE_KEYS_KEY = "analysis:_tracked_state_keys"
+
+
+def _clear_old_load_states(current_key: str, suffix: str) -> None:
+    # Only one heavy figure/projection state should live at a time. We track
+    # the keys we create per suffix so eviction is O(1) instead of scanning
+    # all of session_state on every rerun. Every such key is passed through
+    # this function before it is set, so the registry stays authoritative.
+    tracked: dict[str, set[str]] = st.session_state.setdefault(
+        _TRACKED_STATE_KEYS_KEY, {}
+    )
+    for key in tracked.get(suffix, ()):
+        if key != current_key:
+            st.session_state.pop(key, None)
+    tracked[suffix] = {current_key}
+
+
+def _clear_old_figure_states(current_key: str) -> None:
+    _clear_old_load_states(current_key, "_fig_state")
+
+
+def _clear_old_projection_states(current_key: str) -> None:
+    _clear_old_load_states(current_key, "_projection_state")
+
+
+def _store_figure_state(key: str, value: object) -> None:
+    _clear_old_figure_states(key)
+    st.session_state[key] = value
+
+
+def _seed_selectbox_key(
+    *,
+    key: str,
+    remember_key: str,
+    options: list[str],
+    default: str,
+) -> str:
+    value = st.session_state.get(key, st.session_state.get(remember_key, default))
+    if value not in options:
+        value = default
+    return value
+
+
+def _remembered_selectbox(
+    label: str,
+    *,
+    key: str,
+    remember_key: str,
+    options: list[str],
+    default: str,
+    **selectbox_kwargs: object,
+) -> str:
+    selected = _seed_selectbox_key(
+        key=key,
+        remember_key=remember_key,
+        options=options,
+        default=default,
+    )
+    choice = st.selectbox(
+        label,
+        options=options,
+        index=options.index(selected),
+        key=key,
+        **selectbox_kwargs,
+    )
+    st.session_state[remember_key] = choice
+    return choice
+
+
+def _personas_empty_message(variants: list[str]) -> str:
+    if len(variants) > 1:
+        return (
+            "No personas have vectors for all selected variants. "
+            "Pick a single variant or change the source."
+        )
+    return "No personas found for this model and variant."
+
+
+def _remember_multiselect(
+    *,
+    key: str,
+    remember_key: str,
+    options: list[str],
+) -> list[str]:
+    remembered = st.session_state.get(key, st.session_state.get(remember_key, []))
+    if not isinstance(remembered, list):
+        remembered = []
+    return [value for value in remembered if value in options]

tabs/analysis/cosine.py

@@ -0,0 +1,226 @@
+from itertools import combinations
+
+import streamlit as st
+from persona_vectors.extraction import MaskStrategy
+from persona_vectors.plots import plot_layer_similarity
+
+from utils.analysis_sources import Store, available_variants, store_id
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
+
+from tabs.analysis._shared import (
+    _load_variant_vectors,
+    _plotly_chart,
+    _release_vector_memory,
+    _render_save_buttons,
+    _select_artifact_personas,
+)
+from tabs.analysis._state import (
+    _LAST_COSINE_PERSONAS_KEY,
+    CosineSelection,
+    _clear_old_figure_states,
+    _filename,
+    _store_figure_state,
+)
+
+
+def _render_cosine_selection(
+    store: Store,
+    mask_strategy: MaskStrategy,
+) -> CosineSelection | None:
+    variants = available_variants(store, mask_strategy)
+    if len(variants) < 2:
+        st.info("Need at least two variants with saved vectors for cosine comparison.")
+        return None
+
+    with st.expander("Vector selection", expanded=True):
+        col1, col2 = st.columns(2)
+        with col1:
+            variant_a = st.selectbox(
+                "Variant A",
+                options=variants,
+                index=0,
+                format_func=prompt_variant_label,
+                key=widget_key("load", "variant_a", store_id(store)),
+            )
+        with col2:
+            variant_b = st.selectbox(
+                "Variant B",
+                options=variants,
+                index=min(1, len(variants) - 1),
+                format_func=prompt_variant_label,
+                key=widget_key("load", "variant_b", store_id(store)),
+            )
+
+        if variant_a == variant_b:
+            st.warning("Choose two different variants to compare.")
+            return None
+
+        persona_ids = _select_artifact_personas(
+            store,
+            [variant_a, variant_b],
+            mask_strategy,
+            widget_scope=f"cosine:{store_id(store)}",
+            remember_key=_LAST_COSINE_PERSONAS_KEY,
+        )
+    if not persona_ids:
+        return None
+    return CosineSelection(
+        variants=variants,
+        variant_a=variant_a,
+        variant_b=variant_b,
+        persona_ids=persona_ids,
+        persona_key=personas_fingerprint(persona_ids),
+    )
+
+
+def _build_cosine_figures(
+    store: Store,
+    mask_strategy: MaskStrategy,
+    selection: CosineSelection,
+) -> tuple[object, object | None, int, int] | None:
+    variant_sample_cache: dict[str, object] = {}
+
+    def _load_variant(variant: str):
+        if variant not in variant_sample_cache:
+            samples = _load_variant_vectors(
+                store,
+                [variant],
+                mask_strategy,
+                persona_ids=selection.persona_ids,
+            )
+            variant_sample_cache[variant] = samples[variant]
+        return variant_sample_cache[variant]
+
+    try:
+        samples_a = _load_variant(selection.variant_a)
+        samples_b = _load_variant(selection.variant_b)
+    except Exception as exc:
+        st.error(f"Could not load vectors: {exc}")
+        return None
+
+    labels = samples_a.labels
+    display_traces = [
+        (
+            label,
+            samples_a.vectors[index],
+            samples_b.vectors[index],
+        )
+        for index, label in enumerate(labels)
+    ]
+    fig = plot_layer_similarity(
+        display_traces,
+        title=(
+            f"{prompt_variant_label(selection.variant_a)} vs "
+            f"{prompt_variant_label(selection.variant_b)}"
+        ),
+        show=False,
+    )
+
+    pair_traces = []
+    pair_errors = []
+    for left, right in combinations(selection.variants, 2):
+        try:
+            left_samples = _load_variant(left)
+            right_samples = _load_variant(right)
+            pair_traces.append(
+                (
+                    f"{prompt_variant_label(left)} vs {prompt_variant_label(right)}",
+                    left_samples.vectors.mean(dim=0),
+                    right_samples.vectors.mean(dim=0),
+                )
+            )
+        except Exception as exc:
+            pair_errors.append(f"{left} vs {right}: {exc}")
+            continue
+
+    for err in pair_errors:
+        st.warning(f"Skipped pair trace: `{err}`")
+    pair_fig = (
+        plot_layer_similarity(
+            pair_traces,
+            title="Variant-pair cosine similarity averaged over selected personas",
+            show=False,
+        )
+        if pair_traces
+        else None
+    )
+    return fig, pair_fig, len(display_traces), len(pair_traces)
+
+
+def _render_cosine_similarity(
+    store: Store,
+    mask_strategy: MaskStrategy,
+) -> None:
+    selection = _render_cosine_selection(store, mask_strategy)
+    if selection is None:
+        return
+
+    cosine_fig_key = widget_key(
+        "load",
+        "cosine_fig_state",
+        store_id(store),
+        store.model_name,
+        mask_strategy.value,
+        selection.variant_a,
+        selection.variant_b,
+        selection.persona_key,
+    )
+    filename = _filename(
+        "analysis",
+        "cosine",
+        store.model_name,
+        mask_strategy.value,
+        selection.variant_a,
+        selection.variant_b,
+    )
+    pairs_filename = _filename(
+        "analysis",
+        "cosine_pairs",
+        store.model_name,
+        mask_strategy.value,
+        "_".join(selection.variants),
+    )
+    _clear_old_figure_states(cosine_fig_key)
+
+    if st.button(
+        "Compare vectors",
+        type="primary",
+        key=widget_key(
+            "load",
+            "analysis_vectors",
+            store_id(store),
+            store.model_name,
+            mask_strategy.value,
+            selection.variant_a,
+            selection.variant_b,
+            selection.persona_key,
+        ),
+    ):
+        progress = st.progress(0, text="Loading activation vectors…")
+        try:
+            progress.progress(15, text="Loading activation vectors…")
+            figures = _build_cosine_figures(store, mask_strategy, selection)
+            if figures is None:
+                st.session_state.pop(cosine_fig_key, None)
+                return
+            progress.progress(90, text="Storing figure state…")
+            _store_figure_state(cosine_fig_key, figures)
+            progress.progress(100, text="Done.")
+        finally:
+            _release_vector_memory(store, selection.variants)
+            progress.empty()
+
+    if cosine_fig_key in st.session_state:
+        fig, pair_fig, n_traces, n_pair_traces = st.session_state[cosine_fig_key]
+        _plotly_chart(fig)
+        figs = [fig]
+        filenames = [filename]
+        if pair_fig is not None:
+            st.subheader("Variant pairs")
+            _plotly_chart(pair_fig)
+            figs.append(pair_fig)
+            filenames.append(pairs_filename)
+        _render_save_buttons(figs, filenames, "cosine")
+        st.success(f"Loaded {n_traces} personas for cosine comparison.")
+        if pair_fig is not None:
+            st.caption(f"Generated {n_pair_traces} averaged variant-pair trace(s).")

tabs/analysis/dendrogram.py

@@ -0,0 +1,279 @@
+import streamlit as st
+from persona_vectors.extraction import MaskStrategy
+from persona_vectors.plots import plot_persona_dendrogram
+
+from utils.analysis_sources import (
+    Store,
+    available_variants,
+    store_cache_parts,
+    store_id,
+    store_layers_cached,
+)
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
+
+from tabs.analysis._shared import (
+    _load_persona_options,
+    _load_persona_vectors,
+    _plotly_chart,
+    _release_vector_memory,
+    _render_layer_frame_controls,
+    _render_save_buttons,
+    _select_artifact_personas,
+)
+from tabs.analysis._state import (
+    _DEFAULT_PERSONA_LIMITS,
+    PersonaOptions,
+    _clear_old_figure_states,
+    _filename,
+    _persona_names_state_key,
+    _personas_empty_message,
+    _store_figure_state,
+)
+
+_LAST_DENDRO_PERSONAS_KEY = "analysis:last_personas:dendro"
+_DENDRO_LINKAGE_OPTIONS = ["ward", "complete", "average", "single"]
+
+
+def _render_persona_select_controls(
+    options: PersonaOptions,
+    widget_scope: str,
+) -> list[str]:
+    select_key = widget_key("load", "persona_select", widget_scope)
+    assistant_key = widget_key("load", "persona_select_assistant", widget_scope)
+
+    label_map = {
+        pid: f"{options.persona_names.get(pid, pid)} ({pid})"
+        for pid in options.regular_ids
+    }
+    sorted_labels = sorted(label_map.values())
+    selected_labels = st.multiselect(
+        "Select personas",
+        options=sorted_labels,
+        key=select_key,
+        placeholder="Search and select personas...",
+    )
+    label_to_id = {v: k for k, v in label_map.items()}
+    selected_ids = [label_to_id[lbl] for lbl in selected_labels]
+
+    if options.assistant_id is not None:
+        include_assistant = st.checkbox(
+            "Include Assistant persona",
+            key=assistant_key,
+        )
+        if include_assistant:
+            selected_ids.append(options.assistant_id)
+
+    st.session_state[_persona_names_state_key(widget_scope)] = dict(
+        options.persona_names
+    )
+
+    if not selected_ids:
+        st.info("Select at least one persona.")
+
+    return selected_ids
+
+
+def _render_dendrogram_analysis(
+    store: Store,
+    mask_strategy: MaskStrategy,
+) -> None:
+    variants = available_variants(store, mask_strategy)
+    if not variants:
+        st.info("No variants with saved vectors for this model.")
+        return
+
+    with st.expander("Variant selection", expanded=True):
+        col1, col2 = st.columns(2)
+        default_a = "biography" if "biography" in variants else variants[0]
+        default_b_idx = (
+            variants.index("templated")
+            if "templated" in variants
+            else min(1, len(variants) - 1)
+        )
+        with col1:
+            variant_a = st.selectbox(
+                "Variant A",
+                options=variants,
+                index=variants.index(default_a),
+                format_func=prompt_variant_label,
+                key=widget_key("load", "dendro_variant_a", store_id(store)),
+            )
+        with col2:
+            variant_b = st.selectbox(
+                "Variant B",
+                options=variants,
+                index=default_b_idx,
+                format_func=prompt_variant_label,
+                key=widget_key("load", "dendro_variant_b", store_id(store)),
+            )
+
+    shared_variants = list(dict.fromkeys([variant_a, variant_b]))
+
+    select_specific = st.toggle(
+        "Select specific personas",
+        value=False,
+        key=widget_key("load", "dendro_select_mode", store_id(store)),
+        help="Search and select specific personas instead of using the first N.",
+    )
+
+    if select_specific:
+        empty_message = _personas_empty_message(shared_variants)
+        options = _load_persona_options(
+            store,
+            shared_variants,
+            mask_strategy,
+            empty_message=empty_message,
+        )
+        if options is None:
+            st.session_state.pop(
+                _persona_names_state_key(f"dendro:{store_id(store)}"), None
+            )
+            return
+        persona_ids = _render_persona_select_controls(
+            options,
+            widget_scope=f"dendro:{store_id(store)}",
+        )
+        if not persona_ids:
+            return
+    else:
+        persona_ids = _select_artifact_personas(
+            store,
+            shared_variants,
+            mask_strategy,
+            widget_scope=f"dendro:{store_id(store)}",
+            remember_key=_LAST_DENDRO_PERSONAS_KEY,
+            default_count_limit=_DEFAULT_PERSONA_LIMITS["dendro"],
+        )
+        if not persona_ids:
+            return
+
+    col_opts1, col_opts2 = st.columns(2)
+    with col_opts1:
+        layered_mode = st.toggle(
+            "Per-layer animated",
+            value=False,
+            key=widget_key("load", "dendro_layered", store_id(store)),
+            help="Animated dendrogram with one frame per layer instead of averaging all layers.",
+        )
+    with col_opts2:
+        linkage = st.selectbox(
+            "Linkage",
+            options=_DENDRO_LINKAGE_OPTIONS,
+            index=0,
+            key=widget_key("load", "dendro_linkage", store_id(store)),
+        )
+
+    selected_layers: list[int] | None = None
+    if layered_mode:
+        source, location, model_name = store_cache_parts(store)
+        layer_options = store_layers_cached(
+            source,
+            location,
+            model_name,
+            mask_strategy.value,
+            tuple(shared_variants),
+            tuple(persona_ids),
+        )
+        if not layer_options:
+            st.info("No shared layers are available for the selected personas.")
+            return
+        selected_layers = _render_layer_frame_controls(store, "dendro", layer_options)
+
+    persona_key = personas_fingerprint(persona_ids)
+    fig_key = widget_key(
+        "load",
+        "dendro_fig_state",
+        store_id(store),
+        store.model_name,
+        mask_strategy.value,
+        variant_a,
+        variant_b,
+        persona_key,
+        str(layered_mode),
+        linkage,
+        "_".join(map(str, selected_layers or [])),
+    )
+    _clear_old_figure_states(fig_key)
+
+    if st.button(
+        "Generate dendrograms",
+        type="primary",
+        key=widget_key(
+            "load", "dendro_btn", store_id(store), variant_a, variant_b, persona_key
+        ),
+    ):
+        progress = st.progress(0, text="Loading first variant vectors…")
+        try:
+            progress.progress(15, text="Loading first variant vectors…")
+            samples_a = _load_persona_vectors(
+                store,
+                variant_a,
+                mask_strategy,
+                persona_ids,
+            )
+            progress.progress(40, text="Building first dendrogram…")
+            fig_a = plot_persona_dendrogram(
+                samples_a,
+                layered=layered_mode,
+                layers=selected_layers,
+                linkage=linkage,
+                title=f"Dendrogram — {prompt_variant_label(variant_a)}",
+            )
+            fig_a.update_layout(height=750)
+            del samples_a
+            fig_b = None
+            if variant_a != variant_b:
+                progress.progress(60, text="Loading second variant vectors…")
+                samples_b = _load_persona_vectors(
+                    store,
+                    variant_b,
+                    mask_strategy,
+                    persona_ids,
+                )
+                progress.progress(75, text="Building second dendrogram…")
+                fig_b = plot_persona_dendrogram(
+                    samples_b,
+                    layered=layered_mode,
+                    layers=selected_layers,
+                    linkage=linkage,
+                    title=f"Dendrogram — {prompt_variant_label(variant_b)}",
+                )
+                fig_b.update_layout(height=750)
+                del samples_b
+            progress.progress(90, text="Storing figure state…")
+            _store_figure_state(
+                fig_key,
+                (fig_a, fig_b, len(persona_ids), variant_a, variant_b),
+            )
+            progress.progress(100, text="Done.")
+        except Exception as exc:
+            st.error(f"Could not build dendrogram: {exc}")
+            st.session_state.pop(fig_key, None)
+        finally:
+            _release_vector_memory(store, shared_variants)
+            progress.empty()
+
+    if fig_key in st.session_state:
+        fig_a, fig_b, n_personas, va, vb = st.session_state[fig_key]
+        if fig_b is not None:
+            col_a, col_b = st.columns(2)
+            with col_a:
+                st.subheader(prompt_variant_label(va))
+                _plotly_chart(fig_a)
+            with col_b:
+                st.subheader(prompt_variant_label(vb))
+                _plotly_chart(fig_b)
+        else:
+            _plotly_chart(fig_a)
+
+        figs = [fig_a] + ([fig_b] if fig_b else [])
+        filenames = [
+            _filename("dendro", store.model_name, mask_strategy.value, va),
+            *(
+                [_filename("dendro", store.model_name, mask_strategy.value, vb)]
+                if fig_b
+                else []
+            ),
+        ]
+        _render_save_buttons(figs, filenames, "dendro")
+        st.success(f"Generated dendrogram(s) for {n_personas} persona(s).")

tabs/analysis/layered.py

@@ -0,0 +1,563 @@
+from collections.abc import Callable
+
+import plotly.graph_objects as go
+import streamlit as st
+from persona_vectors.attributes import (
+    attribute_color_kwargs,
+    attribute_display_label,
+)
+from persona_vectors.extraction import MaskStrategy
+from persona_vectors.plots import (
+    build_layered_figure,
+    build_pair_similarity_figure,
+    build_similarity_figures,
+    prepare_layered_projection_data,
+)
+
+from utils.analysis_metadata import (
+    synth_persona_attribute_names,
+    synth_persona_dataset_cached,
+)
+from utils.analysis_sources import Store, store_id
+from utils.helpers import personas_fingerprint, prompt_variant_label, widget_key
+
+from tabs.analysis._shared import (
+    _gray_out_unselected_personas,
+    _load_persona_vectors,
+    _plotly_chart,
+    _release_vector_memory,
+    _render_save_buttons,
+    _select_single_variant_samples,
+)
+from tabs.analysis._state import (
+    _CLUSTER_MODES,
+    _DEFAULT_GRAPH_NEIGHBORS,
+    _LAST_PROJECTION_ATTRIBUTE_KEY,
+    _LAST_PROJECTION_CLUSTER_K_KEY,
+    _LAST_PROJECTION_CLUSTER_MODE_KEY,
+    _LAST_PROJECTION_COLOR_MODE_KEY,
+    _LAST_PROJECTION_HIGHLIGHTS_KEY,
+    _LAST_PROJECTION_PERSONAS_KEY,
+    _LAST_PROJECTION_VARIANT_KEY,
+    _LAST_SIMILARITY_VARIANT_KEY,
+    _MAX_ATTRIBUTE_CATEGORIES,
+    _MAX_PAIR_TRAJECTORY_TRACES,
+    _MAX_SIMILARITY_CELLS,
+    _PROJECTION_COLOR_MODES,
+    _PROJECTION_KINDS,
+    LayeredFigureStateKeys,
+    ProjectionColorConfig,
+    _clear_old_figure_states,
+    _clear_old_projection_states,
+    _highlight_persona_groups,
+    _persona_display_label,
+    _persona_names_state_key,
+    _remember_multiselect,
+    _remembered_selectbox,
+    _store_figure_state,
+)
+
+
+def _render_pair_trajectory_control(
+    *,
+    enabled: bool,
+    persona_count: int,
+    scope: str,
+    store: Store,
+) -> bool:
+    if not enabled:
+        return False
+    pair_count = persona_count * (persona_count - 1) // 2
+    if pair_count > _MAX_PAIR_TRAJECTORY_TRACES:
+        st.caption(
+            "Pair trajectories hidden because this selection would create "
+            f"{pair_count:,} Plotly traces."
+        )
+        return False
+    return st.checkbox(
+        "Pair trajectories",
+        value=False,
+        key=widget_key("load", "pair_trajectories", scope, store_id(store)),
+        help="Adds one line per persona pair. Keep this off for larger selections.",
+    )
+
+
+def _validate_layered_figure_size(
+    figure_kind: str,
+    persona_count: int,
+    selected_layers: list[int],
+) -> bool:
+    if figure_kind != "similarity":
+        return True
+    similarity_cells = persona_count * persona_count * len(selected_layers)
+    if similarity_cells <= _MAX_SIMILARITY_CELLS:
+        return True
+    st.error(
+        "Reduce personas or layer frames before generating the similarity "
+        f"matrix ({similarity_cells:,} cells selected)."
+    )
+    return False
+
+
+def _render_projection_color_config(
+    store: Store,
+    scope: str,
+    persona_ids: list[str],
+) -> ProjectionColorConfig | None:
+    widget_scope = f"{scope}:{store_id(store)}"
+    persona_key = personas_fingerprint(persona_ids)
+    persona_names = st.session_state.get(
+        _persona_names_state_key(widget_scope),
+        {},
+    )
+    color_mode_key = widget_key("load", "color_mode", scope, store_id(store))
+    color_mode = _remembered_selectbox(
+        "Color by",
+        key=color_mode_key,
+        remember_key=_LAST_PROJECTION_COLOR_MODE_KEY,
+        options=_PROJECTION_COLOR_MODES,
+        default="Persona",
+    )
+    if color_mode == "K-means clusters":
+        max_clusters = min(10, len(persona_ids))
+        if max_clusters < 2:
+            st.info("Select at least two personas to use K-means coloring.")
+            return None
+        cluster_key = widget_key("load", "cluster_k", scope, store_id(store))
+        default_clusters = min(3, len(persona_ids))
+        if cluster_key not in st.session_state:
+            st.session_state[cluster_key] = min(
+                max(
+                    int(
+                        st.session_state.get(
+                            _LAST_PROJECTION_CLUSTER_K_KEY,
+                            default_clusters,
+                        )
+                    ),
+                    2,
+                ),
+                max_clusters,
+            )
+        n_clusters = st.slider(
+            "K (clusters)",
+            min_value=2,
+            max_value=max_clusters,
+            key=cluster_key,
+        )
+        mode_key = widget_key("load", "cluster_mode", scope, store_id(store))
+        mode_options = list(_CLUSTER_MODES)
+        mode_label = _remembered_selectbox(
+            "Cluster fit",
+            key=mode_key,
+            remember_key=_LAST_PROJECTION_CLUSTER_MODE_KEY,
+            options=mode_options,
+            default=mode_options[0],
+            help=(
+                "Mean across layers is the previous behavior. First selected "
+                "layer keeps one fixed clustering from the first frame. Per layer "
+                "recomputes clustering for each animation frame."
+            ),
+        )
+        st.session_state[_LAST_PROJECTION_CLUSTER_K_KEY] = n_clusters
+        return ProjectionColorConfig(
+            color_mode=color_mode,
+            n_clusters=n_clusters,
+            cluster_mode=_CLUSTER_MODES[mode_label],
+        )
+
+    if color_mode == "Persona attribute":
+        persona_dataset = synth_persona_dataset_cached()
+        attribute_options = list(synth_persona_attribute_names())
+        if not attribute_options:
+            st.info("No persona attributes are available for this dataset.")
+            return None
+        default_attribute = (
+            attribute_options.index("sex") if "sex" in attribute_options else 0
+        )
+        attribute_key = widget_key("load", "attribute", scope, store_id(store))
+        attribute_name = _remembered_selectbox(
+            "Attribute",
+            key=attribute_key,
+            remember_key=_LAST_PROJECTION_ATTRIBUTE_KEY,
+            options=attribute_options,
+            default=attribute_options[default_attribute],
+            format_func=lambda name: attribute_display_label(persona_dataset, name),
+        )
+        info = persona_dataset.attribute_info(attribute_name)
+        if info.get("high_cardinality"):
+            st.caption(
+                "High-cardinality categorical attributes are grouped to the "
+                f"top {_MAX_ATTRIBUTE_CATEGORIES} values plus Other."
+            )
+        return ProjectionColorConfig(
+            color_mode=color_mode,
+            attribute_name=attribute_name,
+        )
+
+    highlight_persona_ids: tuple[str, ...] = ()
+    if persona_ids:
+        highlight_key = widget_key(
+            "load", "persona_highlight", scope, store_id(store), persona_key
+        )
+        highlighted = st.multiselect(
+            "Highlight personas",
+            options=persona_ids,
+            default=_remember_multiselect(
+                key=highlight_key,
+                remember_key=_LAST_PROJECTION_HIGHLIGHTS_KEY,
+                options=persona_ids,
+            ),
+            format_func=lambda persona_id: _persona_display_label(
+                persona_names, persona_id
+            ),
+            key=highlight_key,
+            help=(
+                "Select a few personas to keep their default colors while the rest "
+                "are grayed out."
+            ),
+        )
+        highlight_persona_ids = tuple(highlighted)
+        st.session_state[_LAST_PROJECTION_HIGHLIGHTS_KEY] = list(highlighted)
+
+    highlight_persona_key = (
+        personas_fingerprint(highlight_persona_ids) if highlight_persona_ids else ""
+    )
+
+    return ProjectionColorConfig(
+        color_mode=color_mode,
+        highlight_persona_ids=highlight_persona_ids,
+        highlight_persona_key=highlight_persona_key,
+    )
+
+
+def _layered_figure_state_keys(
+    store: Store,
+    mask_strategy: MaskStrategy,
+    *,
+    scope: str,
+    figure_kind: str,
+    n_components: int,
+    color_config: ProjectionColorConfig,
+    variant: str,
+    persona_key: str,
+    selected_layers: list[int],
+    pair_trajectories: bool,
+) -> LayeredFigureStateKeys:
+    layer_key = "_".join(map(str, selected_layers))
+    figure_key = widget_key(
+        "load",
+        f"{scope}_fig_state",
+        store_id(store),
+        store.model_name,
+        mask_strategy.value,
+        figure_kind,
+        str(n_components),
+        color_config.color_mode,
+        str(color_config.attribute_name),
+        str(color_config.n_clusters),
+        str(color_config.cluster_mode),
+        str(color_config.highlight_persona_key),
+        variant,
+        "persona_vector",
+        persona_key,
+        layer_key,
+        str(pair_trajectories),
+    )
+    if figure_kind not in _PROJECTION_KINDS:
+        return LayeredFigureStateKeys(figure=figure_key)
+
+    graph_overlay = figure_kind == "isomap"
+    projection_key = widget_key(
+        "load",
+        f"{scope}_projection_state",
+        store_id(store),
+        store.model_name,
+        mask_strategy.value,
+        figure_kind,
+        str(n_components),
+        str(graph_overlay),
+        str(_DEFAULT_GRAPH_NEIGHBORS),
+        variant,
+        "persona_vector",
+        persona_key,
+        layer_key,
+    )
+    return LayeredFigureStateKeys(figure=figure_key, projection=projection_key)
+
+
+def _projection_build_kwargs(
+    samples,
+    *,
+    figure_kind: str,
+    selected_layers: list[int],
+    n_components: int,
+    color_config: ProjectionColorConfig,
+    persona_ids: list[str],
+    persona_names: dict[str, str],
+    projection_key: str | None,
+) -> dict:
+    if figure_kind not in _PROJECTION_KINDS:
+        return {}
+
+    graph_overlay = figure_kind == "isomap"
+    build_kwargs = {
+        "n_components": n_components,
+        "graph_overlay": graph_overlay,
+        "graph_n_neighbors": _DEFAULT_GRAPH_NEIGHBORS,
+    }
+    if color_config.n_clusters is not None:
+        build_kwargs["n_clusters"] = color_config.n_clusters
+        build_kwargs["cluster_mode"] = color_config.cluster_mode
+    if projection_key is not None:
+        projection_data = st.session_state.get(projection_key)
+        if projection_data is None:
+            projection_data = prepare_layered_projection_data(
+                samples,
+                figure_kind,
+                layers=selected_layers,
+                n_components=n_components,
+                graph_overlay=graph_overlay,
+                graph_n_neighbors=_DEFAULT_GRAPH_NEIGHBORS,
+            )
+            st.session_state[projection_key] = projection_data
+        build_kwargs["projection_data"] = projection_data
+    if color_config.attribute_name is not None:
+        build_kwargs.update(
+            attribute_color_kwargs(
+                synth_persona_dataset_cached(),
+                color_config.attribute_name,
+                persona_ids,
+                max_categories=_MAX_ATTRIBUTE_CATEGORIES,
+            )
+        )
+    if color_config.color_mode == "Persona" and color_config.highlight_persona_ids:
+        groups = _highlight_persona_groups(
+            persona_ids,
+            persona_names,
+            color_config.highlight_persona_ids,
+        )
+        if groups is not None:
+            build_kwargs["groups"] = groups
+    return build_kwargs
+
+
+def _build_layered_analysis_figures(
+    samples,
+    *,
+    figure_kind: str,
+    selected_layers: list[int],
+    variant: str,
+    title_fn: Callable[[str], str],
+    pair_trajectories: bool,
+    build_kwargs: dict,
+) -> tuple[go.Figure, go.Figure | None]:
+    if figure_kind == "similarity" and pair_trajectories:
+        return build_similarity_figures(
+            samples,
+            layers=selected_layers,
+            title=title_fn(variant),
+            pair_title=(
+                "Pair similarity trajectories - "
+                f"{prompt_variant_label(variant)} - persona vectors"
+            ),
+        )
+
+    main_fig = build_layered_figure(
+        samples,
+        figure_kind,
+        layers=selected_layers,
+        title=title_fn(variant),
+        **build_kwargs,
+    )
+    if figure_kind == "isomap":
+        _add_isomap_connection_toggle(main_fig)
+    if figure_kind in _PROJECTION_KINDS:
+        main_fig.update_layout(height=700)
+    extra_fig = (
+        build_pair_similarity_figure(
+            samples,
+            layers=selected_layers,
+            title=(
+                "Pair similarity trajectories - "
+                f"{prompt_variant_label(variant)} - persona vectors"
+            ),
+        )
+        if pair_trajectories
+        else None
+    )
+    return main_fig, extra_fig
+
+
+def _add_isomap_connection_toggle(fig: go.Figure) -> None:
+    """Add an in-plot control for the Isomap kNN graph trace."""
+    if not fig.data or fig.data[0].name != "kNN graph":
+        return
+
+    existing_menus = tuple(fig.layout.updatemenus or ())
+    fig.update_layout(
+        updatemenus=existing_menus
+        + (
+            dict(
+                type="buttons",
+                direction="left",
+                active=0,
+                showactive=False,
+                x=0,
+                xanchor="left",
+                y=1.16,
+                yanchor="top",
+                pad=dict(t=0, r=10),
+                buttons=[
+                    dict(
+                        label="Show connections",
+                        method="restyle",
+                        args=[{"visible": True}, [0]],
+                    ),
+                    dict(
+                        label="Hide connections",
+                        method="restyle",
+                        args=[{"visible": False}, [0]],
+                    ),
+                ],
+            ),
+        ),
+    )
+
+
+def _render_layered_figure_analysis(
+    store: Store,
+    mask_strategy: MaskStrategy,
+    *,
+    scope: str,
+    figure_kind: str,
+    button_label: str,
+    title_fn: Callable[[str], str],
+    include_pair_trajectories: bool = False,
+    n_components: int = 2,
+    remember_key: str = _LAST_PROJECTION_PERSONAS_KEY,
+    default_count_limit: int = 500,
+) -> None:
+    """Render a single-variant layered analysis: select → button → figure(s).
+
+    Used for similarity matrix, PCA, and UMAP. Set ``include_pair_trajectories``
+    to add the pair-similarity-trajectory figure (similarity matrix only).
+    """
+    selected = _select_single_variant_samples(
+        store,
+        mask_strategy,
+        scope,
+        remember_key=remember_key,
+        variant_remember_key=(
+            _LAST_PROJECTION_VARIANT_KEY
+            if figure_kind in _PROJECTION_KINDS
+            else _LAST_SIMILARITY_VARIANT_KEY
+        ),
+        default_count_limit=default_count_limit,
+    )
+    if selected is None:
+        return
+    variant, persona_ids, persona_key, selected_layers = selected
+
+    pair_trajectories = _render_pair_trajectory_control(
+        enabled=include_pair_trajectories,
+        persona_count=len(persona_ids),
+        scope=scope,
+        store=store,
+    )
+    if not _validate_layered_figure_size(
+        figure_kind, len(persona_ids), selected_layers
+    ):
+        return
+
+    color_config = ProjectionColorConfig()
+    if figure_kind in _PROJECTION_KINDS:
+        color_config = _render_projection_color_config(store, scope, persona_ids)
+        if color_config is None:
+            return
+
+    state_keys = _layered_figure_state_keys(
+        store,
+        mask_strategy,
+        scope=scope,
+        figure_kind=figure_kind,
+        n_components=n_components,
+        color_config=color_config,
+        variant=variant,
+        persona_key=persona_key,
+        selected_layers=selected_layers,
+        pair_trajectories=pair_trajectories,
+    )
+    if state_keys.projection is not None:
+        _clear_old_projection_states(state_keys.projection)
+    filename = scope
+    _clear_old_figure_states(state_keys.figure)
+    persona_names = st.session_state.get(
+        _persona_names_state_key(f"{scope}:{store_id(store)}"),
+        {},
+    )
+
+    if st.button(button_label, type="primary"):
+        build_label = {
+            "umap": "Computing UMAP projections…",
+            "pca": "Computing PCA projections…",
+            "isomap": "Computing Isomap projections…",
+            "similarity": "Computing similarity matrices…",
+        }.get(figure_kind, "Building figure…")
+        progress = st.progress(0, text="Loading activation vectors…")
+        try:
+            progress.progress(15, text="Loading activation vectors…")
+            samples = _load_persona_vectors(
+                store,
+                variant,
+                mask_strategy,
+                persona_ids,
+            )
+            progress.progress(55, text=build_label)
+            build_kwargs = _projection_build_kwargs(
+                samples,
+                figure_kind=figure_kind,
+                selected_layers=selected_layers,
+                n_components=n_components,
+                color_config=color_config,
+                persona_ids=persona_ids,
+                persona_names=persona_names,
+                projection_key=state_keys.projection,
+            )
+            main_fig, extra_fig = _build_layered_analysis_figures(
+                samples,
+                figure_kind=figure_kind,
+                selected_layers=selected_layers,
+                variant=variant,
+                title_fn=title_fn,
+                pair_trajectories=pair_trajectories,
+                build_kwargs=build_kwargs,
+            )
+            if (
+                color_config.color_mode == "Persona"
+                and color_config.highlight_persona_ids
+            ):
+                _gray_out_unselected_personas(main_fig)
+            progress.progress(90, text="Storing figure state…")
+            n_samples = samples.vectors.shape[0]
+            del samples
+            _store_figure_state(state_keys.figure, (main_fig, extra_fig, n_samples))
+            progress.progress(100, text="Done.")
+        except Exception as exc:
+            st.error(f"Could not build figure: {exc}")
+            st.session_state.pop(state_keys.figure, None)
+        finally:
+            _release_vector_memory(store, [variant])
+            progress.empty()
+
+    if state_keys.figure in st.session_state:
+        main_fig, extra_fig, n_samples = st.session_state[state_keys.figure]
+        _plotly_chart(main_fig)
+        figs = [main_fig]
+        filenames = [filename]
+        if extra_fig is not None:
+            st.subheader("Pair trajectories")
+            _plotly_chart(extra_fig)
+            figs.append(extra_fig)
+            filenames.append(f"{filename}__pair_trajectories")
+        _render_save_buttons(figs, filenames, scope)
+        st.success(f"Loaded {n_samples} samples.")

tabs/analysis_core.py

@@ -1,28 +1,8 @@
-import gc
-from collections.abc import Callable
-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
-from persona_vectors.attributes import (
-    DEFAULT_MAX_ATTRIBUTE_CATEGORIES,
-    attribute_color_kwargs,
-    attribute_display_label,
-)
 from persona_vectors.extraction import MaskStrategy
-from persona_vectors.plots import (
-    build_layered_figure,
-    build_pair_similarity_figure,
-    build_similarity_figures,
-    plot_layer_similarity,
-    plot_persona_dendrogram,
-    prepare_layered_projection_data,
-    save_plot_html,
-)
 
 from utils.analysis_sources import (
     DEFAULT_COMPARE_MODEL,
@@ -32,1611 +12,27 @@ from utils.analysis_sources import (
     SOURCES,
     Store,
     activation_store_cached,
-    available_variants,
     hub_models_by_mask_strategy,
-    load_persona_vectors_cached,
-    load_variant_vectors_cached,
     local_model_matches,
     local_model_options_cached,
-    persona_names_cached,
-    personas_cached,
-    release_hf_store_cache,
-    store_cache_parts,
-    store_id,
-    store_layers_cached,
-)
-from utils.analysis_metadata import (
-    synth_persona_attribute_names,
-    synth_persona_dataset_cached,
 )
-from utils.controls import render_mask_strategy_select
 from utils.helpers import (
     ANALYSIS_HELP_TEXT,
     ANALYSIS_MODES,
-    personas_fingerprint,
     prompt_variant_label,
-    slugify,
     widget_key,
 )
-from utils.theme import active_base, style_plotly_layer_controls
-
-
-def _filename(*parts: str) -> str:
-    return "__".join(slugify(part) for part in parts if part)
-
-
-# Keep analysis-tab selection state separate so projection defaults do not
-# overwrite cosine similarity defaults.
-_LAST_COSINE_PERSONAS_KEY = "analysis:last_personas:cosine"
-_LAST_PROJECTION_PERSONAS_KEY = "analysis:last_personas:projection"
-_LAST_SIMILARITY_PERSONAS_KEY = "analysis:last_personas:similarity"
-_LAST_MASK_STRATEGY_KEY = "analysis:last_mask_strategy"
-_LAST_SOURCE_KEY = "analysis:last_source"
-_LAST_PROJECTION_VARIANT_KEY = "analysis:last_projection_variant"
-_LAST_SIMILARITY_VARIANT_KEY = "analysis:last_similarity_variant"
-_LAST_PROJECTION_COLOR_MODE_KEY = "analysis:last_projection_color_mode"
-_LAST_PROJECTION_ATTRIBUTE_KEY = "analysis:last_projection_attribute"
-_LAST_PROJECTION_CLUSTER_K_KEY = "analysis:last_projection_cluster_k"
-_LAST_PROJECTION_CLUSTER_MODE_KEY = "analysis:last_projection_cluster_mode"
-_LAST_PROJECTION_HIGHLIGHTS_KEY = "analysis:last_projection_highlights"
-_LAST_PROJECTION_DIMS_KEY = "analysis:last_projection_dims"
-_LAST_LAYER_FRAMES_KEY = "analysis:last_layer_frames"
-
-_DEFAULT_LAYER_FRAMES = 16
-_DEFAULT_PERSONA_LIMITS = {
-    "similarity": 120,
-    "pca": 500,
-    "umap": 500,
-    "isomap": 500,
-    "dendro": 160,
-}
-_MAX_SIMILARITY_CELLS = 4_000_000
-_MAX_PAIR_TRAJECTORY_TRACES = 500
-_DEFAULT_GRAPH_NEIGHBORS = 5
-_PROJECTION_KINDS = {"pca", "umap", "isomap"}
-_CLUSTER_MODES = {
-    "Mean across layers": "mean_across_layers",
-    "First selected layer": "first_layer",
-    "Per layer": "per_layer",
-}
-_PROJECTION_COLOR_MODES = ["Persona", "K-means clusters", "Persona attribute"]
-_MAX_ATTRIBUTE_CATEGORIES = DEFAULT_MAX_ATTRIBUTE_CATEGORIES
-
-
-def _is_assistant_persona(persona_id: str, persona_name: str | None = None) -> bool:
-    persona_id_normalized = persona_id.strip().lower()
-    persona_name_normalized = (persona_name or "").strip().lower()
-    return (
-        persona_id_normalized in {"assistant", BASELINE_PERSONA_ID.lower()}
-        or persona_name_normalized == "assistant"
-    )
-
-
-@dataclass(frozen=True)
-class CosineSelection:
-    variants: list[str]
-    variant_a: str
-    variant_b: str
-    persona_ids: list[str]
-    persona_key: str
-
-
-@dataclass(frozen=True)
-class PersonaOptions:
-    regular_ids: list[str]
-    assistant_id: str | None
-    persona_names: dict[str, str]
-
-
-@dataclass(frozen=True)
-class ProjectionColorConfig:
-    color_mode: str = "Persona"
-    n_clusters: int | None = None
-    cluster_mode: str | None = None
-    attribute_name: str | None = None
-    highlight_persona_ids: tuple[str, ...] = ()
-    highlight_persona_key: str = ""
-
-
-@dataclass(frozen=True)
-class LayeredFigureStateKeys:
-    figure: str
-    projection: str | None = None
-
-
-_HIGHLIGHT_OTHER_LABEL = "Other"
-_HIGHLIGHT_OTHER_COLOR = "rgba(148, 163, 184, 0.35)"
-
-
-def _persona_names_state_key(widget_scope: str) -> str:
-    return widget_key("load", "persona_names", widget_scope)
-
-
-def _persona_display_label(persona_names: dict[str, str], persona_id: str) -> str:
-    name = persona_names.get(persona_id, persona_id)
-    return f"{name} ({persona_id})" if name != persona_id else persona_id
-
-
-def _highlight_persona_groups(
-    persona_ids: list[str],
-    persona_names: dict[str, str],
-    highlight_persona_ids: tuple[str, ...],
-) -> list[str] | None:
-    if not highlight_persona_ids:
-        return None
-
-    highlighted = set(highlight_persona_ids)
-    return [
-        (
-            _persona_display_label(persona_names, persona_id)
-            if persona_id in highlighted
-            else _HIGHLIGHT_OTHER_LABEL
-        )
-        for persona_id in persona_ids
-    ]
-
-
-def _sequence_to_list(value: object) -> list[object] | None:
-    if value is None or isinstance(value, (str, bytes)):
-        return None
-    if isinstance(value, list):
-        return value
-    if isinstance(value, tuple):
-        return list(value)
-    try:
-        return list(value)
-    except TypeError:
-        return None
-
-
-def _gray_out_unselected_personas(fig: go.Figure) -> None:
-    def _gray_trace(trace: object) -> None:
-        marker = getattr(trace, "marker", None)
-        if marker is None:
-            return
-
-        colors = _sequence_to_list(getattr(marker, "color", None))
-        labels = _sequence_to_list(getattr(trace, "customdata", None))
-        if colors is not None and labels is not None and len(colors) == len(labels):
-            trace.marker.color = [
-                (
-                    _HIGHLIGHT_OTHER_COLOR
-                    if str(label) == _HIGHLIGHT_OTHER_LABEL
-                    else color
-                )
-                for label, color in zip(labels, colors, strict=True)
-            ]
-            return
-
-        if getattr(trace, "name", None) == _HIGHLIGHT_OTHER_LABEL:
-            trace.marker.color = _HIGHLIGHT_OTHER_COLOR
-            trace.opacity = 0.28
-
-    for trace in fig.data:
-        _gray_trace(trace)
-    for frame in fig.frames:
-        for trace in frame.data:
-            _gray_trace(trace)
-
-
-def _layers_for_variant(
-    store: Store,
-    variant: str,
-    persona_ids: list[str],
-    mask_strategy: MaskStrategy,
-) -> list[int]:
-    source, location, model_name = store_cache_parts(store)
-    return store_layers_cached(
-        source,
-        location,
-        model_name,
-        mask_strategy.value,
-        (variant,),
-        tuple(persona_ids),
-    )
-
-
-def _load_persona_vectors(
-    store: Store,
-    variant: str,
-    mask_strategy: MaskStrategy,
-    persona_ids: list[str],
-):
-    source, location, model_name = store_cache_parts(store)
-    return load_persona_vectors_cached(
-        source,
-        location,
-        model_name,
-        mask_strategy.value,
-        variant,
-        tuple(persona_ids),
-    )
-
-
-def _load_variant_vectors(
-    store: Store,
-    variants: list[str] | tuple[str, ...],
-    mask_strategy: MaskStrategy,
-    persona_ids: list[str],
-):
-    source, location, model_name = store_cache_parts(store)
-    return load_variant_vectors_cached(
-        source,
-        location,
-        model_name,
-        mask_strategy.value,
-        tuple(variants),
-        tuple(persona_ids),
-    )
-
-
-def _clear_old_load_states(current_key: str, suffix: str) -> None:
-    for key in list(st.session_state):
-        if key == current_key or not isinstance(key, str):
-            continue
-        parts = key.split("::", 2)
-        if len(parts) >= 2 and parts[0] == "load" and parts[1].endswith(suffix):
-            st.session_state.pop(key, None)
-
-
-def _clear_old_figure_states(current_key: str) -> None:
-    _clear_old_load_states(current_key, "_fig_state")
-
-
-def _clear_old_projection_states(current_key: str) -> None:
-    _clear_old_load_states(current_key, "_projection_state")
-
-
-def _store_figure_state(key: str, value: object) -> None:
-    _clear_old_figure_states(key)
-    st.session_state[key] = value
-
-
-def _seed_selectbox_key(
-    *,
-    key: str,
-    remember_key: str,
-    options: list[str],
-    default: str,
-) -> str:
-    value = st.session_state.get(key, st.session_state.get(remember_key, default))
-    if value not in options:
-        value = default
-    return value
-
-
-def _remember_multiselect(
-    *,
-    key: str,
-    remember_key: str,
-    options: list[str],
-) -> list[str]:
-    remembered = st.session_state.get(key, st.session_state.get(remember_key, []))
-    if not isinstance(remembered, list):
-        remembered = []
-    return [value for value in remembered if value in options]
-
-
-def _release_vector_memory(store: Store, variants: list[str] | tuple[str, ...]) -> None:
-    release_hf_store_cache(store, variants)
-    gc.collect()
-
-
-def _evenly_spaced_layers(layers: list[int], max_count: int) -> list[int]:
-    if max_count >= len(layers):
-        return layers
-    if max_count <= 1:
-        return [layers[0]]
-
-    last = len(layers) - 1
-    indices = [round(i * last / (max_count - 1)) for i in range(max_count)]
-    return [layers[index] for index in dict.fromkeys(indices)]
-
-
-def _render_layer_frame_controls(
-    store: Store,
-    scope: str,
-    layers: list[int],
-) -> list[int]:
-    if len(layers) <= _DEFAULT_LAYER_FRAMES:
-        st.caption(f"Using all {len(layers)} available layer(s).")
-        return layers
-
-    frame_count = st.slider(
-        "Layer frames",
-        min_value=2,
-        max_value=len(layers),
-        value=min(
-            max(
-                int(
-                    st.session_state.get(
-                        _LAST_LAYER_FRAMES_KEY,
-                        _DEFAULT_LAYER_FRAMES,
-                    )
-                ),
-                2,
-            ),
-            len(layers),
-        ),
-        key=widget_key("load", "layer_frames", scope, store_id(store)),
-        help="Limit animated Plotly frames to keep browser and RAM usage bounded.",
-    )
-    st.session_state[_LAST_LAYER_FRAMES_KEY] = frame_count
-    selected = _evenly_spaced_layers(layers, frame_count)
-    st.caption(f"Using {len(selected)} of {len(layers)} layers.")
-    return selected
 
-
-def _load_persona_options(
-    store: Store,
-    variants: list[str],
-    mask_strategy: MaskStrategy,
-    *,
-    empty_message: str,
-) -> PersonaOptions | None:
-    source, location, model_name = store_cache_parts(store)
-    variant_key = tuple(variants)
-    persona_ids = personas_cached(
-        source,
-        location,
-        model_name,
-        mask_strategy.value,
-        variant_key,
-        include_baseline=True,
-    )
-    if not persona_ids:
-        st.info(empty_message)
-        return None
-
-    persona_names = persona_names_cached(
-        source,
-        location,
-        model_name,
-        mask_strategy.value,
-        variant_key,
-        tuple(persona_ids),
-    )
-    assistant_ids = [
-        persona_id
-        for persona_id in persona_ids
-        if _is_assistant_persona(persona_id, persona_names.get(persona_id))
-    ]
-    assistant_id = next(
-        (
-            persona_id
-            for persona_id in assistant_ids
-            if persona_id == BASELINE_PERSONA_ID
-        ),
-        assistant_ids[0] if assistant_ids else None,
-    )
-    regular_ids = [
-        persona_id for persona_id in persona_ids if persona_id not in assistant_ids
-    ]
-    if not regular_ids and assistant_id is None:
-        st.info("No personas found for this model and variant.")
-        return None
-    return PersonaOptions(
-        regular_ids=regular_ids,
-        assistant_id=assistant_id,
-        persona_names=persona_names,
-    )
-
-
-def _seed_persona_memory(
-    remember_key: str,
-    options: PersonaOptions,
-    *,
-    default_all: bool,
-    default_count_limit: int | None = None,
-) -> tuple[int, bool]:
-    remembered_count_key = f"{remember_key}:count"
-    remembered_assistant_key = f"{remember_key}:include_assistant"
-    legacy_ids = st.session_state.get(remember_key, [])
-    if isinstance(legacy_ids, list) and legacy_ids:
-        st.session_state.setdefault(
-            remembered_count_key,
-            sum(persona_id in options.regular_ids for persona_id in legacy_ids),
-        )
-        st.session_state.setdefault(
-            remembered_assistant_key,
-            options.assistant_id in legacy_ids,
-        )
-
-    if default_count_limit is not None:
-        default_count = min(default_count_limit, len(options.regular_ids))
-    elif default_all:
-        default_count = len(options.regular_ids)
-    else:
-        default_count = min(1, len(options.regular_ids))
-    remembered_count = int(st.session_state.get(remembered_count_key, default_count))
-    persona_count = min(max(remembered_count, 0), len(options.regular_ids))
-    include_assistant = bool(st.session_state.get(remembered_assistant_key, False))
-    return persona_count, include_assistant
-
-
-def _render_persona_count_controls(
-    store: Store,
-    variants: list[str],
-    mask_strategy: MaskStrategy,
-    widget_scope: str,
-    options: PersonaOptions,
-    *,
-    default_count: int,
-    include_assistant_default: bool,
-) -> tuple[int, bool]:
-    count_key = widget_key(
-        "load",
-        "persona_count",
-        widget_scope,
-        store.model_name,
-        mask_strategy.value,
-        *variants,
-    )
-    assistant_key = widget_key(
-        "load",
-        "include_assistant",
-        widget_scope,
-        store.model_name,
-        mask_strategy.value,
-        *variants,
-    )
-
-    if options.regular_ids:
-        persona_count = st.slider(
-            "Personas",
-            min_value=0 if options.assistant_id is not None else 1,
-            max_value=len(options.regular_ids),
-            value=default_count,
-            key=count_key,
-            help="Use the first N available non-assistant personas.",
-        )
-    else:
-        persona_count = 0
-        st.caption("No non-assistant personas are available for this selection.")
-    include_assistant = False
-    if options.assistant_id is not None:
-        include_assistant = st.checkbox(
-            "Include Assistant persona",
-            value=include_assistant_default,
-            key=assistant_key,
-        )
-    return persona_count, include_assistant
-
-
-def _select_artifact_personas(
-    store: Store,
-    variants: list[str],
-    mask_strategy: MaskStrategy,
-    *,
-    widget_scope: str,
-    remember_key: str,
-    default_all: bool = False,
-    default_count_limit: int | None = None,
-) -> list[str]:
-    empty_message = (
-        "No personas have vectors for all selected variants. "
-        "Pick a single variant or change the source."
-        if len(variants) > 1
-        else "No personas found for this model and variant."
-    )
-    options = _load_persona_options(
-        store,
-        variants,
-        mask_strategy,
-        empty_message=empty_message,
-    )
-    if options is None:
-        st.session_state.pop(_persona_names_state_key(widget_scope), None)
-        return []
-
-    default_count, include_assistant_default = _seed_persona_memory(
-        remember_key,
-        options,
-        default_all=default_all,
-        default_count_limit=default_count_limit,
-    )
-    persona_count, include_assistant = _render_persona_count_controls(
-        store,
-        variants,
-        mask_strategy,
-        widget_scope,
-        options,
-        default_count=default_count,
-        include_assistant_default=include_assistant_default,
-    )
-
-    persona_ids = options.regular_ids[:persona_count]
-    if include_assistant and options.assistant_id is not None:
-        persona_ids.append(options.assistant_id)
-
-    remembered_count_key = f"{remember_key}:count"
-    remembered_assistant_key = f"{remember_key}:include_assistant"
-    st.session_state[remembered_count_key] = persona_count
-    st.session_state[remembered_assistant_key] = include_assistant
-    st.session_state[remember_key] = persona_ids
-    st.session_state[_persona_names_state_key(widget_scope)] = options.persona_names
-
-    if not persona_ids:
-        st.info("Select at least one persona or include the Assistant persona.")
-        return []
-
-    regular_label = f"{persona_count} persona{'s' if persona_count != 1 else ''}"
-    assistant_label = (
-        " plus Assistant" if include_assistant and options.assistant_id else ""
-    )
-    st.caption(f"Using {regular_label}{assistant_label}.")
-    return persona_ids
-
-
-def _render_save_buttons(
-    figs: list[object],
-    filenames: list[str],
-    key_suffix: str,
-) -> None:
-    """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, strict=True)
-            ]
-            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 = active_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",
-        config={"responsive": True, "displaylogo": False},
-    )
-
-
-def _render_mask_strategy_select(scope: str) -> MaskStrategy:
-    return render_mask_strategy_select(
-        key=widget_key("load", "mask_strategy", scope),
-        last_key=_LAST_MASK_STRATEGY_KEY,
-        help_text="Which extracted activation set to load.",
-    )
-
-
-def _render_cosine_selection(
-    store: Store,
-    mask_strategy: MaskStrategy,
-) -> CosineSelection | None:
-    variants = available_variants(store, mask_strategy)
-    if len(variants) < 2:
-        st.info("Need at least two variants with saved vectors for cosine comparison.")
-        return None
-
-    with st.expander("Vector selection", expanded=True):
-        col1, col2 = st.columns(2)
-        with col1:
-            variant_a = st.selectbox(
-                "Variant A",
-                options=variants,
-                index=0,
-                format_func=prompt_variant_label,
-                key=widget_key("load", "variant_a", store_id(store)),
-            )
-        with col2:
-            variant_b = st.selectbox(
-                "Variant B",
-                options=variants,
-                index=min(1, len(variants) - 1),
-                format_func=prompt_variant_label,
-                key=widget_key("load", "variant_b", store_id(store)),
-            )
-
-        if variant_a == variant_b:
-            st.warning("Choose two different variants to compare.")
-            return None
-
-        persona_ids = _select_artifact_personas(
-            store,
-            [variant_a, variant_b],
-            mask_strategy,
-            widget_scope=f"cosine:{store_id(store)}",
-            remember_key=_LAST_COSINE_PERSONAS_KEY,
-        )
-    if not persona_ids:
-        return None
-    return CosineSelection(
-        variants=variants,
-        variant_a=variant_a,
-        variant_b=variant_b,
-        persona_ids=persona_ids,
-        persona_key=personas_fingerprint(persona_ids),
-    )
-
-
-def _build_cosine_figures(
-    store: Store,
-    mask_strategy: MaskStrategy,
-    selection: CosineSelection,
-) -> tuple[object, object | None, int, int] | None:
-    variant_sample_cache: dict[str, object] = {}
-
-    def _load_variant(variant: str):
-        if variant not in variant_sample_cache:
-            samples = _load_variant_vectors(
-                store,
-                [variant],
-                mask_strategy,
-                persona_ids=selection.persona_ids,
-            )
-            variant_sample_cache[variant] = samples[variant]
-        return variant_sample_cache[variant]
-
-    try:
-        samples_a = _load_variant(selection.variant_a)
-        samples_b = _load_variant(selection.variant_b)
-    except Exception as exc:
-        st.error(f"Could not load vectors: {exc}")
-        return None
-
-    labels = samples_a.labels
-    display_traces = [
-        (
-            label,
-            samples_a.vectors[index],
-            samples_b.vectors[index],
-        )
-        for index, label in enumerate(labels)
-    ]
-    fig = plot_layer_similarity(
-        display_traces,
-        title=(
-            f"{prompt_variant_label(selection.variant_a)} vs "
-            f"{prompt_variant_label(selection.variant_b)}"
-        ),
-        show=False,
-    )
-
-    pair_traces = []
-    pair_errors = []
-    for left, right in combinations(selection.variants, 2):
-        try:
-            left_samples = _load_variant(left)
-            right_samples = _load_variant(right)
-            pair_traces.append(
-                (
-                    f"{prompt_variant_label(left)} vs {prompt_variant_label(right)}",
-                    left_samples.vectors.mean(dim=0),
-                    right_samples.vectors.mean(dim=0),
-                )
-            )
-        except Exception as exc:
-            pair_errors.append(f"{left} vs {right}: {exc}")
-            continue
-
-    for err in pair_errors:
-        st.warning(f"Skipped pair trace: `{err}`")
-    pair_fig = (
-        plot_layer_similarity(
-            pair_traces,
-            title="Variant-pair cosine similarity averaged over selected personas",
-            show=False,
-        )
-        if pair_traces
-        else None
-    )
-    return fig, pair_fig, len(display_traces), len(pair_traces)
-
-
-def _render_cosine_similarity(
-    store: Store,
-    mask_strategy: MaskStrategy,
-) -> None:
-    selection = _render_cosine_selection(store, mask_strategy)
-    if selection is None:
-        return
-
-    cosine_fig_key = widget_key(
-        "load",
-        "cosine_fig_state",
-        store_id(store),
-        store.model_name,
-        mask_strategy.value,
-        selection.variant_a,
-        selection.variant_b,
-        selection.persona_key,
-    )
-    filename = _filename(
-        "analysis",
-        "cosine",
-        store.model_name,
-        mask_strategy.value,
-        selection.variant_a,
-        selection.variant_b,
-    )
-    pairs_filename = _filename(
-        "analysis",
-        "cosine_pairs",
-        store.model_name,
-        mask_strategy.value,
-        "_".join(selection.variants),
-    )
-    _clear_old_figure_states(cosine_fig_key)
-
-    if st.button(
-        "Compare vectors",
-        type="primary",
-        key=widget_key(
-            "load",
-            "analysis_vectors",
-            store_id(store),
-            store.model_name,
-            mask_strategy.value,
-            selection.variant_a,
-            selection.variant_b,
-            selection.persona_key,
-        ),
-    ):
-        progress = st.progress(0, text="Loading activation vectors…")
-        try:
-            progress.progress(15, text="Loading activation vectors…")
-            figures = _build_cosine_figures(store, mask_strategy, selection)
-            if figures is None:
-                st.session_state.pop(cosine_fig_key, None)
-                return
-            progress.progress(90, text="Storing figure state…")
-            _store_figure_state(cosine_fig_key, figures)
-            progress.progress(100, text="Done.")
-        finally:
-            _release_vector_memory(store, selection.variants)
-            progress.empty()
-
-    if cosine_fig_key in st.session_state:
-        fig, pair_fig, n_traces, n_pair_traces = st.session_state[cosine_fig_key]
-        _plotly_chart(fig)
-        figs = [fig]
-        filenames = [filename]
-        if pair_fig is not None:
-            st.subheader("Variant pairs")
-            _plotly_chart(pair_fig)
-            figs.append(pair_fig)
-            filenames.append(pairs_filename)
-        _render_save_buttons(figs, filenames, "cosine")
-        st.success(f"Loaded {n_traces} personas for cosine comparison.")
-        if pair_fig is not None:
-            st.caption(f"Generated {n_pair_traces} averaged variant-pair trace(s).")
-
-
-def _select_single_variant_samples(
-    store: Store,
-    mask_strategy: MaskStrategy,
-    scope: str,
-    *,
-    remember_key: str,
-    variant_remember_key: str,
-    default_count_limit: int,
-) -> tuple[str, list[str], str, list[int]] | None:
-    variants = available_variants(store, mask_strategy)
-    if not variants:
-        st.info("No variants with saved vectors for this model.")
-        return None
-    variant_key = widget_key("load", "variant", scope, store_id(store))
-    default_variant = "biography" if "biography" in variants else variants[0]
-    selected_variant = _seed_selectbox_key(
-        key=variant_key,
-        remember_key=variant_remember_key,
-        options=variants,
-        default=default_variant,
-    )
-    variant = st.selectbox(
-        "Variant",
-        options=variants,
-        index=variants.index(selected_variant),
-        format_func=prompt_variant_label,
-        key=variant_key,
-    )
-    st.session_state[variant_remember_key] = variant
-    persona_ids = _select_artifact_personas(
-        store,
-        [variant],
-        mask_strategy,
-        widget_scope=f"{scope}:{store_id(store)}",
-        remember_key=remember_key,
-        default_count_limit=default_count_limit,
-    )
-    if not persona_ids:
-        return None
-
-    persona_key = personas_fingerprint(persona_ids)
-    layer_options = _layers_for_variant(store, variant, persona_ids, mask_strategy)
-    if not layer_options:
-        st.info("No shared layers are available for the selected personas.")
-        return None
-
-    selected_layers = _render_layer_frame_controls(store, scope, layer_options)
-    return variant, persona_ids, persona_key, selected_layers
-
-
-def _render_pair_trajectory_control(
-    *,
-    enabled: bool,
-    persona_count: int,
-    scope: str,
-    store: Store,
-) -> bool:
-    if not enabled:
-        return False
-    pair_count = persona_count * (persona_count - 1) // 2
-    if pair_count > _MAX_PAIR_TRAJECTORY_TRACES:
-        st.caption(
-            "Pair trajectories hidden because this selection would create "
-            f"{pair_count:,} Plotly traces."
-        )
-        return False
-    return st.checkbox(
-        "Pair trajectories",
-        value=False,
-        key=widget_key("load", "pair_trajectories", scope, store_id(store)),
-        help="Adds one line per persona pair. Keep this off for larger selections.",
-    )
-
-
-def _validate_layered_figure_size(
-    figure_kind: str,
-    persona_count: int,
-    selected_layers: list[int],
-) -> bool:
-    if figure_kind != "similarity":
-        return True
-    similarity_cells = persona_count * persona_count * len(selected_layers)
-    if similarity_cells <= _MAX_SIMILARITY_CELLS:
-        return True
-    st.error(
-        "Reduce personas or layer frames before generating the similarity "
-        f"matrix ({similarity_cells:,} cells selected)."
-    )
-    return False
-
-
-def _render_projection_color_config(
-    store: Store,
-    scope: str,
-    persona_ids: list[str],
-) -> ProjectionColorConfig | None:
-    widget_scope = f"{scope}:{store_id(store)}"
-    persona_key = personas_fingerprint(persona_ids)
-    persona_names = st.session_state.get(
-        _persona_names_state_key(widget_scope),
-        {},
-    )
-    color_mode_key = widget_key("load", "color_mode", scope, store_id(store))
-    selected_color_mode = _seed_selectbox_key(
-        key=color_mode_key,
-        remember_key=_LAST_PROJECTION_COLOR_MODE_KEY,
-        options=_PROJECTION_COLOR_MODES,
-        default="Persona",
-    )
-    color_mode = st.selectbox(
-        "Color by",
-        options=_PROJECTION_COLOR_MODES,
-        index=_PROJECTION_COLOR_MODES.index(selected_color_mode),
-        key=color_mode_key,
-    )
-    st.session_state[_LAST_PROJECTION_COLOR_MODE_KEY] = color_mode
-    if color_mode == "K-means clusters":
-        max_clusters = min(10, len(persona_ids))
-        if max_clusters < 2:
-            st.info("Select at least two personas to use K-means coloring.")
-            return None
-        cluster_key = widget_key("load", "cluster_k", scope, store_id(store))
-        default_clusters = min(3, len(persona_ids))
-        if cluster_key not in st.session_state:
-            st.session_state[cluster_key] = min(
-                max(
-                    int(
-                        st.session_state.get(
-                            _LAST_PROJECTION_CLUSTER_K_KEY,
-                            default_clusters,
-                        )
-                    ),
-                    2,
-                ),
-                max_clusters,
-            )
-        n_clusters = st.slider(
-            "K (clusters)",
-            min_value=2,
-            max_value=max_clusters,
-            key=cluster_key,
-        )
-        mode_key = widget_key("load", "cluster_mode", scope, store_id(store))
-        mode_options = list(_CLUSTER_MODES)
-        selected_mode = _seed_selectbox_key(
-            key=mode_key,
-            remember_key=_LAST_PROJECTION_CLUSTER_MODE_KEY,
-            options=mode_options,
-            default=mode_options[0],
-        )
-        mode_label = st.selectbox(
-            "Cluster fit",
-            options=mode_options,
-            index=mode_options.index(selected_mode),
-            key=mode_key,
-            help=(
-                "Mean across layers is the previous behavior. First selected "
-                "layer keeps one fixed clustering from the first frame. Per layer "
-                "recomputes clustering for each animation frame."
-            ),
-        )
-        st.session_state[_LAST_PROJECTION_CLUSTER_K_KEY] = n_clusters
-        st.session_state[_LAST_PROJECTION_CLUSTER_MODE_KEY] = mode_label
-        return ProjectionColorConfig(
-            color_mode=color_mode,
-            n_clusters=n_clusters,
-            cluster_mode=_CLUSTER_MODES[mode_label],
-        )
-
-    if color_mode == "Persona attribute":
-        persona_dataset = synth_persona_dataset_cached()
-        attribute_options = list(synth_persona_attribute_names())
-        if not attribute_options:
-            st.info("No persona attributes are available for this dataset.")
-            return None
-        default_attribute = (
-            attribute_options.index("sex") if "sex" in attribute_options else 0
-        )
-        attribute_key = widget_key("load", "attribute", scope, store_id(store))
-        selected_attribute = _seed_selectbox_key(
-            key=attribute_key,
-            remember_key=_LAST_PROJECTION_ATTRIBUTE_KEY,
-            options=attribute_options,
-            default=attribute_options[default_attribute],
-        )
-        attribute_name = st.selectbox(
-            "Attribute",
-            options=attribute_options,
-            index=attribute_options.index(selected_attribute),
-            format_func=lambda name: attribute_display_label(persona_dataset, name),
-            key=attribute_key,
-        )
-        st.session_state[_LAST_PROJECTION_ATTRIBUTE_KEY] = attribute_name
-        info = persona_dataset.attribute_info(attribute_name)
-        if info.get("high_cardinality"):
-            st.caption(
-                "High-cardinality categorical attributes are grouped to the "
-                f"top {_MAX_ATTRIBUTE_CATEGORIES} values plus Other."
-            )
-        return ProjectionColorConfig(
-            color_mode=color_mode,
-            attribute_name=attribute_name,
-        )
-
-    highlight_persona_ids: tuple[str, ...] = ()
-    if persona_ids:
-        highlight_key = widget_key(
-            "load", "persona_highlight", scope, store_id(store), persona_key
-        )
-        highlighted = st.multiselect(
-            "Highlight personas",
-            options=persona_ids,
-            default=_remember_multiselect(
-                key=highlight_key,
-                remember_key=_LAST_PROJECTION_HIGHLIGHTS_KEY,
-                options=persona_ids,
-            ),
-            format_func=lambda persona_id: _persona_display_label(
-                persona_names, persona_id
-            ),
-            key=highlight_key,
-            help=(
-                "Select a few personas to keep their default colors while the rest "
-                "are grayed out."
-            ),
-        )
-        highlight_persona_ids = tuple(highlighted)
-        st.session_state[_LAST_PROJECTION_HIGHLIGHTS_KEY] = list(highlighted)
-
-    highlight_persona_key = (
-        personas_fingerprint(highlight_persona_ids) if highlight_persona_ids else ""
-    )
-
-    return ProjectionColorConfig(
-        color_mode=color_mode,
-        highlight_persona_ids=highlight_persona_ids,
-        highlight_persona_key=highlight_persona_key,
-    )
-
-
-def _layered_figure_state_keys(
-    store: Store,
-    mask_strategy: MaskStrategy,
-    *,
-    scope: str,
-    figure_kind: str,
-    n_components: int,
-    color_config: ProjectionColorConfig,
-    variant: str,
-    persona_key: str,
-    selected_layers: list[int],
-    pair_trajectories: bool,
-) -> LayeredFigureStateKeys:
-    layer_key = "_".join(map(str, selected_layers))
-    figure_key = widget_key(
-        "load",
-        f"{scope}_fig_state",
-        store_id(store),
-        store.model_name,
-        mask_strategy.value,
-        figure_kind,
-        str(n_components),
-        color_config.color_mode,
-        str(color_config.attribute_name),
-        str(color_config.n_clusters),
-        str(color_config.cluster_mode),
-        str(color_config.highlight_persona_key),
-        variant,
-        "persona_vector",
-        persona_key,
-        layer_key,
-        str(pair_trajectories),
-    )
-    if figure_kind not in _PROJECTION_KINDS:
-        return LayeredFigureStateKeys(figure=figure_key)
-
-    graph_overlay = figure_kind == "isomap"
-    projection_key = widget_key(
-        "load",
-        f"{scope}_projection_state",
-        store_id(store),
-        store.model_name,
-        mask_strategy.value,
-        figure_kind,
-        str(n_components),
-        str(graph_overlay),
-        str(_DEFAULT_GRAPH_NEIGHBORS),
-        variant,
-        "persona_vector",
-        persona_key,
-        layer_key,
-    )
-    return LayeredFigureStateKeys(figure=figure_key, projection=projection_key)
-
-
-def _projection_build_kwargs(
-    samples,
-    *,
-    figure_kind: str,
-    selected_layers: list[int],
-    n_components: int,
-    color_config: ProjectionColorConfig,
-    persona_ids: list[str],
-    persona_names: dict[str, str],
-    projection_key: str | None,
-) -> dict:
-    if figure_kind not in _PROJECTION_KINDS:
-        return {}
-
-    graph_overlay = figure_kind == "isomap"
-    build_kwargs = {
-        "n_components": n_components,
-        "graph_overlay": graph_overlay,
-        "graph_n_neighbors": _DEFAULT_GRAPH_NEIGHBORS,
-    }
-    if color_config.n_clusters is not None:
-        build_kwargs["n_clusters"] = color_config.n_clusters
-        build_kwargs["cluster_mode"] = color_config.cluster_mode
-    if projection_key is not None:
-        projection_data = st.session_state.get(projection_key)
-        if projection_data is None:
-            projection_data = prepare_layered_projection_data(
-                samples,
-                figure_kind,
-                layers=selected_layers,
-                n_components=n_components,
-                graph_overlay=graph_overlay,
-                graph_n_neighbors=_DEFAULT_GRAPH_NEIGHBORS,
-            )
-            st.session_state[projection_key] = projection_data
-        build_kwargs["projection_data"] = projection_data
-    if color_config.attribute_name is not None:
-        build_kwargs.update(
-            attribute_color_kwargs(
-                synth_persona_dataset_cached(),
-                color_config.attribute_name,
-                persona_ids,
-                max_categories=_MAX_ATTRIBUTE_CATEGORIES,
-            )
-        )
-    if color_config.color_mode == "Persona" and color_config.highlight_persona_ids:
-        groups = _highlight_persona_groups(
-            persona_ids,
-            persona_names,
-            color_config.highlight_persona_ids,
-        )
-        if groups is not None:
-            build_kwargs["groups"] = groups
-    return build_kwargs
-
-
-def _build_layered_analysis_figures(
-    samples,
-    *,
-    figure_kind: str,
-    selected_layers: list[int],
-    variant: str,
-    title_fn: Callable[[str], str],
-    pair_trajectories: bool,
-    build_kwargs: dict,
-) -> tuple[go.Figure, go.Figure | None]:
-    if figure_kind == "similarity" and pair_trajectories:
-        return build_similarity_figures(
-            samples,
-            layers=selected_layers,
-            title=title_fn(variant),
-            pair_title=(
-                "Pair similarity trajectories - "
-                f"{prompt_variant_label(variant)} - persona vectors"
-            ),
-        )
-
-    main_fig = build_layered_figure(
-        samples,
-        figure_kind,
-        layers=selected_layers,
-        title=title_fn(variant),
-        **build_kwargs,
-    )
-    if figure_kind == "isomap":
-        _add_isomap_connection_toggle(main_fig)
-    if figure_kind in _PROJECTION_KINDS:
-        main_fig.update_layout(height=700)
-    extra_fig = (
-        build_pair_similarity_figure(
-            samples,
-            layers=selected_layers,
-            title=(
-                "Pair similarity trajectories - "
-                f"{prompt_variant_label(variant)} - persona vectors"
-            ),
-        )
-        if pair_trajectories
-        else None
-    )
-    return main_fig, extra_fig
-
-
-def _add_isomap_connection_toggle(fig: go.Figure) -> None:
-    """Add an in-plot control for the Isomap kNN graph trace."""
-    if not fig.data or fig.data[0].name != "kNN graph":
-        return
-
-    existing_menus = tuple(fig.layout.updatemenus or ())
-    fig.update_layout(
-        updatemenus=existing_menus
-        + (
-            dict(
-                type="buttons",
-                direction="left",
-                active=0,
-                showactive=False,
-                x=0,
-                xanchor="left",
-                y=1.16,
-                yanchor="top",
-                pad=dict(t=0, r=10),
-                buttons=[
-                    dict(
-                        label="Show connections",
-                        method="restyle",
-                        args=[{"visible": True}, [0]],
-                    ),
-                    dict(
-                        label="Hide connections",
-                        method="restyle",
-                        args=[{"visible": False}, [0]],
-                    ),
-                ],
-            ),
-        ),
-    )
-
-
-def _render_layered_figure_analysis(
-    store: Store,
-    mask_strategy: MaskStrategy,
-    *,
-    scope: str,
-    figure_kind: str,
-    button_label: str,
-    title_fn: Callable[[str], str],
-    include_pair_trajectories: bool = False,
-    n_components: int = 2,
-    remember_key: str = _LAST_PROJECTION_PERSONAS_KEY,
-    default_count_limit: int = 500,
-) -> None:
-    """Render a single-variant layered analysis: select → button → figure(s).
-
-    Used for similarity matrix, PCA, and UMAP. Set ``include_pair_trajectories``
-    to add the pair-similarity-trajectory figure (similarity matrix only).
-    """
-    selected = _select_single_variant_samples(
-        store,
-        mask_strategy,
-        scope,
-        remember_key=remember_key,
-        variant_remember_key=(
-            _LAST_PROJECTION_VARIANT_KEY
-            if figure_kind in _PROJECTION_KINDS
-            else _LAST_SIMILARITY_VARIANT_KEY
-        ),
-        default_count_limit=default_count_limit,
-    )
-    if selected is None:
-        return
-    variant, persona_ids, persona_key, selected_layers = selected
-
-    pair_trajectories = _render_pair_trajectory_control(
-        enabled=include_pair_trajectories,
-        persona_count=len(persona_ids),
-        scope=scope,
-        store=store,
-    )
-    if not _validate_layered_figure_size(
-        figure_kind, len(persona_ids), selected_layers
-    ):
-        return
-
-    color_config = ProjectionColorConfig()
-    if figure_kind in _PROJECTION_KINDS:
-        color_config = _render_projection_color_config(store, scope, persona_ids)
-        if color_config is None:
-            return
-
-    state_keys = _layered_figure_state_keys(
-        store,
-        mask_strategy,
-        scope=scope,
-        figure_kind=figure_kind,
-        n_components=n_components,
-        color_config=color_config,
-        variant=variant,
-        persona_key=persona_key,
-        selected_layers=selected_layers,
-        pair_trajectories=pair_trajectories,
-    )
-    if state_keys.projection is not None:
-        _clear_old_projection_states(state_keys.projection)
-    filename = scope
-    _clear_old_figure_states(state_keys.figure)
-    persona_names = st.session_state.get(
-        _persona_names_state_key(f"{scope}:{store_id(store)}"),
-        {},
-    )
-
-    if st.button(button_label, type="primary"):
-        build_label = {
-            "umap": "Computing UMAP projections…",
-            "pca": "Computing PCA projections…",
-            "isomap": "Computing Isomap projections…",
-            "similarity": "Computing similarity matrices…",
-        }.get(figure_kind, "Building figure…")
-        progress = st.progress(0, text="Loading activation vectors…")
-        try:
-            progress.progress(15, text="Loading activation vectors…")
-            samples = _load_persona_vectors(
-                store,
-                variant,
-                mask_strategy,
-                persona_ids,
-            )
-            progress.progress(55, text=build_label)
-            build_kwargs = _projection_build_kwargs(
-                samples,
-                figure_kind=figure_kind,
-                selected_layers=selected_layers,
-                n_components=n_components,
-                color_config=color_config,
-                persona_ids=persona_ids,
-                persona_names=persona_names,
-                projection_key=state_keys.projection,
-            )
-            main_fig, extra_fig = _build_layered_analysis_figures(
-                samples,
-                figure_kind=figure_kind,
-                selected_layers=selected_layers,
-                variant=variant,
-                title_fn=title_fn,
-                pair_trajectories=pair_trajectories,
-                build_kwargs=build_kwargs,
-            )
-            if (
-                color_config.color_mode == "Persona"
-                and color_config.highlight_persona_ids
-            ):
-                _gray_out_unselected_personas(main_fig)
-            progress.progress(90, text="Storing figure state…")
-            n_samples = samples.vectors.shape[0]
-            del samples
-            _store_figure_state(state_keys.figure, (main_fig, extra_fig, n_samples))
-            progress.progress(100, text="Done.")
-        except Exception as exc:
-            st.error(f"Could not build figure: {exc}")
-            st.session_state.pop(state_keys.figure, None)
-        finally:
-            _release_vector_memory(store, [variant])
-            progress.empty()
-
-    if state_keys.figure in st.session_state:
-        main_fig, extra_fig, n_samples = st.session_state[state_keys.figure]
-        _plotly_chart(main_fig)
-        figs = [main_fig]
-        filenames = [filename]
-        if extra_fig is not None:
-            st.subheader("Pair trajectories")
-            _plotly_chart(extra_fig)
-            figs.append(extra_fig)
-            filenames.append(f"{filename}__pair_trajectories")
-        _render_save_buttons(figs, filenames, scope)
-        st.success(f"Loaded {n_samples} samples.")
-
-
-_LAST_DENDRO_PERSONAS_KEY = "analysis:last_personas:dendro"
-_DENDRO_LINKAGE_OPTIONS = ["ward", "complete", "average", "single"]
-
-
-def _render_persona_select_controls(
-    options: PersonaOptions,
-    widget_scope: str,
-) -> list[str]:
-    select_key = widget_key("load", "persona_select", widget_scope)
-    assistant_key = widget_key("load", "persona_select_assistant", widget_scope)
-
-    label_map = {
-        pid: f"{options.persona_names.get(pid, pid)} ({pid})"
-        for pid in options.regular_ids
-    }
-    sorted_labels = sorted(label_map.values())
-    selected_labels = st.multiselect(
-        "Select personas",
-        options=sorted_labels,
-        key=select_key,
-        placeholder="Search and select personas...",
-    )
-    label_to_id = {v: k for k, v in label_map.items()}
-    selected_ids = [label_to_id[lbl] for lbl in selected_labels]
-
-    if options.assistant_id is not None:
-        include_assistant = st.checkbox(
-            "Include Assistant persona",
-            key=assistant_key,
-        )
-        if include_assistant:
-            selected_ids.append(options.assistant_id)
-
-    st.session_state[_persona_names_state_key(widget_scope)] = dict(
-        options.persona_names
-    )
-
-    if not selected_ids:
-        st.info("Select at least one persona.")
-
-    return selected_ids
-
-
-def _render_dendrogram_analysis(
-    store: Store,
-    mask_strategy: MaskStrategy,
-) -> None:
-    variants = available_variants(store, mask_strategy)
-    if not variants:
-        st.info("No variants with saved vectors for this model.")
-        return
-
-    with st.expander("Variant selection", expanded=True):
-        col1, col2 = st.columns(2)
-        default_a = "biography" if "biography" in variants else variants[0]
-        default_b_idx = (
-            variants.index("templated")
-            if "templated" in variants
-            else min(1, len(variants) - 1)
-        )
-        with col1:
-            variant_a = st.selectbox(
-                "Variant A",
-                options=variants,
-                index=variants.index(default_a),
-                format_func=prompt_variant_label,
-                key=widget_key("load", "dendro_variant_a", store_id(store)),
-            )
-        with col2:
-            variant_b = st.selectbox(
-                "Variant B",
-                options=variants,
-                index=default_b_idx,
-                format_func=prompt_variant_label,
-                key=widget_key("load", "dendro_variant_b", store_id(store)),
-            )
-
-    shared_variants = list(dict.fromkeys([variant_a, variant_b]))
-
-    select_specific = st.toggle(
-        "Select specific personas",
-        value=False,
-        key=widget_key("load", "dendro_select_mode", store_id(store)),
-        help="Search and select specific personas instead of using the first N.",
-    )
-
-    if select_specific:
-        empty_message = (
-            "No personas have vectors for all selected variants. "
-            "Pick a single variant or change the source."
-            if len(shared_variants) > 1
-            else "No personas found for this model and variant."
-        )
-        options = _load_persona_options(
-            store,
-            shared_variants,
-            mask_strategy,
-            empty_message=empty_message,
-        )
-        if options is None:
-            st.session_state.pop(
-                _persona_names_state_key(f"dendro:{store_id(store)}"), None
-            )
-            return
-        persona_ids = _render_persona_select_controls(
-            options,
-            widget_scope=f"dendro:{store_id(store)}",
-        )
-        if not persona_ids:
-            return
-    else:
-        persona_ids = _select_artifact_personas(
-            store,
-            shared_variants,
-            mask_strategy,
-            widget_scope=f"dendro:{store_id(store)}",
-            remember_key=_LAST_DENDRO_PERSONAS_KEY,
-            default_count_limit=_DEFAULT_PERSONA_LIMITS["dendro"],
-        )
-        if not persona_ids:
-            return
-
-    col_opts1, col_opts2 = st.columns(2)
-    with col_opts1:
-        layered_mode = st.toggle(
-            "Per-layer animated",
-            value=False,
-            key=widget_key("load", "dendro_layered", store_id(store)),
-            help="Animated dendrogram with one frame per layer instead of averaging all layers.",
-        )
-    with col_opts2:
-        linkage = st.selectbox(
-            "Linkage",
-            options=_DENDRO_LINKAGE_OPTIONS,
-            index=0,
-            key=widget_key("load", "dendro_linkage", store_id(store)),
-        )
-
-    selected_layers: list[int] | None = None
-    if layered_mode:
-        source, location, model_name = store_cache_parts(store)
-        layer_options = store_layers_cached(
-            source,
-            location,
-            model_name,
-            mask_strategy.value,
-            tuple(shared_variants),
-            tuple(persona_ids),
-        )
-        if not layer_options:
-            st.info("No shared layers are available for the selected personas.")
-            return
-        selected_layers = _render_layer_frame_controls(store, "dendro", layer_options)
-
-    persona_key = personas_fingerprint(persona_ids)
-    fig_key = widget_key(
-        "load",
-        "dendro_fig_state",
-        store_id(store),
-        store.model_name,
-        mask_strategy.value,
-        variant_a,
-        variant_b,
-        persona_key,
-        str(layered_mode),
-        linkage,
-        "_".join(map(str, selected_layers or [])),
-    )
-    _clear_old_figure_states(fig_key)
-
-    if st.button(
-        "Generate dendrograms",
-        type="primary",
-        key=widget_key(
-            "load", "dendro_btn", store_id(store), variant_a, variant_b, persona_key
-        ),
-    ):
-        progress = st.progress(0, text="Loading first variant vectors…")
-        try:
-            progress.progress(15, text="Loading first variant vectors…")
-            samples_a = _load_persona_vectors(
-                store,
-                variant_a,
-                mask_strategy,
-                persona_ids,
-            )
-            progress.progress(40, text="Building first dendrogram…")
-            fig_a = plot_persona_dendrogram(
-                samples_a,
-                layered=layered_mode,
-                layers=selected_layers,
-                linkage=linkage,
-                title=f"Dendrogram — {prompt_variant_label(variant_a)}",
-            )
-            fig_a.update_layout(height=750)
-            del samples_a
-            fig_b = None
-            if variant_a != variant_b:
-                progress.progress(60, text="Loading second variant vectors…")
-                samples_b = _load_persona_vectors(
-                    store,
-                    variant_b,
-                    mask_strategy,
-                    persona_ids,
-                )
-                progress.progress(75, text="Building second dendrogram…")
-                fig_b = plot_persona_dendrogram(
-                    samples_b,
-                    layered=layered_mode,
-                    layers=selected_layers,
-                    linkage=linkage,
-                    title=f"Dendrogram — {prompt_variant_label(variant_b)}",
-                )
-                fig_b.update_layout(height=750)
-                del samples_b
-            progress.progress(90, text="Storing figure state…")
-            _store_figure_state(
-                fig_key,
-                (fig_a, fig_b, len(persona_ids), variant_a, variant_b),
-            )
-            progress.progress(100, text="Done.")
-        except Exception as exc:
-            st.error(f"Could not build dendrogram: {exc}")
-            st.session_state.pop(fig_key, None)
-        finally:
-            _release_vector_memory(store, shared_variants)
-            progress.empty()
-
-    if fig_key in st.session_state:
-        fig_a, fig_b, n_personas, va, vb = st.session_state[fig_key]
-        if fig_b is not None:
-            col_a, col_b = st.columns(2)
-            with col_a:
-                st.subheader(prompt_variant_label(va))
-                _plotly_chart(fig_a)
-            with col_b:
-                st.subheader(prompt_variant_label(vb))
-                _plotly_chart(fig_b)
-        else:
-            _plotly_chart(fig_a)
-
-        figs = [fig_a] + ([fig_b] if fig_b else [])
-        filenames = [
-            _filename("dendro", store.model_name, mask_strategy.value, va),
-            *(
-                [_filename("dendro", store.model_name, mask_strategy.value, vb)]
-                if fig_b
-                else []
-            ),
-        ]
-        _render_save_buttons(figs, filenames, "dendro")
-        st.success(f"Generated dendrogram(s) for {n_personas} persona(s).")
+from tabs.analysis._shared import _render_mask_strategy_select
+from tabs.analysis._state import (
+    _DEFAULT_PERSONA_LIMITS,
+    _LAST_PROJECTION_DIMS_KEY,
+    _LAST_SIMILARITY_PERSONAS_KEY,
+    _LAST_SOURCE_KEY,
+)
+from tabs.analysis.cosine import _render_cosine_similarity
+from tabs.analysis.dendrogram import _render_dendrogram_analysis
+from tabs.analysis.layered import _render_layered_figure_analysis
 
 
 def _render_source_select() -> str:

tabs/probe.py

@@ -322,8 +322,8 @@ def _cached_sweep(
     if inputs.n_pca_components is not None:
         # Always overlay the compressed sweep against full activations.
         rows_by_label = {
-            f"pca{inputs.n_pca_components}": _sweep(inputs.n_pca_components),
             "full": _sweep(None),
+            f"pca{inputs.n_pca_components}": _sweep(inputs.n_pca_components),
         }
     else:
         rows_by_label = {"full": _sweep(None)}

tests/test_probes.py

@@ -0,0 +1,198 @@
+"""Regression tests for utils.probes.
+
+Covers the probe-artifact filename parser (both naming conventions) and the
+two correctness fixes:
+
+* ``_normalize_batch`` applies PCA independently of the scaler (previously the
+  PCA branch was unreachable when no scaler was present).
+* ``LoadedProbe.run`` predicts class 1 for a single-output probe whose sigmoid
+  score is >= 0.5 (previously it always predicted class 0).
+"""
+
+import pytest
+import torch
+
+from utils.probes import (
+    LoadedProbe,
+    _LinearProbe,
+    _normalize_labels,
+    parse_probe_filename,
+)
+
+
+# --------------------------------------------------------------------------- #
+# parse_probe_filename
+# --------------------------------------------------------------------------- #
+
+
+def test_parse_cognitive_map_filename():
+    meta = parse_probe_filename(
+        "cognitive_map_probe_layer12_lr_pre_reasoning_all_general.pt"
+    )
+    assert meta.layer == 12
+    assert meta.model_type == "lr"
+    assert meta.location == "pre_reasoning"
+    assert meta.scope == "general"
+
+
+def test_parse_persona_probe_dir_without_pca():
+    meta = parse_probe_filename(
+        "google__gemma-3-27b-it/answer_mean/biography/sex/"
+        "logistic_regression_layer20/probe.json"
+    )
+    assert meta.layer == 20
+    assert meta.model_type == "logistic_regression"
+    assert meta.scope is None
+    assert meta.attribute_name == "sex"
+    assert meta.model_name == "google/gemma-3-27b-it"
+
+
+def test_parse_persona_probe_dir_with_pca():
+    meta = parse_probe_filename(
+        "google__gemma-3-27b-it/answer_mean/biography/sex/"
+        "logistic_regression_pca10_layer46/weights.safetensors"
+    )
+    assert meta.layer == 46
+    assert meta.model_type == "logistic_regression"
+    assert meta.scope == "pca10"
+    assert meta.attribute_name == "sex"
+
+
+def test_parse_unknown_filename_falls_back():
+    meta = parse_probe_filename("something_else.bin")
+    assert meta.layer is None
+    assert meta.model_type == "unknown"
+
+
+# --------------------------------------------------------------------------- #
+# _normalize_labels
+# --------------------------------------------------------------------------- #
+
+
+def test_normalize_labels_list_pads_and_truncates():
+    assert _normalize_labels(["a", "b"], 3) == ["a", "b", None]
+    assert _normalize_labels(["a", "b", "c"], 2) == ["a", "b"]
+
+
+def test_normalize_labels_dict_indexes_by_key():
+    assert _normalize_labels({"1": "pos", "0": "neg"}, 2) == ["neg", "pos"]
+
+
+def test_normalize_labels_none():
+    assert _normalize_labels(None, 2) == [None, None]
+
+
+# --------------------------------------------------------------------------- #
+# _normalize_batch — scaler and PCA are applied independently
+# --------------------------------------------------------------------------- #
+
+
+def _probe(model_input_dim: int, **kwargs) -> LoadedProbe:
+    return LoadedProbe(
+        model=_LinearProbe(input_dim=model_input_dim, num_classes=1),
+        input_dim=model_input_dim,
+        labels=[None],
+        model_type="linear",
+        layer=0,
+        location=None,
+        **kwargs,
+    )
+
+
+def test_normalize_batch_noop_without_scaler_or_pca():
+    probe = _probe(3)
+    batch = torch.tensor([[1.0, 2.0, 3.0]])
+    assert torch.equal(probe._normalize_batch(batch), batch)
+
+
+def test_normalize_batch_scaler_only():
+    probe = _probe(
+        3,
+        scaler_mean=torch.ones(3),
+        scaler_std=torch.full((3,), 2.0),
+    )
+    batch = torch.tensor([[3.0, 5.0, 7.0]])
+    out = probe._normalize_batch(batch)
+    torch.testing.assert_close(out, torch.tensor([[1.0, 2.0, 3.0]]))
+
+
+def test_normalize_batch_pca_only_applies_pca():
+    """Regression: PCA must apply even when no scaler is present."""
+    probe = _probe(
+        2,
+        pca_mean=torch.ones(3),
+        pca_components=torch.tensor(
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
+        ),
+    )
+    batch = torch.tensor([[2.0, 4.0, 9.0]])
+    out = probe._normalize_batch(batch)
+    # (batch - pca_mean) @ components.T -> rows [1, 3] selected by components
+    torch.testing.assert_close(out, torch.tensor([[1.0, 3.0]]))
+
+
+def test_normalize_batch_scaler_then_pca():
+    probe = _probe(
+        3,
+        scaler_mean=torch.zeros(3),
+        scaler_std=torch.ones(3),
+        pca_mean=torch.zeros(3),
+        pca_components=torch.eye(3),
+    )
+    batch = torch.tensor([[1.0, 2.0, 3.0]])
+    torch.testing.assert_close(probe._normalize_batch(batch), batch)
+
+
+def test_normalize_batch_scaler_shape_mismatch_raises():
+    probe = _probe(
+        3,
+        scaler_mean=torch.ones(5),
+        scaler_std=torch.ones(5),
+    )
+    with pytest.raises(ValueError, match="scaler shape"):
+        probe._normalize_batch(torch.zeros(1, 3))
+
+
+def test_normalize_batch_pca_shape_mismatch_raises():
+    probe = _probe(
+        2,
+        pca_mean=torch.ones(5),
+        pca_components=torch.zeros(2, 5),
+    )
+    with pytest.raises(ValueError, match="PCA mean shape"):
+        probe._normalize_batch(torch.zeros(1, 3))
+
+
+# --------------------------------------------------------------------------- #
+# LoadedProbe.run — single-output prediction
+# --------------------------------------------------------------------------- #
+
+
+def _single_output_probe(weight: list[float], bias: float) -> LoadedProbe:
+    model = _LinearProbe(input_dim=len(weight), num_classes=1)
+    with torch.no_grad():
+        model.linear.weight.copy_(torch.tensor([weight]))
+        model.linear.bias.copy_(torch.tensor([bias]))
+    return LoadedProbe(
+        model=model,
+        input_dim=len(weight),
+        labels=["neg", "pos"],
+        model_type="linear",
+        layer=0,
+        location=None,
+    )
+
+
+def test_run_single_output_predicts_positive_when_score_high():
+    """Regression: single-output probe must predict class 1 when sigmoid >= 0.5."""
+    probe = _single_output_probe(weight=[1.0, 1.0], bias=5.0)
+    result = probe.run(torch.tensor([1.0, 1.0]))
+    assert result.predicted_index == 1
+    assert result.predicted_label == "pos"
+
+
+def test_run_single_output_predicts_negative_when_score_low():
+    probe = _single_output_probe(weight=[1.0, 1.0], bias=-5.0)
+    result = probe.run(torch.tensor([1.0, 1.0]))
+    assert result.predicted_index == 0
+    assert result.predicted_label == "neg"

utils/probes.py

@@ -154,7 +154,9 @@ class LoadedProbe:
             probs = probs.unsqueeze(0)
 
         predicted_index = (
-            0 if probs.numel() == 1 else int(torch.argmax(probs).item())
+            int(probs.item() >= 0.5)
+            if probs.numel() == 1
+            else int(torch.argmax(probs).item())
         )
         predicted_label = (
             self.labels[predicted_index]
@@ -208,28 +210,27 @@ class LoadedProbe:
         return logits, probs
 
     def _normalize_batch(self, batch: torch.Tensor) -> torch.Tensor:
-        if self.scaler_mean is None or self.scaler_std is None:
-            return batch
-        mean = self.scaler_mean.to(dtype=torch.float32)
-        std = self.scaler_std.to(dtype=torch.float32)
-        if mean.ndim != 1 or std.ndim != 1 or mean.shape[0] != batch.shape[1]:
-            raise ValueError(
-                "Probe scaler shape does not match activation hidden size: "
-                f"mean={tuple(mean.shape)} std={tuple(std.shape)} "
-                f"batch={tuple(batch.shape)}"
-            )
-        safe_std = torch.where(std == 0, torch.ones_like(std), std)
-        batch = (batch - mean) / safe_std
-        if self.pca_mean is None or self.pca_components is None:
-            return batch
-        pca_mean = self.pca_mean.to(dtype=torch.float32)
-        components = self.pca_components.to(dtype=torch.float32)
-        if pca_mean.ndim != 1 or pca_mean.shape[0] != batch.shape[1]:
-            raise ValueError(
-                "Probe PCA mean shape does not match activation hidden size: "
-                f"mean={tuple(pca_mean.shape)} batch={tuple(batch.shape)}"
-            )
-        return (batch - pca_mean) @ components.T
+        if self.scaler_mean is not None and self.scaler_std is not None:
+            mean = self.scaler_mean.to(dtype=torch.float32)
+            std = self.scaler_std.to(dtype=torch.float32)
+            if mean.ndim != 1 or std.ndim != 1 or mean.shape[0] != batch.shape[1]:
+                raise ValueError(
+                    "Probe scaler shape does not match activation hidden size: "
+                    f"mean={tuple(mean.shape)} std={tuple(std.shape)} "
+                    f"batch={tuple(batch.shape)}"
+                )
+            safe_std = torch.where(std == 0, torch.ones_like(std), std)
+            batch = (batch - mean) / safe_std
+        if self.pca_mean is not None and self.pca_components is not None:
+            pca_mean = self.pca_mean.to(dtype=torch.float32)
+            components = self.pca_components.to(dtype=torch.float32)
+            if pca_mean.ndim != 1 or pca_mean.shape[0] != batch.shape[1]:
+                raise ValueError(
+                    "Probe PCA mean shape does not match activation hidden size: "
+                    f"mean={tuple(pca_mean.shape)} batch={tuple(batch.shape)}"
+                )
+            batch = (batch - pca_mean) @ components.T
+        return batch
 
 
 def model_probe_dir_name(model_name: str) -> str:

uv.lock

@@ -343,14 +343,14 @@ wheels = [
 
 [[package]]
 name = "click"
-version = "8.3.3"
+version = "8.4.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "colorama", marker = "sys_platform == 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/bb/63/f9e1ea081ce35720d8b92acde70daaedace594dc93b693c869e0d5910718/click-8.3.3.tar.gz", hash = "sha256:398329ad4837b2ff7cbe1dd166a4c0f8900c3ca3a218de04466f38f6497f18a2", size = 328061, upload-time = "2026-04-22T15:11:27.506Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/23/e4/796662cd90cf80e3a363c99db2b88e0e394b988a575f60a17e16440cd011/click-8.4.0.tar.gz", hash = "sha256:638f1338fe1235c8f4e008e4a8a254fb5c5fbdcbb40ece3c9142ebb78e792973", size = 350843, upload-time = "2026-05-17T00:47:58.425Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl", hash = "sha256:a2bf429bb3033c89fa4936ffb35d5cb471e3719e1f3c8a7c3fff0b8314305613", size = 110502, upload-time = "2026-04-22T15:11:25.044Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/ae/8e92f8058baf87f6c7d86ee7e457668690195cc77efedb8d3797a06e3940/click-8.4.0-py3-none-any.whl", hash = "sha256:40c50b7c6c6adac2823d411041ec84f3f103f1b280d5e9ce0d7f998995832f81", size = 116147, upload-time = "2026-05-17T00:47:56.842Z" },
 ]
 
 [[package]]
@@ -464,11 +464,11 @@ wheels = [
 
 [[package]]
 name = "decorator"
-version = "5.2.1"
+version = "5.3.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/5c/50/a39dd7ab407e93978dfa07d109b7d633e37958c89f30cbcec061b77b3ebc/decorator-5.3.0.tar.gz", hash = "sha256:95fda3122972c847cf0ff7e0ce2829bf25136f2526b627b3da85b60ca5f485c0", size = 58431, upload-time = "2026-05-17T06:59:57.258Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/6f/f8d0bba4dc2a69817d74f640d504650241ebf2f9f7263426f1b953b344d4/decorator-5.3.0-py3-none-any.whl", hash = "sha256:f8c2d71ede92f073144ddd7f3e9fbbc3bd0f2f29522c9d75ee648d66553834f4", size = 11104, upload-time = "2026-05-17T06:59:54.676Z" },
 ]
 
 [[package]]
@@ -752,6 +752,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d2/23/408243171aa9aaba178d3e2559159c24c1171a641aa83b67bdd3394ead8e/idna-3.15-py3-none-any.whl", hash = "sha256:048adeaf8c2d788c40fee287673ccaa74c24ffd8dcf09ffa555a2fbb59f10ac8", size = 72340, upload-time = "2026-05-12T22:45:55.733Z" },
 ]
 
+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
+]
+
 [[package]]
 name = "ipython"
 version = "9.13.0"
@@ -1589,6 +1598,11 @@ dependencies = [
     { name = "streamlit" },
 ]
 
+[package.dev-dependencies]
+dev = [
+    { name = "pytest" },
+]
+
 [package.metadata]
 requires-dist = [
     { name = "catppuccin", specifier = ">=2.5.0" },
@@ -1601,6 +1615,9 @@ requires-dist = [
     { name = "streamlit", specifier = ">=1.44.0" },
 ]
 
+[package.metadata.requires-dev]
+dev = [{ name = "pytest", specifier = ">=9.0.3" }]
+
 [[package]]
 name = "persona-vectors"
 version = "0.8.2"
@@ -1730,6 +1747,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/90/ad/cba91b3bcf04073e4d1655a5c1710ef3f457f56f7d1b79dcc3d72f4dd912/plotly-6.7.0-py3-none-any.whl", hash = "sha256:ac8aca1c25c663a59b5b9140a549264a5badde2e057d79b8c772ae2920e32ff0", size = 9898444, upload-time = "2026-04-09T20:36:39.812Z" },
 ]
 
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
+]
+
 [[package]]
 name = "prompt-toolkit"
 version = "3.0.52"
@@ -2068,6 +2094,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b2/e6/94145d714402fd5ade00b5661f2d0ab981219e07f7db9bfa16786cdb9c04/pynndescent-0.6.0-py3-none-any.whl", hash = "sha256:dc8c74844e4c7f5cbd1e0cd6909da86fdc789e6ff4997336e344779c3d5538ef", size = 73511, upload-time = "2026-01-08T21:29:57.306Z" },
 ]
 
+[[package]]
+name = "pytest"
+version = "9.0.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"