https://embed-goodfire-attr-graph.netlify.app/attribution_graph/init-cg-node-connections.js

1window.initCgNodeConnections = function({visState, renderAll, data, cgSel}){
2  
3  var nodeConnectionsSel = cgSel.select('.node-connections')
4  var headerSel = null
5  var clickedNode = null
6  var featureSel = null
7
8  function render() {
9    nodeConnectionsSel.html('').st({display: 'flex', flexDirection: 'column'})
10    clickedNode = data.nodes.find(e => e.nodeId === visState.clickedId)    
11    headerSel = nodeConnectionsSel.append('div.header-top-row.section-title').st({marginBottom: 20}).datum(clickedNode)
12    
13    if (!clickedNode) return headerSel.text('Click a feature on the left for details')
14    
15    addHeaderRow(headerSel) 
16
17    var types = [
18      { id: 'input', title: 'Input Features'},
19      { id: 'output', title: 'Output Features' }
20    ]
21    types.forEach(type => type.sections = ['Positive', 'Negative'].map(title => {
22      var nodes = data.nodes.filter(d => {
23        var weight = type.id === 'input' ? d.tmpClickedSourceLink?.weight : d.tmpClickedTargetLink?.weight
24        return title == 'Positive' ? weight > 0 : weight < 0
25      })
26      nodes = d3.sort(nodes, d => -(type.id === 'input' ? d.tmpClickedSourceLink?.pctInput : d.tmpClickedTargetLink?.pctInput))
27      return {title, nodes}
28    }))
29
30    var typesSel = nodeConnectionsSel.append('div.connections')
31      .st({flex: '1 0 auto', display: 'flex', overflow: 'hidden', gap: '20px'})
32      .appendMany('div.features', types)
33      .classed('output', d => d.id === 'output')
34      .classed('input', d => d.id === 'input')
35
36    typesSel.append('div.section-title').text(d => d.title)
37      .st({marginBottom: 0})
38    // addInputSumTable(typesSel.filter(d => d.id == 'input').append('div.sum-table').append('div.section'))
39
40    var featuresContainerSel = typesSel.append('div.effects')
41
42    var sectionSel = featuresContainerSel.appendMany('div.section', d => d.sections)
43    featureSel = sectionSel.appendMany('div.feature-row', d => d.nodes)
44      .classed('clicked', e => e.featureId == visState.clickedId)
45
46    classPinned()
47    classHidden()
48
49    typesSel.each(function(type) {
50      d3.select(this).selectAll('.feature-row')
51        .call(utilCg.renderFeatureRow, visState, renderAll, type.id === 'input' ? 'tmpClickedSourceLink' : 'tmpClickedTargetLink')
52    })
53  }
54  
55  function addHeaderRow(headerSel){
56    if (!clickedNode) return 
57    
58    headerSel.append('text')
59      .text(clickedNode.feature_type == 'cross layer transcoder' ? 'F#' + d3.format('08')(clickedNode.feature) : ' ')
60      .st({display: 'inline-block', marginRight: 5, 'font-variant-numeric': 'tabular-nums', width: 82})
61    headerSel.append('span.feature-icon').text(utilCg.featureTypeToText(clickedNode.feature_type))
62    headerSel.append('span.feature-title').text(clickedNode.ppClerp)
63    
64    // add cmd-click toggle to title
65    headerSel.on('click', ev => {
66      utilCg.clickFeature(visState, renderAll, clickedNode, ev.metaKey || ev.ctrlKey)
67
68      if (visState.clickedId) return
69      // double render to toggle on hoveredId, could expose more of utilCg.clickFeature to prevent
70      visState.hoveredId = clickedNode.featureId
71      renderAll.hoveredId()
72    })
73  }
74
75  function addInputSumTable(sumSel){
76    return // TODO: turn back on? 
77    
78    var clickedNode = data.nodes.idToNode[visState.clickedId]
79    if (!clickedNode) return
80    var inputSum = d3.sum(clickedNode.sourceLinks, d => Math.abs(d.weight))
81
82    var tableSel = sumSel.append('table')
83    tableSel.appendMany('th', ['% Feat', '% Err', '% Emb']).text(d => d)
84
85    var trSel = tableSel.appendMany('tr.data', [
86      {str: '←', links: clickedNode.sourceLinks.filter(d => d.sourceNode.ctx_idx != clickedNode.ctx_idx)},
87      {str: '↓', links: clickedNode.sourceLinks.filter(d => d.sourceNode.ctx_idx == clickedNode.ctx_idx)}
88    ])
89
90    trSel.append('td').text(d => d.str).at({title: d => d.str == '↓' ? 'cur token' : 'prev token'})
91    trSel.appendMany('td', d => {
92        var rv = [
93          d.links.filter(e => !e.sourceNode.isError && e.sourceNode.feature_type != 'embedding'),
94          d.links.filter(e => e.sourceNode.isError),
95          d.links.filter(e => e.sourceNode.feature_type == 'embedding'),
96        ]
97
98        if (rv.flat().length != d.links.length) console.error("Non-feature/error/embedding node present")
99        return rv
100      })
101      .html(links => {
102        var pos = d3.sum(links.filter(d => d.weight > 0), d => d.weight)/inputSum
103        var neg = d3.sum(links.filter(d => d.weight < 0), d => d.weight)/inputSum
104
105        var outStr = `+${d3.format('.2f')(pos)}&#160;&#160;&#160;−${d3.format('.2f')(neg)}`
106        return outStr.replaceAll('0.', '.').replaceAll('−−', '−')
107      })
108  }
109
110  renderAll.hClerpUpdate.fns.push(render)
111  renderAll.hoveredId.fns['nodeconnections'] = () => featureSel?.classed('hovered', e => e.featureId == visState.hoveredId)
112  renderAll.pinnedIds.fns['nodeconnections'] = classPinned
113  renderAll.hiddenIds.fns['nodeconnections'] = classHidden
114  renderAll.clickedId.fns['nodeconnections'] = render
115  
116
117  function classPinned(){
118    var pinnedIdSet = new Set(visState.pinnedIds)
119    featureSel?.classed('pinned', d => pinnedIdSet.has(d.nodeId))
120    headerSel?.classed('pinned', d => d && d.nodeId && visState.pinnedIds.includes(d.nodeId))
121      .select('span.feature-title').text(d => d.ppClerp)
122  }
123
124  function classHidden(){
125    var hiddenIdSet = new Set(visState.hiddenIds)
126    featureSel?.classed('hidden', d => hiddenIdSet.has(d.featureId))
127  }
128}
129
130window.init?.()