File size: 3,906 Bytes
4c1c394
62166b9
 
fb9c7be
4c1c394
 
62166b9
 
 
 
4c1c394
 
2f91961
 
4c1c394
62166b9
4c1c394
 
fb9c7be
dceeb0b
fb9c7be
4c1c394
 
62166b9
4c1c394
 
 
62166b9
 
 
 
 
 
 
 
 
 
 
4c1c394
 
 
 
 
 
 
 
 
fb9c7be
4c1c394
 
 
62166b9
fb9c7be
4c1c394
62166b9
4c1c394
 
 
 
62166b9
4c1c394
62166b9
 
4c1c394
 
62166b9
 
4c1c394
2f91961
 
4c1c394
 
62166b9
4c1c394
62166b9
b7739d5
 
 
 
 
fb9c7be
 
b7739d5
 
fb9c7be
4c1c394
 
62166b9
 
4c1c394
62166b9
4c1c394
2f91961
 
 
 
 
 
 
 
 
 
 
 
 
4c1c394
 
 
62166b9
 
4c1c394
 
 
62166b9
 
 
 
 
 
 
fb9c7be
62166b9
 
 
 
 
 
4c1c394
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""
Feature detail panel: stats header, top NSD MEI grid, cortical profile.

Pure display logic lives in explorer.feature_logic.

Exports:
  update_feature_display(feat)
  select_and_display(feat)
  status_div, stats_div, top_heatmap_div, brain_div
  add_steer_btn
"""

import threading

from bokeh.io import curdoc
from bokeh.models import Button, Div

from ..state import _S, active_ds
from ..brain import HAS_DYNADIFF, render_cortical_profile
from ..feature_logic import build_stats_html, build_top_mei_items
from ..rendering import status_html, make_image_grid_html
from .. import widgets

N_DISPLAY = 9

# ---------- Divs ----------

status_div      = Div(text=status_html('idle', 'Select a feature to begin.'), width=680)
stats_div       = Div(text="<h3>Select a feature to explore it</h3>", width=530)
top_heatmap_div = Div(text="", width=680)
brain_div       = Div(text="", width=680)

add_steer_btn = Button(
    label="+ Add to Steer",
    button_type="success",
    width=140,
    visible=HAS_DYNADIFF,
)


# ---------- Core display logic ----------

def update_feature_display(feature_idx):
    feat = int(feature_idx)
    _S.render_token += 1
    my_token = _S.render_token

    ds = active_ds()
    freq_val = ds['feature_frequency'][feat].item()
    feat_name = ds['feature_names'].get(feat, "")
    auto_name = ds['auto_interp_names'].get(feat, "")

    stats_div.text = build_stats_html(feat, freq_val, feat_name, auto_name)

    # Sync name-editing widget
    from .feature_list import name_input
    name_input.value = feat_name

    if freq_val == 0:
        status_div.text = status_html('dead', f'Feature {feat} never activated.')
        brain_div.text = render_cortical_profile(feat)
        top_heatmap_div.text = ""
        add_steer_btn.visible = False
        return

    add_steer_btn.visible = HAS_DYNADIFF
    status_div.text = status_html('loading', '&#x23F3; Loading...')

    doc = curdoc()

    def _render():
        if _S.render_token != my_token:
            return

        _ds = active_ds()
        try:
            zp = int(widgets.zoom_slider.value)
            ha = widgets.heatmap_alpha_slider.value
        except Exception:
            zp, ha = 16, 1.0
        top_infos, top_img_is, subset_label = build_top_mei_items(
            feat, _ds, n_display=N_DISPLAY,
            zoom_patches=zp,
            heatmap_alpha=ha,
        )

        top_heatmap_div.text = make_image_grid_html(
            top_infos, f"Top MEIs (feature {feat}){subset_label}",
            img_indices=top_img_is, cols=3)

        status_div.text = status_html('idle', '')

        # Render cortical profile in background thread so the UI stays responsive
        brain_div.text = ""

        def _render_brain():
            html = render_cortical_profile(feat)
            def _update_brain():
                if _S.render_token == my_token:
                    brain_div.text = html
            doc.add_next_tick_callback(_update_brain)

        threading.Thread(target=_render_brain, daemon=True).start()

    doc.add_next_tick_callback(_render)


def select_and_display(feat: int):
    """Display feature detail and highlight it."""
    widgets.feature_input.value = str(feat)
    update_feature_display(feat)


# ---------- Add-to-Steer callback ----------

def _on_add_steer():
    try:
        feat = int(widgets.feature_input.value)
    except ValueError:
        return
    from .steering import add_feature
    add_feature(feat)

add_steer_btn.on_click(_on_add_steer)


# ---------- Re-render on zoom/alpha changes ----------

def _rerender_current(attr, old, new):
    try:
        feat = int(widgets.feature_input.value)
        if 0 <= feat < active_ds()['d_model']:
            update_feature_display(feat)
    except ValueError:
        pass

widgets.zoom_slider.on_change('value', _rerender_current)
widgets.heatmap_alpha_slider.on_change('value', _rerender_current)