Spaces:
Running
Running
| import { useMemo } from "react"; | |
| import { formatNumber } from "./formula"; | |
| import { edgePath, layoutGraph, OUTPUT, SLOT_OFFSET, type PlacedEdge, type PlacedShape } from "./layout"; | |
| import MathText from "./MathText"; | |
| import StepFlow, { TOKEN_BASELINE } from "./StepFlow"; | |
| import type { Flow, Timeline } from "./Timeline"; | |
| import type { Graph, State } from "./types"; | |
| // Small graphs keep their size rather than filling the plot | |
| const MIN_VIEW_WIDTH = 560; | |
| const MIN_VIEW_HEIGHT = 240; | |
| // The baseline of an edge's name, relative to the edge | |
| const ABOVE_EDGE = -5; | |
| function SlotText({ x, y, number }: { x: number; y: number; number: number }) { | |
| return ( | |
| <text x={x} y={y + TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs"> | |
| {formatNumber(number)} | |
| </text> | |
| ); | |
| } | |
| // A shape with its name or operator in the middle, its value above and its derivative below | |
| function Shape({ shape, current, value, derivative }: { shape: PlacedShape; current: boolean; value: number | undefined; derivative: number | undefined }) { | |
| const fill = current ? "fill-marker-soft" : shape.kind === "op" ? "fill-white" : "fill-paper"; | |
| return ( | |
| <g data-node={shape.id}> | |
| {shape.kind === "op" ? ( | |
| <ellipse cx={shape.x} cy={shape.y} rx={shape.width / 2} ry={shape.height / 2} strokeWidth={1.5} className={`stroke-ink ${fill}`} /> | |
| ) : ( | |
| <rect | |
| x={shape.x - shape.width / 2} | |
| y={shape.y - shape.height / 2} | |
| width={shape.width} | |
| height={shape.height} | |
| strokeWidth={1.5} | |
| className={`stroke-ink ${fill}`} | |
| /> | |
| )} | |
| <text x={shape.x} y={shape.y + 5} textAnchor="middle" className="fill-ink font-mono text-sm font-medium"> | |
| <MathText text={shape.label} svg /> | |
| </text> | |
| {value !== undefined && <SlotText x={shape.x} y={shape.y - SLOT_OFFSET} number={value} />} | |
| {derivative !== undefined && <SlotText x={shape.x} y={shape.y + SLOT_OFFSET} number={derivative} />} | |
| </g> | |
| ); | |
| } | |
| export default function ComputationGraph({ | |
| graph, | |
| state, | |
| flow, | |
| timeline, | |
| }: { | |
| graph: Graph; | |
| state: State | undefined; | |
| flow: Flow | null; | |
| timeline: Timeline; | |
| }) { | |
| const layout = useMemo(() => layoutGraph(graph), [graph]); | |
| const viewWidth = Math.max(layout.width, MIN_VIEW_WIDTH); | |
| const viewHeight = Math.max(layout.height, MIN_VIEW_HEIGHT); | |
| const step = state?.step; | |
| const backward = step?.phase === "backward"; | |
| // The edges into the current node, or the edge from f when the backward pass starts | |
| const isActive = (edge: PlacedEdge) => | |
| step !== undefined && | |
| (backward && step.partials.length === 0 ? edge.to === OUTPUT : edge.to === step.node && edge.to !== OUTPUT); | |
| // Drawn last, so they sit on top | |
| const edges = [...layout.edges.filter((edge) => !isActive(edge)), ...layout.edges.filter(isActive)]; | |
| return ( | |
| <svg | |
| className="min-h-[240px] w-full flex-1" | |
| viewBox={`${(layout.width - viewWidth) / 2} ${(layout.height - viewHeight) / 2} ${viewWidth} ${viewHeight}`} | |
| role="img" | |
| aria-label="Computational graph of the function" | |
| > | |
| <defs> | |
| <marker id="arrowhead" viewBox="0 0 8 8" refX="7" refY="4" markerWidth="8" markerHeight="8" orient="auto"> | |
| <path d="M 0 0 L 8 4 L 0 8 z" className="fill-ink-muted" /> | |
| </marker> | |
| <marker id="arrowhead-active" viewBox="0 0 8 8" refX="7" refY="4" markerWidth="6" markerHeight="6" orient="auto-start-reverse"> | |
| <path d="M 0 0 L 8 4 L 0 8 z" className="fill-ink" /> | |
| </marker> | |
| </defs> | |
| {edges.map((edge) => { | |
| const active = isActive(edge); | |
| // The backward pass flows against the edge, so its arrow points back | |
| const flipped = active && backward; | |
| return ( | |
| <path | |
| key={edge.id} | |
| d={edgePath(edge)} | |
| fill="none" | |
| strokeWidth={active ? 3 : 1.5} | |
| markerStart={flipped ? "url(#arrowhead-active)" : undefined} | |
| markerEnd={flipped ? undefined : active ? "url(#arrowhead-active)" : "url(#arrowhead)"} | |
| className={active ? "stroke-ink" : "stroke-ink-muted"} | |
| /> | |
| ); | |
| })} | |
| {layout.shapes.map((shape) => ( | |
| <Shape | |
| key={shape.id} | |
| shape={shape} | |
| current={shape.id === step?.node} | |
| // f shows what the output node holds | |
| value={state?.values[shape.id === OUTPUT ? graph.output : shape.id]} | |
| derivative={state?.derivatives[shape.id === OUTPUT ? graph.output : shape.id]} | |
| /> | |
| ))} | |
| {layout.edges.map((edge) => ( | |
| <text | |
| key={edge.id} | |
| data-edge={edge.from} | |
| x={edge.labelAt.x} | |
| y={edge.labelAt.y + ABOVE_EDGE} | |
| // On a slanted edge, the name moves to the side the edge leaves clear | |
| textAnchor={edge.labelSlope === 0 ? "middle" : edge.labelSlope < 0 ? "end" : "start"} | |
| strokeWidth={4} | |
| paintOrder="stroke" | |
| className="fill-ink stroke-white font-mono text-xs" | |
| > | |
| <MathText text={edge.from} svg /> | |
| </text> | |
| ))} | |
| {flow && <StepFlow key={flow.id} graph={graph} layout={layout} flow={flow} timeline={timeline} />} | |
| </svg> | |
| ); | |
| } | |