Spaces:
Running
Running
Download frontends/react/src/StepFlow.tsx from nhamid289/autodiff-dev: direct link, hf CLI and curl.
- Browser
- Download file 12.1 kB
-
https://huggingface.co/spaces/nhamid289/autodiff-dev/resolve/main/frontends/react/src/StepFlow.tsx
- Command line
-
hf download hf://spaces/nhamid289/autodiff-dev/frontends/react/src/StepFlow.tsx
-
curl -L -o StepFlow.tsx https://huggingface.co/spaces/nhamid289/autodiff-dev/resolve/main/frontends/react/src/StepFlow.tsx
12.1 kB
| import { useMemo, useSyncExternalStore } from "react"; | |
| import { formatNumber, nodeArgs } from "./formula"; | |
| import { edgePath, expressionOf, OUTPUT, SLOT_OFFSET, type GraphLayout, type PlacedEdge, type PlacedShape } from "./layout"; | |
| import { | |
| APPLY_MS, FADE_AT, IN_AT, IN_MS, LAND_AT, OUT_MS, PRODUCT_AT, RESULT_AT, | |
| SEED_HOLD_MS, SEED_TRAVEL_MS, SETTLE_AT, SETTLE_MS, SUM_AT, TRAVEL_MS, | |
| } from "./flowTiming"; | |
| import type { Graph, GraphNode, Step } from "./types"; | |
| import type { Flow, Timeline } from "./Timeline"; | |
| const POP_MS = 150; | |
| const CHAR_WIDTH = 7.2; | |
| const TOKEN_HEIGHT = 20; | |
| // A token's text baseline below its centre; a slot's text uses it too, so a settling token lands exactly on it | |
| export const TOKEN_BASELINE = 4; | |
| const SVG_NS = "http://www.w3.org/2000/svg"; | |
| function easeInOut(t: number): number { | |
| return t < 0.5 ? 2 * t * t : 1 - (1 - t) ** 2 * 2; | |
| } | |
| function progress(elapsed: number, from: number, duration: number): number { | |
| return Math.min(Math.max((elapsed - from) / duration, 0), 1); | |
| } | |
| function shapeOf(layout: GraphLayout, id: string): PlacedShape { | |
| return layout.shapes.find((s) => s.id === id)!; | |
| } | |
| const valueSlot = (shape: PlacedShape) => ({ x: shape.x, y: shape.y - SLOT_OFFSET }); | |
| const derivativeSlot = (shape: PlacedShape) => ({ x: shape.x, y: shape.y + SLOT_OFFSET }); | |
| // A path for a token: along an edge, joining it from one point and leaving it for another | |
| function travelPath(edge: PlacedEdge, start: { x: number; y: number }, end: { x: number; y: number }) { | |
| const path = document.createElementNS(SVG_NS, "path"); | |
| path.setAttribute("d", `${edgePath(edge).replace("M", `M ${start.x} ${start.y} L`)} L ${end.x} ${end.y}`); | |
| return { path, length: path.getTotalLength() }; | |
| } | |
| // Milliseconds since the animation started, from the timeline's playhead | |
| function useElapsed({ timeline, flow }: FlowProps): number { | |
| const time = useSyncExternalStore(timeline.subscribeTime, timeline.getTime); | |
| return Math.max(time - flow.start, 0); | |
| } | |
| interface TokenProps { | |
| x: number; | |
| y: number; | |
| text: string; | |
| scale?: number; | |
| strong?: boolean; | |
| // How much of the pill shows; at 0, only the text is left | |
| outline?: number; | |
| // Hide what's behind the pill, such as a number that's being replaced | |
| backdrop?: boolean; | |
| } | |
| // A number or calculation in a pill, centred on a point | |
| function Token({ x, y, text, scale = 1, strong = false, outline = 1, backdrop = false }: TokenProps) { | |
| const width = text.length * CHAR_WIDTH + 12; | |
| const pill = { x: -width / 2, y: -TOKEN_HEIGHT / 2, width, height: TOKEN_HEIGHT, rx: TOKEN_HEIGHT / 2 }; | |
| return ( | |
| <g transform={`translate(${x} ${y}) scale(${scale})`}> | |
| {backdrop && <rect {...pill} className="fill-white" />} | |
| <rect {...pill} strokeWidth={1.5} opacity={outline} className={`stroke-ink ${strong ? "fill-marker" : "fill-marker-soft"}`} /> | |
| {/* As the pill fades, the text gets the same outline as a number resting in its slot */} | |
| <text y={TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} strokeOpacity={1 - outline} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs"> | |
| {text} | |
| </text> | |
| </g> | |
| ); | |
| } | |
| // Copies of a step's arguments flow along their edges and combine at the node | |
| function ForwardFlow(props: FlowProps) { | |
| const { graph, layout, flow } = props; | |
| const { step, before } = flow; | |
| const node = graph.nodes.find((n) => n.id === step.node)!; | |
| const shape = shapeOf(layout, step.node); | |
| const forwardMode = step.derivative !== undefined; | |
| const { tokens, toOutput, tangents, tangentToOutput } = useMemo(() => { | |
| const target = shapeOf(layout, step.node); | |
| const outputEdge = step.node === graph.output ? layout.edges.find((edge) => edge.to === OUTPUT) : undefined; | |
| return { | |
| tokens: layout.edges | |
| .filter((edge) => edge.to === step.node) | |
| .map((edge) => ({ id: edge.id, text: formatNumber(before.values[edge.from]), ...travelPath(edge, valueSlot(shapeOf(layout, edge.from)), target) })), | |
| // The last result also carries on to f | |
| toOutput: outputEdge && travelPath(outputEdge, valueSlot(target), valueSlot(shapeOf(layout, OUTPUT))), | |
| // Forward mode's tangents follow the same way | |
| tangents: step.derivative === undefined ? undefined : derivativeMoves(graph, layout, step, before.derivatives), | |
| tangentToOutput: outputEdge && travelPath(outputEdge, derivativeSlot(target), derivativeSlot(shapeOf(layout, OUTPUT))), | |
| }; | |
| }, [graph, layout, step, before]); | |
| const elapsed = useElapsed(props); | |
| const pop = (from: number) => 0.7 + 0.3 * progress(elapsed, from, POP_MS); | |
| const calculation = expressionOf(node, (arg) => (arg in before.values ? formatNumber(before.values[arg]) : arg)); | |
| const pulse = (elapsed - TRAVEL_MS) / APPLY_MS; | |
| const settled = easeInOut(progress(elapsed, SETTLE_AT, SETTLE_MS)); | |
| const result = formatNumber(step.value); | |
| // How far the tangents have got, once they've set off | |
| const t = elapsed - SETTLE_AT; | |
| const valuePhase = elapsed < TRAVEL_MS ? "travel" : elapsed < RESULT_AT ? "apply" : elapsed < SETTLE_AT ? "result" : "settle"; | |
| return ( | |
| <g pointerEvents="none" data-flow={forwardMode && t >= 0 ? `tangent-${derivativePhase(t)}` : valuePhase}> | |
| {elapsed < TRAVEL_MS && | |
| tokens.map(({ id, path, length, text }) => { | |
| const { x, y } = path.getPointAtLength(length * easeInOut(elapsed / TRAVEL_MS)); | |
| return <Token key={id} x={x} y={y} text={text} />; | |
| })} | |
| {pulse >= 0 && pulse < 1 && ( | |
| <ellipse | |
| cx={shape.x} | |
| cy={shape.y} | |
| rx={(shape.width / 2) * (1 + pulse * 0.4)} | |
| ry={(shape.height / 2) * (1 + pulse * 0.4)} | |
| opacity={1 - pulse} | |
| fill="none" | |
| strokeWidth={2} | |
| className="stroke-ink" | |
| /> | |
| )} | |
| {elapsed >= TRAVEL_MS && elapsed < RESULT_AT && <Token x={shape.x} y={shape.y} text={calculation} scale={pop(TRAVEL_MS)} />} | |
| {elapsed >= RESULT_AT && ( | |
| <Token x={shape.x} y={shape.y - SLOT_OFFSET * settled} text={result} scale={pop(RESULT_AT)} strong outline={1 - settled} /> | |
| )} | |
| {toOutput && elapsed >= SETTLE_AT && (() => { | |
| const { x, y } = toOutput.path.getPointAtLength(toOutput.length * settled); | |
| return <Token x={x} y={y} text={result} strong outline={1 - settled ** 4} />; | |
| })()} | |
| {tangents && t >= 0 && <DerivativeTokens {...tangents} t={t} />} | |
| {tangents && tangentToOutput && t >= FADE_AT && (() => { | |
| const moved = easeInOut(progress(t, FADE_AT, SETTLE_MS)); | |
| const { x, y } = tangentToOutput.path.getPointAtLength(tangentToOutput.length * moved); | |
| return <Token x={x} y={y} text={formatNumber(step.derivative!)} strong outline={1 - moved ** 4} />; | |
| })()} | |
| </g> | |
| ); | |
| } | |
| // The seed leaves f, keeping a copy there, and settles below the output node | |
| function SeedFlow(props: FlowProps) { | |
| const { layout, flow } = props; | |
| const output = shapeOf(layout, OUTPUT); | |
| const { path, length } = useMemo(() => { | |
| const edge = layout.edges.find((e) => e.to === OUTPUT)!; | |
| return travelPath(edge, derivativeSlot(shapeOf(layout, edge.from)), derivativeSlot(shapeOf(layout, OUTPUT))); | |
| }, [layout]); | |
| const elapsed = useElapsed(props); | |
| const moved = easeInOut(progress(elapsed, SEED_HOLD_MS, SEED_TRAVEL_MS)); | |
| const { x, y } = path.getPointAtLength(length * (1 - moved)); | |
| const text = formatNumber(flow.step.derivative!); | |
| return ( | |
| <g pointerEvents="none" data-flow={elapsed < SEED_HOLD_MS ? "seed" : elapsed < SEED_HOLD_MS + SEED_TRAVEL_MS ? "travel" : "settle"}> | |
| <text x={output.x} y={output.y + SLOT_OFFSET + TOKEN_BASELINE} textAnchor="middle" strokeWidth={3} paintOrder="stroke" className="fill-ink stroke-white font-mono text-xs"> | |
| {text} | |
| </text> | |
| <Token x={x} y={y} text={text} strong outline={1 - progress(elapsed, SEED_HOLD_MS + SEED_TRAVEL_MS, SETTLE_MS)} /> | |
| </g> | |
| ); | |
| } | |
| function sumText(terms: number[]): string { | |
| return terms.map(formatNumber).join(" + ").replaceAll("+ -", "− "); | |
| } | |
| // A derivative's moves along one step's edges: one token per partial, and where they land | |
| function derivativeMoves(graph: Graph, layout: GraphLayout, step: Step, derivatives: Record<string, number>) { | |
| const node = graph.nodes.find((n) => n.id === step.node) as GraphNode & { kind: "op" }; | |
| const forward = step.phase === "forward"; | |
| const args = nodeArgs(graph, node); | |
| const tokens = step.partials.map((p, n) => { | |
| const edge = layout.edges.find((e) => e.id === `${p.arg}-${step.node}-${args[n].index}`)!; | |
| // A tangent comes from the argument; an adjoint comes from the node | |
| const start = forward ? derivatives[p.arg] : step.derivative!; | |
| return { | |
| id: edge.id, | |
| start: formatNumber(start), | |
| multiply: `${formatNumber(start)} · ${formatNumber(p.partial)}`, | |
| product: formatNumber(p.amount), | |
| reversed: !forward, | |
| ...travelPath(edge, derivativeSlot(shapeOf(layout, p.arg)), derivativeSlot(shapeOf(layout, step.node))), | |
| }; | |
| }); | |
| // Forward: everything adds up at the node. Backward: at each argument, on top of what it has. | |
| const amounts = (arg: string) => step.partials.filter((p) => p.arg === arg).map((p) => p.amount); | |
| const landing = (id: string, terms: number[]) => ({ | |
| id, | |
| ...derivativeSlot(shapeOf(layout, id)), | |
| sum: sumText(terms), | |
| total: formatNumber(terms.reduce((a, b) => a + b, 0)), | |
| }); | |
| const landings = forward | |
| ? [landing(step.node, step.partials.map((p) => p.amount))] | |
| : [...new Set(step.partials.map((p) => p.arg))].map((arg) => | |
| landing(arg, [...(arg in derivatives ? [derivatives[arg]] : []), ...amounts(arg)]), | |
| ); | |
| return { tokens, landings }; | |
| } | |
| function derivativePhase(t: number): string { | |
| return t < OUT_MS ? "out" : t < PRODUCT_AT ? "multiply" : t < IN_AT ? "product" : t < LAND_AT ? "in" : t < SUM_AT ? "land" : "sum"; | |
| } | |
| // t is the time since the derivative set off | |
| function DerivativeTokens({ tokens, landings, t }: ReturnType<typeof derivativeMoves> & { t: number }) { | |
| const pop = (from: number) => 0.7 + 0.3 * progress(t, from, POP_MS); | |
| const phase = derivativePhase(t); | |
| // How far along the edge a token has come, from where it set off (0) to where it lands (1) | |
| const along = 0.5 * easeInOut(progress(t, 0, OUT_MS)) + 0.5 * easeInOut(progress(t, IN_AT, IN_MS)); | |
| return ( | |
| <> | |
| {t < LAND_AT && | |
| tokens.map(({ id, path, length, start, multiply, product, reversed }) => { | |
| const { x, y } = path.getPointAtLength(length * (reversed ? 1 - along : along)); | |
| if (phase === "multiply") { | |
| return <Token key={id} x={x} y={y} text={multiply} scale={pop(OUT_MS)} />; | |
| } | |
| return phase === "out" ? <Token key={id} x={x} y={y} text={start} /> : <Token key={id} x={x} y={y} text={product} strong />; | |
| })} | |
| {t >= LAND_AT && | |
| landings.map(({ id, x, y, sum, total }) => | |
| t < SUM_AT ? ( | |
| <Token key={id} x={x} y={y} text={sum} scale={pop(LAND_AT)} backdrop /> | |
| ) : ( | |
| <Token key={id} x={x} y={y} text={total} strong outline={1 - progress(t, FADE_AT, SETTLE_MS)} backdrop /> | |
| ), | |
| )} | |
| </> | |
| ); | |
| } | |
| // The node's adjoint leaves along each edge, picks up the edge's partial on the way, and adds to the argument's adjoint | |
| function BackwardFlow(props: FlowProps) { | |
| const { graph, layout, flow } = props; | |
| const { step, before } = flow; | |
| const moves = useMemo(() => derivativeMoves(graph, layout, step, before.derivatives), [graph, layout, step, before]); | |
| const t = useElapsed(props); | |
| return ( | |
| <g pointerEvents="none" data-flow={derivativePhase(t)}> | |
| <DerivativeTokens {...moves} t={t} /> | |
| </g> | |
| ); | |
| } | |
| interface FlowProps { | |
| graph: Graph; | |
| layout: GraphLayout; | |
| flow: Flow; | |
| timeline: Timeline; | |
| } | |
| export default function StepFlow(props: FlowProps) { | |
| const { step } = props.flow; | |
| if (step.phase === "forward") { | |
| return <ForwardFlow {...props} />; | |
| } | |
| return step.partials.length === 0 ? <SeedFlow {...props} /> : <BackwardFlow {...props} />; | |
| } | |