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

1window.initCgLinkGraph = function({visState, renderAll, data, cgSel}){
2  var {nodes, links, metadata} = data
3
4  var c = d3.conventions({
5    sel: cgSel.select('.link-graph').html(''),
6    margin: {left: visState.isHideLayer ? 0 : 30, bottom: 85},
7    layers: 'sccccs',
8  })
9  
10  c.svgBot = c.layers[0]
11  var allCtx = {
12    allLinks: c.layers[1],
13    pinnedLinks: c.layers[2],
14    bgLinks: c.layers[3],
15    clickedLinks: c.layers[4]
16  }
17  c.svg = c.layers[5]
18
19  // Count max number of nodes at each context to create a polylinear x scale
20  var earliestCtxWithNodes = d3.min(nodes, d => d.ctx_idx)
21  var cumsum = 0
22  var ctxCounts = d3.range(d3.max(nodes, d => d.ctx_idx) + 1).map(ctx_idx => {
23    if (ctx_idx >= earliestCtxWithNodes) {
24      var group = nodes.filter(d => d.ctx_idx === ctx_idx)
25      var maxCount = d3.max([1, d3.max(d3.nestBy(group, d => d.streamIdx), e => e.length)])
26      cumsum += maxCount
27    }
28    return {ctx_idx, maxCount, cumsum}
29  })
30
31  var xDomain = [-1].concat(ctxCounts.map(d => d.ctx_idx))
32  var xRange = [0].concat(ctxCounts.map(d => d.cumsum * c.width / cumsum))
33  c.x = d3.scaleLinear().domain(xDomain.map(d => d + 1)).range(xRange)
34  
35  const layers = 14// 8 // 19 // TODO: dynamically set layers instead of hardcode
36
37  var yNumTicks= visState.isHideLayer ? data.byStream.length : layers 
38  c.y = d3.scaleBand(d3.range(yNumTicks), [c.height, 0])
39
40  c.yAxis = d3.axisLeft(c.y)
41    .tickValues(d3.range(yNumTicks))
42    .tickFormat(i => {
43      if (i % 2) return
44      
45      return i == layers - 1 ? 'Lgt' : i == 0 ? 'Emb' : 'L' + i
46      var label = data.byStream[i][0].layerLocationLabel
47      var layer = +label.replace('L', '')
48      return isFinite(layer) && layer % 2 ? '' : label
49    })
50  
51  c.svgBot.append('rect').at({width: c.width, height: c.height, fill: '#F5F4EE'})
52  c.svgBot.append('g').appendMany('rect', [0, yNumTicks - 1])
53    .at({width: c.width, height: c.y.bandwidth(), y: c.y, fill: '#F0EEE7'})
54  
55  c.svgBot.append('g').appendMany('path', d3.range(-1, yNumTicks - 1))
56    .translate(d => [0, c.y(d + 1)])
57    .at({d: `M0,0H${c.width}`, stroke: 'white', strokeWidth: .5})
58  
59  c.drawAxis(c.svgBot)
60  c.svgBot.select('.x').remove()
61  c.svgBot.selectAll('.y line').remove()
62  if (visState.isHideLayer) c.svgBot.select('.y').remove()
63  
64  // Spread nodes across each context 
65  // d.width is the total amount of px space in each column
66  ctxCounts.forEach((d, i) => d.width = c.x(d.ctx_idx + 1) - c.x(ctxCounts[i].ctx_idx))
67  
68  // if default to 8px padding right, if pad right to center singletons 
69  var padR = Math.min(8, d3.min(ctxCounts.slice(1), d => d.width/2)) + 0
70  
71  // find the tightest spacing between nodes and use for all ctx (but don't go below 20)
72  ctxCounts.forEach(d => d.minS = (d.width - padR)/d.maxCount)
73  var overallS = Math.max(20, d3.min(ctxCounts, d => d.minS))
74
75  // apply to nodes
76  d3.nestBy(nodes, d => [d.ctx_idx, d.streamIdx].join('-')).forEach(ctxLayer => {
77    var ctxWidth = c.x(ctxLayer[0].ctx_idx + 1) - c.x(ctxLayer[0].ctx_idx) - padR
78    var s = Math.min(overallS, ctxWidth/ctxLayer.length)
79    
80    // sorting by pinned stacks all the links on top of each other
81    // ctxLayer = d3.sort(ctxLayer, d => visState.pinnedIds.includes(d.nodeId) ? -1 : 1)
82    ctxLayer = d3.sort(ctxLayer, d => -d.logitPct)
83    ctxLayer.forEach((d, i) => {
84      d.xOffset = d.feature_type === 'logit' ? ctxWidth - (padR/2 + i*s) : ctxWidth - (padR/2 + i*s)
85      d.yOffset = 0
86    })
87  })
88  nodes.forEach(d => d.pos = [
89    c.x(d.ctx_idx) + d.xOffset, 
90    c.y(d.streamIdx) + c.y.bandwidth()/2 + d.yOffset
91  ])
92
93  
94  // hover poitns
95  var maxHoverDistance = 30
96  c.sel
97    .on('mousemove', (ev) => {
98      if (ev.shiftKey) return
99      var [mouseX, mouseY] = d3.pointer(ev)
100      var [closestNode, closestDistance] = findClosestPoint(mouseX - c.margin.left, mouseY - c.margin.top, nodes)
101      if (closestDistance > maxHoverDistance) {
102        utilCg.unHoverFeature(visState, renderAll)
103        utilCg.hideTooltip()
104      } else if (visState.hoveredId !== closestNode) {
105        utilCg.hoverFeature(visState, renderAll, closestNode)
106        utilCg.showTooltip(ev, closestNode)
107      }
108    })
109    .on('mouseleave', (ev) => {
110      if (ev.shiftKey) return
111      utilCg.unHoverFeature(visState, renderAll)
112      utilCg.hideTooltip()
113    })
114    .on('click', (ev) => {
115      var [mouseX, mouseY] = d3.pointer(ev)
116      var [closestNode, closestDistance] = findClosestPoint(mouseX - c.margin.left, mou

116seY - c.margin.top, nodes)
117      if (closestDistance > maxHoverDistance) {
118        visState.clickedId = null
119        visState.clickedCtxIdx = null
120        renderAll.clickedId()
121      } else {
122        utilCg.clickFeature(visState, renderAll, closestNode, ev.metaKey || ev.ctrlKey)
123      }
124    })
125
126  function findClosestPoint(mouseX, mouseY, points) {
127    if (points.length === 0) return null
128
129    let closestPoint = points[0]
130    let closestDistance = distance(mouseX, mouseY, closestPoint.pos[0], closestPoint.pos[1])
131
132    for (let i = 1; i < points.length; i++){
133      var point = points[i]
134      var dist = distance(mouseX, mouseY, point.pos[0], point.pos[1])
135      if (dist < closestDistance){
136        closestPoint = point
137        closestDistance = dist
138      }
139    }
140    return [closestPoint, closestDistance]
141
142    function distance(x1, y1, x2, y2) {
143      return Math.sqrt(Math.pow(x2 - x1, 2) + Math.pow(y2 - y1, 2))
144    }
145  }
146
147  // set up dom
148  var nodeSel = c.svg.appendMany('text.node', nodes)
149    .translate(d => d.pos)
150    .text(d => utilCg.featureTypeToText(d.feature_type))
151    .at({
152      fontSize: 9,
153      fill: d => d.nodeColor,
154      stroke: '#000',
155      strokeWidth: .5,
156      textAnchor: 'middle',
157      dominantBaseline: 'central',
158    })
159    // .call(utilCg.addFeatureTooltip)
160    // .call(utilCg.addFeatureEvents(visState, renderAll, ev => ev.shiftKey))
161
162  var hoverSel = c.svg.appendMany('circle', nodes)
163    .translate(d => d.pos)
164    .at({r: 6, cy: .5, stroke: '#f0f', strokeWidth: 2, strokeDasharray: '2 2', fill: 'none', display: 'xnone', pointEvents: 'none'})
165
166  links.forEach(d => {
167    var [x1, y1] = d.sourceNode.pos
168    var [x2, y2] = d.targetNode.pos
169    d.pathStr = `M${x1},${y1}L${x2},${y2}`
170  })
171
172  
173  function drawLinks(links, ctx, strokeWidthOffset=0, colorOverride){
174    ctx.clearRect(-c.margin.left, -c.margin.top, c.totalWidth, c.totalHeight)
175    d3.sort(links, d =>

175 d.strokeWidth).forEach(d => {
176      ctx.beginPath()
177      ctx.moveTo(d.sourceNode.pos[0], d.sourceNode.pos[1])
178      ctx.lineTo(d.targetNode.pos[0], d.targetNode.pos[1])
179      ctx.strokeStyle = colorOverride || d.color
180      ctx.lineWidth = d.strokeWidth + strokeWidthOffset
181      ctx.stroke()
182    })
183  }
184
185  function filterLinks(featureIds){
186    var filteredLinks = []
187
188    featureIds.forEach(nodeId => {
189      nodes.filter(n => n.nodeId == nodeId).forEach(node => {
190        if (visState.linkType == 'input' || visState.linkType == 'either') {
191          Array.prototype.push.apply(filteredLinks, node.sourceLinks)
192        }
193        if (visState.linkType == 'output' || visState.linkType == 'either') {
194          Array.prototype.push.apply(filteredLinks, node.targetLinks)
195        }
196        if (visState.linkType == 'both') {
197          Array.prototype.push.apply(filteredLinks, node.sourceLinks.filter(
198            link => visState.pinnedIds.includes(link.sourceNode.nodeId)
199          ))
200          Array.prototype.push.apply(filteredLinks, node.targetLinks.filter(
201            link => visState.pinnedIds.includes(link.targetNode.nodeId)
202          ))
203        }
204      })
205    })
206
207    return filteredLinks
208  }
209
210  drawLinks(links, allCtx.allLinks, 0, 'rgba(0,0,0,.05)')
211  // renderAll.isShowAllLinks.fns['linkGraph'] = () => c.sel.select('canvas').st({display: visState.isShowAllLinks ? '' : 'none'})
212
213  function renderPinnedIds(){
214    drawLinks(visState.clickedId ? [] : filterLinks(visState.pinnedIds), allCtx.pinnedLinks)
215    nodeSel.classed('pinned', d => visState.pinnedIds.includes(d.nodeId))
216  }
217  renderAll.pinnedIds.fns['linkGraph'] = renderPinnedIds
218
219  function renderHiddenIds(){
220    var hiddenIdSet = new Set(visState.hiddenIds)
221    nodeSel.classed('hidden', d => hiddenIdSet.has(d.featureId))
222  }
223  renderAll.hiddenIds.fns['linkGraph'] = renderHiddenIds
224
225  function renderClicked(){
226    var clickedLinks = []
227    // if (visState.clickedId) {
228    //   clickedLinks = links.filter(link => 
229    //     link.sourceNode.nodeId === visState.clickedId || 
230    //     link.targetNode.nodeId === visState.clickedId
231    //   )
232    // } 
233
234    drawLinks(clickedLinks, allCtx.bgLinks, .05, '#000')
235    drawLinks(clickedLinks, allCtx.clickedLinks)
236    nodeSel.classed('clicked', e => e.nodeId === visState.clickedId)
237    
238    drawLinks(visState.clickedId ? [] : filterLinks(visState.pinnedIds), allCtx.pinnedLinks)
239
240
241    nodeSel
242      .at({fill: '#fff'})
243      .filter(d => d.tmpClickedLink?.tmpColor)
244      .at({fill: d => d.tmpClickedLink.tmpColor})
245      .raise()
246  }
247
248  renderAll.clickedId.fns['linkGraph'] = renderClicked
249  renderAll.linkType.fns['linkGraph'] = () => {
250    renderPinnedIds()
251    renderClicked()
252  }
253  renderAll.hoveredId.fns['linkGraph'] = () => {
254    hoverSel.st({display: e => e.featureId == visState.hoveredId ? '' : 'none'})
255  }
256
257  // Add x axis text/lines
258  var promptTicks = data.metadata.prompt_tokens.slice(earliestCtxWithNodes).map((token, i) =>{
259    var ctx_idx = i + earliestCtxWithNodes 
260    var mNodes = nodes.filter(d => d.ctx_idx == ctx_idx)
261    var hasEmbed = mNodes.some(d => d.feature_type == 'embedding')
262    return {token, ctx_idx, mNodes, hasEmbed}
263  })
264
265  var xTickSel = c.svgBot.appendMany('g.prompt-token', promptTicks)
266    .translate(d => [c.x(d.ctx_idx + 1), c.height])
267  
268  xTickSel.append('path').at({d: `M0,0v${-c.height}`, stroke: '#fff',strokeWidth: 1})
269  xTickSel.filter(d => d.hasEmbed).append('path').at({
270    stroke: '#B0AEA6',
271    d: `M-${padR + 3.5},${-c.y.bandwidth()/2 + 6}V${8}`,
272  })
273  
274  xTickSel.filter(d => d.hasEmbed).append('g').translate([-12, 8])
275    .append('text').text(d => d.token) 
276    .at({
277      x: -5,
278      y: 2,
279      textAnchor: 'end',
280      transform: 'rotate(-45)',
281      dominantBaseline: 'middle',
282      fontSize: 12,
283      // fontSize: (d, i) => c.x(i+1) - c.x(i) < 15 ? 9 : 14,
284    })
285  
286  var logitTickSel = c.svgBot.append('g.axis').appendMany('g', nodes.filter(d => d.feature_type == 'logit'))
287    .translate(d => d.pos)
288  logitTickSel.append('path').at({
289    stroke: '#B0AEA6',
290    d: `M0,${-6}V${-c.y.bandwidth()/2 - 6}`,
291  })
292  logitTickSel.append('g').translate([-5, -c.y.bandwidth()/2 - 8])
293    .append('text').text(d => d.logitToken) 
294    .at({
295      x: 5,
296      y: 2,
297      textAnchor: 'start',
298      transform: 'rotate(-45)',
299      dominantBaseline: 'middle',
300      fontSize: 12,
301      // fontSize: (d, i) => c.x(i+1) - c.x(i) < 15 ? 9 : 14,
302    })
303
304
305  utilCg.addPinnedClickedGradient(c.svg)
306}
307
308window.init?.()