Spaces:
Running
Running
File size: 5,293 Bytes
b9499ab | 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 | 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>
);
}
|