autodiff-dev / frontends /react /src /StepFlow.tsx
nhamid289's picture
Deploy autodiff ff35e9c
b9499ab verified
Raw History Blame Contribute Delete
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} />;
}